diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index cecb277cb844a..0f52362d0d125 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(_, _) ) } @@ -1918,4 +1942,321 @@ 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` 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::([ + 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::([ + 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::(); + assert_eq!(result.len(), 2); + let g0 = result.value(0); + let g0 = g0.as_primitive::(); + 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::(); + 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::( + repeat_with(|| Some(vec![Some(1_i32); 10])).take(N), + ); + let orderings = Int64Array::from((0..N as i64).collect::>()); + let group_indices: Vec = (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::(); + assert_eq!(result.len(), GROUPS); + for g in 0..GROUPS { + let winner = result.value(g); + let winner = winner.as_primitive::(); + 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` 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 `Arc` reference to any past source batch is retained by the + /// accumulator — every batch's `Weak` handle reports + /// `strong_count == 0` once the local `Arc` is dropped. + /// 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; + + // Keep a `Weak` handle to each batch's payload array. After we drop + // the local `Arc`s, `strong_count == 0` means nothing in the + // accumulator is still holding a reference to that batch. + // 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::( + (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::>(), + ); + let group_indices: Vec = + (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 = Arc::new(values); + let orderings_arc: Arc = 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::(); + 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::(); + 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 cd7114bf04f9c..d99b4f6ecc6da 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>, + 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 { + 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::>(); + 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::>() + 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` 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::(); + assert_eq!(result.len(), 2); + let g0 = result.value(0); + let g0 = g0.as_any().downcast_ref::().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::().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::().unwrap(); + assert_eq!(out.len(), 2); + + let out_id = out.column(0).as_any().downcast_ref::().unwrap(); + let out_name = out + .column(1) + .as_any() + .downcast_ref::() + .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 = + 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::(); + assert_eq!(out.len(), 2); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().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::().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::() + .expect("emitted FixedSizeListArray"); + assert_eq!(out.len(), 2); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().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::().unwrap(); + assert_eq!(g1.value(0), 5); + assert_eq!(g1.value(1), 6); + Ok(()) + } + + #[test] + fn test_generic_value_state_map() -> Result<()> { + // Map 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 = + 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::().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::(); + assert_eq!(head.len(), 2); + let h0 = head.value(0); + let h0 = h0.as_any().downcast_ref::().unwrap(); + assert_eq!(h0.value(0), "a"); + let h1 = head.value(1); + let h1 = h1.as_any().downcast_ref::().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::(); + assert_eq!(tail.len(), 1); + let t0 = tail.value(0); + let t0 = t0.as_any().downcast_ref::().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 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::(); + 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::().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::(); + assert_eq!(out.len(), 1); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().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(()) + } }