Skip to content

Commit

Permalink
perf: Don't rechunk in align_ if arrays are aligned (pola-rs#16850)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jun 10, 2024
1 parent 96ef044 commit eaedf74
Showing 1 changed file with 21 additions and 0 deletions.
21 changes: 21 additions & 0 deletions crates/polars-core/src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -783,7 +783,18 @@ where
)
};
match (left.chunks.len(), right.chunks.len()) {
// All chunks are equal length
(1, 1) => (Cow::Borrowed(left), Cow::Borrowed(right)),
// All chunks are equal length
(a, b)
if a == b
&& left
.chunk_lengths()
.zip(right.chunk_lengths())
.all(|(l, r)| l == r) =>
{
(Cow::Borrowed(left), Cow::Borrowed(right))
},
(_, 1) => {
assert();
(
Expand Down Expand Up @@ -902,6 +913,16 @@ where
Cow::Owned(c.match_chunks(a.chunk_lengths())),
)
},
(len_a, len_b, len_c)
if len_a == len_b
&& len_b == len_c
&& a.chunk_lengths()
.zip(b.chunk_lengths())
.zip(c.chunk_lengths())
.all(|((a, b), c)| a == b && b == c) =>
{
(Cow::Borrowed(a), Cow::Borrowed(b), Cow::Borrowed(c))
},
_ => {
// could optimize to choose to rechunk a primitive and not a string or list type
let a = a.rechunk();
Expand Down

0 comments on commit eaedf74

Please sign in to comment.