[ray] Ray Data의 차세대 데이터 소스 API: DataSourceV2 설계 및 최적화 전략

PR 링크: ray-project/ray#61615 상태: Merged | 변경: +608 / -1

들어가며

Ray Data는 분산 데이터 처리를 위한 강력한 프레임워크입니다. 하지만 기존의 데이터 소스 추상화는 다양한 데이터 소스(파일, DB, 데이터 레이크 등)의 특수한 최적화 기법을 통합하는 데 한계가 있었습니다. 이번에 도입된 DataSourceV2는 데이터 소스 읽기 과정을 Indexer, Scanner, Reader로 명확히 분리하고, Mixin을 통해 필터 푸시다운(Filter Pushdown)이나 컬럼 프루닝(Column Pruning) 같은 최적화를 선언적으로 지원하도록 설계되었습니다.

코드 분석

1. DataSourceV2의 핵심 추상화

DataSourceV2는 데이터 소스의 진입점입니다. 기존의 모호했던 InputBucket이라는 용어를 InputSplit으로 변경하여 데이터 분할의 의미를 명확히 했습니다.

# Before
InputBucket = TypeVar("InputBucket")

# After
InputSplit = TypeVar("InputSplit")

DataSourceV2 클래스는 데이터 소스의 카테고리를 정의하여 실행 엔진이 적절한 최적화 전략을 선택할 수 있게 합니다.

@DeveloperAPI
class DataSourceV2(ABC, Generic[InputSplit]):
    def __init__(self, name: str, category: DatasourceCategory):
        self._name = name
        self._category = category
    @abstractmethod
    def create_scanner(self, schema: pa.Schema, ...) -> Scanner[InputSplit]:
        ...

2. 논리적 최적화 Mixin

가장 강력한 변화는 logical_optimizers.py의 도입입니다. 이제 각 데이터 소스는 자신이 지원하는 최적화 기능을 Mixin을 통해 명시적으로 선언할 수 있습니다.

@DeveloperAPI
class SupportsColumnPruning(ABC):
    @abstractmethod
    def prune_columns(self, columns: List[str]) -> "Scanner[InputSplit]":
        ...

이를 통해 실행 엔진은 런타임에 isinstance(scanner, SupportsColumnPruning)를 확인하여 자동으로 컬럼 프루닝을 적용할 수 있습니다.

3. Reader와 Scanner의 분리

Scanner는 읽기 계획을 담당하고, Reader는 실제 워커에서 데이터를 읽는 역할을 수행합니다. FileReader는 PyArrow의 Dataset API를 활용하여 효율적인 읽기를 보장합니다.

왜 이게 좋은가

1. 모듈화(Modularity): 인덱싱, 계획 수립, 실행이 분리되어 유지보수가 용이합니다.
2. 선언적 최적화: SupportsFilterPushdown과 같은 인터페이스를 통해 데이터 소스별로 최적화 기능을 쉽게 추가할 수 있습니다.
3. 성능 향상: 컬럼 프루닝과 필터 푸시다운을 통해 불필요한 I/O를 획기적으로 줄일 수 있습니다. 특히 Parquet와 같은 컬럼형 포맷에서 전체 데이터셋을 읽지 않고 필요한 부분만 읽어오는 최적화가 가능해집니다.

리뷰어 피드백 반영

리뷰 과정에서 InputBucket이라는 명칭이 다소 모호하다는 지적이 있었고, 이를 InputSplit으로 변경함으로써 데이터 분할(Partitioning)의 의미를 더 명확히 했습니다. 또한, 모든 API에 @DeveloperAPI 어노테이션을 추가하여 안정성 수준을 명시했습니다.

결론

DataSourceV2는 Ray Data가 단순한 데이터 로더를 넘어, 고도로 최적화된 분산 쿼리 엔진으로 진화하고 있음을 보여줍니다. 향후 다양한 데이터 소스(Iceberg, Delta Lake 등)를 이 프레임워크 위에서 구현하면 일관된 성능 최적화 혜택을 누릴 수 있을 것입니다.

참고 자료

• https://arrow.apache.org/docs/python/dataset.html
• https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.