fix: Allow nested `is_in()` in `when()/then()` for full-streaming #20052

nameexhaustion · 2024-11-28T13:44:17Z

This PR also does some refactoring to consolidate and improve the way we identify expressions as elementwise.

Replacing `streamable` with `is_elementwise`

A lot of places where we currently use the streamable terminology actually needed to use an even stricter elementwise requirement (e.g. filter / slice pushdown. streamable expressions include all elementwise expressions and more. For example, filter() / explode() expressions are streamable, but they are not elementwise.

Where this matters is if the result of a streamable expression needs to be projected next to other result columns, which is almost always the case. The existing in-memory engine and evaluators cannot do this properly as it projects the result columns independently within every chunk, which could lead to height mismatches. The only place that properly supports this is the new-streaming engine, where the ZipNode performs the projection properly.

nameexhaustion · 2024-11-28T14:27:35Z

crates/polars-plan/src/plans/aexpr/mod.rs

-    pub(crate) fn groups_sensitive(&self) -> bool {
+
+    /// Checks whether this expression is elementwise. This only checks the top level expression.
+    pub(crate) fn is_elementwise_top_level(&self) -> bool {


is_elementwise_top_level() replaces groups_sensitive(). It is also changed to consider Explode and Filter as non-elementwise

nameexhaustion · 2024-11-28T14:30:54Z

crates/polars-plan/src/plans/aexpr/utils.rs

-}
+/// Checks if the top-level expression node is elementwise. If this is the case, then `stack` will
+/// be extended further with any nested expression nodes.
+pub fn is_elementwise(stack: &mut Vec<Node>, ae: &AExpr, expr_arena: &Arena<AExpr>) -> bool {


Base is_elementwise() function - recursion is done through a nodes stack which allows the caller to inspect the exprs.

nameexhaustion · 2024-11-28T14:32:05Z

crates/polars-plan/src/plans/aexpr/utils.rs

-        }
+                if matches!(expr_arena.get(rhs), AExpr::Literal { .. }) {
+                    stack.push_node(input.first().unwrap().node());
+                    return;


This is the fix for the when(A.is_in(B)) issue - we don't traverse into RHS literals of is_in() during recursive checking. We had this already in predicate pushdown but it has been moved here.

nameexhaustion · 2024-11-28T14:51:22Z

crates/polars-plan/src/plans/aexpr/mod.rs

+                ..
+            } => {
+                assert!(options.is_elementwise());
+                opts.strict


This rule was moved from new-streaming below

nameexhaustion · 2024-11-28T14:56:37Z

crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs

+                if cfg!(debug_assertions) {
+                    for v in acc_predicates.values() {
+                        let ae = expr_arena.get(v.node());
+                        assert!(permits_filter_pushdown(


drive-by - predicates that hit here should already satisfy these requirements - otherwise they should have been caught at the Filter node

nameexhaustion · 2024-11-28T15:57:19Z

crates/polars-plan/src/plans/aexpr/utils.rs

-                    options.collect_groups,
-                    ApplyOptions::ElementWise | ApplyOptions::ApplyList
-                ),
-                Context::Aggregation => matches!(options.collect_groups, ApplyOptions::ElementWise),


This function was never called with Context::Aggregation

nameexhaustion · 2024-11-28T15:58:11Z

crates/polars-plan/src/plans/options.rs

-                .contains(FunctionFlags::CHANGES_LENGTH | FunctionFlags::RETURNS_SCALAR)
+        matches!(
+            self.collect_groups,
+            ApplyOptions::ElementWise | ApplyOptions::ApplyList


I don't fully understand this, but we apparently consider ApplyList as elementwise according to the test suite

I think we should see when we hit this. I believe only in a few python udf cases. I think we would never hit the streaming engine with this.

Yes, this is hit by map_elements() / map_batches(.., agg_list=True). This also hits the streaming engines because FunctionOptions::is_elementwise() gets used by all of the engines during physical plan creation / IR lowering (indirectly through fn is_elementwise(_rec)).

I don't think ApplyList is strictly elementwise, but both the in-memory and existing streaming engine currently expect it to be identified as such, so I've made it the default here. For the new-streaming engine I added an override below.

codecov · 2024-11-28T17:37:09Z

Codecov Report

Attention: Patch coverage is 84.18605% with 34 lines in your changes missing coverage. Please review.

Project coverage is 79.52%. Comparing base (44ddbc2) to head (5c45bf1).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/polars-stream/src/physical_plan/lower_expr.rs	0.00%	23 Missing ⚠️
crates/polars-plan/src/plans/aexpr/utils.rs	91.22%	5 Missing ⚠️
...rs-plan/src/plans/optimizer/slice_pushdown_expr.rs	0.00%	3 Missing ⚠️
crates/polars-plan/src/plans/lit.rs	75.00%	1 Missing ⚠️
...es/polars-plan/src/plans/optimizer/cse/cse_expr.rs	50.00%	1 Missing ⚠️
crates/polars-stream/src/physical_plan/lower_ir.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #20052      +/-   ##
==========================================
- Coverage   79.52%   79.52%   -0.01%     
==========================================
  Files        1563     1563              
  Lines      217194   217173      -21     
  Branches     2464     2464              
==========================================
- Hits       172729   172701      -28     
- Misses      43905    43912       +7     
  Partials      560      560

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

ritchie46 · 2024-11-29T07:51:36Z

crates/polars-plan/src/plans/aexpr/mod.rs

+            Alias(_, _) | BinaryExpr { .. } | Column(_) | Ternary { .. } | Cast { .. } => true,
+
+            Agg { .. }
+            | Explode(_)


Filter, gather, sort and explode don't seem elementwise. They are group sensitive but not elementwise. Maybe this is an unlucky name?

crates/polars-plan/src/plans/aexpr/utils.rs

ritchie46 · 2024-11-29T07:56:45Z

crates/polars-plan/src/plans/options.rs

-                .contains(FunctionFlags::CHANGES_LENGTH | FunctionFlags::RETURNS_SCALAR)
+        matches!(
+            self.collect_groups,
+            ApplyOptions::ElementWise | ApplyOptions::ApplyList


I think we should see when we hit this. I believe only in a few python udf cases. I think we would never hit the streaming engine with this.

github-actions bot added fix Bug fix python Related to Python Polars rust Related to Rust Polars labels Nov 28, 2024

nameexhaustion changed the title ~~fix: Allow nested is_in() in when()/then() for streaming sink~~ fix: Allow nested is_in() in when()/then() for full-streaming Nov 28, 2024

nameexhaustion commented Nov 28, 2024

View reviewed changes

nameexhaustion force-pushed the is-elementwise-streamable branch from 8adc9dd to 55ec6b6 Compare November 28, 2024 16:22

nameexhaustion force-pushed the is-elementwise-streamable branch from a57466f to 66c30d4 Compare November 29, 2024 05:59

nameexhaustion marked this pull request as ready for review November 29, 2024 06:09

nameexhaustion requested review from ritchie46, c-peters, alexander-beedie, MarcoGorelli, reswqa and orlp as code owners November 29, 2024 06:09

nameexhaustion marked this pull request as draft November 29, 2024 06:10

nameexhaustion marked this pull request as ready for review November 29, 2024 06:17

ritchie46 reviewed Nov 29, 2024

View reviewed changes

nameexhaustion added 3 commits November 29, 2024 20:43

c

ac95383

c

ec06d38

c

208fca5

nameexhaustion marked this pull request as draft November 29, 2024 09:43

unitvec

5c45bf1

nameexhaustion force-pushed the is-elementwise-streamable branch from 66c30d4 to 5c45bf1 Compare November 29, 2024 09:55

nameexhaustion marked this pull request as ready for review November 29, 2024 10:19

ritchie46 merged commit ca8c1ef into pola-rs:main Nov 29, 2024
25 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: Allow nested `is_in()` in `when()/then()` for full-streaming #20052

fix: Allow nested `is_in()` in `when()/then()` for full-streaming #20052

nameexhaustion commented Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 28, 2024

nameexhaustion Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 28, 2024

nameexhaustion Nov 28, 2024

nameexhaustion Nov 28, 2024

nameexhaustion Nov 28, 2024 •

edited

Loading

ritchie46 Nov 29, 2024

nameexhaustion Nov 29, 2024 •

edited

Loading

codecov bot commented Nov 28, 2024 •

edited

Loading

ritchie46 Nov 29, 2024 •

edited

Loading

ritchie46 Nov 29, 2024

fix: Allow nested is_in() in when()/then() for full-streaming #20052

fix: Allow nested is_in() in when()/then() for full-streaming #20052

Conversation

nameexhaustion commented Nov 28, 2024 • edited Loading

Replacing streamable with is_elementwise

nameexhaustion Nov 28, 2024 • edited Loading

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024 • edited Loading

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024

Choose a reason for hiding this comment

nameexhaustion Nov 28, 2024 • edited Loading

Choose a reason for hiding this comment

ritchie46 Nov 29, 2024

Choose a reason for hiding this comment

nameexhaustion Nov 29, 2024 • edited Loading

Choose a reason for hiding this comment

codecov bot commented Nov 28, 2024 • edited Loading

Codecov Report

ritchie46 Nov 29, 2024 • edited Loading

Choose a reason for hiding this comment

ritchie46 Nov 29, 2024

Choose a reason for hiding this comment

fix: Allow nested `is_in()` in `when()/then()` for full-streaming #20052

fix: Allow nested `is_in()` in `when()/then()` for full-streaming #20052

nameexhaustion commented Nov 28, 2024 •

edited

Loading

Replacing `streamable` with `is_elementwise`

nameexhaustion Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 28, 2024 •

edited

Loading

nameexhaustion Nov 29, 2024 •

edited

Loading

codecov bot commented Nov 28, 2024 •

edited

Loading

ritchie46 Nov 29, 2024 •

edited

Loading