feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator

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<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 `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::<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(())
+    }
 }