More Polars plan optimizations for TPC-DS - #22395
Merged
rapids-bot[bot] merged 4 commits intoMay 19, 2026
Merged
Conversation
Matt711
added
Performance
github-actions
Bot
added
Python
Matt711
marked this pull request as ready for review
May 13, 2026 18:11
📝 WalkthroughSummary by CodeRabbit
WalkthroughAcross 16 TPC-DS benchmark queries, the Polars implementations shift from "join full dimension/fact tables then apply filters" to "pre-filter and project dimensions early, use semi-joins to reduce cardinality, derive key pairs, and reorder aggregation stages for earlier row reduction". No exported signatures change; all modifications are internal optimization refactors. ChangesPolars Query Optimization Pattern
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.py`:
- Line 121: The ranking uses Polars' ordinal method which gives unique ranks and
diverges from SQL RANK() semantics; update both occurrences where
pl.col("avg_profit").rank(method="ordinal").alias("rnk") is used (and any
subsequent filters like rnk < 11) to use method="min" instead so tied avg_profit
values receive the same rank with gaps (matching DuckDB/SQL RANK behavior).
In `@python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.py`:
- Line 177: The aggregation pl.col("ext_sales_price").sum().alias("sales_amt")
returns 0 for all-null groups in Polars, which differs from SQL/DuckDB; replace
this with a conditional aggregation that checks
pl.col("ext_sales_price").count() > 0 and only returns the sum when count>0,
otherwise returns None, mirroring the null-sum handling used in q1/q49 so that
the "sales_amt" column matches SQL semantics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4b334632-2668-4f86-af10-2ff2b1dba11b
📒 Files selected for processing (19)
python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py
Matt711
force-pushed
the
imp/pdsds/more-pds-optimizations
branch
from
May 14, 2026 02:02
e94077b to
e1bb3be
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py (1)
121-121: ⚡ Quick winDeduplicate the semi-join probe keys.
Same pattern here:
sr_customer_itemis acting as an existence set for two semi joins, so keeping duplicate(sr_customer_sk, sr_item_sk)rows only makes the join builds heavier. A.unique()should make the prefilter cheaper without changing semantics.♻️ Proposed change
- sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"]) + sr_customer_item = store_returns_filtered.select( + ["sr_customer_sk", "sr_item_sk"] + ).unique()🤖 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/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py` at line 121, sr_customer_item currently keeps duplicate (sr_customer_sk, sr_item_sk) rows used as an existence set for two semi-joins; make the prefilter cheaper by deduplicating it. Update the assignment for sr_customer_item (from store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"])) to call .unique() (or the equivalent drop_duplicates()) on the resulting frame so duplicates are removed before using sr_customer_item in the semi-joins.python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py (1)
126-126: ⚡ Quick winDeduplicate the semi-join probe keys.
sr_customer_itemis only used as the RHS ofhow="semi"joins, so duplicate(sr_customer_sk, sr_item_sk)pairs cannot change results but can still bloat both downstream join builds. Adding.unique()here should reduce the work in the hottest part of this rewrite.♻️ Proposed change
- sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"]) + sr_customer_item = store_returns_filtered.select( + ["sr_customer_sk", "sr_item_sk"] + ).unique()🤖 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/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py` at line 126, sr_customer_item currently selects ["sr_customer_sk","sr_item_sk"] from store_returns_filtered but is only used as the RHS of semi-joins (how="semi"), so duplicate (sr_customer_sk, sr_item_sk) pairs only bloat downstream joins; change the creation to deduplicate the probe keys by applying .unique() to the selection (i.e., replace the current store_returns_filtered.select([...]) usage for sr_customer_item with store_returns_filtered.select([...]).unique()) so the semi-join builds operate on distinct keys.
🤖 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/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py`:
- Line 121: sr_customer_item currently keeps duplicate (sr_customer_sk,
sr_item_sk) rows used as an existence set for two semi-joins; make the prefilter
cheaper by deduplicating it. Update the assignment for sr_customer_item (from
store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"])) to call
.unique() (or the equivalent drop_duplicates()) on the resulting frame so
duplicates are removed before using sr_customer_item in the semi-joins.
In `@python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py`:
- Line 126: sr_customer_item currently selects ["sr_customer_sk","sr_item_sk"]
from store_returns_filtered but is only used as the RHS of semi-joins
(how="semi"), so duplicate (sr_customer_sk, sr_item_sk) pairs only bloat
downstream joins; change the creation to deduplicate the probe keys by applying
.unique() to the selection (i.e., replace the current
store_returns_filtered.select([...]) usage for sr_customer_item with
store_returns_filtered.select([...]).unique()) so the semi-join builds operate
on distinct keys.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 800a9d03-24c5-465a-8998-3ee5af1eaa20
📒 Files selected for processing (19)
python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.pypython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py
🚧 Files skipped from review as they are similar to previous changes (17)
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.py
- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.py
Member
Author
|
/ok to test f1a800d |
Member
Author
|
/ok to test 6a216d1 |
… imp/pdsds/more-pds-optimizations
Matt711
force-pushed
the
imp/pdsds/more-pds-optimizations
branch
from
May 19, 2026 15:02
6a216d1 to
ed3e928
Compare
Member
Author
|
/ok to test ed3e928 |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py (1)
180-195:⚠️ Potential issue | 🟠 Major | ⚡ Quick winJoin to
date_dimmay produce duplicate rows perd_week_seq.
date_dimcontains ~7 rows perd_week_seq(one per day). Sincewswscshas one row perd_week_seq, joining it to the fulldate_dimond_week_seqwill duplicate each aggregated row ~7 times. The subsequent.filter(pl.col("d_year") == year)doesn't deduplicate.To preserve semantics, join against a deduplicated mapping:
Proposed fix
+ # One row per week: pick the year of that week (any day in the week has the same d_year). + week_to_year = date_dim.select(["d_week_seq", "d_year"]).unique(subset=["d_week_seq"]) + # Step 3: Create year data (y subquery equivalent) y_year = ( - wswscs.join(date_dim, left_on="d_week_seq", right_on="d_week_seq") + wswscs.join(week_to_year, on="d_week_seq") .filter(pl.col("d_year") == year) .select( [ pl.col("d_week_seq").alias("d_week_seq1"),The same issue applies to
z_year_plus_1at lines 197-212.🤖 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/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py` around lines 180 - 195, The join between wswscs and date_dim (used to build y_year and z_year_plus_1) currently multiplies rows because date_dim has multiple rows per d_week_seq; instead create a deduplicated mapping of date_dim keyed by d_week_seq with the d_year (e.g., select d_week_seq and d_year and call unique()/drop_duplicates()) and join wswscs to that deduped Date mapping before filtering by pl.col("d_year") == year (and similarly for the z_year_plus_1 path), ensuring you reference the variables/expressions y_year, z_year_plus_1, wswscs, date_dim and the join key d_week_seq when making the change.
🤖 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.
Outside diff comments:
In `@python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py`:
- Around line 180-195: The join between wswscs and date_dim (used to build
y_year and z_year_plus_1) currently multiplies rows because date_dim has
multiple rows per d_week_seq; instead create a deduplicated mapping of date_dim
keyed by d_week_seq with the d_year (e.g., select d_week_seq and d_year and call
unique()/drop_duplicates()) and join wswscs to that deduped Date mapping before
filtering by pl.col("d_year") == year (and similarly for the z_year_plus_1
path), ensuring you reference the variables/expressions y_year, z_year_plus_1,
wswscs, date_dim and the join key d_week_seq when making the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 515ffcbe-e4e2-4528-afa8-400ffb783f8c
📒 Files selected for processing (19)
python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q14.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q17.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q23.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q25.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q29.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q43.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q44.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q52.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q53.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q55.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q63.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q8.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q88.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q9.pypython/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q98.py
TomAugspurger
approved these changes
May 19, 2026
TomAugspurger
left a comment
TomAugspurger
left a comment
Contributor
There was a problem hiding this comment.
I haven't run these to confirm that validation still passes, but I'll assume you have ;)
One general question: Your PR description has a nice categorization of the different optimizations here. Do we want to include comments labeling the different optimizations in the source code (with the idea being it'll be easier to back out individual by-hand optimizations as they're implemented in polars)? Or is that not necessary?
| "d_date_sk" | ||
| ) | ||
|
|
||
| # store_returns has [6] partitions — at the broadcast limit. Filter it to Q1-Q3 dates |
Contributor
There was a problem hiding this comment.
Is this (6 partitions, near the broadcast limit) always true, or does that depend on specific configuration settings, GPU size, dataset size, etc?
Member
Author
That might be a good follow up. We're adding unoptimized queries too. And in a lot of cases it's easy to compare the two. But some of the optimizations are complex enough that's it's worth adding an explanatory comment. |
Member
Author
|
/merge |
rapids-bot
Bot
merged commit May 19, 2026
691d638
into
NVIDIA:release/26.06
187 of 197 checks passed
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Hand-tune
polars_implfor 19 TPC-DS benchmark queries inpython/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/. Each rewrite preserves query semantics and only changes how the polars LazyFrame is constructed;duckdb_implis unchanged.The optimizations apply a small set of recurring patterns that the polars optimizer does not (yet) perform automatically:
date_dim,item,store, etc. by literal predicates (year, quarter, month window, category/class/brand) before any join, so the join builds smaller hash tables.store_returns(customer, item) pairs) as semi-join probes against the fact tables, shrinking them before the expensive joins.select(...)only the columns each table contributes before joining, instead of relying on the planner to prune them later.Test plan
Checklist