[Ray Data] PyArrow 스키마 해싱 방식 개선으로 대규모 데이터셋 성능 향상

PR 링크: ray-project/ray#62108 상태: Merged | 변경: +33 / -20

들어가며

Ray Data에서 RefBundle은 데이터 블록의 참조를 관리하는 핵심 자료구조입니다. 이 구조체의 __eq__와 __hash__에서 PyArrow 스키마의 전체 동등성 비교와 해싱을 수행하고 있었는데, 대규모 스키마에서는 이 연산이 상당히 비쌌습니다. 또한 이미지 데이터셋처럼 수천 개의 입력 파일이 있는 경우, input_files 목록이 DatasetStats로 전파되어 불필요한 메모리를 소비하고 있었습니다.

핵심 코드 분석

스키마 동등성 비교: 값 비교에서 객체 ID 비교로

Before:

def __eq__(self, other: "RefBundle"):
    return (
        self.blocks == other.blocks
        and self.slices == other.slices
        and self.schema == other.schema  # 전체 스키마 값 비교
        and self.owns_blocks == other.owns_blocks
    )

After:

def __eq__(self, other: "RefBundle"):
    return (
        self.blocks == other.blocks
        and self.slices == other.slices
        and self.schema is other.schema  # 객체 동일성 비교
        and self.owns_blocks == other.owns_blocks
    )

스키마 해싱: 값 해싱에서 객체 ID 해싱으로

Before:

def __hash__(self) -> int:
    return hash((
        *self.blocks,
        *self.slices,
        _make_hashable_schema(self.schema) if self.schema is not None else None,
        self.owns_blocks,
    ))

After:

def __hash__(self) -> int:
    return hash((
        *[b for b, _ in self.blocks],  # 블록 참조만 해싱
        *self.slices,
        id(self.schema),  # 객체 ID로 해싱
        self.owns_blocks,
    ))

입력 파일 목록 제거

# DatasetStats에 전달 시 input_files를 제거
stats = DatasetStats(
    metadata={
        "Read": [
            replace(read_task.metadata, input_files=None)
            for read_task in read_tasks
        ]
    },
)

왜 이게 좋은가

• 해싱 비용 O(1)로 감소: _make_hashable_schema는 스키마의 모든 이름과 타입을 튜플로 변환했지만, id()는 상수 시간입니다.
• 대규모 스키마에서 효과적: 이미지 데이터셋처럼 컬럼이 많은 스키마에서 동등성 비교와 해싱 비용이 크게 줄어듭니다.
• 메모리 절감: 수천 개의 입력 파일 경로 문자열이 DatasetStats에 복사되지 않습니다.
• 안전한 가정: RefBundle의 스키마가 동일한 파이프라인 내에서 같은 객체를 공유한다는 가정은, Ray Data의 실제 사용 패턴과 일치합니다.

참고 자료

• PR #62108 - Revisiting hashing method of Pyarrow schemas to improve perf