feat: Add per partition sort and finish callback to sinks (#22789)

Author: coastalwhite

crates/polars-expr/src/reduce/mod.rs

@@ -15,6 +15,7 @@ use std::marker::PhantomData;
 use arrow::array::{Array, PrimitiveArray, StaticArray};
 use arrow::bitmap::{Bitmap, BitmapBuilder, MutableBitmap};
 pub use convert::into_reduction;
+pub use min_max::{new_max_reduction, new_min_reduction};
 use polars_core::prelude::*;
 
 use crate::EvictIdx;

crates/polars-lazy/src/frame/mod.rs

@@ -1081,6 +1081,7 @@ impl LazyFrame {
     /// final result doesn't fit into memory. This methods will return an error if the query cannot
     /// be completely done in a streaming fashion.
     #[cfg(feature = "parquet")]
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_parquet_partitioned(
         self,
         base_path: Arc<PathBuf>,
@@ -1089,6 +1090,8 @@ impl LazyFrame {
         options: ParquetWriteOptions,
         cloud_options: Option<polars_io::cloud::CloudOptions>,
         sink_options: SinkOptions,
+        per_partition_sort_by: Option<Vec<SortColumn>>,
+        finish_callback: Option<SinkFinishCallback>,
     ) -> PolarsResult<Self> {
         self.sink(SinkType::Partition(PartitionSinkType {
             base_path,
@@ -1097,13 +1100,16 @@ impl LazyFrame {
             variant,
             file_type: FileType::Parquet(options),
             cloud_options,
+            per_partition_sort_by,
+            finish_callback,
         }))
     }
 
     /// Stream a query result into an ipc/arrow file in a partitioned manner. This is useful if the
     /// final result doesn't fit into memory. This methods will return an error if the query cannot
     /// be completely done in a streaming fashion.
     #[cfg(feature = "ipc")]
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_ipc_partitioned(
         self,
         base_path: Arc<PathBuf>,
@@ -1112,6 +1118,8 @@ impl LazyFrame {
         options: IpcWriterOptions,
         cloud_options: Option<polars_io::cloud::CloudOptions>,
         sink_options: SinkOptions,
+        per_partition_sort_by: Option<Vec<SortColumn>>,
+        finish_callback: Option<SinkFinishCallback>,
     ) -> PolarsResult<Self> {
         self.sink(SinkType::Partition(PartitionSinkType {
             base_path,
@@ -1120,13 +1128,16 @@ impl LazyFrame {
             variant,
             file_type: FileType::Ipc(options),
             cloud_options,
+            per_partition_sort_by,
+            finish_callback,
         }))
     }
 
     /// Stream a query result into an csv file in a partitioned manner. This is useful if the final
     /// result doesn't fit into memory. This methods will return an error if the query cannot be
     /// completely done in a streaming fashion.
     #[cfg(feature = "csv")]
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_csv_partitioned(
         self,
         base_path: Arc<PathBuf>,
@@ -1135,6 +1146,8 @@ impl LazyFrame {
         options: CsvWriterOptions,
         cloud_options: Option<polars_io::cloud::CloudOptions>,
         sink_options: SinkOptions,
+        per_partition_sort_by: Option<Vec<SortColumn>>,
+        finish_callback: Option<SinkFinishCallback>,
     ) -> PolarsResult<Self> {
         self.sink(SinkType::Partition(PartitionSinkType {
             base_path,
@@ -1143,13 +1156,16 @@ impl LazyFrame {
             variant,
             file_type: FileType::Csv(options),
             cloud_options,
+            per_partition_sort_by,
+            finish_callback,
         }))
     }
 
     /// Stream a query result into a JSON file in a partitioned manner. This is useful if the final
     /// result doesn't fit into memory. This methods will return an error if the query cannot be
     /// completely done in a streaming fashion.
     #[cfg(feature = "json")]
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_json_partitioned(
         self,
         base_path: Arc<PathBuf>,
@@ -1158,6 +1174,8 @@ impl LazyFrame {
         options: JsonWriterOptions,
         cloud_options: Option<polars_io::cloud::CloudOptions>,
         sink_options: SinkOptions,
+        per_partition_sort_by: Option<Vec<SortColumn>>,
+        finish_callback: Option<SinkFinishCallback>,
     ) -> PolarsResult<Self> {
         self.sink(SinkType::Partition(PartitionSinkType {
             base_path,
@@ -1166,6 +1184,8 @@ impl LazyFrame {
             variant,
             file_type: FileType::Json(options),
             cloud_options,
+            per_partition_sort_by,
+            finish_callback,
         }))
     }
 

crates/polars-plan/src/dsl/options/sink.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use polars_core::error::PolarsResult;
+use polars_core::frame::DataFrame;
 use polars_core::prelude::DataType;
 use polars_core::scalar::Scalar;
 use polars_io::cloud::CloudOptions;
@@ -99,6 +100,13 @@ impl SinkTarget {
             )),
         }
     }
+
+    pub fn to_display_string(&self) -> String {
+        match self {
+            Self::Path(p) => p.display().to_string(),
+            Self::Dyn(_) => "dynamic-target".to_string(),
+        }
+    }
 }
 
 impl fmt::Debug for SinkTarget {
@@ -260,6 +268,47 @@ pub enum PartitionTargetCallback {
     Python(polars_utils::python_function::PythonFunction),
 }
 
+#[cfg_attr(feature = "python", pyo3::pyclass)]
+pub struct SinkWritten {
+    pub file_idx: usize,
+    pub part_idx: usize,
+    pub in_part_idx: usize,
+    pub keys: Vec<PartitionTargetContextKey>,
+    pub file_path: PathBuf,
+    pub full_path: PathBuf,
+    pub num_rows: usize,
+    pub file_size: usize,
+    pub gathered: Option<DataFrame>,
+}
+
+#[cfg_attr(feature = "python", pyo3::pyclass)]
+pub struct SinkFinishContext {
+    pub written: Vec<SinkWritten>,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum SinkFinishCallback {
+    Rust(SpecialEq<Arc<dyn Fn(DataFrame) -> PolarsResult<()> + Send + Sync>>),
+    #[cfg(feature = "python")]
+    Python(polars_utils::python_function::PythonFunction),
+}
+
+impl SinkFinishCallback {
+    pub fn call(&self, df: DataFrame) -> PolarsResult<()> {
+        match self {
+            Self::Rust(f) => f(df),
+            #[cfg(feature = "python")]
+            Self::Python(f) => pyo3::Python::with_gil(|py| {
+                let converter =
+                    polars_utils::python_convert_registry::get_python_convert_registry();
+                let df = (converter.to_py.df)(Box::new(df) as Box<dyn std::any::Any>)?;
+                f.call1(py, (df,))?;
+                PolarsResult::Ok(())
+            }),
+        }
+    }
+}
+
 impl PartitionTargetCallback {
     pub fn call(&self, ctx: PartitionTargetContext) -> PolarsResult<SinkTarget> {
         match self {
@@ -277,6 +326,60 @@ impl PartitionTargetCallback {
     }
 }
 
+#[cfg(feature = "serde")]
+impl serde::Serialize for SinkFinishCallback {
+    fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        use serde::ser::Error;
+
+        #[cfg(feature = "python")]
+        if let Self::Python(v) = self {
+            return v.serialize(_serializer);
+        }
+
+        Err(S::Error::custom(format!("cannot serialize {self:?}")))
+    }
+}
+
+#[cfg(feature = "serde")]
+impl<'de> serde::Deserialize<'de> for SinkFinishCallback {
+    fn deserialize<D>(_deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        #[cfg(feature = "python")]
+        {
+            Ok(Self::Python(
+                polars_utils::python_function::PythonFunction::deserialize(_deserializer)?,
+            ))
+        }
+        #[cfg(not(feature = "python"))]
+        {
+            use serde::de::Error;
+            Err(D::Error::custom(
+                "cannot deserialize PartitionOutputCallback",
+            ))
+        }
+    }
+}
+
+#[cfg(feature = "dsl-schema")]
+impl schemars::JsonSchema for SinkFinishCallback {
+    fn schema_name() -> String {
+        "PartitionTargetCallback".to_owned()
+    }
+
+    fn schema_id() -> std::borrow::Cow<'static, str> {
+        std::borrow::Cow::Borrowed(concat!(module_path!(), "::", "SinkFinishCallback"))
+    }
+
+    fn json_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema {
+        Vec::<u8>::json_schema(generator)
+    }
+}
+
 #[cfg(feature = "serde")]
 impl<'de> serde::Deserialize<'de> for PartitionTargetCallback {
     fn deserialize<D>(_deserializer: D) -> Result<Self, D::Error>
@@ -331,6 +434,23 @@ impl schemars::JsonSchema for PartitionTargetCallback {
     }
 }
 
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
+#[derive(Clone, Debug, PartialEq)]
+pub struct SortColumn {
+    pub expr: Expr,
+    pub descending: bool,
+    pub nulls_last: bool,
+}
+
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Clone, Debug, PartialEq)]
+pub struct SortColumnIR {
+    pub expr: ExprIR,
+    pub descending: bool,
+    pub nulls_last: bool,
+}
+
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
 #[derive(Clone, Debug, PartialEq)]
@@ -341,6 +461,8 @@ pub struct PartitionSinkType {
     pub sink_options: SinkOptions,
     pub variant: PartitionVariant,
     pub cloud_options: Option<polars_io::cloud::CloudOptions>,
+    pub per_partition_sort_by: Option<Vec<SortColumn>>,
+    pub finish_callback: Option<SinkFinishCallback>,
 }
 
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -352,6 +474,8 @@ pub struct PartitionSinkTypeIR {
     pub sink_options: SinkOptions,
     pub variant: PartitionVariantIR,
     pub cloud_options: Option<polars_io::cloud::CloudOptions>,
+    pub per_partition_sort_by: Option<Vec<SortColumnIR>>,
+    pub finish_callback: Option<SinkFinishCallback>,
 }
 
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -392,23 +516,44 @@ pub enum PartitionVariantIR {
     },
 }
 
+#[cfg(feature = "cse")]
 impl SinkTypeIR {
-    #[cfg(feature = "cse")]
     pub(crate) fn traverse_and_hash<H: Hasher>(&self, expr_arena: &Arena<AExpr>, state: &mut H) {
         std::mem::discriminant(self).hash(state);
         match self {
             Self::Memory => {},
             Self::File(f) => f.hash(state),
-            Self::Partition(f) => {
-                f.file_type.hash(state);
-                f.sink_options.hash(state);
-                f.variant.traverse_and_hash(expr_arena, state);
-                f.cloud_options.hash(state);
-            },
+            Self::Partition(f) => f.traverse_and_hash(expr_arena, state),
+        }
+    }
+}
+
+#[cfg(feature = "cse")]
+impl PartitionSinkTypeIR {
+    pub(crate) fn traverse_and_hash<H: Hasher>(&self, expr_arena: &Arena<AExpr>, state: &mut H) {
+        self.file_type.hash(state);
+        self.sink_options.hash(state);
+        self.variant.traverse_and_hash(expr_arena, state);
+        self.cloud_options.hash(state);
+        std::mem::discriminant(&self.per_partition_sort_by).hash(state);
+        if let Some(v) = &self.per_partition_sort_by {
+            v.len().hash(state);
+            for v in v {
+                v.traverse_and_hash(expr_arena, state);
+            }
         }
     }
 }
 
+#[cfg(feature = "cse")]
+impl SortColumnIR {
+    pub(crate) fn traverse_and_hash<H: Hasher>(&self, expr_arena: &Arena<AExpr>, state: &mut H) {
+        self.expr.traverse_and_hash(expr_arena, state);
+        self.descending.hash(state);
+        self.nulls_last.hash(state);
+    }
+}
+
 impl PartitionVariantIR {
     #[cfg(feature = "cse")]
     pub(crate) fn traverse_and_hash<H: Hasher>(&self, expr_arena: &Arena<AExpr>, state: &mut H) {

crates/polars-plan/src/plans/conversion/dsl_to_ir.rs

@@ -1186,6 +1186,25 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
                         },
                     },
                     cloud_options: f.cloud_options,
+                    per_partition_sort_by: match f.per_partition_sort_by {
+                        None => None,
+                        Some(sort_by) => Some(
+                            sort_by
+                                .into_iter()
+                                .map(|s| {
+                                    let expr = to_expr_ir(s.expr, ctxt.expr_arena)?;
+                                    ctxt.conversion_optimizer
+                                        .push_scratch(expr.node(), ctxt.expr_arena);
+                                    Ok(SortColumnIR {
+                                        expr,
+                                        descending: s.descending,
+                                        nulls_last: s.nulls_last,
+                                    })
+                                })
+                                .collect::<PolarsResult<Vec<_>>>()?,
+                        ),
+                    },
+                    finish_callback: f.finish_callback,
                 }),
             };
 

crates/polars-plan/src/plans/conversion/mod.rs

@@ -302,6 +302,17 @@ impl IR {
                             },
                         },
                         cloud_options: f.cloud_options,
+                        per_partition_sort_by: f.per_partition_sort_by.map(|sort_by| {
+                            sort_by
+                                .into_iter()
+                                .map(|s| SortColumn {
+                                    expr: s.expr.to_expr(expr_arena),
+                                    descending: s.descending,
+                                    nulls_last: s.descending,
+                                })
+                                .collect()
+                        }),
+                        finish_callback: f.finish_callback,
                     }),
                 };
                 DslPlan::Sink { input, payload }