[ray] Ray Data의 Join 너머 Predicate Pushdown: AND 분해와 OR에서 단측 조건 유도하기

PR 링크: ray-project/ray#67122 상태: Merged | 변경: +1200 / -194

들어가며

Ray Data의 logical optimizer는 Filter를 가능한 한 데이터 소스 쪽으로 내리는 predicate pushdown을 수행합니다. 그런데 Join 위에 놓인 Filter는 predicate 전체를 하나의 덩어리로 취급했습니다. 그 결과 다음과 같은 문제가 있었습니다.

• filter(left_pred & right_pred)처럼 양쪽 컬럼을 섞은 AND는 어느 한쪽에도 "전부" 속하지 않으므로 아무것도 push되지 않았습니다.
• 한쪽 조건만 가진 filter 두 개를 연달아 걸어도, optimizer가 이를 fuse하면 위와 같은 모양이 되어 역시 push되지 않았습니다.
• TPC-H q7의 (n1 == 'FRANCE' & n2 == 'GERMANY') | (n1 == 'GERMANY' & n2 == 'FRANCE') 같은 disjunction은 논리적으로 n1 in (FRANCE, GERMANY)와 n2 in (FRANCE, GERMANY)를 함축하지만, 이 사실을 활용하지 못해 nation 스캔이 25행을 모두 읽었습니다.

덧붙여 side 라우팅이 컬럼 이름만으로 이루어져, RIGHT_OUTER + right suffix 조합에서 col("val") == 1이 잘못된 입력으로 내려가 틀린 결과를 내는 기존 버그도 있었습니다. 이 PR은 이 세 가지를 함께 해결합니다.

코드 분석

1. pushdown_utils.py: predicate chain 유틸리티

핵심 빌딩 블록은 AND/OR 체인을 평탄화하고 다시 접는 두 함수입니다.

def flatten_predicate_chain(predicate: Expr, op: Operation) -> List[Expr]:
    operands: List[Expr] = []
    stack = [predicate]
    while stack:
        expr = stack.pop()
        if isinstance(expr, BinaryExpr) and expr.op == op:
            stack.append(expr.right)
            stack.append(expr.left)
        else:
            operands.append(expr)
    return operands

def combine_predicate_chain(predicates: List[Expr], op: Operation) -> Optional[Expr]:
    if not predicates:
        return None
    return functools.reduce(lambda left, right: BinaryExpr(op, left, right), predicates)

재귀 대신 명시적 stack을 사용해 매우 긴 AND 체인에서도 recursion limit에 걸리지 않으며, right를 먼저 push해 좌→우 순서를 보존합니다. 이 유틸 덕분에 기존 _split_predicate_by_columns가 크게 단순해졌습니다.

Before — 재귀적으로 좌/우를 나누고 combine_predicates로 다시 합침:

if isinstance(predicate, BinaryExpr) and predicate.op == Operation.AND:
    left_result = _split_predicate_by_columns(predicate.left, partition_columns)
    right_result = _split_predicate_by_columns(predicate.right, partition_columns)
    return _SplitPredicateResult(
        data_predicate=combine_predicates(
            left_result.data_predicate, right_result.data_predicate
        ),
        ...
    )

After — conjunct 목록으로 평탄화한 뒤 한 번의 루프로 분류:

for conjunct in flatten_predicate_chain(predicate, Operation.AND):
    referenced_cols = set(get_column_references(conjunct))
    if not (referenced_cols & partition_columns):
        data.append(conjunct)
    elif referenced_cols <= partition_columns:
        partition.append(conjunct)
    else:
        residual.append(conjunct)

이는 리뷰어 owenowenisme의 "이 함수들을 pushdown_utils.py로 옮기고 비슷한 함수(_combine_with_and 등)를 정리하자"는 피드백이 반영된 결과로, 같은 "conjunct 분해 → 분류 → 재조립" 패턴을 join pushdown과 partition pruning이 공유하게 되었습니다. append_unique_predicate는 structurally_equals로 중복 predicate가 쌓이지 않도록 합니다.

2. join_operator.py: 이름 기반에서 컬럼 소유권 기반으로

Before — 컬럼 이름 집합의 포함 관계로 "하나의" side를 결정:

can_evaluate_on_left = predicate_columns.issubset(
    left_columns
) or predicate_columns.issubset(left_join_keys)
if can_evaluate_on_left and can_push_left:
    return JoinSide.LEFT

이 방식은 출력에서 val_r로 바뀐 right 컬럼과 left의 val을 구분하지 못합니다. 그래서 RIGHT_OUTER 조인에서 출력의 val(실제로는 left 소유)에 대한 조건이 right 쪽 val로 내려가는 오류가 생겼습니다. 또한 반환값이 단일 side라 inner join의 공유 key 조건을 양쪽에 동시에 내릴 수 없었습니다.

After — 출력 컬럼 → 입력 컬럼 매핑을 물리 조인의 naming 규칙을 그대로 미러링해 계산합니다.

collisions = set(left_columns) & (set(right_columns) - right_keys)
left_renamed = collisions - left_keys if self.left_columns_suffix else set()
right_renamed = collisions if self.right_columns_suffix else set()

left = {
    output_name(c, left_renamed, self.left_columns_suffix): c
    for c in left_columns
}
right = {
    output_name(c, right_renamed, self.right_columns_suffix): c
    for c in right_columns
}
emitted_by_both = (left.keys() & right.keys()) - shared_keys

