perf: fast path to generate group idxs for vanilla int_range in group_by_dynamic #19932

jvdd · 2024-11-22T21:10:49Z

Performance optimization of group_by_dynamic when passing a vanillaint_range (i.e., start=0, step=1) as index_column.
If we know that the index column for the dynamic group by is an int range, we can generate the group indices (as there is a fixed step between the index values) and can thus avoid the slow group_by_windows function.

Any feedback to improve this PR is welcome :)

Further enhancements (that could? be done):

support non-vanilla int-ranges (i.e., step != 1 and/or start !=0)
do not materialize the int_range (as indices can be generated)

Using the updated code, I observe ~10x performance improvements on my machine

…_by_dynamic

jvdd · 2024-11-22T21:11:49Z

crates/polars-lazy/src/frame/mod.rs

+                    Expr::Function {
+                        input, function, ..
+                    } => {
+                        matches!(function, FunctionExpr::Range(f) if f.to_string() == "int_range")


A bit awkward that I have to check against the .to_string method.

You don't have to:

matches!( function, FunctionExpr::Range(RangeFunction::IntRange { .. }) );

I think we should do this check during the IR::conversion, in resolve_groupby here:

polars/crates/polars-plan/src/plans/conversion/dsl_to_ir.rs

Line 1013 in 05f2abb

} else if let Some(options) = _options.dynamic.as_ref() {

Sure! I am just unsure how I should check in resolve_group_by whether the index of the group_by_dynamic is an int_range. Can I get the expression for this index (conveniently) from the expr_arena?

Ah, right. I see that we first do a with_columns here. We should store the index column on line 1155 and do the with_column rewrite during IR conversion.

I understand it's a bit more than you anticipated. I can do the pre-work for that later if you like?

I suppose this involves passing the index_column Expr from the lazy_frame.group_by_dynamic method to the dsl_to_ir functions? However, the best approach to do this is not clear to me..

Would thus be very helpful if you could do the pre-work for this! :)

jvdd · 2024-11-22T21:12:53Z

crates/polars-time/src/group_by/dynamic.rs

@@ -34,6 +34,7 @@ pub struct DynamicGroupOptions {
    pub include_boundaries: bool,
    pub closed_window: ClosedWindow,
    pub start_by: StartBy,
+    pub int_range: bool, // Whether the index column is a range column


Perhaps this should be private (as we do not want users to set this "flag")

No, we don't.

We can add a new DynamicGroupOptionsIR where we add private flags set by the IR.

codecov · 2024-11-23T07:14:26Z

Codecov Report

Attention: Patch coverage is 17.54386% with 94 lines in your changes missing coverage. Please review.

Project coverage is 79.47%. Comparing base (414d883) to head (f98b8fc).
Report is 37 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/polars-time/src/windows/group_by.rs	0.00%	81 Missing ⚠️
crates/polars-time/src/group_by/dynamic.rs	60.00%	10 Missing ⚠️
crates/polars-lazy/src/frame/mod.rs	62.50%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main   #19932       +/-   ##
===========================================
+ Coverage   59.28%   79.47%   +20.18%     
===========================================
  Files        1555     1562        +7     
  Lines      216180   217047      +867     
  Branches     2456     2459        +3     
===========================================
+ Hits       128155   172488    +44333     
+ Misses      87467    44000    -43467     
- Partials      558      559        +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ritchie46

Nice find! I've left some comments.

ritchie46 · 2024-11-23T13:56:34Z

crates/polars-lazy/src/frame/mod.rs

+                    Expr::Function {
+                        input, function, ..
+                    } => {
+                        matches!(function, FunctionExpr::Range(f) if f.to_string() == "int_range")


You don't have to:

matches!( function, FunctionExpr::Range(RangeFunction::IntRange { .. }) );

ritchie46 · 2024-11-23T14:01:23Z

crates/polars-lazy/src/frame/mod.rs

+                    Expr::Function {
+                        input, function, ..
+                    } => {
+                        matches!(function, FunctionExpr::Range(f) if f.to_string() == "int_range")


I think we should do this check during the IR::conversion, in resolve_groupby here:

polars/crates/polars-plan/src/plans/conversion/dsl_to_ir.rs

Line 1013 in 05f2abb

} else if let Some(options) = _options.dynamic.as_ref() {

ritchie46 · 2024-11-23T14:02:29Z

crates/polars-time/src/group_by/dynamic.rs

@@ -34,6 +34,7 @@ pub struct DynamicGroupOptions {
    pub include_boundaries: bool,
    pub closed_window: ClosedWindow,
    pub start_by: StartBy,
+    pub int_range: bool, // Whether the index column is a range column


No, we don't.

We can add a new DynamicGroupOptionsIR where we add private flags set by the IR.

ritchie46 · 2024-11-23T14:03:22Z

crates/polars-time/src/group_by/dynamic.rs

-                options.start_by,
-            );
+
+            let vanilla_start_step = (ts[0] == 0) && (ts[1] - ts[0] == 1);


Can we move this to a separate function so that implementation is separate from dispatch?

Moved the logic to a separate function in the group_by.rs file

crates/polars-time/src/group_by/dynamic.rs

perf: fast path to generate group idxs for vanilla int_range in group…

dbae112

…_by_dynamic

jvdd requested review from ritchie46, orlp, c-peters and MarcoGorelli as code owners November 22, 2024 21:10

github-actions bot added the title needs formatting label Nov 22, 2024

jvdd commented Nov 22, 2024

View reviewed changes

jvdd added 3 commits November 22, 2024 22:23

🧹 cleanup match condition for vanilla int_range in impl_group_by_dynamic

4ada18e

fix: use IdxSize instead of u32

b75a3b4

fix: deal with clippy warnings

8a14bb4

ritchie46 reviewed Nov 23, 2024

View reviewed changes

jvdd added 2 commits November 26, 2024 09:06

move group by dynamic int_rage logic to group_by_windows_ir

2502983

fix: deal with clippy warnings

f98b8fc

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

perf: fast path to generate group idxs for vanilla int_range in group_by_dynamic #19932

perf: fast path to generate group idxs for vanilla int_range in group_by_dynamic #19932

jvdd commented Nov 22, 2024

jvdd Nov 22, 2024

ritchie46 Nov 23, 2024

ritchie46 Nov 23, 2024

jvdd Nov 24, 2024

ritchie46 Nov 24, 2024

jvdd Nov 27, 2024

jvdd Nov 22, 2024

ritchie46 Nov 23, 2024

codecov bot commented Nov 23, 2024 •

edited

Loading

ritchie46 left a comment

ritchie46 Nov 23, 2024

ritchie46 Nov 23, 2024

ritchie46 Nov 23, 2024

ritchie46 Nov 23, 2024

jvdd Nov 28, 2024

perf: fast path to generate group idxs for vanilla int_range in group_by_dynamic #19932

Are you sure you want to change the base?

perf: fast path to generate group idxs for vanilla int_range in group_by_dynamic #19932

Conversation

jvdd commented Nov 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 23, 2024 • edited Loading

Codecov Report

ritchie46 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 23, 2024 •

edited

Loading