feat: Add maintain order argument on implode (#26782)

Author: orlp

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

@@ -211,18 +211,20 @@ impl MutableBooleanArray {
                     validity.extend_constant(additional, true);
                 }
             },
-            None => {
-                self.values.extend_constant(additional, false);
-                if let Some(validity) = self.validity.as_mut() {
-                    validity.extend_constant(additional, false)
-                } else {
-                    self.init_validity();
-                    self.validity
-                        .as_mut()
-                        .unwrap()
-                        .extend_constant(additional, false)
-                };
-            },
+            None => self.extend_null(additional),
+        };
+    }
+
+    pub fn extend_null(&mut self, additional: usize) {
+        self.values.extend_constant(additional, false);
+        if let Some(validity) = self.validity.as_mut() {
+            validity.extend_constant(additional, false)
+        } else {
+            self.init_validity();
+            self.validity
+                .as_mut()
+                .unwrap()
+                .extend_constant(additional, false)
         };
     }
 

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

@@ -171,17 +171,21 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
                 validity.extend_constant(additional, true)
             }
         } else {
-            if let Some(validity) = &mut self.validity {
-                validity.extend_constant(additional, false)
-            } else {
-                let mut validity = MutableBitmap::with_capacity(self.values.capacity());
-                validity.extend_constant(self.len(), true);
-                validity.extend_constant(additional, false);
-                self.validity = Some(validity)
-            }
-            self.values
-                .resize(self.values.len() + additional, T::default());
+            self.extend_null(additional);
+        }
+    }
+
+    pub fn extend_null(&mut self, additional: usize) {
+        if let Some(validity) = &mut self.validity {
+            validity.extend_constant(additional, false)
+        } else {
+            let mut validity = MutableBitmap::with_capacity(self.values.capacity());
+            validity.extend_constant(self.len(), true);
+            validity.extend_constant(additional, false);
+            self.validity = Some(validity)
         }
+        self.values
+            .resize(self.values.len() + additional, T::default());
     }
 
     /// Extends the [`MutablePrimitiveArray`] from an iterator of trusted len.

crates/polars-core/src/chunked_array/builder/list/anonymous.rs

@@ -1,110 +1,5 @@
 use super::*;
 
