Implement more unary polars.Expressions - #23037
Closed
mroeschke wants to merge 18 commits into
Closed
Conversation
Truncate numeric values toward zero to the requested number of decimal places. Integer inputs are returned unchanged; float inputs are scaled by 10^decimals (when nonzero), truncated via copy_if_else(floor(x), ceil(x), x >= 0), and unscaled.
Drop NaN values (keeping nulls) via pylibcudf.stream_compaction.drop_nans. Non-floating-point inputs cannot contain NaN and are returned unchanged.
rechunk is an identity passthrough on the GPU since cudf columns are always stored as a single contiguous chunk.
Find insertion indices for one or more needles in a sorted column via pylibcudf.search.lower_bound (side any/left) or upper_bound (side right), selecting ascending or descending order from the descending flag. Indices are cast to the unsigned output dtype.
Return the first index where the column equals the search value (or null if absent). A NULL_EQUALS comparison builds the match mask so a null search value matches null entries; the matching positions of an index sequence are filtered and the first is returned, cast to the unsigned output dtype.
Map approx_n_unique to an exact distinct count via pylibcudf.reduce.distinct_count with NullPolicy.INCLUDE and NanPolicy.NAN_IS_VALID, matching Polars' n_unique semantics where null is counted once and all NaNs compare equal. The result is returned as a length-1 unsigned column.
mroeschke
added
improvement
github-actions
Bot
added
Python
mroeschke
changed the title
📝 WalkthroughWalkthrough
ChangesNew UnaryFunction Operations
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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/tests/expressions/test_index_of.py`:
- Around line 12-31: The current test only covers literal needles for
pl.col("a").index_of, so add cases that use an expression-valued needle to
exercise the UnaryFunction.do_evaluate path that evaluates value_expr and
obj_scalar. Extend test_index_of with at least one query where the needle is
another column or derived expression, and keep using assert_gpu_result_equal on
the resulting LazyFrame so regressions in the non-literal index_of branch are
caught.
In `@python/cudf_polars/tests/expressions/test_search_sorted.py`:
- Around line 14-41: The current search_sorted tests cover only normal non-null
arrays, so add cases in test_search_sorted, test_search_sorted_descending, or
nearby helpers that exercise empty inputs, single-element inputs, and
null-containing inputs to validate boundary and nullable behavior. Reuse the
existing pl.col("a").search_sorted(...) queries and assert_gpu_result_equal so
the GPU path is checked against CPU for these edge cases, especially around
UnaryFunction.do_evaluate where null ordering is hard-coded.
🪄 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: c3d8345d-f780-4709-b5cb-aec917b5a0a7
📒 Files selected for processing (7)
python/cudf_polars/cudf_polars/dsl/expressions/unary.pypython/cudf_polars/tests/expressions/test_approx_n_unique.pypython/cudf_polars/tests/expressions/test_drop_nans.pypython/cudf_polars/tests/expressions/test_index_of.pypython/cudf_polars/tests/expressions/test_rechunk.pypython/cudf_polars/tests/expressions/test_search_sorted.pypython/cudf_polars/tests/expressions/test_truncate.py
3 tasks
3 tasks
Comment on lines
+386
to
+388
| elif self.name == "rechunk": | ||
| (column,) = (child.evaluate(df, context=context) for child in self.children) | ||
| return column |
Contributor
There was a problem hiding this comment.
Just for my own understanding, does this actually do anything related to rechunking?
This API from polars seems very tied to its execution model, so we might not have something similar in cudf-polars, which is fine. Implementing this just to avoid fallback seems fine.
Contributor
Author
There was a problem hiding this comment.
Just for my own understanding, does this actually do anything related to rechunking?
It does not. Yeah this essentially is implemented as a no-op not to fallback to CPU
Comment on lines
+389
to
+390
| elif self.name == "approx_n_unique": | ||
| (column,) = (child.evaluate(df, context=context) for child in self.children) |
Contributor
There was a problem hiding this comment.
I see https://docs.rapids.ai/api/libcudf/stable/classcudf_1_1approx__distinct__count. Would that be appropriate to use here? (no bindings in pylibcudf though).
Contributor
There was a problem hiding this comment.
Makes sense since it's "approx"
Also FYI @wence- since he mentioned using cudf::approx_distinct_count for deciding when to use bloom filters
Comment on lines
+391
to
+396
| count = plc.reduce.distinct_count( | ||
| column.obj, | ||
| plc.types.NullPolicy.INCLUDE, | ||
| plc.types.NanPolicy.NAN_IS_VALID, | ||
| stream=df.stream, | ||
| ) |
Contributor
There was a problem hiding this comment.
This seems like something that'll need special handling for partitioned data. Is that happening anywhere? Maybe there's a fallback to single partition?
| ), | ||
| dtype=self.dtype, | ||
| ) | ||
| elif self.name == "search_sorted": |
Contributor
There was a problem hiding this comment.
Same question about multi-partition here, and for index_of.
| def test_search_sorted_nulls( | ||
| engine: pl.GPUEngine, side: Literal["any", "left", "right"] | ||
| ) -> None: | ||
| lf = pl.LazyFrame({"a": [None, 1, 2, 2, 4]}) |
Contributor
There was a problem hiding this comment.
Maybe parametrize this over data with nulls at the start, middle, and end?
| value = value_expr.evaluate(df, context=context).obj_scalar( | ||
| stream=df.stream | ||
| ) | ||
| py_value = value.to_py(stream=df.stream) |
Contributor
There was a problem hiding this comment.
Will this block df.stream entirely? I guess we need to, so that the isinstance(py_value, float) can run, but maybe we can perform that check in some other way? Does libcudf have a way to check if some scalar is nan?
Oh, but I guess then that would need to block so that we can evaluate the if condition. So really, we'd need to push that if/else into some CUDA expressions...
Contributor
Author
|
Closing as discussed offline to open separate PRs (will take into account reviews left here) |
rapids-bot Bot
pushed a commit
that referenced
this pull request
Jul 8, 2026
I used a more barebones prompt that resembles this PR to implement #23015 and #23037, but I think it could be generally be useful to have when implementing reviewing expression support in cudf_polars. Some of the guidelines is probably more suitable for generally cudf_polars development, but this skill is mainly a starting point that can be iterated on/split out Authors: - Matthew Roeschke (https://github.com/mroeschke) Approvers: - Tom Augspurger (https://github.com/TomAugspurger) URL: #23078
rapids-bot Bot
pushed a commit
that referenced
this pull request
Jul 9, 2026
xref #23151 Prior review #23037 (comment) https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.rechunk.html Authors: - Matthew Roeschke (https://github.com/mroeschke) Approvers: - Matthew Murray (https://github.com/Matt711) URL: #23192
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
4 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
xref #23015
Implements the following unary expressions
Expr.truncateExpr.drop_nansExpr.rechunkExpr.search_sortedExpr.index_ofExpr.approx_n_uniqueChecklist