From 2649f42efd39312307b862ec006c577613f55318 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 20 Jul 2026 11:22:36 +0200
Subject: [PATCH 1/7] improve array_agg() performance

---
 Cargo.lock                                    |   1 +
 datafusion/functions-aggregate/Cargo.toml     |   1 +
 .../functions-aggregate/src/array_agg.rs      | 291 +++++++++++++-----
 3 files changed, 216 insertions(+), 77 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f00c931f15032..a88f6a236375c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2264,6 +2264,7 @@ dependencies = [
  "datafusion-physical-expr",
  "datafusion-physical-expr-common",
  "half",
+ "hashbrown 0.17.1",
  "log",
  "num-traits",
  "rand 0.9.4",
diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml
index c1b992a6d89b0..865df24e31f17 100644
--- a/datafusion/functions-aggregate/Cargo.toml
+++ b/datafusion/functions-aggregate/Cargo.toml
@@ -51,6 +51,7 @@ datafusion-macros = { workspace = true }
 datafusion-physical-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 half = { workspace = true }
+hashbrown = { workspace = true }
 log = { workspace = true }
 num-traits = { workspace = true }
 
diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 1937f17973950..66b3cb2735f02 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -18,7 +18,7 @@
 //! `ARRAY_AGG` aggregate implementation: [`ArrayAgg`]
 
 use std::cmp::Ordering;
-use std::collections::{HashMap, VecDeque};
+use std::collections::VecDeque;
 use std::mem::{size_of, size_of_val, take};
 use std::sync::Arc;
 
@@ -28,6 +28,7 @@ use arrow::array::{
 };
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::{SortOptions, filter};
+use arrow::row::{Row, RowConverter, Rows, SortField};
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
 
 use datafusion_common::cast::as_list_array;
@@ -47,7 +48,10 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls:
 use datafusion_functions_aggregate_common::merge_arrays::merge_ordered_arrays;
 use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity;
 use datafusion_functions_aggregate_common::utils::ordering_fields;
+use datafusion_common::hash_utils::{RandomState, create_hashes};
+use datafusion_common::utils::proxy::HashTableAllocExt;
 use datafusion_macros::user_doc;
+use hashbrown::hash_table::HashTable;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 
 make_udaf_expr_and_func!(
@@ -815,10 +819,27 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator {
 
 #[derive(Debug)]
 pub struct DistinctArrayAggAccumulator {
-    // Value → live refcount. Multiset state lets `retract_batch` correctly
-    // drop a duplicate occurrence while keeping the key alive if other
-    // copies remain in the current window frame.
-    values: HashMap<ScalarValue, u64>,
+    /// Contiguous buffer of one encoded row per distinct value ever inserted.
+    /// Dead slots (where `counts[i] == 0`) may exist after `retract_batch`.
+    group_values: Option<Rows>,
+    /// Live refcount per group index. `counts[i]` is how many times the value
+    /// at `group_values.row(i)` is currently present in the window frame.
+    counts: Vec<u64>,
+    /// Hash table storing `(hash, group_index)`. Only contains live entries
+    /// (those whose count is > 0). Evicted on `retract_batch` when count
+    /// drops to zero.
+    map: HashTable<(u64, usize)>,
+    /// Heap size of `map` in bytes, tracked for `size()` reporting.
+    map_size: usize,
+    /// Temporary buffer for encoding an incoming batch; reused across calls.
+    rows_buffer: Option<Rows>,
+    /// Reused buffer for batch hashes.
+    hashes_buffer: Vec<u64>,
+    /// Random state used by `create_hashes`.
+    random_state: RandomState,
+    /// Lazily initialised converter — set on the first `update_batch` call
+    /// using the actual runtime column type.
+    converter: Option<RowConverter>,
     datatype: DataType,
     sort_options: Option<SortOptions>,
     ignore_nulls: bool,
@@ -831,12 +852,35 @@ impl DistinctArrayAggAccumulator {
         ignore_nulls: bool,
     ) -> Result<Self> {
         Ok(Self {
-            values: HashMap::new(),
+            group_values: None,
+            counts: Vec::new(),
+            map: HashTable::new(),
+            map_size: 0,
+            rows_buffer: None,
+            hashes_buffer: Vec::new(),
+            random_state: RandomState::default(),
+            converter: None,
             datatype: datatype.clone(),
             sort_options,
             ignore_nulls,
         })
     }
+
+    /// Lazily initialises the `RowConverter`, `rows_buffer`, and `group_values`
+    /// on the first call, using the actual runtime column type.
+    fn ensure_converter(&mut self, data_type: &DataType) -> Result<()> {
+        if self.converter.is_none() {
+            let sort_field = match self.sort_options {
+                Some(opts) => SortField::new_with_options(data_type.clone(), opts),
+                None => SortField::new(data_type.clone()),
+            };
+            let converter = RowConverter::new(vec![sort_field])?;
+            self.rows_buffer = Some(converter.empty_rows(0, 0));
+            self.group_values = Some(converter.empty_rows(0, 0));
+            self.converter = Some(converter);
+        }
+        Ok(())
+    }
 }
 
 impl Accumulator for DistinctArrayAggAccumulator {
@@ -850,22 +894,72 @@ impl Accumulator for DistinctArrayAggAccumulator {
         }
 
         let val = &values[0];
-        let nulls = if self.ignore_nulls {
-            val.logical_nulls()
+
+        // Filter nulls out upfront when ignore_nulls is set so they are
+        // never inserted into the dedup state.
+        let filtered;
+        let col: &ArrayRef = if self.ignore_nulls {
+            if let Some(nulls) = val.logical_nulls() {
+                if nulls.null_count() > 0 {
+                    let mask: BooleanArray = nulls.iter().map(Some).collect();
+                    filtered = filter(val.as_ref(), &mask)?;
+                    &filtered
+                } else {
+                    val
+                }
+            } else {
+                val
+            }
         } else {
-            None
+            val
         };
 
-        let nulls = nulls.as_ref();
-        if nulls.is_none_or(|nulls| nulls.null_count() < val.len()) {
-            for i in 0..val.len() {
-                if nulls.is_none_or(|nulls| nulls.is_valid(i)) {
-                    let key = ScalarValue::try_from_array(val, i)?.compacted();
-                    *self.values.entry(key).or_insert(0) += 1;
+        if col.is_empty() {
+            return Ok(());
+        }
+
+        self.ensure_converter(col.data_type())?;
+
+        // Encode the entire incoming batch into rows_buffer in one pass.
+        let converter = self.converter.as_ref().unwrap();
+        let rows_buffer = self.rows_buffer.as_mut().unwrap();
+        let group_values = self.group_values.as_mut().unwrap();
+        rows_buffer.clear();
+        converter.append(rows_buffer, std::slice::from_ref(col))?;
+
+        // Pre-compute all hashes for the batch in one SIMD-friendly pass.
+        self.hashes_buffer.clear();
+        self.hashes_buffer.resize(col.len(), 0);
+        create_hashes(
+            std::slice::from_ref(col),
+            &self.random_state,
+            &mut self.hashes_buffer,
+        )?;
+
+        for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() {
+            let row = rows_buffer.row(row_idx);
+            let entry = self.map.find_mut(hash, |&(h, group_idx)| {
+                h == hash && group_values.row(group_idx) == row
+            });
+            match entry {
+                Some((_, group_idx)) => {
+                    // Already known: just increment the live refcount.
+                    self.counts[*group_idx] += 1;
+                }
+                None => {
+                    // New distinct value: append to the contiguous buffer and
+                    // record it in the map.
+                    let new_group_idx = group_values.num_rows();
+                    group_values.push(row);
+                    self.counts.push(1);
+                    self.map.insert_accounted(
+                        (hash, new_group_idx),
+                        |&(h, _)| h,
+                        &mut self.map_size,
+                    );
                 }
             }
         }
-
         Ok(())
     }
 
@@ -876,12 +970,9 @@ impl Accumulator for DistinctArrayAggAccumulator {
 
         assert_eq_or_internal_err!(states.len(), 1, "expects single state");
 
-        // The DISTINCT state schema is `List<value>` — partial accumulators
-        // ship the set of values they saw, not multiplicities. Re-ingesting
-        // each element here makes the merged counts represent "partitions
-        // that emitted this value," which is fine because `evaluate` only
-        // reads keys. Refcount semantics for retract are only valid within
-        // a single accumulator instance (window execution).
+        // The DISTINCT state is `List<value>`. Partial accumulators ship the
+        // set of live values, not multiplicities. Re-ingesting them here is
+        // correct: `evaluate` reads only the map keys.
         states[0]
             .as_list::<i32>()
             .iter()
@@ -890,37 +981,38 @@ impl Accumulator for DistinctArrayAggAccumulator {
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
-        let mut values: Vec<ScalarValue> = self.values.keys().cloned().collect();
-        if values.is_empty() {
+        if self.map.is_empty() {
             return Ok(ScalarValue::new_null_list(self.datatype.clone(), true, 1));
         }
 
-        if let Some(opts) = self.sort_options {
-            let mut delayed_cmp_err = Ok(());
-            values.sort_by(|a, b| {
-                if a.is_null() {
-                    return match opts.nulls_first {
-                        true => Ordering::Less,
-                        false => Ordering::Greater,
-                    };
-                }
-                if b.is_null() {
-                    return match opts.nulls_first {
-                        true => Ordering::Greater,
-                        false => Ordering::Less,
-                    };
-                }
-                match opts.descending {
-                    true => b.try_cmp(a),
-                    false => a.try_cmp(b),
-                }
-                .unwrap_or_else(|err| {
-                    delayed_cmp_err = Err(err);
-                    Ordering::Equal
-                })
-            });
-            delayed_cmp_err?;
-        };
+        let group_values = self
+            .group_values
+            .as_ref()
+            .expect("group_values must be set when map is non-empty");
+        let converter = self
+            .converter
+            .as_ref()
+            .expect("converter must be set when map is non-empty");
+
+        // Collect the group indices of all live entries.
+        let mut live_indices: Vec<usize> =
+            self.map.iter().map(|&(_, group_idx)| group_idx).collect();
+
+        // If ORDER BY was specified, the RowConverter bakes the sort direction
+        // into the row bytes, so lexicographic sort gives the correct order.
+        if self.sort_options.is_some() {
+            live_indices
+                .sort_unstable_by(|&a, &b| group_values.row(a).cmp(&group_values.row(b)));
+        }
+
+        // Decode the selected rows back into an Arrow array.
+        let rows: Vec<Row<'_>> =
+            live_indices.iter().map(|&i| group_values.row(i)).collect();
+        let arrays = converter.convert_rows(rows)?;
+
+        let values: Vec<ScalarValue> = (0..arrays[0].len())
+            .map(|i| ScalarValue::try_from_array(arrays[0].as_ref(), i))
+            .collect::<Result<_>>()?;
 
         let arr = ScalarValue::new_list(&values, &self.datatype, true);
         Ok(ScalarValue::List(arr))
@@ -934,33 +1026,74 @@ impl Accumulator for DistinctArrayAggAccumulator {
         assert_eq_or_internal_err!(values.len(), 1, "expects single batch");
 
         let val = &values[0];
-        let nulls = if self.ignore_nulls {
-            val.logical_nulls()
+
+        // Mirror the null-filtering logic from update_batch so we only
+        // retract values that were actually inserted.
+        let filtered;
+        let col: &ArrayRef = if self.ignore_nulls {
+            if let Some(nulls) = val.logical_nulls() {
+                if nulls.null_count() > 0 {
+                    let mask: BooleanArray = nulls.iter().map(Some).collect();
+                    filtered = filter(val.as_ref(), &mask)?;
+                    &filtered
+                } else {
+                    val
+                }
+            } else {
+                val
+            }
         } else {
-            None
+            val
         };
-        let nulls = nulls.as_ref();
 
-        for i in 0..val.len() {
-            if nulls.is_some_and(|nulls| !nulls.is_valid(i)) {
-                continue;
-            }
-            let key = ScalarValue::try_from_array(val, i)?;
-            match self.values.get_mut(&key) {
-                Some(count) => {
-                    *count -= 1;
-                    if *count == 0 {
-                        self.values.remove(&key);
-                    }
-                }
-                None => {
+        if col.is_empty() {
+            return Ok(());
+        }
+
+        let converter = self
+            .converter
+            .as_ref()
+            .expect("retract_batch called before update_batch");
+        let rows_buffer = self
+            .rows_buffer
+            .as_mut()
+            .expect("rows_buffer must be initialised");
+        let group_values = self
+            .group_values
+            .as_ref()
+            .expect("group_values must be initialised");
+
+        rows_buffer.clear();
+        converter.append(rows_buffer, std::slice::from_ref(col))?;
+
+        self.hashes_buffer.clear();
+        self.hashes_buffer.resize(col.len(), 0);
+        create_hashes(
+            std::slice::from_ref(col),
+            &self.random_state,
+            &mut self.hashes_buffer,
+        )?;
+
+        for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() {
+            let row = rows_buffer.row(row_idx);
+            match self.map.find_entry(hash, |&(h, group_idx)| {
+                h == hash && group_values.row(group_idx) == row
+            }) {
+                Err(_) => {
                     return internal_err!(
-                        "DistinctArrayAggAccumulator::retract_batch: value not present in state"
+                        "DistinctArrayAggAccumulator::retract_batch: \
+                         value not present in state"
                     );
                 }
+                Ok(occupied) => {
+                    let (_, group_idx) = *occupied.get();
+                    self.counts[group_idx] -= 1;
+                    if self.counts[group_idx] == 0 {
+                        occupied.remove();
+                    }
+                }
             }
         }
-
         Ok(())
     }
 
@@ -969,12 +1102,14 @@ impl Accumulator for DistinctArrayAggAccumulator {
     }
 
     fn size(&self) -> usize {
-        size_of_val(self) + ScalarValue::size_of_hashmap(&self.values)
-            - size_of_val(&self.values)
-            + self.datatype.size()
-            - size_of_val(&self.datatype)
-            - size_of_val(&self.sort_options)
-            + size_of::<Option<SortOptions>>()
+        size_of_val(self)
+            + self.group_values.as_ref().map(|r| r.size()).unwrap_or(0)
+            + self.rows_buffer.as_ref().map(|r| r.size()).unwrap_or(0)
+            + self.converter.as_ref().map(|c| c.size()).unwrap_or(0)
+            + self.map_size
+            + self.counts.capacity() * size_of::<u64>()
+            + self.hashes_buffer.capacity() * size_of::<u64>()
+            + self.datatype.size() - size_of_val(&self.datatype)
     }
 }
 
@@ -1546,8 +1681,10 @@ mod tests {
         acc2.update_batch(&[string_list_data([vec!["e", "f", "g"]])])?;
         acc1 = merge(acc1, acc2)?;
 
-        // without compaction, the size is 16684
-        assert_eq!(acc1.size(), 1684);
+        // The GroupValuesRows-based implementation uses a contiguous Rows
+        // buffer + HashTable instead of individual ScalarValue allocations,
+        // so the reported size differs from the previous implementation (was 1684).
+        assert_eq!(acc1.size(), 2268);
 
         Ok(())
     }

From 1cd130563b4f8ac948dd9ee08429c6e17a1b7007 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 20 Jul 2026 13:19:12 +0200
Subject: [PATCH 2/7] cargo fmt failing

---
 datafusion/functions-aggregate/src/array_agg.rs | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 66b3cb2735f02..8cd9097e3144f 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -28,10 +28,12 @@ use arrow::array::{
 };
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::{SortOptions, filter};
-use arrow::row::{Row, RowConverter, Rows, SortField};
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
+use arrow::row::{Row, RowConverter, Rows, SortField};
 
 use datafusion_common::cast::as_list_array;
+use datafusion_common::hash_utils::{RandomState, create_hashes};
+use datafusion_common::utils::proxy::HashTableAllocExt;
 use datafusion_common::utils::{
     SingleRowListArrayBuilder, compare_rows, get_row_at_idx, take_function_args,
 };
@@ -48,11 +50,9 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls:
 use datafusion_functions_aggregate_common::merge_arrays::merge_ordered_arrays;
 use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity;
 use datafusion_functions_aggregate_common::utils::ordering_fields;
-use datafusion_common::hash_utils::{RandomState, create_hashes};
-use datafusion_common::utils::proxy::HashTableAllocExt;
 use datafusion_macros::user_doc;
-use hashbrown::hash_table::HashTable;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
+use hashbrown::hash_table::HashTable;
 
 make_udaf_expr_and_func!(
     ArrayAgg,
@@ -1109,7 +1109,8 @@ impl Accumulator for DistinctArrayAggAccumulator {
             + self.map_size
             + self.counts.capacity() * size_of::<u64>()
             + self.hashes_buffer.capacity() * size_of::<u64>()
-            + self.datatype.size() - size_of_val(&self.datatype)
+            + self.datatype.size()
+            - size_of_val(&self.datatype)
     }
 }
 

From bcb69faf28c894d645c4661f16101ade1ee18363 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 27 Jul 2026 15:52:29 +0200
Subject: [PATCH 3/7] cleanup comments

---
 datafusion/functions-aggregate/src/array_agg.rs | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 8cd9097e3144f..d7286ecfaa431 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -970,9 +970,7 @@ impl Accumulator for DistinctArrayAggAccumulator {
 
         assert_eq_or_internal_err!(states.len(), 1, "expects single state");
 
-        // The DISTINCT state is `List<value>`. Partial accumulators ship the
-        // set of live values, not multiplicities. Re-ingesting them here is
-        // correct: `evaluate` reads only the map keys.
+        // The DISTINCT state is `List<value>`.
         states[0]
             .as_list::<i32>()
             .iter()
@@ -1682,9 +1680,6 @@ mod tests {
         acc2.update_batch(&[string_list_data([vec!["e", "f", "g"]])])?;
         acc1 = merge(acc1, acc2)?;
 
-        // The GroupValuesRows-based implementation uses a contiguous Rows
-        // buffer + HashTable instead of individual ScalarValue allocations,
-        // so the reported size differs from the previous implementation (was 1684).
         assert_eq!(acc1.size(), 2268);
 
         Ok(())

From df327a674dd0fda2e0d9068c15150c91bb77d453 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 27 Jul 2026 16:10:23 +0200
Subject: [PATCH 4/7] add an implicit state machine

---
 .../functions-aggregate/src/array_agg.rs      | 98 ++++++++++---------
 1 file changed, 54 insertions(+), 44 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index d7286ecfaa431..2769fb2620df2 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -817,13 +817,29 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator {
     }
 }
 
+/// Resources that are allocated lazily on the first `update_batch` call,
+/// once the concrete runtime Arrow type is known.
+///
+/// Grouping all three fields together makes the "either all present or all
+/// absent" invariant explicit in the type system, replacing the scattered
+/// `.expect()` calls that would otherwise be needed.
 #[derive(Debug)]
-pub struct DistinctArrayAggAccumulator {
+struct DistinctState {
+    /// Converts Arrow arrays to/from the comparable row format.
+    converter: RowConverter,
     /// Contiguous buffer of one encoded row per distinct value ever inserted.
     /// Dead slots (where `counts[i] == 0`) may exist after `retract_batch`.
-    group_values: Option<Rows>,
+    group_values: Rows,
+    /// Temporary buffer for encoding an incoming batch; reused across calls.
+    rows_buffer: Rows,
+}
+
+#[derive(Debug)]
+pub struct DistinctArrayAggAccumulator {
+    /// Lazily allocated on the first `update_batch`; `None` until then.
+    state: Option<DistinctState>,
     /// Live refcount per group index. `counts[i]` is how many times the value
-    /// at `group_values.row(i)` is currently present in the window frame.
+    /// at `state.group_values.row(i)` is currently present in the window frame.
     counts: Vec<u64>,
     /// Hash table storing `(hash, group_index)`. Only contains live entries
     /// (those whose count is > 0). Evicted on `retract_batch` when count
@@ -831,15 +847,10 @@ pub struct DistinctArrayAggAccumulator {
     map: HashTable<(u64, usize)>,
     /// Heap size of `map` in bytes, tracked for `size()` reporting.
     map_size: usize,
-    /// Temporary buffer for encoding an incoming batch; reused across calls.
-    rows_buffer: Option<Rows>,
     /// Reused buffer for batch hashes.
     hashes_buffer: Vec<u64>,
     /// Random state used by `create_hashes`.
     random_state: RandomState,
-    /// Lazily initialised converter — set on the first `update_batch` call
-    /// using the actual runtime column type.
-    converter: Option<RowConverter>,
     datatype: DataType,
     sort_options: Option<SortOptions>,
     ignore_nulls: bool,
@@ -852,32 +863,34 @@ impl DistinctArrayAggAccumulator {
         ignore_nulls: bool,
     ) -> Result<Self> {
         Ok(Self {
-            group_values: None,
+            state: None,
             counts: Vec::new(),
             map: HashTable::new(),
             map_size: 0,
-            rows_buffer: None,
             hashes_buffer: Vec::new(),
             random_state: RandomState::default(),
-            converter: None,
             datatype: datatype.clone(),
             sort_options,
             ignore_nulls,
         })
     }
 
-    /// Lazily initialises the `RowConverter`, `rows_buffer`, and `group_values`
-    /// on the first call, using the actual runtime column type.
-    fn ensure_converter(&mut self, data_type: &DataType) -> Result<()> {
-        if self.converter.is_none() {
+    /// Lazily initialises the `DistinctState` on the first call, using the
+    /// actual runtime column type.
+    fn ensure_state(&mut self, data_type: &DataType) -> Result<()> {
+        if self.state.is_none() {
             let sort_field = match self.sort_options {
                 Some(opts) => SortField::new_with_options(data_type.clone(), opts),
                 None => SortField::new(data_type.clone()),
             };
             let converter = RowConverter::new(vec![sort_field])?;
-            self.rows_buffer = Some(converter.empty_rows(0, 0));
-            self.group_values = Some(converter.empty_rows(0, 0));
-            self.converter = Some(converter);
+            let group_values = converter.empty_rows(0, 0);
+            let rows_buffer = converter.empty_rows(0, 0);
+            self.state = Some(DistinctState {
+                converter,
+                group_values,
+                rows_buffer,
+            });
         }
         Ok(())
     }
@@ -918,12 +931,14 @@ impl Accumulator for DistinctArrayAggAccumulator {
             return Ok(());
         }
 
-        self.ensure_converter(col.data_type())?;
+        self.ensure_state(col.data_type())?;
 
         // Encode the entire incoming batch into rows_buffer in one pass.
-        let converter = self.converter.as_ref().unwrap();
-        let rows_buffer = self.rows_buffer.as_mut().unwrap();
-        let group_values = self.group_values.as_mut().unwrap();
+        let DistinctState {
+            converter,
+            rows_buffer,
+            group_values,
+        } = self.state.as_mut().unwrap();
         rows_buffer.clear();
         converter.append(rows_buffer, std::slice::from_ref(col))?;
 
@@ -983,14 +998,11 @@ impl Accumulator for DistinctArrayAggAccumulator {
             return Ok(ScalarValue::new_null_list(self.datatype.clone(), true, 1));
         }
 
-        let group_values = self
-            .group_values
-            .as_ref()
-            .expect("group_values must be set when map is non-empty");
-        let converter = self
-            .converter
-            .as_ref()
-            .expect("converter must be set when map is non-empty");
+        let DistinctState {
+            group_values,
+            converter,
+            ..
+        } = self.state.as_ref().expect("state must be set when map is non-empty");
 
         // Collect the group indices of all live entries.
         let mut live_indices: Vec<usize> =
@@ -1048,18 +1060,14 @@ impl Accumulator for DistinctArrayAggAccumulator {
             return Ok(());
         }
 
-        let converter = self
-            .converter
-            .as_ref()
-            .expect("retract_batch called before update_batch");
-        let rows_buffer = self
-            .rows_buffer
+        let DistinctState {
+            converter,
+            rows_buffer,
+            group_values,
+        } = self
+            .state
             .as_mut()
-            .expect("rows_buffer must be initialised");
-        let group_values = self
-            .group_values
-            .as_ref()
-            .expect("group_values must be initialised");
+            .expect("retract_batch called before update_batch");
 
         rows_buffer.clear();
         converter.append(rows_buffer, std::slice::from_ref(col))?;
@@ -1101,9 +1109,11 @@ impl Accumulator for DistinctArrayAggAccumulator {
 
     fn size(&self) -> usize {
         size_of_val(self)
-            + self.group_values.as_ref().map(|r| r.size()).unwrap_or(0)
-            + self.rows_buffer.as_ref().map(|r| r.size()).unwrap_or(0)
-            + self.converter.as_ref().map(|c| c.size()).unwrap_or(0)
+            + self
+                .state
+                .as_ref()
+                .map(|s| s.group_values.size() + s.rows_buffer.size() + s.converter.size())
+                .unwrap_or(0)
             + self.map_size
             + self.counts.capacity() * size_of::<u64>()
             + self.hashes_buffer.capacity() * size_of::<u64>()

From d70f66dda4af45d577a023bb16099af589e1d028 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 27 Jul 2026 17:33:18 +0200
Subject: [PATCH 5/7] fix dictionaries

---
 .../functions-aggregate/src/array_agg.rs      | 59 +++++++++++++++++--
 1 file changed, 54 insertions(+), 5 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 2769fb2620df2..826341902f167 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -27,7 +27,7 @@ use arrow::array::{
     UInt32Array, new_empty_array,
 };
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
-use arrow::compute::{SortOptions, filter};
+use arrow::compute::{SortOptions, cast, filter};
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
 use arrow::row::{Row, RowConverter, Rows, SortField};
 
@@ -1002,7 +1002,10 @@ impl Accumulator for DistinctArrayAggAccumulator {
             group_values,
             converter,
             ..
-        } = self.state.as_ref().expect("state must be set when map is non-empty");
+        } = self
+            .state
+            .as_ref()
+            .expect("state must be set when map is non-empty");
 
         // Collect the group indices of all live entries.
         let mut live_indices: Vec<usize> =
@@ -1020,8 +1023,19 @@ impl Accumulator for DistinctArrayAggAccumulator {
             live_indices.iter().map(|&i| group_values.row(i)).collect();
         let arrays = converter.convert_rows(rows)?;
 
-        let values: Vec<ScalarValue> = (0..arrays[0].len())
-            .map(|i| ScalarValue::try_from_array(arrays[0].as_ref(), i))
+        // `convert_rows` always returns the physical (non-dictionary) type.
+        // Cast back to the declared logical type when they differ so that
+        // e.g. Dictionary columns round-trip correctly through the RowConverter.
+        let decoded = if arrays[0].data_type() != &self.datatype
+            && matches!(self.datatype, DataType::Dictionary(_, _))
+        {
+            cast(arrays[0].as_ref(), &self.datatype)?
+        } else {
+            Arc::clone(&arrays[0])
+        };
+
+        let values: Vec<ScalarValue> = (0..decoded.len())
+            .map(|i| ScalarValue::try_from_array(decoded.as_ref(), i))
             .collect::<Result<_>>()?;
 
         let arr = ScalarValue::new_list(&values, &self.datatype, true);
@@ -1112,7 +1126,9 @@ impl Accumulator for DistinctArrayAggAccumulator {
             + self
                 .state
                 .as_ref()
-                .map(|s| s.group_values.size() + s.rows_buffer.size() + s.converter.size())
+                .map(|s| {
+                    s.group_values.size() + s.rows_buffer.size() + s.converter.size()
+                })
                 .unwrap_or(0)
             + self.map_size
             + self.counts.capacity() * size_of::<u64>()
@@ -2826,6 +2842,39 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn distinct_array_agg_dictionary_preserves_type() -> Result<()> {
+        use arrow::array::{DictionaryArray, Int32Array, StringArray};
+
+        // Dictionary(Int32, Utf8) input with duplicates.
+        let keys = Int32Array::from(vec![0, 1, 0, 2, 1]); // "a", "b", "a", "c", "b"
+        let values = StringArray::from(vec!["a", "b", "c"]);
+        let dict: ArrayRef = Arc::new(DictionaryArray::new(keys, Arc::new(values)));
+
+        let datatype =
+            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+        let mut acc = DistinctArrayAggAccumulator::try_new(&datatype, None, false)?;
+        acc.update_batch(&[dict])?;
+
+        let result = acc.evaluate()?;
+        let ScalarValue::List(arr) = &result else {
+            panic!("expected ScalarValue::List, got {result:?}");
+        };
+
+        // The element type of the returned list must stay Dictionary(Int32, Utf8),
+        // not be silently widened to Utf8.
+        assert_eq!(
+            arr.values().data_type(),
+            &datatype,
+            "element type must be Dictionary(Int32, Utf8), got {}",
+            arr.values().data_type()
+        );
+
+        // There should be exactly 3 distinct values.
+        assert_eq!(arr.value(0).len(), 3);
+        Ok(())
+    }
+
     #[test]
     fn distinct_array_agg_date32_deduplicates() -> Result<()> {
         use arrow::array::Date32Array;

From 8eb193125b7864517e1fe10229885b5c26b68e71 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Mon, 27 Jul 2026 18:10:39 +0200
Subject: [PATCH 6/7] fix unbound memory usage

---
 .../functions-aggregate/src/array_agg.rs      | 143 ++++++++++++++----
 1 file changed, 114 insertions(+), 29 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 826341902f167..1691af134d31f 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -29,7 +29,7 @@ use arrow::array::{
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::{SortOptions, cast, filter};
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
-use arrow::row::{Row, RowConverter, Rows, SortField};
+use arrow::row::{OwnedRow, Row, RowConverter, Rows, SortField};
 
 use datafusion_common::cast::as_list_array;
 use datafusion_common::hash_utils::{RandomState, create_hashes};
@@ -827,9 +827,16 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator {
 struct DistinctState {
     /// Converts Arrow arrays to/from the comparable row format.
     converter: RowConverter,
-    /// Contiguous buffer of one encoded row per distinct value ever inserted.
-    /// Dead slots (where `counts[i] == 0`) may exist after `retract_batch`.
-    group_values: Rows,
+    /// One owned encoded row per live distinct value, indexed by group index.
+    /// Compacted via swap-remove on eviction so there are never dead slots.
+    group_rows: Vec<OwnedRow>,
+    /// Live refcount per group index. `counts[i]` is how many times the value
+    /// at `group_rows[i]` is currently present in the window frame.
+    counts: Vec<u64>,
+    /// Hash of the encoded row at group index `i`, kept in sync with
+    /// `group_rows` and `counts`. Needed to patch the map on swap-remove
+    /// eviction without re-encoding the moved row.
+    row_hashes: Vec<u64>,
     /// Temporary buffer for encoding an incoming batch; reused across calls.
     rows_buffer: Rows,
 }
@@ -838,9 +845,6 @@ struct DistinctState {
 pub struct DistinctArrayAggAccumulator {
     /// Lazily allocated on the first `update_batch`; `None` until then.
     state: Option<DistinctState>,
-    /// Live refcount per group index. `counts[i]` is how many times the value
-    /// at `state.group_values.row(i)` is currently present in the window frame.
-    counts: Vec<u64>,
     /// Hash table storing `(hash, group_index)`. Only contains live entries
     /// (those whose count is > 0). Evicted on `retract_batch` when count
     /// drops to zero.
@@ -864,7 +868,6 @@ impl DistinctArrayAggAccumulator {
     ) -> Result<Self> {
         Ok(Self {
             state: None,
-            counts: Vec::new(),
             map: HashTable::new(),
             map_size: 0,
             hashes_buffer: Vec::new(),
@@ -884,11 +887,12 @@ impl DistinctArrayAggAccumulator {
                 None => SortField::new(data_type.clone()),
             };
             let converter = RowConverter::new(vec![sort_field])?;
-            let group_values = converter.empty_rows(0, 0);
             let rows_buffer = converter.empty_rows(0, 0);
             self.state = Some(DistinctState {
                 converter,
-                group_values,
+                group_rows: Vec::new(),
+                counts: Vec::new(),
+                row_hashes: Vec::new(),
                 rows_buffer,
             });
         }
@@ -936,8 +940,10 @@ impl Accumulator for DistinctArrayAggAccumulator {
         // Encode the entire incoming batch into rows_buffer in one pass.
         let DistinctState {
             converter,
+            group_rows,
+            counts,
+            row_hashes,
             rows_buffer,
-            group_values,
         } = self.state.as_mut().unwrap();
         rows_buffer.clear();
         converter.append(rows_buffer, std::slice::from_ref(col))?;
@@ -954,19 +960,19 @@ impl Accumulator for DistinctArrayAggAccumulator {
         for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() {
             let row = rows_buffer.row(row_idx);
             let entry = self.map.find_mut(hash, |&(h, group_idx)| {
-                h == hash && group_values.row(group_idx) == row
+                h == hash && group_rows[group_idx].row() == row
             });
             match entry {
                 Some((_, group_idx)) => {
                     // Already known: just increment the live refcount.
-                    self.counts[*group_idx] += 1;
+                    counts[*group_idx] += 1;
                 }
                 None => {
-                    // New distinct value: append to the contiguous buffer and
-                    // record it in the map.
-                    let new_group_idx = group_values.num_rows();
-                    group_values.push(row);
-                    self.counts.push(1);
+                    // New distinct value: own the encoded row, record it.
+                    let new_group_idx = group_rows.len();
+                    group_rows.push(row.owned());
+                    counts.push(1);
+                    row_hashes.push(hash);
                     self.map.insert_accounted(
                         (hash, new_group_idx),
                         |&(h, _)| h,
@@ -999,8 +1005,8 @@ impl Accumulator for DistinctArrayAggAccumulator {
         }
 
         let DistinctState {
-            group_values,
             converter,
+            group_rows,
             ..
         } = self
             .state
@@ -1015,12 +1021,12 @@ impl Accumulator for DistinctArrayAggAccumulator {
         // into the row bytes, so lexicographic sort gives the correct order.
         if self.sort_options.is_some() {
             live_indices
-                .sort_unstable_by(|&a, &b| group_values.row(a).cmp(&group_values.row(b)));
+                .sort_unstable_by(|&a, &b| group_rows[a].row().cmp(&group_rows[b].row()));
         }
 
         // Decode the selected rows back into an Arrow array.
         let rows: Vec<Row<'_>> =
-            live_indices.iter().map(|&i| group_values.row(i)).collect();
+            live_indices.iter().map(|&i| group_rows[i].row()).collect();
         let arrays = converter.convert_rows(rows)?;
 
         // `convert_rows` always returns the physical (non-dictionary) type.
@@ -1076,8 +1082,10 @@ impl Accumulator for DistinctArrayAggAccumulator {
 
         let DistinctState {
             converter,
+            group_rows,
+            counts,
+            row_hashes,
             rows_buffer,
-            group_values,
         } = self
             .state
             .as_mut()
@@ -1097,7 +1105,7 @@ impl Accumulator for DistinctArrayAggAccumulator {
         for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() {
             let row = rows_buffer.row(row_idx);
             match self.map.find_entry(hash, |&(h, group_idx)| {
-                h == hash && group_values.row(group_idx) == row
+                h == hash && group_rows[group_idx].row() == row
             }) {
                 Err(_) => {
                     return internal_err!(
@@ -1106,10 +1114,32 @@ impl Accumulator for DistinctArrayAggAccumulator {
                     );
                 }
                 Ok(occupied) => {
-                    let (_, group_idx) = *occupied.get();
-                    self.counts[group_idx] -= 1;
-                    if self.counts[group_idx] == 0 {
+                    let (_, dead_idx) = *occupied.get();
+                    counts[dead_idx] -= 1;
+                    if counts[dead_idx] == 0 {
                         occupied.remove();
+                        // Compact via swap-remove: move the last slot into the
+                        // dead slot so group_rows / counts / row_hashes stay
+                        // dense with no dead entries.
+                        let last_idx = group_rows.len() - 1;
+                        if dead_idx != last_idx {
+                            // Patch the map entry that points to last_idx so
+                            // it points to dead_idx instead.
+                            let last_hash = row_hashes[last_idx];
+                            self.map
+                                .find_mut(last_hash, |&(_, idx)| idx == last_idx)
+                                .ok_or_else(|| {
+                                    datafusion_common::internal_datafusion_err!(
+                                        "DistinctArrayAggAccumulator: map is missing \
+                                         group index {last_idx} during swap-remove \
+                                         compaction"
+                                    )
+                                })?
+                                .1 = dead_idx;
+                        }
+                        group_rows.swap_remove(dead_idx);
+                        counts.swap_remove(dead_idx);
+                        row_hashes.swap_remove(dead_idx);
                     }
                 }
             }
@@ -1127,11 +1157,18 @@ impl Accumulator for DistinctArrayAggAccumulator {
                 .state
                 .as_ref()
                 .map(|s| {
-                    s.group_values.size() + s.rows_buffer.size() + s.converter.size()
+                    s.group_rows
+                        .iter()
+                        .map(|r| r.row().data().len())
+                        .sum::<usize>()
+                        + s.group_rows.capacity() * size_of::<OwnedRow>()
+                        + s.counts.capacity() * size_of::<u64>()
+                        + s.row_hashes.capacity() * size_of::<u64>()
+                        + s.rows_buffer.size()
+                        + s.converter.size()
                 })
                 .unwrap_or(0)
             + self.map_size
-            + self.counts.capacity() * size_of::<u64>()
             + self.hashes_buffer.capacity() * size_of::<u64>()
             + self.datatype.size()
             - size_of_val(&self.datatype)
@@ -1706,7 +1743,7 @@ mod tests {
         acc2.update_batch(&[string_list_data([vec!["e", "f", "g"]])])?;
         acc1 = merge(acc1, acc2)?;
 
-        assert_eq!(acc1.size(), 2268);
+        assert_eq!(acc1.size(), 2274);
 
         Ok(())
     }
@@ -2905,4 +2942,52 @@ mod tests {
         assert_eq!(values, vec![100i32, 200, 300, 400]);
         Ok(())
     }
+
+    #[test]
+    fn distinct_retract_memory_is_bounded() -> Result<()> {
+        use arrow::array::Int64Array;
+
+        // Emulates a sliding window where each value enters and immediately
+        // leaves. Only CARDINALITY distinct values are ever live at once;
+        // memory must not grow with the number of rows processed.
+        const CARDINALITY: i64 = 10;
+        const WARMUP_ROWS: i64 = 1_000;
+        const EXTRA_ROWS: i64 = 20_000;
+
+        let mut acc =
+            DistinctArrayAggAccumulator::try_new(&DataType::Int64, None, false)?;
+
+        let slide = |acc: &mut DistinctArrayAggAccumulator, rows: i64| -> Result<()> {
+            for i in 0..rows {
+                let value: ArrayRef = Arc::new(Int64Array::from(vec![i % CARDINALITY]));
+                acc.update_batch(std::slice::from_ref(&value))?;
+                acc.retract_batch(std::slice::from_ref(&value))?;
+            }
+            Ok(())
+        };
+
+        // Let every buffer reach its steady state before taking a baseline.
+        slide(&mut acc, WARMUP_ROWS)?;
+        let baseline = acc.size();
+
+        slide(&mut acc, EXTRA_ROWS)?;
+        let grown = acc.size();
+
+        assert!(
+            grown <= 2 * baseline,
+            "size() must not grow with the number of retracted rows: \
+             {baseline} bytes after {WARMUP_ROWS} rows, \
+             {grown} bytes after {} rows",
+            WARMUP_ROWS + EXTRA_ROWS
+        );
+
+        // Everything was retracted so evaluate must return null.
+        let result = acc.evaluate()?;
+        assert!(
+            matches!(&result, ScalarValue::List(arr) if arr.is_null(0)),
+            "expected null list after retracting every row, got {result:?}"
+        );
+
+        Ok(())
+    }
 }

From 51fb11f805e751a5d6fb5967498732eecdf95636 Mon Sep 17 00:00:00 2001
From: Fred Thomas <frederic.thomas@datadoghq.com>
Date: Wed, 29 Jul 2026 09:39:47 +0200
Subject: [PATCH 7/7] fix dictionary handling

---
 .../functions-aggregate/src/array_agg.rs      | 26 ++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index 17a8fa8e660b0..cfacd771968c2 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -855,6 +855,24 @@ pub struct DistinctArrayAggAccumulator {
     ignore_nulls: bool,
 }
 
+/// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type.
+///
+/// `RowConverter` always decodes to the physical (non-dictionary) type, so a
+/// cast back to the declared logical type is required when this is true.
+fn datatype_contains_dictionary(dt: &DataType) -> bool {
+    match dt {
+        DataType::Dictionary(_, _) => true,
+        DataType::List(f)
+        | DataType::LargeList(f)
+        | DataType::FixedSizeList(f, _)
+        | DataType::Map(f, _) => datatype_contains_dictionary(f.data_type()),
+        DataType::Struct(fields) => fields
+            .iter()
+            .any(|f| datatype_contains_dictionary(f.data_type())),
+        _ => false,
+    }
+}
+
 impl DistinctArrayAggAccumulator {
     pub fn try_new(
         datatype: &DataType,
@@ -1025,10 +1043,12 @@ impl Accumulator for DistinctArrayAggAccumulator {
         let arrays = converter.convert_rows(rows)?;
 
         // `convert_rows` always returns the physical (non-dictionary) type.
-        // Cast back to the declared logical type when they differ so that
-        // e.g. Dictionary columns round-trip correctly through the RowConverter.
+        // Cast back to the declared logical type when they differ AND the
+        // declared type contains a Dictionary somewhere (directly or nested
+        // inside a Struct, List, etc.) — that is the only case where
+        // RowConverter strips the logical type.
         let decoded = if arrays[0].data_type() != &self.datatype
-            && matches!(self.datatype, DataType::Dictionary(_, _))
+            && datatype_contains_dictionary(&self.datatype)
         {
             cast(arrays[0].as_ref(), &self.datatype)?
         } else {