[ray] Ray Data의 sort().limit(k)를 Top-N 실행으로 바꾸기: 전체 셔플 없이 상위 k개만 고르는 법

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

들어가며

ds.sort("key").limit(100)은 의미상으로는 "상위 100개만 달라"는 요청입니다. 그런데 기존 Ray Data(SHUFFLE_V2 경로)는 이 요청을 받아도 전체 정렬부터 수행했습니다.

1. SortSamplingOp가 boundary를 샘플링하고
2. SortShuffleMapOp가 모든 row를 정렬한 뒤 파티션으로 나누고
3. ShuffleReduceOp가 파티션을 병합한 다음
4. 마지막으로 LimitOperator가 앞에서 k개를 자릅니다.

Limit이 소비를 일찍 멈출 수는 있지만, map 단계에서 모든 입력 row는 이미 정렬되고 셔플됩니다. k가 100이고 입력이 수억 row라면 대부분의 작업이 버려지는 셈입니다.

이 PR(#66831 해결)은 TopNRewriteRule을 추가해 조건에 맞는 Sort + Limit plan을 Top-N 실행으로 바꿉니다.

• LocalTopN: 각 block에서 최대 k개의 후보만 남깁니다.
• TopNMergeOperator: 후보 batch를 받아 running Top-N과 병합하는 작업을 반복하고, 입력이 끝나면 정렬된 결과를 내보냅니다.

기능은 private이고 기본값은 off입니다. 적용 조건은 다음과 같습니다.

• k > 0
• custom boundary가 없을 것
• k가 설정 가능한 최대값(기본 10,000) 이하일 것

코드 분석

1. arrow_block.py — block 단위 Top-N

기존 ArrowBlockAccessor에는 block 전체를 정렬하는 sort()만 있었습니다.

# Before
def sort(self, sort_key: "SortKey") -> Block:
    ...
    return sort(self._table, sort_key)

PR은 여기에 top_n()을 추가했습니다. 핵심 경로는 PyArrow의 select_k_unstable입니다.

# After
try:
    indices = pac.select_k_unstable(
        self._table,
        k=result_size,
        sort_keys=sort_key.to_arrow_sort_args(),
    )
except pyarrow.ArrowNotImplementedError:
    return sort_and_take()

return self.take(indices)

select_k_unstable은 전체를 정렬하지 않고 상위 k개의 index만 뽑습니다. 비용은 대략 O(n log k)이므로 k가 n보다 훨씬 작을 때 full sort(O(n log n))보다 유리합니다. 리뷰어(owenowenisme)가 "select_k_unstable을 확인해 block 전체 정렬을 피하라"고 제안했고, 그 피드백이 반영된 부분입니다.

흥미로운 부분은 null/NaN 처리입니다.

primary_column = self._table[sort_key.get_columns()[0]]
# SelectK excludes rows whose primary key is null or NaN, while a full
# sort places them at the end. Fall back if the result must include them.
excluded_primary_values = primary_column.null_count

if pyarrow.types.is_floating(primary_column.type):
    nan_mask = pac.is_nan(primary_column)
    nan_count = pac.sum(nan_mask).as_py()
    excluded_primary_values += nan_count or 0

if excluded_primary_values > self._table.num_rows - result_size:
    return sort_and_take()

select_k는 primary key가 null 또는 NaN인 row를 결과에서 아예 빼버립니다. 반면 full sort는 이런 row를 끝에 배치합니다. 유효 값의 개수가 k보다 적다면 결과에 null/NaN row가 포함되어야 하는데, 이 경우를 미리 세어 fallback합니다. 최적화가 의미(semantics)를 바꾸지 않도록 지키는 장치입니다.

fallback 경로인 sort_and_take도 신경을 썼습니다.

# Sort only the keys to avoid materializing the full payload.
keys = self._table.select(sort_key.get_columns())
...
indices = pac.sort_indices(keys, sort_keys=sort_key.to_arrow_sort_args())
return self.take(indices.slice(0, result_size))

전체 테이블을 정렬하지 않습니다. key 컬럼만 sort_indices로 정렬한 다음 상위 result_size개 index로 take합니다. 128바이트 payload 같은 wide row에서 불필요한 복사를 피할 수 있습니다.

Polars 정렬 경로에서는 __top_n_row_index 컬럼을 덧붙여 원래 row 위치를 보존합니다. 마지막에는 polars.from_arrow(result).to_arrow()로 Polars의 schema 변환을 그대로 재현해 기존 sort 결과와 타입을 맞춥니다.

리뷰 중 Polars 정렬 결과 불일치 문제가 제기되었습니다. 이는 기존부터 있던 문제이고, 결과도 실제로는 정렬되어 있어 "false alarm"으로 정리되었습니다.

2. top_n_operator.py — 점진적 병합 연산자

초기 버전은 AllToAllOperator 기반이었고, 모든 후보를 한 번에 fetch해 병합했습니다. 리뷰에서 두 가지 지적이 나왔습니다.

• AllToAllOperator는 곧 deprecated된다. 따라서 PhysicalOperator 기반의 native operator로 만들어야 한다.
• 후보를 한 번에 모으지 말고 부분적으로 가져와 running Top-N을 유지해야 한다.

최종 구현은 이 두 지적을 반영했습니다.

_MERGE_BATCH_BYTES = 128 * 1024 * 1024
_MERGE_BATCH_BLOCKS = 32
_MERGE_BATCH_WAIT_S = 0.1

@dataclass
class _MergeState:
    task: Optional[DataOpTask] = None
    inputs: List[RefBundle] = field(default_factory=list)
    output: Optional[RefBundle] = None

동시에 실행되는 merge task는 하나뿐입니다. 이전 merge 결과(output)는 다음 batch와 함께 다시 입력으로 들어갑니다.

if self._merge.output is not None:
    state = self._merge.output
    self._merge.output = None
    self._metrics.on_input_dequeued(state, input_index=0)
    self._merge.inputs.append(state)

이 구조 덕분에 merge 단계의 peak 메모리는 "batch 하나 + running Top-N"으로 제한됩니다. 전체 후보 수에는 영향을 받지 않습니다.

batch를 언제 flush할지는 _batch_ready가 결정합니다.

return (
    self._inputs_complete
    or self._input_buffer.num_blocks() >= _MERGE_BATCH_BLOCKS
    or self._input_buffer.estimate_size_bytes() >= _MERGE_BATCH_BYTES
    or monotonic() - self._pending_since >= _MERGE_BATCH_WAIT_S
)

flush 조건은 32 block, 128 MiB, 100 ms 중 하나라도 충족되는 경우입니다. 시간 조건은 입력이 느리게 들어올 때 partial batch가 무한정 대기하는 것을 막습니다.

can_add_input은 버퍼가 임계치 미만일 때만 입력을 받아 backpressure 역할을 합니다. 주석대로 bundle 하나만큼 초과할 수 있으므로 hard limit은 아닙니다.

merge task는 DataOpTask로 감쌉니다. 그 결과 resource accounting, backpressure, input cleanup을 기존 executor 인프라에 그대로 맡길 수 있습니다.

왜 이게 좋은가

성능 수치

PR 본문 벤치마크(SplitMix64 key, 정렬 후 limit) 결과입니다.

조건	기존	Top-N	속도 향상
1M / 128 blocks / k=100	3.838s	1.429s	2.69x
5M / 128 blocks / k=100	4.949s	1.602s	3.09x

worker USS 메모리는 약 50% 감소했습니다.

리뷰어 rayo5o는 32-CPU 단일 노드에서 warm-up 후 5회 median으로 측정했습니다.

조건	기존	Top-N	속도 향상
10M / 2048 blocks / k=1000	61.18s	9.26s	6.6x

추가 RSS는 5.6 GiB에서 284 MiB로 줄었습니다.

점진적 병합으로 바뀐 버전에서는 결과가 더 좋아졌습니다. 100M row + 128B payload 기준으로 다음과 같습니다.

항목	sort+limit	Top-N
시간	21.4s	3.5s
extra worker 메모리	16.3 GiB	0.5 GiB

리뷰가 바꾼 설계

• 런타임 실패 제거: 초기 버전은 plan 시점에 max_limit을 검사하고, 런타임에 다시 merge budget을 검사했습니다. 그래서 k=10000에 256 block(후보 2.56M row)이면 ValueError가 발생했습니다. 점진적 병합 도입 후 해당 케이스가 정상 완료됩니다.
• stats 누락 수정: TopNMerge의 wall time이 3us로 표시되던 문제가 고쳐졌습니다.
• max_limit 기준 논의: Spark(topKSortFallbackThreshold), DuckDB(limit 5,000과 입력 대비 0.7% 비율), DataFusion의 접근 방식이 비교되었습니다. 그 결과 고정 임계치보다 후보량 기반 budget이 적절하다는 방향이 제시되었습니다.
• 남은 trade-off: block 수가 적은 소규모 입력(32 blocks)에서는 오히려 약간 느립니다(0.87x). 또 object store peak가 6.2 GiB에서 9.0 GiB로 늘어난 사례가 있습니다. 기본값을 off로 둔 판단이 합리적인 이유입니다.

일반적 교훈

1. 연산자 조합을 의미 단위로 재작성하라. Sort + Limit은 각각 최적화할 대상이 아니라, 하나의 Top-N 문제로 보아야 합니다.
2. 후보를 먼저 줄이고 나중에 모아라. 각 파티션에서 k개만 남기면 셔플 데이터량이 O(N)에서 O(blocks × k)로 줄어듭니다.
3. 최적화가 semantics를 깨지 않는지 확인하라. null/NaN 처리 차이처럼 API의 미묘한 동작을 짚고 fallback을 두어야 합니다.
4. 상태를 bounded하게 유지하라. running Top-N과 batch 단위 병합은 전체 입력 크기와 무관하게 메모리를 묶어 둡니다.

참고 자료

• pyarrow.compute.select_k_unstable — block 단위 Top-N 선택에 사용된 함수
• pyarrow.compute.sort_indices — fallback 경로에서 key만 정렬할 때 사용
• pyarrow.compute.is_nan — NaN 개수 계산에 사용
• ray.data.Dataset.sort — 재작성 대상 연산
• ray.data.Dataset.limit — 재작성 대상 연산

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