Skip to content
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.

Sign up for GitHub

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: Implement sorted flags for struct series #21290

Merged
merged 4 commits into from
Feb 19, 2025
Merged

feat: Implement sorted flags for struct series #21290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1c5e133
Remove duplicate `StructChunked::arg_sort`
JakubValtar Feb 11, 2025
f4cd8c7
Get and set struct series flags
JakubValtar Feb 14, 2025
91e385f
Add a test
JakubValtar Feb 14, 2025
0cee9a7
Fix formatting
JakubValtar Feb 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 9 additions & 16 deletions crates/polars-core/src/chunked_array/ops/sort/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ use compare_inner::NonNull;
use rayon::prelude::*;
pub use slice::*;

use crate::chunked_array::ops::row_encode::_get_rows_encoded_ca;
use crate::prelude::compare_inner::TotalOrdInner;
use crate::prelude::sort::arg_sort_multiple::*;
use crate::prelude::*;
Expand Down Expand Up @@ -645,26 +644,20 @@ impl ChunkSort<BinaryOffsetType> for BinaryOffsetChunked {
}
}

#[cfg(feature = "dtype-struct")]
impl StructChunked {
pub(crate) fn arg_sort(&self, options: SortOptions) -> IdxCa {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why was this removed? Was it unused?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems like this method was left over when ChunkSort was implemented for structs in #17522. A method with the same name, signature and implementation was added in ChunkSort impl and now some calls resolve to this one and some to the ChunkSort one. I removed it, because it is duplicate code, confusing to read, and a potential footgun if the implementations are not kept in sync.

This is the situation:

impl StructChunked {
    // <<< A >>> only used in `sort_with` below and one other place
    pub(crate) fn arg_sort(&self, options: SortOptions) -> IdxCa {
        let bin = _get_rows_encoded_ca(...).unwrap(); // same impl as B, just `get_row_encoded` inlined
        bin.arg_sort(Default::default())
    }
}

impl ChunkSort<StructType> for StructChunked {
    fn sort_with(&self, mut options: SortOptions) -> ChunkedArray<StructType> {
        ...
        let idx = self.arg_sort(options); // Looks like it calls B, but actually calls A
        ...
    }

    fn sort(&self, descending: bool) -> ChunkedArray<StructType> {
        self.sort_with(...) // Also calls A through `sort_with`
    }

    // <<< B >>> used in `T: ChunkSort` generic contexts
    fn arg_sort(&self, options: SortOptions) -> IdxCa {
        let bin = self.get_row_encoded(options).unwrap();
        bin.arg_sort(Default::default())
    }
}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, now I see.

let bin = _get_rows_encoded_ca(
self.name().clone(),
&[self.clone().into_column()],
&[options.descending],
&[options.nulls_last],
)
.unwrap();
bin.arg_sort(Default::default())
}
}

#[cfg(feature = "dtype-struct")]
impl ChunkSort<StructType> for StructChunked {
fn sort_with(&self, mut options: SortOptions) -> ChunkedArray<StructType> {
options.multithreaded &= POOL.current_num_threads() > 1;
let idx = self.arg_sort(options);
unsafe { self.take_unchecked(&idx) }
let mut out = unsafe { self.take_unchecked(&idx) };

let s = if options.descending {
IsSorted::Descending
} else {
IsSorted::Ascending
};
out.set_sorted_flag(s);
out
}

fn sort(&self, descending: bool) -> ChunkedArray<StructType> {
Expand Down
6 changes: 4 additions & 2 deletions crates/polars-core/src/series/implementations/struct_.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,12 @@ impl PrivateSeries for SeriesWrap<StructChunked> {
}

fn _get_flags(&self) -> StatisticsFlags {
StatisticsFlags::empty()
self.0.get_flags()
}

fn _set_flags(&mut self, _flags: StatisticsFlags) {}
fn _set_flags(&mut self, flags: StatisticsFlags) {
self.0.set_flags(flags);
}

// TODO! remove this. Very slow. Asof join should use row-encoding.
unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
Expand Down
10 changes: 10 additions & 0 deletions py-polars/tests/unit/operations/test_is_sorted.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,3 +427,13 @@ def test_is_sorted_chunked_select() -> None:
def test_is_sorted_arithmetic_overflow_14106() -> None:
s = pl.Series([0, 200], dtype=pl.UInt8).sort()
assert not (s + 200).is_sorted()


def test_is_sorted_struct() -> None:
s = pl.Series("a", [{"x": 3}, {"x": 1}, {"x": 2}]).sort()
assert s.flags["SORTED_ASC"]
assert not s.flags["SORTED_DESC"]

s = s.sort(descending=True)
assert s.flags["SORTED_DESC"]
assert not s.flags["SORTED_ASC"]
Loading