-pub struct AnonymousListBuilder<'a> {
-    name: PlSmallStr,
-    builder: AnonymousBuilder<'a>,
-    fast_explode: bool,
-    inner_dtype: Option<DataType>,
-}
-
-impl Default for AnonymousListBuilder<'_> {
-    fn default() -> Self {
-        Self::new(PlSmallStr::EMPTY, 0, None)
-    }
-}
-
-impl<'a> AnonymousListBuilder<'a> {
-    pub fn new(name: PlSmallStr, capacity: usize, inner_dtype: Option<DataType>) -> Self {
-        Self {
-            name,
-            builder: AnonymousBuilder::new(capacity),
-            fast_explode: true,
-            inner_dtype,
-        }
-    }
-
-    pub fn append_opt_series(&mut self, opt_s: Option<&'a Series>) -> PolarsResult<()> {
-        match opt_s {
-            Some(s) => return self.append_series(s),
-            None => {
-                self.append_null();
-            },
-        }
-        Ok(())
-    }
-
-    pub fn append_opt_array(&mut self, opt_s: Option<&'a dyn Array>) {
-        match opt_s {
-            Some(s) => self.append_array(s),
-            None => {
-                self.append_null();
-            },
-        }
-    }
-
-    pub fn append_array(&mut self, arr: &'a dyn Array) {
-        self.builder.push(arr)
-    }
-
-    #[inline]
-    pub fn append_null(&mut self) {
-        self.fast_explode = false;
-        self.builder.push_null();
-    }
-
-    #[inline]
-    pub fn append_empty(&mut self) {
-        self.fast_explode = false;
-        self.builder.push_empty()
-    }
-
-    pub fn append_series(&mut self, s: &'a Series) -> PolarsResult<()> {
-        match (s.dtype(), &self.inner_dtype) {
-            (DataType::Null, _) => {},
-            (dt, None) => self.inner_dtype = Some(dt.clone()),
-            (dt, Some(set_dt)) => {
-                polars_bail!(ComputeError: "dtypes don't match, got {}, expected: {}", dt.pretty_format(), set_dt.pretty_format());
-            },
-        }
-        if s.is_empty() {
-            self.append_empty();
-        } else {
-            self.builder.push_multiple(s.chunks());
-        }
-        Ok(())
-    }
-
-    pub fn finish(&mut self) -> ListChunked {
-        // Don't use self from here on out.
-        let slf = std::mem::take(self);
-        if slf.builder.is_empty() {
-            ListChunked::full_null_with_dtype(
-                slf.name.clone(),
-                0,
-                &slf.inner_dtype.unwrap_or(DataType::Null),
-            )
-        } else {
-            let inner_dtype_physical = self
-                .inner_dtype
-                .as_ref()
-                .map(|dt| dt.to_physical().to_arrow(CompatLevel::newest()));
-            let arr = slf.builder.finish(inner_dtype_physical.as_ref()).unwrap();
-
-            let list_dtype_logical = match &self.inner_dtype {
-                None => DataType::from_arrow_dtype(arr.dtype()),
-                Some(dt) => DataType::List(Box::new(dt.clone())),
-            };
-
-            let mut ca = ListChunked::with_chunk(PlSmallStr::EMPTY, arr);
-            if slf.fast_explode {
-                ca.set_fast_explode();
-            }
-            ca.field = Arc::new(Field::new(slf.name, list_dtype_logical));
-            ca
-        }
-    }
-}
-
 pub struct AnonymousOwnedListBuilder {
     name: PlSmallStr,
     builder: AnonymousBuilder<'static>,
@@ -121,6 +16,10 @@ impl Default for AnonymousOwnedListBuilder {
 
 impl ListBuilderTrait for AnonymousOwnedListBuilder {
     fn append_series(&mut self, s: &Series) -> PolarsResult<()> {
+        self.append_owned_series(s.clone())
+    }
+
+    fn append_owned_series(&mut self, s: Series) -> PolarsResult<()> {
         match (s.dtype(), &self.inner_dtype) {
             (DataType::Null, _) => {},
             (dt, None) => self.inner_dtype = Some(dt.clone()),
@@ -135,8 +34,8 @@ impl ListBuilderTrait for AnonymousOwnedListBuilder {
                 self.builder
                     .push_multiple(&*(s.chunks().as_ref() as *const [ArrayRef]));
             }
-            // This make sure that the underlying ArrayRef's are not dropped.
-            self.owned.push(s.clone());
+            // This ensures that the underlying ArrayRef's are not dropped.
+            self.owned.push(s);
         }
         Ok(())
     }

crates/polars-core/src/chunked_array/builder/list/mod.rs

@@ -25,6 +25,9 @@ pub trait ListBuilderTrait {
         Ok(())
     }
     fn append_series(&mut self, s: &Series) -> PolarsResult<()>;
+    fn append_owned_series(&mut self, s: Series) -> PolarsResult<()> {
+        self.append_series(&s)
+    }
     fn append_null(&mut self);
 
     fn field(&self) -> &Field {

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

@@ -742,44 +742,6 @@ impl<'a> GroupBy<'a> {
         DataFrame::new_infer_height(cols)
     }
 
-    /// Aggregate the groups of the group_by operation into lists.
-    ///
-    /// # Example
-    ///
-    /// ```rust
-    /// # use polars_core::prelude::*;
-    /// fn example(df: DataFrame) -> PolarsResult<DataFrame> {
-    ///     // GroupBy and aggregate to Lists
-    ///     df.group_by(["date"])?.select(["temp"]).agg_list()
-    /// }
-    /// ```
-    /// Returns:
-    ///
-    /// ```text
-    /// +------------+------------------------+
-    /// | date       | temp_agg_list          |
-    /// | ---        | ---                    |
-    /// | Date       | list [i32]             |
-    /// +============+========================+
-    /// | 2020-08-23 | "[Some(9)]"            |
-    /// +------------+------------------------+
-    /// | 2020-08-22 | "[Some(7), Some(1)]"   |
-    /// +------------+------------------------+
-    /// | 2020-08-21 | "[Some(20), Some(10)]" |
-    /// +------------+------------------------+
-    /// ```
-    #[deprecated(since = "0.24.1", note = "use polars.lazy aggregations")]
-    pub fn agg_list(&self) -> PolarsResult<DataFrame> {
-        let (mut cols, agg_cols) = self.prepare_agg()?;
-        for agg_col in agg_cols {
-            let new_name = fmt_group_by_column(agg_col.name().as_str(), GroupByMethod::Implode);
-            let mut agg = unsafe { agg_col.agg_list(&self.groups) };
-            agg.rename(new_name);
-            cols.push(agg);
-        }
-        DataFrame::new_infer_height(cols)
-    }
-
     fn prepare_apply(&self) -> PolarsResult<DataFrame> {
         if let Some(agg) = &self.selected_agg {
             if agg.is_empty() {
@@ -912,7 +874,7 @@ pub enum GroupByMethod {
     NUnique,
     Quantile(f64, QuantileMethod),
     Count { include_nulls: bool },
-    Implode,
+    Implode { maintain_order: bool },
     Std(u8),
     Var(u8),
     ArgMin,
@@ -939,7 +901,7 @@ impl Display for GroupByMethod {
             NUnique => "n_unique",
             Quantile(_, _) => "quantile",
             Count { .. } => "count",
-            Implode => "list",
+            Implode { .. } => "implode",
             Std(_) => "std",
             Var(_) => "var",
             ArgMin => "arg_min",
@@ -968,7 +930,7 @@ pub fn fmt_group_by_column(name: &str, method: GroupByMethod) -> PlSmallStr {
         Groups => PlSmallStr::from_static("groups"),
         NUnique => format_pl_smallstr!("{name}_n_unique"),
         Count { .. } => format_pl_smallstr!("{name}_count"),
-        Implode => format_pl_smallstr!("{name}_agg_list"),
+        Implode { .. } => format_pl_smallstr!("{name}_agg_list"),
         Quantile(quantile, _interpol) => format_pl_smallstr!("{name}_quantile_{quantile:.2}"),
         Std(_) => format_pl_smallstr!("{name}_agg_std"),
         Var(_) => format_pl_smallstr!("{name}_agg_var"),

crates/polars-core/src/series/proptest.rs

@@ -7,10 +7,12 @@ use arrow::bitmap::bitmask::nth_set_bit_u32;
 use polars_dtype::categorical::{Categories, FrozenCategories};
 use proptest::prelude::*;
 
-use crate::chunked_array::builder::AnonymousListBuilder;
+use crate::chunked_array::builder::AnonymousOwnedListBuilder;
 #[cfg(feature = "dtype-categorical")]
 use crate::chunked_array::builder::CategoricalChunkedBuilder;
-use crate::prelude::{Int32Chunked, Int64Chunked, Int128Chunked, NamedFrom, Series, TimeUnit};
+use crate::prelude::{
+    Int32Chunked, Int64Chunked, Int128Chunked, ListBuilderTrait, NamedFrom, Series, TimeUnit,
+};
 #[cfg(feature = "dtype-struct")]
 use crate::series::StructChunked;
 use crate::series::from::IntoSeries;
@@ -393,7 +395,7 @@ fn series_list_strategy(
 ) -> impl Strategy<Value = Series> {
     inner.prop_flat_map(move |sample_series| {
         series_length_range.clone().prop_map(move |num_lists| {
-            let mut builder = AnonymousListBuilder::new(
+            let mut builder = AnonymousOwnedListBuilder::new(
                 next_column_name().into(),
                 num_lists,
                 Some(sample_series.dtype().clone()),
@@ -417,7 +419,7 @@ fn series_array_strategy(
         series_length_range.clone().prop_map(move |num_arrays| {
             let width = sample_series.len();
 
-            let mut builder = AnonymousListBuilder::new(
+            let mut builder = AnonymousOwnedListBuilder::new(
                 next_column_name().into(),
                 num_arrays,
                 Some(sample_series.dtype().clone()),

crates/polars-expr/src/expressions/aggregation.rs

@@ -155,7 +155,7 @@ impl PhysicalExpr for AggregationExpr {
 
                 Ok(IdxCa::from_slice(s.name().clone(), &[count as IdxSize]).into_column())
             },
-            GroupByMethod::Implode => s.implode().map(|ca| ca.into_column()),
+            GroupByMethod::Implode { maintain_order: _ } => s.implode().map(|ca| ca.into_column()),
             GroupByMethod::Std(ddof) => s
                 .std_reduce(ddof)
                 .map(|sc| sc.into_column(s.name().clone())),
@@ -414,11 +414,13 @@ impl PhysicalExpr for AggregationExpr {
                     let agg_s = s.agg_n_unique(&groups);
                     AggregatedScalar(agg_s.with_name(keep_name))
                 },
-                GroupByMethod::Implode => AggregatedScalar(match ac.agg_state() {
-                    AggState::LiteralScalar(_) => unreachable!(), // handled above
-                    AggState::AggregatedScalar(c) => c.as_list().into_column(),
-                    AggState::NotAggregated(_) | AggState::AggregatedList(_) => ac.aggregated(),
-                }),
+                GroupByMethod::Implode { maintain_order: _ } => {
+                    AggregatedScalar(match ac.agg_state() {
+                        AggState::LiteralScalar(_) => unreachable!(), // handled above
+                        AggState::AggregatedScalar(c) => c.as_list().into_column(),
+                        AggState::NotAggregated(_) | AggState::AggregatedList(_) => ac.aggregated(),
+                    })
+                },
                 GroupByMethod::Groups => {
                     let mut column: ListChunked = ac.groups().as_list_chunked();
                     column.rename(keep_name);

crates/polars-expr/src/expressions/window.rs

@@ -237,7 +237,7 @@ impl WindowExpr {
                 let mut finishes_list = false;
                 for e in &**function {
                     match e {
-                        Expr::Agg(AggExpr::Implode(_)) => {
+                        Expr::Agg(AggExpr::Implode { .. }) => {
                             finishes_list = true;
                         },
                         Expr::Alias(_, _) => {},

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

@@ -13,6 +13,7 @@ use crate::reduce::bitwise::{
 use crate::reduce::count::{CountReduce, NullCountReduce};
 use crate::reduce::first_last::{new_first_reduction, new_item_reduction, new_last_reduction};
 use crate::reduce::first_last_nonnull::{new_first_nonnull_reduction, new_last_nonnull_reduction};
+use crate::reduce::implode::new_unordered_implode_reduction;
 use crate::reduce::len::LenReduce;
 use crate::reduce::mean::new_mean_reduction;
 use crate::reduce::min_max::{new_max_reduction, new_min_reduction};
@@ -68,10 +69,14 @@ pub fn into_reduction(
                 let count = Box::new(CountReduce::new(*include_nulls)) as Box<_>;
                 (count, *input)
             },
+            IRAggExpr::Implode {
+                input,
+                maintain_order: false,
+            } => (new_unordered_implode_reduction(get_dt(*input)?), *input),
             IRAggExpr::Quantile { .. } => todo!(),
             IRAggExpr::Median(_) => todo!(),
             IRAggExpr::NUnique(_) => todo!(),
-            IRAggExpr::Implode(_) => todo!(),
+            IRAggExpr::Implode { .. } => todo!(),
             IRAggExpr::AggGroups(_) => todo!(),
         },
         AExpr::Len => {