feat: Improved list arithmetic support #19162

nameexhaustion · 2024-10-09T10:48:43Z

Introduces numeric list kernels that operate directly on the list offsets and leaf arrays.

Support for
- List<->List operations where both lists must only have 1 level of nesting
- List<->Numeric operations, where the list can have any level of nesting
Adds support for floor divide

Notes

I haven't added dedicated codepaths for when the outer validity is NULL to reduce codegen - I've instead forced the outer validity to be materialized
Have not added feature gates yet

Fixes #19010
Fixes #19025

nameexhaustion · 2024-10-09T10:51:40Z

crates/polars-core/src/series/arithmetic/borrowed.rs

+    let mut new_left_dtype = left_dtype.cast_leaf(leaf_super_dtype.clone());
+    let mut new_right_dtype = right_dtype.cast_leaf(leaf_super_dtype);
+
+    // Cast List<->Array to List


We currently already cast list<->array to list as a supertype from #12016

nameexhaustion · 2024-10-09T10:57:22Z

crates/polars-core/src/series/series_trait.rs

@@ -327,6 +328,26 @@ pub trait SeriesTrait:
    /// Aggregate all chunks to a contiguous array of memory.
    fn rechunk(&self) -> Series;

+    fn rechunk_validity(&self) -> Option<Bitmap> {


Copied from impl<T> ChunkedArray<T> so that it can be used without downcasting

nameexhaustion · 2024-10-09T10:59:19Z

crates/polars-plan/src/plans/conversion/type_coercion/binary.rs

@@ -47,55 +47,6 @@ fn is_cat_str_binary(type_left: &DataType, type_right: &DataType) -> bool {
    }
 }

