refactor(rust): Support for named/anonymous aggregations (#25118)

Author: ritchie46

crates/polars-arrow/src/array/fixed_size_list/mutable.rs

@@ -149,6 +149,15 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
             validity.shrink_to_fit()
         }
     }
+
+    pub fn freeze(mut self) -> FixedSizeListArray {
+        FixedSizeListArray::new(
+            self.dtype,
+            self.length,
+            self.values.as_box(),
+            self.validity.map(|b| b.freeze()),
+        )
+    }
 }
 
 impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {

crates/polars-expr/src/planner.rs

@@ -634,5 +634,8 @@ fn create_physical_expr_inner(
                 false,
             )))
         },
+        AnonymousStreamingAgg { .. } => {
+            polars_bail!(ComputeError: "anonymous agg not supported in in-memory engine")
+        },
     }
 }

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

@@ -143,6 +143,20 @@ pub fn into_reduction(
                 _ => unreachable!(),
             }
         },
+        AExpr::AnonymousStreamingAgg {
+            input: inner_exprs,
+            fmt_str: _,
+            function,
+        } => {
+            let ann_agg = function.materialize()?;
+            assert!(inner_exprs.len() == 1);
+            let input = inner_exprs[0].node();
+            let reduction = ann_agg.as_any();
+            let reduction = reduction
+                .downcast_ref::<Box<dyn GroupedReduction>>()
+                .unwrap();
+            (reduction.new_empty(), input)
+        },
         _ => unreachable!(),
     };
     Ok(out)

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

@@ -219,7 +219,7 @@ pub struct VecGroupedReduction<R: Reducer> {
 }
 
 impl<R: Reducer> VecGroupedReduction<R> {
-    fn new(in_dtype: DataType, reducer: R) -> Self {
+    pub fn new(in_dtype: DataType, reducer: R) -> Self {
         Self {
             values: Vec::new(),
             evicted_values: Vec::new(),
@@ -486,7 +486,7 @@ where
 }
 
 #[derive(Clone)]
-struct NullGroupedReduction {
+pub struct NullGroupedReduction {
     num_groups: IdxSize,
     num_evictions: IdxSize,
     output: Scalar,

crates/polars-plan/src/dsl/expr/anonymous/agg.rs

@@ -0,0 +1,62 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use polars_core::prelude::Field;
+use polars_core::schema::Schema;
+use polars_error::{PolarsResult, feature_gated};
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+use super::SpecialEq;
+use crate::dsl::LazySerde;
+
+pub trait AnonymousStreamingAgg: Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+
+    fn get_field(&self, input_schema: &Schema, fields: &[Field]) -> PolarsResult<Field>;
+}
+
+pub type OpaqueStreamingAgg = LazySerde<SpecialEq<Arc<dyn AnonymousStreamingAgg>>>;
+
+impl OpaqueStreamingAgg {
+    pub fn materialize(&self) -> PolarsResult<SpecialEq<Arc<dyn AnonymousStreamingAgg>>> {
+        match self {
+            Self::Deserialized(t) => Ok(t.clone()),
+            Self::Named {
+                name,
+                payload,
+                value,
+            } => feature_gated!("serde", {
+                use super::named_serde::NAMED_SERDE_REGISTRY_EXPR;
+                match value {
+                    Some(v) => Ok(v.clone()),
+                    None => Ok(SpecialEq::new(
+                        NAMED_SERDE_REGISTRY_EXPR
+                            .read()
+                            .unwrap()
+                            .as_ref()
+                            .expect("NAMED EXPR REGISTRY NOT SET")
+                            .get_agg(name, payload.as_ref().unwrap())?
+                            .expect("NAMED AGG NOT FOUND"),
+                    )),
+                }
+            }),
+            Self::Bytes(_b) => {
+                feature_gated!("serde", {
+                    use crate::dsl::anonymous::serde_expr;
+                    serde_expr::deserialize_anon_agg(_b.as_ref()).map(SpecialEq::new)
+                })
+            },
+        }
+    }
+}
+
+#[cfg(feature = "ir_serde")]
+impl Serialize for SpecialEq<Arc<dyn AnonymousStreamingAgg>> {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        unreachable!("should not be hit")
+    }
+}

crates/polars-plan/src/dsl/expr/anonymous/expr.rs

@@ -139,7 +139,7 @@ where
 }
 
 pub type OpaqueColumnUdf = LazySerde<SpecialEq<Arc<dyn AnonymousColumnsUdf>>>;
-pub(crate) fn new_column_udf<F: AnonymousColumnsUdf + 'static>(func: F) -> OpaqueColumnUdf {
+pub fn new_column_udf<F: AnonymousColumnsUdf + 'static>(func: F) -> OpaqueColumnUdf {
     LazySerde::Deserialized(SpecialEq::new(Arc::new(func)))
 }
 

crates/polars-plan/src/dsl/expr/anonymous/mod.rs

@@ -1,5 +1,9 @@
+mod agg;
 mod expr;
+
+pub use agg::{AnonymousStreamingAgg, OpaqueStreamingAgg};
 pub use expr::*;
+
 #[cfg(feature = "dsl-schema")]
 mod json_schema;
 #[cfg(feature = "serde")]

crates/polars-plan/src/dsl/expr/anonymous/named_serde.rs

@@ -1,11 +1,26 @@
 use std::sync::{Arc, LazyLock, RwLock};
 
+use polars_error::PolarsResult;
+
 use super::AnonymousColumnsUdf;
+use super::agg::AnonymousStreamingAgg;
 
 // Can be used to have named anonymous functions.
 // The receiver must have implemented this registry and map the names to the proper UDFs.
 pub trait ExprRegistry: Sync + Send {
-    fn get_function(&self, name: &str, payload: &[u8]) -> Option<Arc<dyn AnonymousColumnsUdf>>;
+    #[allow(unused)]
+    fn get_function(&self, name: &str, payload: &[u8]) -> Option<Arc<dyn AnonymousColumnsUdf>> {
+        None
+    }
+
+    #[allow(unused)]
+    fn get_agg(
+        &self,
+        name: &str,
+        payload: &[u8],
+    ) -> PolarsResult<Option<Arc<dyn AnonymousStreamingAgg>>> {
+        Ok(None)
+    }
 }
 
 pub(super) static NAMED_SERDE_REGISTRY_EXPR: LazyLock<RwLock<Option<Arc<dyn ExprRegistry>>>> =

crates/polars-plan/src/dsl/expr/anonymous/serde_expr.rs

@@ -119,6 +119,35 @@ impl<'a> Deserialize<'a> for SpecialEq<Arc<dyn AnonymousColumnsUdf>> {
     }
 }
 