포인트는 세 가지입니다.

• suffix가 붙는 충돌 payload 컬럼을 정확히 추적해, push 시 predicate를 입력 이름으로 rebind합니다(val_r → val). Project의 rename을 통과할 때와 같은 방식입니다.
• 이름이 같은 join key(shared_keys)는 양쪽 값이 동일하므로 양쪽 모두에 push 가능합니다.
• 어느 쪽 소유인지 모호한 이름(emitted_by_both)은 claim하지 않아 filter가 join 위에 남습니다 — 안전한 쪽으로 실패합니다.

허용 여부는 join type 규칙만 남겨 allows_predicate_pushdown_to(side)로 분리했습니다. "어떤 컬럼이 어디서 왔나"와 "이 side에 내려도 의미가 보존되나"라는 두 관심사가 깔끔히 분리된 것입니다.

흥미롭게도 이 설계는 리뷰에서 다듬어졌습니다. 초기 구현은 이름 충돌 시 소유권을 unknown으로 처리했는데, 리뷰어가 "suffix는 이 op의 필드이므로 bare-name 소유권은 계산 가능하다. L.join(R, join_type=\"inner\", on=\"k\", right_suffix=\"_r\").filter(col(\"val\") == 1)은 left로 push되어야 한다"고 지적했고, 그 결과 현재의 suffix 미러링 로직과 test_join_colliding_payload_column_routes_by_suffix 테스트가 추가되었습니다.

3. DeriveJoinSidePredicates: OR에서 약한 단측 조건 유도

(A1 & B1) | (A2 & B2)가 참이면 A1 | A2도 반드시 참입니다(A가 left 컬럼만 참조할 때). 새 rule은 residual OR의 각 branch에서 한쪽 side 조건을 뽑아 == | == 형태로 묶어 push합니다. 원래 OR는 cross 조합(FRANCE-FRANCE 등)을 걸러내기 위해 join 위에 그대로 남습니다. IN 대신 == | ==를 쓰는 이유는 non-equality branch도 동일하게 다루기 위함입니다.

원래 OR가 남아 있으므로 반복 적용하면 같은 predicate를 계속 push할 위험이 있습니다. 그래서 이 rule은 optimizer loop가 아니라 LogicalOptimizer._post_optimize에서 한 번만 실행되고, _is_enforced_below로 read의 get_current_predicate()까지 내려가 이미 적용된 조건은 건너뜁니다. 이를 위해 ReadFiles.get_current_predicate가 partition pruning predicate도 포함하도록 수정되었습니다. 리뷰 과정에서 Filter에 "이미 유도된 predicate" 필드를 두어 무한 pushdown을 막는 방향으로 로직이 더 단순화되었습니다.

왜 이게 좋은가

TPC-H q7(sf1, 단일 머신 20 CPU)을 SQL 그대로(수동 nation 필터 없이) 실행한 결과입니다.

	master	PR
nation scan 출력 행	25	2
마지막 join 입력 행	6,001,215	478,523
마지막 shuffle map 시간	1.14 s	0.10 s
wall time	~380 s	~374 s

마지막 join으로 들어가는 행이 약 12.5배 줄었고 shuffle map 시간은 1/10 수준입니다. wall time 차이가 작은 것은 ~257초가 lineitem의 단일 머신 S3 읽기이기 때문이며, read가 병렬화되는 클러스터 sf100+ 환경에서는 shuffle 절감이 그대로 드러날 것으로 기대됩니다. 유도된 predicate가 select_columns와 두 번의 rename_columns를 거쳐 ReadFilesParquetV2까지 rebind되어 내려간 점도 인상적입니다.

일반적 교훈

1. predicate는 덩어리가 아니라 conjunct 집합이다. AND를 평탄화해 각 conjunct를 독립적으로 라우팅하는 것은 거의 모든 쿼리 엔진의 기본기입니다.
2. OR는 push할 수 없어도 "함축"은 push할 수 있다. 약한(weaker) 조건을 유도해 내리고 원 조건을 위에 남기는 것은 정확성을 해치지 않으면서 데이터량을 줄이는 고전적 기법입니다.
3. 이름이 아니라 소유권(lineage)으로 라우팅하라. 이름 기반 판단은 suffix/rename이 끼는 순간 틀린 결과를 냅니다. 물리 연산자의 naming 규칙을 logical 계층에서 그대로 미러링하고, 모호하면 push하지 않는 보수적 기본값이 중요합니다.
4. 비멱등 rewrite는 수렴 조건을 설계하라. 원 predicate를 남기는 rule은 "이미 적용되었는가" 검사가 없으면 무한 반복합니다.

테스트 측면에서도 리뷰어가 dag_str 비교는 scanner의 predicate가 보이지 않고 60자 이상에서 잘릴 수 있어 신뢰할 수 없다고 지적하며, plan을 post-order로 순회해 Filter.predicate_expr와 get_current_predicate()를 structurally_equals로 비교하는 helper를 제안한 점도 optimizer 테스트 작성 시 참고할 만한 조언입니다.

참고 자료

• Ray Data Joins — 이 PR이 pushdown을 확장한 Dataset join 가이드
• ray.data.Dataset.filter — join 위로 push되는 Filter를 만드는 API
• ray.data.Dataset.join — suffix와 join type 규칙이 정의된 join API
• functools.reduce — combine_predicate_chain에서 predicate 체인을 접는 데 사용

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