feat(rust, python): out of core sort on multiple columns (#7244)

polars/polars-core/src/chunked_array/ops/mod.rs

@@ -40,7 +40,7 @@ mod reverse;
 pub(crate) mod rolling_window;
 mod set;
 mod shift;
-pub(crate) mod sort;
+pub mod sort;
 pub(crate) mod take;
 pub(crate) mod unique;
 #[cfg(feature = "zip_with")]

polars/polars-core/src/chunked_array/ops/sort/arg_sort_multiple.rs

@@ -1,4 +1,5 @@
 use polars_arrow::data_types::IsFloat;
+use polars_row::{convert_columns, RowsEncoded, SortField};
 use polars_utils::iter::EnumerateIdxTrait;
 
 use super::*;
@@ -67,45 +68,54 @@ pub(crate) fn arg_sort_multiple_impl<T: PartialOrd + Send + IsFloat + Copy>(
     Ok(ca.into_inner())
 }
 
-pub(crate) fn argsort_multiple_row_fmt(
+pub fn _get_rows_encoded_compat_array(by: &Series) -> PolarsResult<ArrayRef> {
+    let by = convert_sort_column_multi_sort(by, true)?;
+    let by = by.rechunk();
+
+    let out = match by.dtype() {
+        #[cfg(feature = "dtype-categorical")]
+        DataType::Categorical(_) => {
+            let ca = by.categorical().unwrap();
+            if ca.use_lexical_sort() {
+                by.to_arrow(0)
+            } else {
+                ca.logical().chunks[0].clone()
+            }
+        }
+        _ => by.to_arrow(0),
+    };
+    Ok(out)
+}
+
+pub fn _get_rows_encoded(
     by: &[Series],
-    mut descending: Vec<bool>,
+    descending: &[bool],
     nulls_last: bool,
-    parallel: bool,
-) -> PolarsResult<IdxCa> {
-    use polars_row::{convert_columns, SortField};
-    broadcast_descending(by.len(), &mut descending);
-
+) -> PolarsResult<RowsEncoded> {
+    debug_assert_eq!(by.len(), descending.len());
     let mut cols = Vec::with_capacity(by.len());
     let mut fields = Vec::with_capacity(by.len());
-
-    debug_assert_eq!(by.len(), descending.len());
     for (by, descending) in by.iter().zip(descending) {
-        let by = convert_sort_column_multi_sort(by, true)?;
-        let by = by.rechunk();
-
-        let arr = match by.dtype() {
-            #[cfg(feature = "dtype-categorical")]
-            DataType::Categorical(_) => {
-                let ca = by.categorical().unwrap();
-                if ca.use_lexical_sort() {
-                    by.to_arrow(0)
-                } else {
-                    ca.logical().chunks[0].clone()
-                }
-            }
-            _ => by.to_arrow(0),
-        };
-        let data_type = arr.data_type().clone();
+        let arr = _get_rows_encoded_compat_array(by)?;
 
         cols.push(arr);
         fields.push(SortField {
-            descending,
+            descending: *descending,
             nulls_last,
-            data_type,
         })
     }
-    let rows_encoded = convert_columns(&cols, fields);
+    Ok(convert_columns(&cols, &fields))
+}
+
+pub(crate) fn argsort_multiple_row_fmt(
+    by: &[Series],
+    mut descending: Vec<bool>,
+    nulls_last: bool,
+    parallel: bool,
+) -> PolarsResult<IdxCa> {
+    _broadcast_descending(by.len(), &mut descending);
+
+    let rows_encoded = _get_rows_encoded(by, &descending, nulls_last)?;
     let mut items: Vec<_> = rows_encoded.iter().enumerate_idx().collect();
 
     if parallel {

polars/polars-core/src/chunked_array/ops/sort/mod.rs

@@ -1,6 +1,6 @@
 mod arg_sort;
 #[cfg(feature = "sort_multiple")]
-mod arg_sort_multiple;
+pub mod arg_sort_multiple;
 #[cfg(feature = "dtype-categorical")]
 mod categorical;
 
@@ -742,7 +742,7 @@ pub(crate) fn convert_sort_column_multi_sort(
     Ok(out)
 }
 
-pub(super) fn broadcast_descending(n_cols: usize, descending: &mut Vec<bool>) {
+pub fn _broadcast_descending(n_cols: usize, descending: &mut Vec<bool>) {
     if n_cols > descending.len() && descending.len() == 1 {
         while n_cols != descending.len() {
             descending.push(descending[0]);
@@ -765,7 +765,7 @@ pub(crate) fn prepare_arg_sort(
     let first = columns.remove(0);
 
     // broadcast ordering
-    broadcast_descending(n_cols, &mut descending);
+    _broadcast_descending(n_cols, &mut descending);
     Ok((first, columns, descending))
 }
 

polars/polars-core/src/frame/mod.rs

@@ -1160,6 +1160,16 @@ impl DataFrame {
         inner(self, series)
     }
 
+    /// Adds a column to the `DataFrame` without doing any checks
+    /// on length or duplicates.
+    ///
+    /// # Safety
+    /// The caller must ensure `column.len() == self.height()` .
+    pub unsafe fn with_column_unchecked(&mut self, column: Series) -> &mut Self {
+        self.get_columns_mut().push(column);
+        self
+    }
+
     fn add_column_by_schema(&mut self, s: Series, schema: &Schema) -> PolarsResult<()> {
         let name = s.name();
         if let Some((idx, _, _)) = schema.get_full(name) {

polars/polars-lazy/polars-pipe/Cargo.toml

@@ -18,6 +18,7 @@ polars-core = { version = "0.27.2", path = "../../polars-core", features = ["laz
 polars-io = { version = "0.27.2", path = "../../polars-io", default-features = false, features = ["ipc"] }
 polars-ops = { version = "0.27.2", path = "../../polars-ops", features = ["search_sorted"] }
 polars-plan = { version = "0.27.2", path = "../polars-plan", default-features = false, features = ["compile"] }
+polars-row = { version = "0.27.2", path = "../../polars-row" }
 polars-utils = { version = "0.27.2", path = "../../polars-utils", features = ["sysinfo"] }
 rayon.workspace = true
 smartstring = { version = "1" }

polars/polars-lazy/polars-pipe/src/executors/sinks/file_sink.rs

@@ -180,7 +180,7 @@ impl Sink for FilesSink {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, _other: Box<dyn Sink>) {
+    fn combine(&mut self, _other: &mut dyn Sink) {
         // already synchronized
     }
 

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs

@@ -309,7 +309,7 @@ impl Sink for GenericGroupbySink {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         // don't parallel this as this is already done in parallel.
 
         let other = other.as_any().downcast_ref::<Self>().unwrap();

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/primitive.rs

@@ -293,7 +293,7 @@ where
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         // don't parallel this as this is already done in parallel.
         let other = other.as_any().downcast_ref::<Self>().unwrap();
 

polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/string.rs

@@ -252,7 +252,7 @@ impl Sink for Utf8GroupbySink {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         // don't parallel this as this is already done in parallel.
 
         let other = other.as_any().downcast_ref::<Self>().unwrap();

polars/polars-lazy/polars-pipe/src/executors/sinks/joins/cross.rs

@@ -34,7 +34,7 @@ impl Sink for CrossJoin {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         let other = other.as_any().downcast_mut::<Self>().unwrap();
         let other_chunks = std::mem::take(&mut other.chunks);
         self.chunks.extend(other_chunks.into_iter());

polars/polars-lazy/polars-pipe/src/executors/sinks/joins/generic_build.rs

@@ -244,7 +244,7 @@ impl Sink for GenericBuild {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         if self.is_empty() {
             let other = other.as_any().downcast_mut::<Self>().unwrap();
             if !other.is_empty() {

polars/polars-lazy/polars-pipe/src/executors/sinks/ordered.rs

@@ -31,7 +31,7 @@ impl Sink for OrderedSink {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         let other = other.as_any().downcast_ref::<OrderedSink>().unwrap();
         self.chunks.extend_from_slice(&other.chunks);
         self.sort();

polars/polars-lazy/polars-pipe/src/executors/sinks/slice.rs

@@ -75,7 +75,7 @@ impl Sink for SliceSink {
         }
     }
 
-    fn combine(&mut self, _other: Box<dyn Sink>) {
+    fn combine(&mut self, _other: &mut dyn Sink) {
         // no-op
     }
 

polars/polars-lazy/polars-pipe/src/executors/sinks/sort/mod.rs

@@ -1,6 +1,8 @@
 mod io;
 mod ooc;
 mod sink;
+mod sink_multiple;
 mod source;
 
 pub(crate) use sink::SortSink;
+pub(crate) use sink_multiple::SortSinkMultiple;

polars/polars-lazy/polars-pipe/src/executors/sinks/sort/sink.rs

@@ -5,8 +5,9 @@ use std::sync::{Arc, Mutex};
 
 use polars_core::error::PolarsResult;
 use polars_core::frame::DataFrame;
-use polars_core::prelude::{AnyValue, SchemaRef, Series};
+use polars_core::prelude::{AnyValue, SchemaRef, Series, SortOptions};
 use polars_core::utils::accumulate_dataframes_vertical_unchecked;
+use polars_plan::prelude::SortArguments;
 use polars_utils::atomic::SyncCounter;
 use polars_utils::sys::MEMINFO;
 
@@ -26,19 +27,13 @@ pub struct SortSink {
     io_thread: Arc<Mutex<Option<IOThread>>>,
     // location in the dataframe of the columns to sort by
     sort_idx: usize,
-    descending: bool,
-    slice: Option<(i64, usize)>,
+    sort_args: SortArguments,
     // sampled values so we can find the distribution.
     dist_sample: Vec<AnyValue<'static>>,
 }
 
 impl SortSink {
-    pub(crate) fn new(
-        sort_idx: usize,
-        descending: bool,
-        schema: SchemaRef,
-        slice: Option<(i64, usize)>,
-    ) -> Self {
+    pub(crate) fn new(sort_idx: usize, sort_args: SortArguments, schema: SchemaRef) -> Self {
         // for testing purposes
         let ooc = std::env::var("POLARS_FORCE_OOC_SORT").is_ok();
 
@@ -50,8 +45,7 @@ impl SortSink {
             ooc,
             io_thread: Default::default(),
             sort_idx,
-            descending,
-            slice,
+            sort_args,
             dist_sample: vec![],
         };
         if ooc {
@@ -129,7 +123,7 @@ impl Sink for SortSink {
         Ok(SinkResult::CanHaveMoreInput)
     }
 
-    fn combine(&mut self, mut other: Box<dyn Sink>) {
+    fn combine(&mut self, other: &mut dyn Sink) {
         let other = other.as_any().downcast_mut::<Self>().unwrap();
         self.chunks.extend(std::mem::take(&mut other.chunks));
         self.ooc |= other.ooc;
@@ -150,9 +144,8 @@ impl Sink for SortSink {
             ooc: self.ooc,
             io_thread: self.io_thread.clone(),
             sort_idx: self.sort_idx,
-            descending: self.descending,
+            sort_args: self.sort_args.clone(),
             dist_sample: vec![],
-            slice: self.slice,
         })
     }
 
@@ -168,15 +161,30 @@ impl Sink for SortSink {
             let io_thread = lock.as_ref().unwrap();
 
             let dist = Series::from_any_values("", &self.dist_sample).unwrap();
-            let dist = dist.sort(self.descending);
+            let dist = dist.sort_with(SortOptions {
+                descending: self.sort_args.descending[0],
+                nulls_last: self.sort_args.nulls_last,
+                multithreaded: true,
+            });
 
             block_thread_until_io_thread_done(io_thread);
 
-            sort_ooc(io_thread, dist, self.sort_idx, self.descending, self.slice)
+            sort_ooc(
+                io_thread,
+                dist,
+                self.sort_idx,
+                self.sort_args.descending[0],
+                self.sort_args.slice,
+            )
         } else {
             let chunks = std::mem::take(&mut self.chunks);
             let df = accumulate_dataframes_vertical_unchecked(chunks);
-            let df = sort_accumulated(df, self.sort_idx, self.descending, self.slice)?;
+            let df = sort_accumulated(
+                df,
+                self.sort_idx,
+                self.sort_args.descending[0],
+                self.sort_args.slice,
+            )?;
             Ok(FinalizedSink::Finished(df))
         }
     }