-fn process_list_arithmetic(


Remove IR checking / modifying logic from type_coercion - I think this should already be handled by get_arithmetic_field, get_truediv_field etc.

nameexhaustion · 2024-10-09T11:01:35Z

py-polars/polars/dataframe/frame.py

@@ -1085,7 +1087,8 @@ def _div(self, other: Any, *, floordiv: bool) -> DataFrame:
            int_casts = [
                col(column).cast(tp)
                for i, (column, tp) in enumerate(self.schema.items())
-                if tp.is_integer() and orig_dtypes[i].is_integer()
+                if tp.is_integer()
+                and (orig_dtypes[i].is_integer() or orig_dtypes[i] == Null)


This is a fix after something I did broke this test case -

polars/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py

Lines 287 to 288 in 2e0438b

assert_frame_equal(df_expected, op(df, None))

assert_series_equal(s_expected, op(s, None))

, I have no idea how it was working before though

py-polars/polars/series/series.py

-            as_float = self._recursive_cast_to_dtype(Float64())
-
-        return as_float._arithmetic(other, "div", "div_<>")
+        return self._arithmetic(other, "div", "div_<>")


nameexhaustion · 2024-10-09T11:21:48Z

py-polars/tests/unit/operations/arithmetic/test_arithmetic.py

-        ),
-    ],
-)
-def test_list_arithmetic_same_size(


I will check these tests back in later to reduce the size of this PR

Maybe just add comment saying "no need to review these functions, they're cut/pasted from test_arithmetic.py"?

Thanks for the tip - I realized I can actually leave those tests here for now instead of removing them

nameexhaustion · 2024-10-09T11:33:03Z

py-polars/tests/unit/operations/arithmetic/test_list_arithmetic.py

+    BROADCAST_SERIES_COMBINATIONS,
+)
+@pytest.mark.parametrize("exec_op", EXEC_OP_COMBINATIONS)
+def test_list_arithmetic_values(


Magical test case parametrization that runs every codepath in BinaryListNumericOpHelper 😄

codecov · 2024-10-09T13:53:04Z

Codecov Report

Attention: Patch coverage is 86.90745% with 116 lines in your changes missing coverage. Please review.

Project coverage is 79.72%. Comparing base (ff10b38) to head (8e7554b).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
...polars-core/src/series/arithmetic/list_borrowed.rs	93.64%	39 Missing ⚠️
crates/polars-compute/src/arithmetic/pl_num.rs	61.61%	38 Missing ⚠️
crates/polars-plan/src/plans/aexpr/schema.rs	71.01%	20 Missing ⚠️
crates/polars-core/src/series/series_trait.rs	22.22%	14 Missing ⚠️
...ates/polars-core/src/series/arithmetic/borrowed.rs	82.14%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #19162      +/-   ##
==========================================
+ Coverage   79.67%   79.72%   +0.04%     
==========================================
  Files        1532     1533       +1     
  Lines      209200   209915     +715     
  Branches     2417     2415       -2     
==========================================
+ Hits       166687   167357     +670     
- Misses      41965    42010      +45     
  Partials      548      548

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

coastalwhite · 2024-10-10T08:50:51Z

crates/polars-arrow/src/offset.rs

+            for o in others {
+                let slc = o.as_slice();
+                l = slc[l].to_usize();
+                r = slc[r].to_usize();


I am not 100% sure, but I feel like this should be r + 1. Might be completely wrong though.

It should be correct as it is - Arrow list offsets are defined as

1st row : offsets[0]..offsets[1] 2nd row : offsets[1]..offsets[2] ..and so on

coastalwhite · 2024-10-10T08:50:57Z

crates/polars-arrow/src/offset.rs

+        for o in &offsets[1..] {
+            let slc = o.as_slice();
+            l = slc[l].to_usize();
+            r = slc[r].to_usize();


coastalwhite · 2024-10-10T08:59:02Z

crates/polars-core/src/series/arithmetic/borrowed.rs

+            | (_, Time)
+            | (_, Date)
+            | (_, Datetime(_, _)) => polars_bail!(opq = div, self.dtype(), rhs.dtype()),
+            _ => match (self.dtype(), rhs.dtype()) {


nit: why is this a nested match?

Updated to remove

crates/polars-core/src/series/arithmetic/list_borrowed.rs

coastalwhite · 2024-10-10T09:11:37Z

crates/polars-core/src/series/arithmetic/list_borrowed.rs

+                // list lengths.
+                let mut mismatch_pos = 0;
+
+                with_match_numeric_list_op!(&self.op, self.swapped, |$OP| {


I kinda feel like this should not be here. Instead, this should be in polars-compute.

I tried to move the entire file, but I ran into some difficulty as polars-compute doesn't have access to the Series struct - I think we can leave it here for now?

I also don't want to split out parts of the logic in this file as it isn't really used anywhere else

coastalwhite · 2024-10-10T09:11:57Z

crates/polars-core/src/series/arithmetic/list_borrowed.rs

+
+                /// Reduce monomorphization
+                #[inline(never)]
+                fn combine_validities_list_to_list_no_broadcast(


I think I want to keep this here for now - it's a very specialized function that's only used in this file. If we need to use it somewhere else later I can move it then

coastalwhite · 2024-10-10T09:14:15Z

crates/polars-core/src/series/ops/reshape.rs

@@ -56,6 +57,25 @@ impl Series {
        }
    }

+    /// TODO: Move this somewhere else?
+    pub fn list_offsets_and_validities_recursive(


This always outputs Vecs with 1 element.

Not necessarily - it returns Vecs with a number of elements corresponding to the nesting level - so e.g. for a List(List(Float64)) -

print((pl.Series([[[1]]]) / pl.Series([1])).dtype) List(List(Float64)) [crates/polars-core/src/series/arithmetic/list_borrowed.rs:93:17] lhs.list_offsets_and_validities_recursive().0 = [ OffsetsBuffer( [ 0, 1, ], ), OffsetsBuffer( [ 0, 1, ], ), ]

itamarst · 2024-10-14T13:10:10Z

Thanks for doing this!

github-actions bot added enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars labels Oct 9, 2024

nameexhaustion commented Oct 9, 2024

View reviewed changes

py-polars/polars/series/series.py

as_float = self._recursive_cast_to_dtype(Float64())

return as_float._arithmetic(other, "div", "div_<>")

return self._arithmetic(other, "div", "div_<>")

This comment was marked as outdated.

Sign in to view

nameexhaustion commented Oct 9, 2024

View reviewed changes

nameexhaustion force-pushed the list-arith branch from 72e5b6e to 12494af Compare October 9, 2024 11:22

nameexhaustion mentioned this pull request Oct 9, 2024

feat(rust, python): allow arithmetic operations between numeric Series and list Series #18901

Closed

nameexhaustion commented Oct 9, 2024

View reviewed changes

nameexhaustion marked this pull request as ready for review October 9, 2024 13:56

nameexhaustion requested review from ritchie46, c-peters, alexander-beedie, MarcoGorelli, reswqa and orlp as code owners October 9, 2024 13:56

nameexhaustion marked this pull request as draft October 10, 2024 04:30

nameexhaustion force-pushed the list-arith branch 2 times, most recently from 3bf866d to 09b0d7f Compare October 10, 2024 04:58

nameexhaustion marked this pull request as ready for review October 10, 2024 05:05

coastalwhite requested changes Oct 10, 2024

View reviewed changes

nameexhaustion marked this pull request as draft October 11, 2024 04:14

nameexhaustion force-pushed the list-arith branch from 09b0d7f to 2ea1524 Compare October 11, 2024 04:21

This comment has been minimized.

Sign in to view

nameexhaustion marked this pull request as ready for review October 11, 2024 05:03

This comment was marked as outdated.

Sign in to view

nameexhaustion added 3 commits October 14, 2024 18:19

c

8d323ed

remove nested branch

7cad3e2

c

8e7554b

nameexhaustion force-pushed the list-arith branch from 1662a10 to 8e7554b Compare October 14, 2024 07:19

ritchie46 merged commit f7c6a05 into pola-rs:main Oct 14, 2024
26 checks passed

This was referenced Oct 14, 2024

Allow element wise computation based on another column in list.eval(pl.element() ... ) #17496

Closed

Simple arithmetic operations on the "list" type columns #8006

Closed

Divide each element in a list by an int #14711

Closed

bike-picket mentioned this pull request Oct 21, 2024

Support pow operation for list dtype #19349

Open

nameexhaustion deleted the list-arith branch October 28, 2024 04:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Improved list arithmetic support #19162

feat: Improved list arithmetic support #19162

nameexhaustion commented Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024

nameexhaustion Oct 9, 2024

nameexhaustion Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024

This comment was marked as outdated.

nameexhaustion Oct 9, 2024 •

edited

Loading

itamarst Oct 9, 2024

nameexhaustion Oct 10, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024 •

edited

Loading

codecov bot commented Oct 9, 2024 •

edited

Loading

coastalwhite Oct 10, 2024

nameexhaustion Oct 11, 2024

coastalwhite Oct 10, 2024

coastalwhite Oct 10, 2024

nameexhaustion Oct 11, 2024

coastalwhite Oct 10, 2024

nameexhaustion Oct 11, 2024 •

edited

Loading

coastalwhite Oct 10, 2024

nameexhaustion Oct 11, 2024

coastalwhite Oct 10, 2024

nameexhaustion Oct 11, 2024

This comment has been minimized.

This comment was marked as outdated.

itamarst commented Oct 14, 2024

	assert_frame_equal(df_expected, op(df, None))
	assert_series_equal(s_expected, op(s, None))

feat: Improved list arithmetic support #19162

feat: Improved list arithmetic support #19162

Conversation

nameexhaustion commented Oct 9, 2024 • edited Loading

Notes

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nameexhaustion Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment was marked as outdated.

nameexhaustion Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nameexhaustion Oct 10, 2024 • edited Loading

Choose a reason for hiding this comment

nameexhaustion Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

codecov bot commented Oct 9, 2024 • edited Loading

Codecov Report

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

nameexhaustion Oct 11, 2024 • edited Loading

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

This comment has been minimized.

This comment was marked as outdated.

itamarst commented Oct 14, 2024

nameexhaustion commented Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 10, 2024 •

edited

Loading

nameexhaustion Oct 9, 2024 •

edited

Loading

codecov bot commented Oct 9, 2024 •

edited

Loading

nameexhaustion Oct 11, 2024 •

edited

Loading