+impl<'a> Deserialize<'a> for SpecialEq<Arc<dyn AnonymousStreamingAgg>> {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'a>,
+    {
+        use serde::de::Error;
+        deserialize_map_bytes(deserializer, |buf| {
+            deserialize_anon_agg(&buf)
+                .map_err(|e| D::Error::custom(format!("{e}")))
+                .map(SpecialEq::new)
+        })?
+    }
+}
+
+pub(super) fn deserialize_anon_agg(buf: &[u8]) -> PolarsResult<Arc<dyn AnonymousStreamingAgg>> {
+    if buf.starts_with(NAMED_SERDE_MAGIC_BYTE_MARK) {
+        let (reg, name, payload) = deserialize_named_registry(buf)?;
+
+        if let Some(func) = reg.get_agg(name, payload)? {
+            Ok(func)
+        } else {
+            let msg = "name not found in named serde registry";
+            polars_bail!(ComputeError: msg)
+        }
+    } else {
+        polars_bail!(ComputeError: "deserialization not supported for this 'opaque' function")
+    }
+}
+
 // Serialize SpecialEq<T>
 
 impl Serialize for SpecialEq<Series> {

crates/polars-plan/src/plans/aexpr/builder.rs

@@ -55,6 +55,18 @@ impl AExprBuilder {
         )
     }
 
+    pub fn map_as_expr_ir<F: Fn(ExprIR, &mut Arena<AExpr>) -> AExpr>(
+        self,
+        mapper: F,
+        arena: &mut Arena<AExpr>,
+    ) -> Self {
+        let eir = ExprIR::from_node(self.node, arena);
+
+        let ae = mapper(eir, arena);
+        let node = arena.add(ae);
+        Self { node }
+    }
+
     pub fn row_encode_unary(
         self,
         variant: RowEncodingVariant,