[Ray] concat_tables의 Happy Path를 최적화하여 동일 스키마 테이블 연결 가속화

PR 링크: ray-project/ray#61315 상태: Merged | 변경: +327 / -54

들어가며

Ray Data의 map_task에서는 여러 블록(PyArrow Table)을 하나로 합칩니다. 최악의 경우 블록들의 스키마가 서로 달라 통합이 필요하지만, 대부분의 경우(happy path) 모든 블록이 동일한 스키마를 가집니다. 기존에는 항상 커스텀 컬럼별 연결 로직을 수행하여 pa.concat_tables 대비 약 25배 느렸습니다. 이 PR은 동일 스키마일 때 네이티브 pa.concat_tables를 직접 사용하는 fast path를 추가합니다.

핵심 코드 분석

Before: 항상 수동 컬럼별 연결

def _concat_cols_with_native_pyarrow_types(
    col_names, blocks, promote_types=False
):
    # 각 컬럼별로 수동으로 ChunkedArray 연결
    # 모든 경우에 대해 동일한 느린 경로 사용
    ...

After: 스키마 동일성에 따른 분기

def _concat_cols_via_concat_tables(
    col_names, blocks, promote_types=False
):
    """동일 스키마 블록에 대해 pa.concat_tables 사용.

    네이티브 타입과 동일한 확장 타입 모두 지원.
    """
    # 선택된 컬럼만 추출하여 pa.concat_tables로 한 번에 연결
    selected = [block.select(col_names) for block in blocks]
    if promote_types:
        combined = pa.concat_tables(selected, promote_options="permissive")
    else:
        combined = pa.concat_tables(selected)
    return {name: combined.column(name) for name in col_names}

확장 타입(텐서, Python 객체)도 모든 블록에서 타입이 동일한 경우 fast path를 사용합니다:

# 같은 타입인 컬럼은 fast path(concat_tables)로 연결
# 다른 타입인 컬럼만 slow path(수동 연결)로 처리
# 마지막으로 두 결과를 합침

왜 이게 좋은가

• 벤치마크 결과: 3000개의 동일 텐서 스키마 테이블에서 기존 25배 느리던 것이 1.5배 수준으로 개선됩니다.
• 점진적 폴백: 모든 컬럼이 동일 타입이면 전부 fast path, 일부만 다르면 혼합 전략, 전부 다르면 기존 slow path를 사용합니다.
• 확장 타입 지원: ArrowTensorType, ArrowPythonObjectType 같은 확장 타입도 타입이 같으면 fast path를 탑니다.
• 상세한 문서화: 모든 내부 함수에 docstring과 예제가 추가되어 복잡한 타입 연결 로직을 이해하기 쉽습니다.

참고 자료

• PyArrow concat_tables
• Ray Data 성능 팁