Fix silent row drops in multi-GPU joins with computed key expressions - #22318
Conversation
github-actions
Bot
added
Python
Workaround for #22317 The issue seems to be that we silently drop rows in multi-rank joins (described more in #22318 (comment)). This is why the sum changes non-deterministically from run-to-run. But the changes in that PR don't seem to be sufficient to validate Q8. Therefore, the change in this PR is a workaround until we figure where else we could be dropping rows. Authors: - Matthew Murray (https://github.com/Matt711) Approvers: - Tom Augspurger (https://github.com/TomAugspurger) URL: #22473
c9b22a0 to
e6a6d89
Compare
e6a6d89 to
7a43ab9
Compare
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a ChangesComputed Join Key Handling
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested Reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cudf_polars/tests/streaming/test_join.py (1)
393-447: ⚡ Quick winCover the symmetric computed-key path (
left_onexpression) in this test.Line 393 currently validates only computed
right_on. Since the planner change applies to both sides, parameterizing this test to also exercise computedleft_onwould better guard against one-sided regressions.Suggested test update
+@pytest.mark.parametrize("computed_side", ["left", "right"]) -def test_join_computed_expr_right_key(streaming_engine_factory) -> None: - """Join on a computed key expression.""" +def test_join_computed_expr_key(streaming_engine_factory, computed_side) -> None: + """Join on a computed key expression from either side.""" @@ - # Now join on a computed key expression. - # This should not silently drop rows across ranks - q = left.join( - right, - left_on="zip_prefix", - right_on=pl.col("full_zip").str.slice(0, 2), - ) + # Now join on a computed key expression. + # This should not silently drop rows across ranks. + if computed_side == "right": + q = left.join( + right, + left_on="zip_prefix", + right_on=pl.col("full_zip").str.slice(0, 2), + ) + else: + q = right.join( + left, + left_on=pl.col("full_zip").str.slice(0, 2), + right_on="zip_prefix", + ) assert_gpu_result_equal(q, engine=engine, check_row_order=False)As per coding guidelines:
python/cudf_polars/**/test_*.py: verify tests compare GPU results against Polars CPU results and cover all supported expression types.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/streaming/test_join.py` around lines 393 - 447, The test test_join_computed_expr_right_key only exercises a computed expression on right_on; update it to also exercise the symmetric case where the computed expression is on left_on (e.g., call left.join(right, left_on=pl.col("zip_prefix").str.slice(0, 2), right_on="full_zip") in addition to the existing join), and ensure both queries (the original q and the new symmetric one) are passed to assert_gpu_result_equal with the same engine and check_row_order=False so the GPU/CPU results for both computed-key paths are compared.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python/cudf_polars/tests/streaming/test_join.py`:
- Around line 393-447: The test test_join_computed_expr_right_key only exercises
a computed expression on right_on; update it to also exercise the symmetric case
where the computed expression is on left_on (e.g., call left.join(right,
left_on=pl.col("zip_prefix").str.slice(0, 2), right_on="full_zip") in addition
to the existing join), and ensure both queries (the original q and the new
symmetric one) are passed to assert_gpu_result_equal with the same engine and
check_row_order=False so the GPU/CPU results for both computed-key paths are
compared.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 610e31f4-4ad3-433d-a6ff-46b178ccbdc8
📒 Files selected for processing (4)
python/cudf_polars/cudf_polars/dsl/utils/naming.pypython/cudf_polars/cudf_polars/streaming/actor_graph/join.pypython/cudf_polars/tests/dsl/test_naming.pypython/cudf_polars/tests/streaming/test_join.py
✅ Files skipped from review due to trivial changes (1)
- python/cudf_polars/tests/dsl/test_naming.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cudf_polars/tests/streaming/test_groupby.py (1)
235-253: ⚡ Quick winAdd one end-to-end streaming regression case for mixed concrete/computed GroupBy keys.
This validates
_key_indicesdirectly, but it doesn’t verify distributed runtime correctness. Please add anassert_gpu_result_equal(...)case with mixed concrete/computed keys under streaming settings that exercise partition-planning decisions.As per coding guidelines "In cudf_polars: verify tests compare GPU results against Polars CPU results and cover all supported expression types."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/tests/streaming/test_groupby.py` around lines 235 - 253, The test only checks _key_indices statically; add an end-to-end streaming regression that runs the GroupBy with mixed concrete/computed keys through the GPU runtime and compares to Polars CPU using assert_gpu_result_equal under streaming settings: construct a small input DataFrame matching the GroupBy schema used in test_groupby_key_indices_concrete_prefix (use GroupBy with NamedExpr("a", Col), NamedExpr("b", Literal), NamedExpr("c", Col)), enable streaming/partitioning flags used elsewhere in tests, execute the GPU pipeline and call assert_gpu_result_equal(...) against the equivalent Polars expression to validate partition-planning/runtime correctness while still keeping the existing static _key_indices assertions. Ensure the test references the same symbols (GroupBy, _key_indices, expr.NamedExpr, ErrorNode, assert_gpu_result_equal) so it exercises mixed concrete/computed keys end-to-end.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python/cudf_polars/tests/streaming/test_groupby.py`:
- Around line 235-253: The test only checks _key_indices statically; add an
end-to-end streaming regression that runs the GroupBy with mixed
concrete/computed keys through the GPU runtime and compares to Polars CPU using
assert_gpu_result_equal under streaming settings: construct a small input
DataFrame matching the GroupBy schema used in
test_groupby_key_indices_concrete_prefix (use GroupBy with NamedExpr("a", Col),
NamedExpr("b", Literal), NamedExpr("c", Col)), enable streaming/partitioning
flags used elsewhere in tests, execute the GPU pipeline and call
assert_gpu_result_equal(...) against the equivalent Polars expression to
validate partition-planning/runtime correctness while still keeping the existing
static _key_indices assertions. Ensure the test references the same symbols
(GroupBy, _key_indices, expr.NamedExpr, ErrorNode, assert_gpu_result_equal) so
it exercises mixed concrete/computed keys end-to-end.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 59c4f45a-d0a9-457b-a77b-b8a7a2cc312b
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.pypython/cudf_polars/tests/streaming/test_groupby.py
|
|
||
| Uses ``hash_partition(input, key_table, ...)`` to support | ||
| non-``Col`` (e.g. expression-derived) shuffle keys. |
There was a problem hiding this comment.
| Uses ``hash_partition(input, key_table, ...)`` to support | |
| non-``Col`` (e.g. expression-derived) shuffle keys. |
But please describe the parameters.
| concrete: list[str] = [] | ||
| for n in names: | ||
| if isinstance(n, str): | ||
| concrete.append(n) | ||
| elif isinstance(n.value, Col): | ||
| concrete.append(n.value.name) | ||
| else: | ||
| break | ||
| return tuple(concrete) |
There was a problem hiding this comment.
return tuple(
n if isinstance(n, str) else n.value.name
for n in itertools.takewhile(
lambda n: isinstance(n, str) or isinstance(n.value, Col),
names
)
)
If we fancy being fancy
| concrete_prefix | ||
| If True, use only the prefix of names corresponding | ||
| to concrete column references. If False (default), | ||
| use all names. |
There was a problem hiding this comment.
Under what circumstances to we not want this?
names_to_indices is used to convert name references into a column indices of a table. So by definition, I think, it can't be used to if the namedexpr isn't referring to a column?
There was a problem hiding this comment.
It depends on whether the schema we are indexing on corresponds the input or the output of the expressions in names.
If the schema references the output DataFrame, then it's fine for the expressions to be non-concrete - The output of the expressions are concrete columns. If the schema references the input DataFrame, then the expression must be concrete.
When we check if the input DataFrame is already partitioned correctly, we must pass in this concrete_prefix=True option.
| assert not _use_pwise_join(executor, partition_info, join_ir) | ||
|
|
||
|
|
||
| def test_join_computed_expr_right_key(streaming_engine_factory) -> None: |
There was a problem hiding this comment.
I can't remember exactly how we run these tests multi-rank. But we should ensure this test is run multi-rank.
There was a problem hiding this comment.
Yes, the Ray variation will run two ranks on the same visible device.
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
Approving assuming the detail test is removed.
|
/merge |
6 participants
Description
When a join key is a computed expression (e.g.
pl.col("ca_zip").str.slice(0, 2)), the join actor incorrectly inferred that the side using that key was already partitioned on it. The join actor derived the column index from the expression's output name, then checked the upstream partitioning metadata against that index. But the data was partitioned on the raw column value, not the derived one. When both sides happened to have the same shuffle modulus from prior joins, the planner chose a chunkwise join and skipped re-shuffling, so matching rows on different ranks were never paired and the result was silently missing rows.The fix passes
Noneas the partitioning metadata for any side whose join key is a computed expression, so the chunkwise path is never taken for those joins.For context: I found this while running TPC-DS Q8 (multi-GPU). It was failing validation ~2-4 out of 15 runs producing wrong row counts.
Closes #22317
Closes #22105
Contributes to #21813
Checklist