[Ray] Parquet 배치 크기를 C++ 32비트 정수 범위로 클램핑하여 OverflowError 수정

PR 링크: ray-project/ray#62242 상태: Merged | 변경: +115 / -0

들어가며

Ray Data가 Parquet 파일을 읽을 때, 행당 바이트 수가 매우 작은 경우(희소하거나 고도로 압축된 데이터) 추정된 배치 크기가 2^31을 초과할 수 있습니다. PyArrow의 ParquetFileFragment.to_batches()는 batch_size를 Cython 레이어를 통해 C의 int(32비트)로 전달하므로, 이 범위를 초과하면 OverflowError: value too large to convert to int가 발생합니다.

핵심 코드 분석

클램핑 함수 추가

_MAX_PYARROW_TO_BATCHES_BATCH_SIZE = 2**31 - 1

def _coerce_pyarrow_fragment_batch_size(batch_size: int) -> int:
    """Clamp batch size for ParquetFileFragment.to_batches to PyArrow's C int range."""
    if batch_size <= 0:
        raise ValueError(f"Batch size must be > 0, got {batch_size}")
    coerced = min(batch_size, _MAX_PYARROW_TO_BATCHES_BATCH_SIZE)
    if coerced != batch_size:
        logger.debug(
            "Clamping Parquet fragment read batch_size from %s to %s "
            "(PyArrow ``to_batches`` requires batch_size in [1, %s]).",
            batch_size, coerced, _MAX_PYARROW_TO_BATCHES_BATCH_SIZE,
        )
    return coerced

호출 지점에 적용

def _read_batches_from(fragment, *, schema, ...):
    # ...
    if to_batches_kwargs.get("batch_size") is not None:
        to_batches_kwargs["batch_size"] = _coerce_pyarrow_fragment_batch_size(
            int(to_batches_kwargs["batch_size"])
        )

int()로 감싸서 numpy.int64 같은 타입도 안전하게 처리합니다.

왜 이게 좋은가

1. 크래시 방지: 희소 데이터셋에서 발생하던 OverflowError를 방어적 클램핑으로 해결합니다.
2. 투명한 동작: 클램핑이 발생하면 debug 로그를 남기되, 정상 범위의 배치 크기에는 아무런 영향을 주지 않습니다.
3. 포괄적 테스트: 2^31, 10^12, 0, 음수, 정상값, numpy.int64 등 다양한 엣지 케이스를 테스트로 커버합니다.
4. 근본 원인 문서화: 상수와 함수에 왜 이 제한이 필요한지(PyArrow Cython C int 제약)를 명확하게 문서화했습니다.

참고 자료

• PyArrow ParquetFileFragment API
• Ray Data GitHub