You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/05 02:55:53 UTC
(datafusion) 01/01: feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator (#23628)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23628-c1366b554994b8ae99bbbe2cfd38abda8bf98cea
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 31ffab1418aae8966dd270b1042c386bc9646950
Author: Qi Zhu <zh...@gmail.com>
AuthorDate: Wed Aug 5 10:55:29 2026 +0800
feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator (#23628)
## Which issue does this PR close?
Closes #23601.
Part of the epic #23600.
## Rationale for this change
`first_value(x ORDER BY o)` / `last_value(x ORDER BY o)` currently fall
back to the per-group `Accumulator` path whenever `x` is a nested type
(`List`, `LargeList`, `Struct`, `Map`, ...), because
[`groups_accumulator_supported()`](https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/first_last.rs)
whitelists only scalar primitives and byte types.
Two consequences we hit:
- The per-group `Accumulator` path calls `ScalarValue::try_from_array`
on every candidate row, and stores `ScalarValue::List` (or
`ScalarValue::Struct`, ...) as an `Arc` slice into the source batch. For
wide payloads this both allocates heavily per row and *pins the source
batch* in memory for every group whose current winner came from it. On
high-cardinality dedup queries (`SELECT DISTINCT ON`, `ROW_NUMBER() =
1`, single-pass `FIRST_VALUE(...ORDER BY)` over a wide `List<Struct>`)
that combination can OOM even when the final output would fit
comfortably.
- This is also what #16620 (`Performance of DISTINCT ON (columns)`)
reports at the SQL layer.
## What changes are included in this PR?
Adds a new `GenericValueState` in `first_last/state.rs` — a
`Vec<Option<ScalarValue>>`-backed `ValueState` — and wires it into
`create_groups_accumulator` for nested types. The state uses
`ScalarValue::compact()` after `try_from_array` so the stored winner is
an owned copy instead of an `Arc` slice into the source batch (otherwise
the fast path would still pin source batches even though the
accumulator's own reported size looks small).
Types now supported by the `GroupsAccumulator` fast path:
- `List`, `LargeList`, `ListView`, `LargeListView`
- `FixedSizeList`
- `Struct`
- `Map`
The existing `PrimitiveValueState` (scalar primitives) and
`BytesValueState` (Utf8 / Binary variants) paths are unchanged.
## Are these changes tested?
Yes.
- 9 unit tests in `first_last/state.rs` covering `List<Utf8>`, `Struct`,
`LargeList`, `FixedSizeList`, `Map`, `EmitTo::First(n)` partial emit,
null handling, size accounting on overwrite and shrink, and a dedicated
regression test
(`test_generic_value_state_compact_releases_parent_batch`) that verifies
`compact()` releases the `Arc` reference to the source batch.
- 2 integration tests in `first_last.rs` exercising the full
`FirstLastGroupsAccumulator<GenericValueState>` for `List<Int32>`: one
functional multi-batch test, and one that asserts the accumulator's
reported `size()` stays proportional to `#groups` rather than `#rows` on
a 10-group × 10 000-row workload.
`cargo fmt` and `cargo clippy -p datafusion-functions-aggregate --lib
--tests -- -D warnings` are clean.
## Are there any user-facing changes?
Yes, but only in the sense that a previously-slow / OOM-prone path
becomes fast for the affected data types. No SQL or API surface changes;
the same queries run without modification and previously-passing tests
continue to pass.
## Follow-ups
- Add a benchmark comparing the two code paths on wide-payload
aggregates so the memory improvement is measurable in-tree.
- The two remaining sub-issues in the epic build on this one:
- #23602 — coalesce peer `FIRST_VALUE(... ORDER BY o)` expressions into
a single struct accumulator.
- #23603 — logical rewrite `Filter(row_number() = 1) →
Aggregate(FIRST_VALUE(... ORDER BY o))`. Depends on this PR to emit the
fast path.
---
datafusion/functions-aggregate/src/first_last.rs | 341 ++++++++++++++-
.../functions-aggregate/src/first_last/state.rs | 459 ++++++++++++++++++++-
.../sqllogictest/test_files/first_last_nested.slt | 80 ++++
3 files changed, 876 insertions(+), 4 deletions(-)
diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs
index ea45e42e84..c56cd73dbe 100644
--- a/datafusion/functions-aggregate/src/first_last.rs
+++ b/datafusion/functions-aggregate/src/first_last.rs
@@ -51,7 +51,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
mod state;
-use state::{BytesValueState, PrimitiveValueState, ValueState};
+use state::{BytesValueState, GenericValueState, PrimitiveValueState, ValueState};
create_func!(FirstValue, first_value_udaf);
create_func!(LastValue, last_value_udaf);
@@ -171,6 +171,23 @@ fn create_groups_accumulator(
BytesValueState::try_new(data_type.clone())?,
),
+ // Nested / composite types fall through to a generic ScalarValue-backed
+ // state. Slower per-batch than the primitive/bytes fast paths but still
+ // avoids the per-row ScalarValue churn of the per-group `Accumulator`
+ // path: winner extraction happens once per group per batch, not once
+ // per candidate row.
+ DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::ListView(_)
+ | DataType::LargeListView(_)
+ | DataType::FixedSizeList(_, _)
+ | DataType::Struct(_)
+ | DataType::Map(_, _) => create_groups_accumulator_helper(
+ args,
+ is_first,
+ GenericValueState::new(data_type.clone()),
+ ),
+
_ => internal_err!(
"GroupsAccumulator not supported for {}({})",
function_name,
@@ -209,6 +226,13 @@ fn groups_accumulator_supported(args: &AccumulatorArgs) -> bool {
| Binary
| LargeBinary
| BinaryView
+ | List(_)
+ | LargeList(_)
+ | ListView(_)
+ | LargeListView(_)
+ | FixedSizeList(_, _)
+ | Struct(_)
+ | Map(_, _)
)
}
@@ -2047,4 +2071,319 @@ mod tests {
Ok(())
}
+
+ /// End-to-end integration test for the nested-type support added to
+ /// [`FirstLastGroupsAccumulator`]: build the accumulator directly with a
+ /// [`GenericValueState`] for `List<Int32>` and verify that winners are
+ /// selected correctly across multiple batches.
+ ///
+ /// Mirrors the shape produced by SQL like:
+ /// ```sql
+ /// SELECT first_value(list_col ORDER BY o DESC) FROM t GROUP BY p
+ /// ```
+ /// which previously fell back to the per-group `Accumulator` path and
+ /// blew up on wide payloads.
+ #[test]
+ fn test_first_group_acc_list_int32() -> Result<()> {
+ let value_type =
+ DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("val", value_type.clone(), true),
+ Field::new("ord", DataType::Int64, true),
+ ]));
+ let sort_keys = [PhysicalSortExpr {
+ expr: col("ord", &schema)?,
+ options: SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ }];
+
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ GenericValueState::new(value_type.clone()),
+ sort_keys.into(),
+ false,
+ &[DataType::Int64],
+ /* pick_first = */ true,
+ )?;
+
+ // Batch 1: four rows across two groups.
+ // Winners (largest ord per group with pick_first=true + DESC):
+ // group 0 -> ord=30 -> [3, 3, 3]
+ // group 1 -> ord=40 -> [4, 4, 4, 4]
+ let values_1 = ListArray::from_iter_primitive::<Int32Type, _, _>([
+ Some(vec![Some(1)]),
+ Some(vec![Some(2), Some(2)]),
+ Some(vec![Some(3), Some(3), Some(3)]),
+ Some(vec![Some(4), Some(4), Some(4), Some(4)]),
+ ]);
+ let orderings_1 = Int64Array::from(vec![10, 20, 30, 40]);
+ group_acc.update_batch(
+ &[
+ Arc::new(values_1) as ArrayRef,
+ Arc::new(orderings_1) as ArrayRef,
+ ],
+ &[0, 1, 0, 1],
+ None,
+ 2,
+ )?;
+
+ // Batch 2: group 0 gets a new winner ord=50 -> [9, 9]; group 1
+ // keeps its previous winner (5 < 40).
+ let values_2 = ListArray::from_iter_primitive::<Int32Type, _, _>([
+ Some(vec![Some(9), Some(9)]),
+ Some(vec![Some(8)]),
+ ]);
+ let orderings_2 = Int64Array::from(vec![50, 5]);
+ group_acc.update_batch(
+ &[
+ Arc::new(values_2) as ArrayRef,
+ Arc::new(orderings_2) as ArrayRef,
+ ],
+ &[0, 1],
+ None,
+ 2,
+ )?;
+
+ let result = group_acc.evaluate(EmitTo::All)?;
+ let result = result.as_list::<i32>();
+ assert_eq!(result.len(), 2);
+ let g0 = result.value(0);
+ let g0 = g0.as_primitive::<Int32Type>();
+ assert_eq!(g0.len(), 2);
+ assert_eq!(g0.value(0), 9);
+ assert_eq!(g0.value(1), 9);
+ let g1 = result.value(1);
+ let g1 = g1.as_primitive::<Int32Type>();
+ assert_eq!(g1.len(), 4);
+ for i in 0..4 {
+ assert_eq!(g1.value(i), 4);
+ }
+ Ok(())
+ }
+
+ /// Regression test for the wide-payload memory blow-up: run the full
+ /// aggregate loop over a batch large enough that the per-group
+ /// `Accumulator` path would have generated N * batch-worth of state
+ /// (via `ScalarValue::List` clones) and verify that the reported
+ /// accumulator size stays proportional to `#groups`, not `#rows`.
+ #[test]
+ fn test_first_group_acc_list_size_bounded_by_groups() -> Result<()> {
+ let value_type =
+ DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("val", value_type.clone(), true),
+ Field::new("ord", DataType::Int64, true),
+ ]));
+ let sort_keys = [PhysicalSortExpr {
+ expr: col("ord", &schema)?,
+ options: SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ }];
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ GenericValueState::new(value_type),
+ sort_keys.into(),
+ false,
+ &[DataType::Int64],
+ true,
+ )?;
+
+ // 10 groups × 10_000 candidate rows per group (100_000 total). Each
+ // list value has ~10 elements. Under the old per-group `Accumulator`
+ // + Arc-slice code path this would pin every batch in memory.
+ const GROUPS: usize = 10;
+ const ROWS_PER_GROUP: usize = 10_000;
+ const N: usize = GROUPS * ROWS_PER_GROUP;
+ let values = ListArray::from_iter_primitive::<Int32Type, _, _>(
+ repeat_with(|| Some(vec![Some(1_i32); 10])).take(N),
+ );
+ let orderings = Int64Array::from((0..N as i64).collect::<Vec<_>>());
+ let group_indices: Vec<usize> = (0..N).map(|i| i % GROUPS).collect();
+
+ group_acc.update_batch(
+ &[
+ Arc::new(values) as ArrayRef,
+ Arc::new(orderings) as ArrayRef,
+ ],
+ &group_indices,
+ None,
+ GROUPS,
+ )?;
+
+ // Sanity: the retained size must be small — well under what a single
+ // input batch worth of list buffers would occupy. The exact number is
+ // implementation-dependent, but should be O(GROUPS * per-list), not
+ // O(N * per-list).
+ let size = group_acc.size();
+ assert!(
+ size < 100_000,
+ "accumulator size {size} bytes is not bounded by #groups (10 groups × ~10 int32 list elements)"
+ );
+
+ // Winner per group is the row with the largest ord — with our layout
+ // that's the last row assigned to each group.
+ let result = group_acc.evaluate(EmitTo::All)?;
+ let result = result.as_list::<i32>();
+ assert_eq!(result.len(), GROUPS);
+ for g in 0..GROUPS {
+ let winner = result.value(g);
+ let winner = winner.as_primitive::<Int32Type>();
+ assert_eq!(winner.len(), 10);
+ for i in 0..10 {
+ assert_eq!(winner.value(i), 1);
+ }
+ }
+ Ok(())
+ }
+
+ /// End-to-end memory-savings regression test.
+ ///
+ /// Streams many independent batches of wide `List<Int32>` payload through
+ /// the accumulator, dropping each source batch immediately after feeding
+ /// it in. The test then verifies three things:
+ ///
+ /// 1. The accumulator still emits the correct winners after every
+ /// source batch has been dropped (proves that stored values are
+ /// owned copies, not `Arc` slices into batches that no longer
+ /// exist).
+ /// 2. No buffer of any past source batch is shared by the emitted
+ /// output — the raw data-buffer pointer of every source batch is
+ /// recorded, and the final output's buffers must not alias any of
+ /// them (proves `compact()` copied the winners into owned memory).
+ /// 3. The accumulator's reported `size()` stays bounded by
+ /// `#groups * per-group-cost`, independent of `#batches * #rows`.
+ ///
+ /// This is the regression test for the wide-payload pinning behaviour
+ /// that motivated this PR.
+ #[test]
+ fn test_first_group_acc_list_no_source_batch_pinning() -> Result<()> {
+ let value_type =
+ DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("val", value_type.clone(), true),
+ Field::new("ord", DataType::Int64, true),
+ ]));
+ let sort_keys = [PhysicalSortExpr {
+ expr: col("ord", &schema)?,
+ options: SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ }];
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ GenericValueState::new(value_type),
+ sort_keys.into(),
+ false,
+ &[DataType::Int64],
+ true,
+ )?;
+
+ const GROUPS: usize = 4;
+ const BATCHES: usize = 50;
+ const ROWS_PER_BATCH: usize = 256;
+
+ // Record the raw pointer of each source batch's Int32 value-data
+ // buffer. If `compact()` did its job, the accumulator's final
+ // output must not share any of these pointers — every winner
+ // value should have been copied into an owned buffer.
+ let mut source_value_ptrs: Vec<*const u8> = Vec::with_capacity(BATCHES);
+
+ // Track the running-max ord we have fed to each group so the test's
+ // "expected winner" oracle matches the accumulator's choice.
+ let mut expected_ord = [i64::MIN; GROUPS];
+ let mut expected_val_repeat = [0_i32; GROUPS];
+
+ for batch in 0..BATCHES {
+ // Each batch's list values are `[batch as i32; group_idx + 1]`
+ // — a distinct payload per (batch, row) so we can verify the
+ // winner by content.
+ let values = ListArray::from_iter_primitive::<Int32Type, _, _>(
+ (0..ROWS_PER_BATCH).map(|i| {
+ let g = i % GROUPS;
+ Some(vec![Some(batch as i32); g + 1])
+ }),
+ );
+ let orderings = Int64Array::from(
+ (0..ROWS_PER_BATCH as i64)
+ .map(|i| batch as i64 * ROWS_PER_BATCH as i64 + i)
+ .collect::<Vec<_>>(),
+ );
+ let group_indices: Vec<usize> =
+ (0..ROWS_PER_BATCH).map(|i| i % GROUPS).collect();
+
+ // Update the oracle: the last row in this batch that hits each
+ // group has the largest ord for that group in this batch.
+ for i in (0..ROWS_PER_BATCH).rev() {
+ let g = i % GROUPS;
+ let ord = batch as i64 * ROWS_PER_BATCH as i64 + i as i64;
+ if ord > expected_ord[g] {
+ expected_ord[g] = ord;
+ expected_val_repeat[g] = batch as i32;
+ }
+ }
+
+ // Capture the raw pointer of this batch's Int32 value-data
+ // buffer *before* handing ownership to the accumulator. Int32
+ // arrays have a single value buffer at index 0.
+ source_value_ptrs.push(values.values().to_data().buffers()[0].as_ptr());
+
+ let values_arc: Arc<dyn Array> = Arc::new(values);
+ let orderings_arc: Arc<dyn Array> = Arc::new(orderings);
+
+ group_acc.update_batch(
+ &[values_arc, orderings_arc],
+ &group_indices,
+ None,
+ GROUPS,
+ )?;
+
+ // Drop happens implicitly at end of scope.
+ }
+
+ // (2) Size is bounded by #groups. The exact number is
+ // implementation-dependent but should be orders of magnitude below
+ // `BATCHES * ROWS_PER_BATCH * per-list-cost` (the amount that would
+ // be retained under the old Arc-slice pinning bug).
+ let size = group_acc.size();
+ assert!(
+ size < 10_000,
+ "accumulator size {size} bytes is not bounded by #groups \
+ (expected O({GROUPS}) not O({BATCHES} * {ROWS_PER_BATCH}))"
+ );
+
+ // (1) Winners are still readable and match the oracle.
+ let result = group_acc.evaluate(EmitTo::All)?;
+ let result_list = result.as_list::<i32>();
+ assert_eq!(result_list.len(), GROUPS);
+ for (g, expected_repeat) in expected_val_repeat.iter().enumerate().take(GROUPS) {
+ let winner = result_list.value(g);
+ let winner = winner.as_primitive::<Int32Type>();
+ assert_eq!(winner.len(), g + 1, "winner list length for group {g}");
+ for i in 0..winner.len() {
+ assert_eq!(
+ winner.value(i),
+ *expected_repeat,
+ "winner payload mismatch for group {g}"
+ );
+ }
+ }
+
+ // (3) The critical byte-level check: the emitted output's Int32
+ // value-data buffer must NOT share a raw pointer with any of the
+ // source batches. If `compact()` were omitted, `list_array.value(i)`
+ // would yield a slice whose backing buffer points into the source
+ // batch — the accumulator would then either pin the batch or emit
+ // an output that shares its buffer.
+ let result_values_ptr = result_list.values().to_data().buffers()[0].as_ptr();
+ for (i, src_ptr) in source_value_ptrs.iter().enumerate() {
+ assert_ne!(
+ *src_ptr, result_values_ptr,
+ "emitted result's Int32 value buffer aliases source batch \
+ {i}'s buffer; compact() is not making an owned copy"
+ );
+ }
+ Ok(())
+ }
}
diff --git a/datafusion/functions-aggregate/src/first_last/state.rs b/datafusion/functions-aggregate/src/first_last/state.rs
index cd7114bf04..d99b4f6ecc 100644
--- a/datafusion/functions-aggregate/src/first_last/state.rs
+++ b/datafusion/functions-aggregate/src/first_last/state.rs
@@ -25,7 +25,7 @@ use arrow::array::{
};
use arrow::buffer::{BooleanBuffer, NullBuffer};
use arrow::datatypes::DataType;
-use datafusion_common::{Result, internal_err};
+use datafusion_common::{Result, ScalarValue, internal_err};
use datafusion_expr::EmitTo;
pub(crate) trait ValueState: Send + Sync {
@@ -290,6 +290,78 @@ impl BytesValueState {
}
}
+/// Fallback state for arbitrary Arrow types (List, LargeList, Struct, Map, ...)
+/// that are not covered by [`PrimitiveValueState`] or [`BytesValueState`].
+///
+/// Stores one [`ScalarValue`] per group. Winners are identified by the
+/// vectorized comparator in the enclosing accumulator, so the per-row
+/// allocation cost of the fallback `Accumulator` path is avoided:
+/// `ScalarValue::try_from_array` is called once per group per batch (at
+/// winner-update time), not once per candidate row.
+pub(crate) struct GenericValueState {
+ vals: Vec<Option<ScalarValue>>,
+ data_type: DataType,
+ /// Cached total heap size of `vals`, updated on each mutation to avoid
+ /// walking the vector on every `size()` call.
+ total_size: usize,
+}
+
+impl GenericValueState {
+ pub(crate) fn new(data_type: DataType) -> Self {
+ Self {
+ vals: vec![],
+ data_type,
+ total_size: 0,
+ }
+ }
+}
+
+impl ValueState for GenericValueState {
+ fn resize(&mut self, new_size: usize) {
+ if new_size < self.vals.len() {
+ for v in self.vals[new_size..].iter().flatten() {
+ self.total_size -= v.size();
+ }
+ }
+ self.vals.resize(new_size, None);
+ }
+
+ fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()> {
+ if let Some(v) = &self.vals[group_idx] {
+ self.total_size -= v.size();
+ }
+ let mut scalar = ScalarValue::try_from_array(array, idx)?;
+ // `try_from_array` for nested types returns Arc slices into the source
+ // batch buffers, so a single stored winner would pin the entire batch
+ // in memory. Compact copies the referenced bytes into an owned buffer
+ // so old batches can be dropped as new ones arrive.
+ scalar.compact();
+ self.total_size += scalar.size();
+ self.vals[group_idx] = Some(scalar);
+ Ok(())
+ }
+
+ fn take(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let taken = emit_to.take_needed(&mut self.vals);
+ let taken_size: usize = taken.iter().flatten().map(|v| v.size()).sum();
+ self.total_size -= taken_size;
+
+ let default = ScalarValue::try_from(&self.data_type)?;
+ let scalars = taken
+ .into_iter()
+ .map(|opt| opt.unwrap_or_else(|| default.clone()))
+ .collect::<Vec<_>>();
+ if scalars.is_empty() {
+ return Ok(arrow::array::new_empty_array(&self.data_type));
+ }
+ ScalarValue::iter_to_array(scalars)
+ }
+
+ fn size(&self) -> usize {
+ self.vals.capacity() * size_of::<Option<ScalarValue>>() + self.total_size
+ }
+}
+
pub(crate) fn take_need(
bool_buf_builder: &mut BooleanBufferBuilder,
emit_to: EmitTo,
@@ -312,9 +384,12 @@ pub(crate) fn take_need(
mod tests {
use super::*;
use arrow::array::{
- BinaryArray, BinaryViewArray, LargeBinaryArray, LargeStringArray, StringArray,
- StringViewArray,
+ Array, BinaryArray, BinaryViewArray, FixedSizeListArray, Int32Array,
+ Int32Builder, LargeBinaryArray, LargeListArray, LargeStringArray, ListBuilder,
+ MapArray, StringArray, StringBuilder, StringViewArray, StructArray,
};
+ use arrow::buffer::{OffsetBuffer, ScalarBuffer};
+ use arrow::datatypes::{DataType, Field, Fields};
#[test]
fn test_bytes_value_state_utf8() -> Result<()> {
@@ -459,4 +534,382 @@ mod tests {
Ok(())
}
+
+ // ---------- GenericValueState (nested types) ----------
+
+ /// Build a `List<Utf8>` array with three rows: `["a"]`, `["b", "c"]`,
+ /// `["d", "e", "f"]`. Used by several tests.
+ fn make_list_utf8_array() -> ArrayRef {
+ let mut builder = ListBuilder::new(StringBuilder::new());
+ builder.values().append_value("a");
+ builder.append(true);
+ builder.values().append_value("b");
+ builder.values().append_value("c");
+ builder.append(true);
+ builder.values().append_value("d");
+ builder.values().append_value("e");
+ builder.values().append_value("f");
+ builder.append(true);
+ Arc::new(builder.finish())
+ }
+
+ #[test]
+ fn test_generic_value_state_list_utf8() -> Result<()> {
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8.clone());
+ state.resize(2);
+
+ let array = make_list_utf8_array();
+
+ // group 0 <- ["a"] ; group 1 <- ["b", "c"]
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+
+ // Overwrite group 0 with the wider ["d", "e", "f"] (size-accounting
+ // must decrement the old value before adding the new one).
+ let size_after_first = state.total_size;
+ state.update(0, &array, 2)?;
+ assert!(
+ state.total_size > 0,
+ "total_size must remain positive after overwrite"
+ );
+ // The overwrite replaced group 0's payload; the delta relative to the
+ // previous state should equal `new.size() - old.size()`. If the caller
+ // forgot to subtract the old size, `total_size` would drift upward.
+ let expected_delta = {
+ let new_scalar = {
+ let mut s = ScalarValue::try_from_array(&array, 2)?;
+ s.compact();
+ s
+ };
+ let old_scalar = {
+ let mut s = ScalarValue::try_from_array(&array, 0)?;
+ s.compact();
+ s
+ };
+ new_scalar.size() as isize - old_scalar.size() as isize
+ };
+ assert_eq!(
+ state.total_size as isize - size_after_first as isize,
+ expected_delta,
+ "size accounting drifted after overwrite"
+ );
+
+ let result = state.take(EmitTo::All)?;
+ let result = result.as_list::<i32>();
+ assert_eq!(result.len(), 2);
+ let g0 = result.value(0);
+ let g0 = g0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(g0.value(0), "d");
+ assert_eq!(g0.value(1), "e");
+ assert_eq!(g0.value(2), "f");
+ let g1 = result.value(1);
+ let g1 = g1.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(g1.value(0), "b");
+ assert_eq!(g1.value(1), "c");
+
+ assert_eq!(state.total_size, 0, "state must be fully drained");
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_struct() -> Result<()> {
+ let fields = Fields::from(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("name", DataType::Utf8, false),
+ ]);
+ let struct_type = DataType::Struct(fields.clone());
+
+ let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ let name = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
+ let struct_array =
+ Arc::new(StructArray::new(fields, vec![id, name], None)) as ArrayRef;
+
+ let mut state = GenericValueState::new(struct_type);
+ state.resize(2);
+ state.update(0, &struct_array, 0)?;
+ state.update(1, &struct_array, 2)?;
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(out.len(), 2);
+
+ let out_id = out.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
+ let out_name = out
+ .column(1)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(out_id.value(0), 1);
+ assert_eq!(out_id.value(1), 3);
+ assert_eq!(out_name.value(0), "a");
+ assert_eq!(out_name.value(1), "c");
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_large_list() -> Result<()> {
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let large_list_type = DataType::LargeList(Arc::clone(&field));
+
+ let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let offsets: OffsetBuffer<i64> =
+ OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 5, 6]));
+ let array = Arc::new(LargeListArray::new(field, offsets, Arc::new(values), None))
+ as ArrayRef;
+
+ let mut state = GenericValueState::new(large_list_type);
+ state.resize(2);
+ state.update(0, &array, 0)?; // [1, 2]
+ state.update(1, &array, 2)?; // [6]
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_list::<i64>();
+ assert_eq!(out.len(), 2);
+ let g0 = out.value(0);
+ let g0 = g0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g0.len(), 2);
+ assert_eq!(g0.value(0), 1);
+ assert_eq!(g0.value(1), 2);
+ let g1 = out.value(1);
+ let g1 = g1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g1.len(), 1);
+ assert_eq!(g1.value(0), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_fixed_size_list() -> Result<()> {
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let fsl_type = DataType::FixedSizeList(Arc::clone(&field), 2);
+
+ let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let array = Arc::new(FixedSizeListArray::new(field, 2, Arc::new(values), None))
+ as ArrayRef;
+
+ let mut state = GenericValueState::new(fsl_type);
+ state.resize(2);
+ state.update(0, &array, 0)?; // [1, 2]
+ state.update(1, &array, 2)?; // [5, 6]
+
+ let out = state.take(EmitTo::All)?;
+ let out = out
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .expect("emitted FixedSizeListArray");
+ assert_eq!(out.len(), 2);
+ let g0 = out.value(0);
+ let g0 = g0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g0.value(0), 1);
+ assert_eq!(g0.value(1), 2);
+ let g1 = out.value(1);
+ let g1 = g1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g1.value(0), 5);
+ assert_eq!(g1.value(1), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_map() -> Result<()> {
+ // Map<Utf8, Int32> with two entries: {"a": 1, "b": 2}, {"c": 3}
+ let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
+ let values = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ let entry_fields = Fields::from(vec![
+ Field::new("keys", DataType::Utf8, false),
+ Field::new("values", DataType::Int32, true),
+ ]);
+ let entries = StructArray::new(entry_fields.clone(), vec![keys, values], None);
+ let offsets: OffsetBuffer<i32> =
+ OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 3]));
+ let map_field =
+ Arc::new(Field::new("entries", DataType::Struct(entry_fields), false));
+ let map_array = Arc::new(MapArray::new(
+ Arc::clone(&map_field),
+ offsets,
+ entries,
+ None,
+ false,
+ )) as ArrayRef;
+ let map_type = DataType::Map(map_field, false);
+
+ let mut state = GenericValueState::new(map_type);
+ state.resize(2);
+ state.update(0, &map_array, 0)?; // {"a": 1, "b": 2}
+ state.update(1, &map_array, 1)?; // {"c": 3}
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_any().downcast_ref::<MapArray>().unwrap();
+ assert_eq!(out.len(), 2);
+ assert_eq!(out.value_length(0), 2);
+ assert_eq!(out.value_length(1), 1);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_emit_first() -> Result<()> {
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8);
+ state.resize(3);
+
+ let array = make_list_utf8_array();
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+ state.update(2, &array, 2)?;
+
+ let after_all_updates = state.total_size;
+ assert!(after_all_updates > 0);
+
+ // Emit the first 2 groups; remaining group 2 stays.
+ let head = state.take(EmitTo::First(2))?;
+ let head = head.as_list::<i32>();
+ assert_eq!(head.len(), 2);
+ let h0 = head.value(0);
+ let h0 = h0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(h0.value(0), "a");
+ let h1 = head.value(1);
+ let h1 = h1.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(h1.value(0), "b");
+ assert_eq!(h1.value(1), "c");
+
+ // After partial emit, `total_size` shrank but is still positive.
+ assert!(state.total_size > 0);
+ assert!(state.total_size < after_all_updates);
+
+ let tail = state.take(EmitTo::All)?;
+ let tail = tail.as_list::<i32>();
+ assert_eq!(tail.len(), 1);
+ let t0 = tail.value(0);
+ let t0 = t0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(t0.value(0), "d");
+ assert_eq!(t0.value(1), "e");
+ assert_eq!(t0.value(2), "f");
+
+ assert_eq!(state.total_size, 0);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_update_null() -> Result<()> {
+ // List<Int32> with rows: [1, 2], NULL
+ let mut builder = ListBuilder::new(Int32Builder::new());
+ builder.values().append_value(1);
+ builder.values().append_value(2);
+ builder.append(true);
+ builder.append(false); // null entry
+ let array: ArrayRef = Arc::new(builder.finish());
+
+ let list_type = array.data_type().clone();
+ let mut state = GenericValueState::new(list_type);
+ state.resize(1);
+
+ // group 0 = [1, 2]
+ state.update(0, &array, 0)?;
+ let size_after_value = state.total_size;
+ assert!(size_after_value > 0);
+
+ // Overwrite group 0 with NULL. The size accounting must subtract the
+ // previous value's size and then add the null-scalar's size; the point
+ // of this test is that `total_size` stays consistent (no drift) and
+ // the null is emitted correctly.
+ state.update(0, &array, 1)?;
+ // Recomputing from scratch must match the cached total_size.
+ let recomputed: usize = state.vals.iter().flatten().map(|v| v.size()).sum();
+ assert_eq!(
+ state.total_size, recomputed,
+ "total_size drifted after null update"
+ );
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_list::<i32>();
+ assert_eq!(out.len(), 1);
+ assert!(out.is_null(0));
+ assert_eq!(state.total_size, 0);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_compact_releases_parent_batch() -> Result<()> {
+ // Regression test for the memory-pinning bug: without compact(),
+ // `ScalarValue::try_from_array` on a List column produces a
+ // ScalarValue whose child values array is an Arrow slice pointing
+ // into the *source* batch's underlying byte buffer. That means the
+ // source batch's memory stays alive until every extracted winner
+ // is dropped, even if the outer ListArray is released. `compact()`
+ // must copy the referenced bytes into a fresh owned buffer.
+ //
+ // Correctly detecting this requires comparing the raw buffer
+ // pointer of the source `Utf8` value-data buffer against the raw
+ // buffer pointer of the stored winner's value-data buffer. Checking
+ // `Arc::strong_count` on the outer `ArrayRef` is not sufficient,
+ // because `list_array.value(idx)` returns a sliced child that keeps
+ // its own Arc chain independent of the outer ListArray.
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8);
+ state.resize(1);
+
+ let array: ArrayRef = make_list_utf8_array();
+
+ // Capture the raw pointer of the *source* Utf8 value-data buffer.
+ // Utf8Array has two buffers: offsets (buffer 0) and value bytes
+ // (buffer 1). Comparing buffer 1 is the direct check for byte
+ // pinning.
+ let source_values_ptr =
+ array.as_list::<i32>().values().to_data().buffers()[1].as_ptr();
+
+ state.update(0, &array, 0)?;
+ drop(array);
+
+ // Directly probe the stored ScalarValue's underlying values buffer.
+ let stored_values_ptr = match state
+ .vals
+ .first()
+ .and_then(|opt| opt.as_ref())
+ .expect("group 0 should have a stored value")
+ {
+ ScalarValue::List(list_arr) => {
+ list_arr.values().to_data().buffers()[1].as_ptr()
+ }
+ other => panic!("expected ScalarValue::List, got {other:?}"),
+ };
+
+ assert_ne!(
+ source_values_ptr, stored_values_ptr,
+ "compact() failed: stored ScalarValue still shares the source \
+ batch's Utf8 value-data buffer, meaning the batch is pinned in \
+ memory even after the outer ArrayRef is dropped"
+ );
+
+ // Data must still be readable from the stored copy.
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_list::<i32>();
+ assert_eq!(out.len(), 1);
+ let g0 = out.value(0);
+ let g0 = g0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(g0.value(0), "a");
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_resize_shrink_recovers_size() -> Result<()> {
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8);
+ state.resize(3);
+
+ let array = make_list_utf8_array();
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+ state.update(2, &array, 2)?;
+ let full_size = state.total_size;
+ assert!(full_size > 0);
+
+ // Shrinking must subtract the dropped groups' sizes from total_size.
+ state.resize(1);
+ assert!(state.total_size > 0);
+ assert!(state.total_size < full_size);
+ Ok(())
+ }
}
diff --git a/datafusion/sqllogictest/test_files/first_last_nested.slt b/datafusion/sqllogictest/test_files/first_last_nested.slt
new file mode 100644
index 0000000000..b96b47f6ab
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/first_last_nested.slt
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# SQL-level coverage for first_value / last_value over nested payloads
+# (Struct, Map) through the grouped `FirstLastGroupsAccumulator` path.
+#
+# The accumulator is unit-tested directly in first_last.rs. These tests
+# add the integration coverage the unit tests cannot: the SLT runner
+# executes with target_partitions = 4, so a GROUP BY drives the two-phase
+# aggregate (Partial -> FinalPartitioned) and the nested intermediate
+# state produced by `state()` is round-tripped back through
+# `merge_batch()` across the partition boundary. Struct and Map otherwise
+# have no SQL-level first_value / last_value coverage (only List does, in
+# array_agg.slt).
+
+########################################
+# Struct payload
+########################################
+
+statement ok
+CREATE TABLE first_last_struct AS VALUES
+ (1, 1, named_struct('a', 10, 'b', 'x')),
+ (1, 2, named_struct('a', 20, 'b', 'y')),
+ (1, 3, named_struct('a', 30, 'b', 'z')),
+ (2, 1, named_struct('a', 40, 'b', 'p')),
+ (2, 2, named_struct('a', 50, 'b', 'q'));
+
+query I??
+select column1, first_value(column3 order by column2), last_value(column3 order by column2)
+from first_last_struct group by column1 order by column1;
+----
+1 {a: 10, b: x} {a: 30, b: z}
+2 {a: 40, b: p} {a: 50, b: q}
+
+# Descending order flips first / last.
+query I??
+select column1, first_value(column3 order by column2 desc), last_value(column3 order by column2 desc)
+from first_last_struct group by column1 order by column1;
+----
+1 {a: 30, b: z} {a: 10, b: x}
+2 {a: 50, b: q} {a: 40, b: p}
+
+statement ok
+drop table first_last_struct;
+
+########################################
+# Map payload
+########################################
+
+statement ok
+CREATE TABLE first_last_map AS VALUES
+ (1, 1, MAP {'k1': 10, 'k2': 20}),
+ (1, 2, MAP {'k3': 30}),
+ (1, 3, MAP {'k4': 40, 'k5': 50}),
+ (2, 1, MAP {'k9': 99}),
+ (2, 2, MAP {'k8': 88, 'k7': 77});
+
+query I??
+select column1, first_value(column3 order by column2), last_value(column3 order by column2)
+from first_last_map group by column1 order by column1;
+----
+1 {k1: 10, k2: 20} {k4: 40, k5: 50}
+2 {k9: 99} {k8: 88, k7: 77}
+
+statement ok
+drop table first_last_map;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org