-
-
Notifications
You must be signed in to change notification settings - Fork 2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: Improved array arithmetic support #19837
Conversation
91f0f5d
to
5bac4c8
Compare
@@ -0,0 +1,820 @@ | |||
use polars_error::{feature_gated, PolarsResult}; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note: This file has mostly the same structure as the one for list arithmetic, but with some array specific adjustments.
|
||
match (&self.op_apply_type, &self.broadcast) { | ||
// Mostly the same as ListNumericOp, however with fixed size list we also have | ||
// (BinaryOpApplyType::ListToPrimitive, Broadcast::Left) as a physical impl. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Arrays have a simpler layout compared to lists, so we can also implement array<->primitive
with the array side being broadcasted without materializing.
} | ||
} | ||
|
||
#[derive(Debug, Clone)] | ||
pub enum NumericListOp { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Factored out this and some other parts of list arithmetic for re-use
// We materialize the list columns with `new_from_index`, as otherwise we'd have to | ||
// implement logic that broadcasts the offsets and validities across multiple levels | ||
// of nesting. But we will re-use the materialized memory to store the result. | ||
(BinaryOpApplyType::ListToList, Broadcast::Left) => { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
drive-by - changed the order of this enum matching to the same order of the comment (expand the diff above to see the comment)
@@ -383,16 +383,3 @@ def test_zero_width_array(fn: str) -> None: | |||
|
|||
df = pl.concat([a.to_frame(), b.to_frame()], how="horizontal") | |||
df.select(c=expr_f(pl.col.a, pl.col.b)) | |||
|
|||
|
|||
def test_elementwise_arithmetic_19682() -> None: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
moved to arithmetic/test_array.py
below
) | ||
@pytest.mark.parametrize("exec_op", EXEC_OP_COMBINATIONS) | ||
@pytest.mark.slow | ||
def test_array_arithmetic_values( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Same parametrized test copied from list arithmetic. It has also extra parametrization to test operations between nested array types (array_side="both3"
).
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #19837 +/- ##
==========================================
+ Coverage 79.33% 79.37% +0.04%
==========================================
Files 1548 1550 +2
Lines 214245 214751 +506
Branches 2460 2460
==========================================
+ Hits 169968 170464 +496
- Misses 43719 43728 +9
- Partials 558 559 +1 ☔ View full report in Codecov by Sentry. |
Table code
Fixes #19356
Fixes #19838