Skip to content

Commit

Permalink
perf: Ensure metadata flags are maintained on vertical parallelization (
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored and atigbadr committed Jul 23, 2024
1 parent c026d93 commit 29c48d7
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
17 changes: 17 additions & 0 deletions crates/polars-core/src/series/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -195,13 +195,30 @@ impl Series {
// TODO! this probably can now be removed, now we don't have special case for structs.
pub fn select_chunk(&self, i: usize) -> Self {
let mut new = self.clear();
let flags = self.get_flags();

let mut new_flags = MetadataFlags::empty();
new_flags.set(
MetadataFlags::SORTED_ASC,
flags.contains(MetadataFlags::SORTED_ASC),
);
new_flags.set(
MetadataFlags::SORTED_DSC,
flags.contains(MetadataFlags::SORTED_DSC),
);
new_flags.set(
MetadataFlags::FAST_EXPLODE_LIST,
flags.contains(MetadataFlags::FAST_EXPLODE_LIST),
);

// Assign mut so we go through arc only once.
let mut_new = new._get_inner_mut();
let chunks = unsafe { mut_new.chunks_mut() };
let chunk = self.chunks()[i].clone();
chunks.clear();
chunks.push(chunk);
mut_new.compute_len();
mut_new._set_flags(new_flags);
new
}

Expand Down
10 changes: 10 additions & 0 deletions py-polars/tests/unit/operations/test_is_sorted.py
Original file line number Diff line number Diff line change
Expand Up @@ -407,3 +407,13 @@ def test_sorted_flag_group_by_dynamic() -> None:

def test_is_sorted_rle_id() -> None:
assert pl.Series([12, 3345, 12, 3, 4, 4, 1, 12]).rle_id().flags["SORTED_ASC"]


def test_is_sorted_chunked_select() -> None:
df = pl.DataFrame({"a": np.ones(14)})

assert (
pl.concat([df, df, df], rechunk=False)
.set_sorted("a")
.select(pl.col("a").alias("b"))
)["b"].flags["SORTED_ASC"]

0 comments on commit 29c48d7

Please sign in to comment.