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/09/06 11:01:48 UTC
(datafusion) branch main updated: feat(physical-plan): `GroupColumn` support for `List` / `LargeList` (#23648)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 722cbf2e7a feat(physical-plan): `GroupColumn` support for `List` / `LargeList` (#23648)
722cbf2e7a is described below
commit 722cbf2e7ad9a43bc6723f7795fe87e0552eb81f
Author: Mikhail Zabaluev <mi...@gmail.com>
AuthorDate: Sun Sep 6 10:52:35 2026 +0000
feat(physical-plan): `GroupColumn` support for `List` / `LargeList` (#23648)
## Which issue does this PR close?
- Part of #22682.
## Rationale for this change
This adds column-native `GroupColumn` implementations for `List<T>` and
`LargeList<T>`, so a `GROUP BY` containing these (along with other
supported data types) no longer falls back from `GroupValuesColumn` to
`GroupValuesRows`.
## What changes are included in this PR?
Cherry-picked @zhuqi-lucas's changes in #22706 that pertained to `List`
/ `LargeList`, and made some optimizations.
## Are these changes tested?
Cherry-picked the unit tests for the `list` module from #22706.
Added supported/unsupported cases for `List` in the
`group_column_supported_type_matches_make_group_column` test.
Benchmarked in a proprietary application using several HashAggregate
operations on 20 columns including two string lists, to about 8%
cumulative improvement.
## Are there any user-facing changes?
No.
---------
Co-authored-by: Qi Zhu <82...@qq.com>
Co-authored-by: Trương Hoàng Long <lo...@gmail.com>
Co-authored-by: Qi Zhu <zh...@gmail.com>
Co-authored-by: Raz Luvaton <16...@users.noreply.github.com>
---
.../aggregates/group_values/multi_group_by/list.rs | 961 +++++++++++++++++++++
.../aggregates/group_values/multi_group_by/mod.rs | 238 ++---
2 files changed, 1102 insertions(+), 97 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs
new file mode 100644
index 0000000000..8c681698be
--- /dev/null
+++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs
@@ -0,0 +1,961 @@
+// 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.
+
+//! [`GroupColumn`] implementation for `List<T>` and `LargeList<T>` group keys.
+//!
+//! Storage:
+//! - `offsets: Vec<O>` of length `outer_len + 1`, with `offsets[0] = 0`. The
+//! children of outer row `i` live at child positions
+//! `[offsets[i] .. offsets[i+1])`.
+//! - `child: Box<dyn GroupColumn>` for the element type.
+//! - `outer_nulls`: outer-row null bitmap.
+//!
+//! Null outer rows are stored with a zero-length range (i.e.
+//! `offsets[i+1] == offsets[i]`). No child slots are pushed for null rows.
+
+use crate::aggregates::group_values::multi_group_by::{GroupColumn, nulls_equal_to};
+use crate::aggregates::group_values::null_builder::NullBufferBuilderExt;
+
+use arrow::array::{
+ Array, ArrayRef, BooleanBufferBuilder, GenericListArray, ListLikeArray,
+ NullBufferBuilder, OffsetSizeTrait, as_generic_list_array,
+};
+use arrow::buffer::{OffsetBuffer, ScalarBuffer};
+use arrow::datatypes::FieldRef;
+use datafusion_common::utils::split_vec_min_alloc;
+use datafusion_common::{Result, internal_datafusion_err};
+use datafusion_execution::memory_pool::proxy::VecAllocExt;
+use datafusion_expr::GroupSelection;
+use std::sync::Arc;
+
+/// A [`GroupColumn`] for `List<T>` (`O = i32`) and `LargeList<T>` (`O = i64`).
+pub struct ListGroupValueBuilder<O: OffsetSizeTrait> {
+ field: FieldRef,
+ offsets: Vec<O>,
+ child: Box<dyn GroupColumn>,
+ outer_nulls: NullBufferBuilder,
+ outer_len: usize,
+}
+
+impl<O: OffsetSizeTrait> ListGroupValueBuilder<O> {
+ pub fn new(field: FieldRef, child: Box<dyn GroupColumn>) -> Self {
+ Self {
+ field,
+ offsets: vec![O::usize_as(0)],
+ child,
+ outer_nulls: NullBufferBuilder::empty(),
+ outer_len: 0,
+ }
+ }
+
+ fn current_end(&self) -> O {
+ *self.offsets.last().expect("offsets is never empty")
+ }
+
+ fn push_offset(&mut self, additional: usize) -> Result<()> {
+ let additional = O::from_usize(additional).ok_or_else(|| {
+ internal_datafusion_err!(
+ "List sublist length {additional} overflows the offset type"
+ )
+ })?;
+ let next = self.current_end().checked_add(&additional).ok_or_else(|| {
+ internal_datafusion_err!(
+ "List offset overflow: current={} additional={}",
+ self.current_end().as_usize(),
+ additional.as_usize()
+ )
+ })?;
+ self.offsets.push(next);
+ Ok(())
+ }
+
+ fn equal_to_typed(
+ &self,
+ lhs_row: usize,
+ array: &GenericListArray<O>,
+ rhs_row: usize,
+ ) -> bool {
+ let lhs_null = self.outer_nulls.is_null(lhs_row);
+ let rhs_null = array.is_null(rhs_row);
+ if let Some(result) = nulls_equal_to(lhs_null, rhs_null) {
+ return result;
+ }
+
+ let rhs_range = array.element_range(rhs_row);
+ let lhs_start = self.offsets[lhs_row].as_usize();
+ let lhs_end = self.offsets[lhs_row + 1].as_usize();
+ let lhs_len = lhs_end - lhs_start;
+ if lhs_len != rhs_range.end - rhs_range.start {
+ return false;
+ }
+ let values = array.values();
+ for j in 0..lhs_len {
+ if !self
+ .child
+ .equal_to(lhs_start + j, values, rhs_range.start + j)
+ {
+ return false;
+ }
+ }
+ true
+ }
+
+ fn append_val_typed(
+ &mut self,
+ array: &GenericListArray<O>,
+ row: usize,
+ ) -> Result<()> {
+ if array.is_null(row) {
+ self.outer_nulls.append_null();
+ // Zero-length range for null outer rows: do not push any child
+ // elements, and the offset for the next row stays the same.
+ let end = self.current_end();
+ self.offsets.push(end);
+ } else {
+ self.outer_nulls.append_non_null();
+ let values = array.values();
+ let rhs_range = array.element_range(row);
+ let rhs_len = rhs_range.end - rhs_range.start;
+ for j in rhs_range {
+ self.child.append_val(values, j)?;
+ }
+ self.push_offset(rhs_len)?;
+ }
+ self.outer_len += 1;
+ Ok(())
+ }
+}
+
+impl<O: OffsetSizeTrait> GroupColumn for ListGroupValueBuilder<O> {
+ fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
+ let other = as_generic_list_array(array);
+ self.equal_to_typed(lhs_row, other, rhs_row)
+ }
+
+ fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> {
+ let other = as_generic_list_array(array);
+ self.append_val_typed(other, row)
+ }
+
+ fn vectorized_equal_to(
+ &self,
+ lhs_rows: &[usize],
+ array: &ArrayRef,
+ rhs_rows: &[usize],
+ equal_to_results: &mut BooleanBufferBuilder,
+ ) {
+ let other = as_generic_list_array(array);
+ for (idx, (&lhs_row, &rhs_row)) in
+ lhs_rows.iter().zip(rhs_rows.iter()).enumerate()
+ {
+ if !equal_to_results.get_bit(idx) {
+ continue;
+ }
+ if !self.equal_to_typed(lhs_row, other, rhs_row) {
+ equal_to_results.set_bit(idx, false);
+ }
+ }
+ }
+
+ fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> {
+ let l = as_generic_list_array(array);
+ for &row in rows {
+ self.append_val_typed(l, row)?;
+ }
+ Ok(())
+ }
+
+ fn len(&self) -> usize {
+ self.outer_len
+ }
+
+ fn size(&self) -> usize {
+ self.offsets.allocated_size()
+ + self.outer_nulls.allocated_size()
+ + self.child.size()
+ }
+
+ fn build(self: Box<Self>) -> ArrayRef {
+ let Self {
+ field,
+ offsets,
+ child,
+ outer_nulls,
+ outer_len: _,
+ } = *self;
+ let outer_nulls = outer_nulls.build();
+ let child_array = child.build();
+ // SAFETY: offsets are constructed monotonically by `push_offset` /
+ // initial `[0]`, and child_array length matches the final offset.
+ let offset_buffer = OffsetBuffer::<O>::new(ScalarBuffer::<O>::from(offsets));
+ Arc::new(GenericListArray::<O>::new(
+ field,
+ offset_buffer,
+ child_array,
+ outer_nulls,
+ ))
+ }
+
+ fn values_preserving(&self, selection: GroupSelection<'_>) -> Result<ArrayRef> {
+ selection.validate_num_groups(self.outer_len)?;
+
+ // Flatten the selected outer rows into the child positions they cover,
+ // rebuilding offsets for the output. Null (and empty) rows contribute a
+ // zero-length range, matching how they are stored.
+ let mut child_indices = Vec::new();
+ let mut offsets = Vec::with_capacity(selection.len() + 1);
+ let mut nulls = NullBufferBuilder::new(selection.len());
+ let mut end = O::usize_as(0);
+ offsets.push(end);
+
+ for row in selection.iter() {
+ nulls.append(!self.outer_nulls.is_null(row));
+ let start = self.offsets[row].as_usize();
+ let row_end = self.offsets[row + 1].as_usize();
+ child_indices.extend(start..row_end);
+ let len = O::from_usize(row_end - start).ok_or_else(|| {
+ internal_datafusion_err!(
+ "List sublist length {} overflows the offset type",
+ row_end - start
+ )
+ })?;
+ end = end.checked_add(&len).ok_or_else(|| {
+ internal_datafusion_err!(
+ "List offset overflow: current={} additional={}",
+ end.as_usize(),
+ len.as_usize()
+ )
+ })?;
+ offsets.push(end);
+ }
+
+ let child_selection =
+ GroupSelection::try_from_indices(&child_indices, self.child.len())?;
+ let child_array = self.child.values_preserving(child_selection)?;
+
+ // SAFETY: offsets start at 0 and grow monotonically by each selected
+ // row's length, and `child_array` has exactly `end` elements.
+ let offset_buffer =
+ unsafe { OffsetBuffer::<O>::new_unchecked(ScalarBuffer::<O>::from(offsets)) };
+ Ok(Arc::new(GenericListArray::<O>::new(
+ Arc::clone(&self.field),
+ offset_buffer,
+ child_array,
+ nulls.build(),
+ )))
+ }
+
+ fn take_n(&mut self, n: usize) -> ArrayRef {
+ // Number of child elements consumed by the first n outer rows.
+ let cut_offset = self.offsets[n];
+ let cut = cut_offset.as_usize();
+
+ // Split off first-n offsets: 0, off[1], ..., off[n-1], top up with off[n].
+ let mut first_n_offsets = split_vec_min_alloc(&mut self.offsets, n);
+ first_n_offsets.push(cut_offset);
+ // Remaining offsets shifted so that what was off[n] becomes 0.
+ self.offsets.iter_mut().for_each(|o| *o = o.sub(cut_offset));
+
+ let first_n_outer_nulls = self.outer_nulls.take_n(n);
+ let first_n_child = self.child.take_n(cut);
+ self.outer_len -= n;
+
+ let offset_buffer =
+ OffsetBuffer::<O>::new(ScalarBuffer::<O>::from(first_n_offsets));
+ Arc::new(GenericListArray::<O>::new(
+ Arc::clone(&self.field),
+ offset_buffer,
+ first_n_child,
+ first_n_outer_nulls,
+ ))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder;
+ use arrow::array::{
+ Int32Array, LargeListArray, ListArray, builder::Int32Builder,
+ builder::LargeListBuilder, builder::ListBuilder,
+ };
+ use arrow::datatypes::{DataType, Field, Int32Type};
+
+ fn child_field() -> FieldRef {
+ Arc::new(Field::new("item", DataType::Int32, true))
+ }
+
+ fn list_array(rows: &[Option<Vec<Option<i32>>>]) -> ArrayRef {
+ let mut b = ListBuilder::new(Int32Builder::new());
+ for row in rows {
+ match row {
+ None => b.append(false),
+ Some(items) => {
+ for v in items {
+ b.values().append_option(*v);
+ }
+ b.append(true);
+ }
+ }
+ }
+ Arc::new(b.finish())
+ }
+
+ fn large_list_array(rows: &[Option<Vec<Option<i32>>>]) -> ArrayRef {
+ let mut b = LargeListBuilder::new(Int32Builder::new());
+ for row in rows {
+ match row {
+ None => b.append(false),
+ Some(items) => {
+ for v in items {
+ b.values().append_option(*v);
+ }
+ b.append(true);
+ }
+ }
+ }
+ Arc::new(b.finish())
+ }
+
+ fn make_child() -> Box<dyn GroupColumn> {
+ Box::new(PrimitiveGroupValueBuilder::<Int32Type, true>::new(
+ DataType::Int32,
+ ))
+ }
+
+ #[test]
+ fn list_append_equal_build_round_trip() {
+ let input = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(3)]),
+ None,
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![]),
+ ]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+ assert_eq!(builder.len(), 5);
+
+ // Verify equality matrix.
+ let probe = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(1), Some(3)]),
+ None,
+ Some(vec![]),
+ ]);
+ assert!(builder.equal_to(0, &probe, 0), "[1,2] == [1,2]");
+ assert!(!builder.equal_to(0, &probe, 1), "[1,2] != [1,3]");
+ assert!(builder.equal_to(2, &probe, 2), "null == null");
+ assert!(!builder.equal_to(0, &probe, 2), "[1,2] != null");
+ assert!(builder.equal_to(4, &probe, 3), "[] == []");
+ assert!(!builder.equal_to(0, &probe, 3), "[1,2] != []");
+
+ let out = Box::new(builder).build();
+ let out_list = out.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(out_list.len(), 5);
+ assert!(out_list.is_null(2));
+ // Verify values for non-null rows.
+ let v0 = out_list.value(0);
+ let v0 = v0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v0.values(), &[1, 2]);
+ let v1 = out_list.value(1);
+ let v1 = v1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v1.values(), &[3]);
+ let v3 = out_list.value(3);
+ let v3 = v3.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v3.values(), &[1, 2]);
+ assert_eq!(out_list.value_length(4), 0);
+ }
+
+ #[test]
+ fn large_list_round_trip() {
+ let input = large_list_array(&[
+ Some(vec![Some(7), Some(8), Some(9)]),
+ None,
+ Some(vec![Some(10)]),
+ ]);
+ let mut builder = ListGroupValueBuilder::<i64>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+ let out = Box::new(builder).build();
+ let out_list = out.as_any().downcast_ref::<LargeListArray>().unwrap();
+ assert_eq!(out_list.len(), 3);
+ assert!(out_list.is_null(1));
+ let v0 = out_list.value(0);
+ let v0 = v0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v0.values(), &[7, 8, 9]);
+ let v2 = out_list.value(2);
+ let v2 = v2.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v2.values(), &[10]);
+ }
+
+ #[test]
+ fn list_take_n_splits_offsets_and_child() {
+ let input = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(3)]),
+ Some(vec![Some(4), Some(5), Some(6)]),
+ ]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+ let first = builder.take_n(2);
+ let first_list = first.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(first_list.len(), 2);
+ let v0 = first_list.value(0);
+ let v0 = v0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v0.values(), &[1, 2]);
+ let v1 = first_list.value(1);
+ let v1 = v1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v1.values(), &[3]);
+
+ // Remaining = row 2 only.
+ assert_eq!(builder.len(), 1);
+ let rest = Box::new(builder).build();
+ let rest_list = rest.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(rest_list.len(), 1);
+ let v = rest_list.value(0);
+ let v = v.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v.values(), &[4, 5, 6]);
+ }
+
+ #[test]
+ fn large_list_of_struct_round_trip_through_new_group_values() {
+ // Composition test: LargeList<Struct<id: Utf8, n: Int32>>, the same
+ // shape as a SEC Form 4 `footnotes` column (`LargeList<Struct<id:
+ // Utf8, description: Utf8>>`). Proves the recursive factory wires
+ // List/Struct together correctly and that intern/emit round-trips
+ // through `GroupValuesColumn` rather than `GroupValuesRows`.
+ use crate::aggregates::group_values::new_group_values;
+ use crate::aggregates::order::GroupOrdering;
+ use arrow::array::{
+ Int32Array, LargeListArray, StringArray, StructArray, builder::Int32Builder,
+ builder::LargeListBuilder, builder::StringBuilder, builder::StructBuilder,
+ };
+ use arrow::datatypes::{Fields, Schema};
+ use datafusion_expr::EmitTo;
+
+ let struct_fields = Fields::from(vec![
+ Field::new("id", DataType::Utf8, true),
+ Field::new("n", DataType::Int32, true),
+ ]);
+ let element_field = Arc::new(Field::new(
+ "element",
+ DataType::Struct(struct_fields.clone()),
+ true,
+ ));
+ let column_field = Field::new(
+ "notes",
+ DataType::LargeList(Arc::clone(&element_field)),
+ true,
+ );
+ let schema = Arc::new(Schema::new(vec![column_field]));
+
+ // Build LargeList<Struct<id: Utf8, n: Int32>> arrays.
+ //
+ // NotesRow models one outer LargeList row of Struct<id, n>:
+ // None -> null outer list
+ // Some(vec![]) -> empty list
+ // Some(vec![(id, n), ...]) -> list with struct entries
+ type NotesRow<'a> = Option<Vec<(Option<&'a str>, Option<i32>)>>;
+ let notes_v = |rows: &[NotesRow<'_>]| -> ArrayRef {
+ let struct_builder = StructBuilder::new(
+ struct_fields.clone(),
+ vec![
+ Box::new(StringBuilder::new()),
+ Box::new(Int32Builder::new()),
+ ],
+ );
+ let mut list_builder = LargeListBuilder::new(struct_builder)
+ .with_field(Arc::clone(&element_field));
+ for row in rows {
+ match row {
+ None => list_builder.append(false),
+ Some(items) => {
+ for (id, n) in items {
+ let s = list_builder.values();
+ s.field_builder::<StringBuilder>(0)
+ .unwrap()
+ .append_option(id.map(|x| x.to_string()));
+ s.field_builder::<Int32Builder>(1)
+ .unwrap()
+ .append_option(*n);
+ s.append(true);
+ }
+ list_builder.append(true);
+ }
+ }
+ }
+ Arc::new(list_builder.finish())
+ };
+
+ let mut gv = new_group_values(schema, &GroupOrdering::None).unwrap();
+
+ // Batch 1: a mix of duplicate / distinct / null lists.
+ let batch1 = notes_v(&[
+ Some(vec![(Some("a"), Some(1))]), // 0
+ Some(vec![(Some("a"), Some(1))]), // dup of 0
+ Some(vec![(Some("a"), Some(2))]), // distinct (n differs)
+ None, // null
+ Some(vec![(Some("a"), Some(1)), (Some("b"), Some(2))]), // distinct (length differs)
+ None, // dup of null
+ Some(vec![]), // empty list
+ ]);
+ let mut groups = Vec::new();
+ gv.intern(&[batch1], &mut groups).unwrap();
+ assert_eq!(
+ groups,
+ vec![0, 0, 1, 2, 3, 2, 4],
+ "composition dedup: same struct list -> same group; null and empty are distinct"
+ );
+
+ // Batch 2: cross-batch dedup including a brand-new struct value.
+ let batch2 = notes_v(&[
+ Some(vec![(Some("a"), Some(1))]), // existing 0
+ Some(vec![(Some("a"), Some(1)), (Some("b"), Some(2))]), // existing 3
+ Some(vec![(Some("z"), Some(99))]), // new 5
+ None, // existing 2
+ ]);
+ let mut groups = Vec::new();
+ gv.intern(&[batch2], &mut groups).unwrap();
+ assert_eq!(groups, vec![0, 3, 5, 2]);
+ assert_eq!(gv.len(), 6, "five distinct lists + one null group");
+
+ // Emit and sanity-check the materialized shape.
+ let out = gv.emit(EmitTo::All).unwrap();
+ let ll = out[0].as_any().downcast_ref::<LargeListArray>().unwrap();
+ assert_eq!(ll.len(), 6);
+ assert!(ll.is_null(2), "row 2 was the null-list group");
+ // row 0 == [(a, 1)]
+ let row0 = ll.value(0);
+ let row0 = row0.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(row0.len(), 1);
+ let ids = row0
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let ns = row0
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ids.value(0), "a");
+ assert_eq!(ns.value(0), 1);
+ // row 3 == [(a, 1), (b, 2)]
+ let row3 = ll.value(3);
+ let row3 = row3.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(row3.len(), 2);
+ // row 5 == [(z, 99)] (the new value from batch 2)
+ let row5 = ll.value(5);
+ let row5 = row5.as_any().downcast_ref::<StructArray>().unwrap();
+ let ids = row5
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let ns = row5
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ids.value(0), "z");
+ assert_eq!(ns.value(0), 99);
+ }
+
+ #[test]
+ fn list_dispatcher_round_trip_through_new_group_values() {
+ use crate::aggregates::group_values::new_group_values;
+ use crate::aggregates::order::GroupOrdering;
+ use arrow::datatypes::Schema;
+ use datafusion_expr::EmitTo;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "tags",
+ DataType::List(child_field()),
+ true,
+ )]));
+ let mut gv = new_group_values(schema, &GroupOrdering::None).unwrap();
+
+ // Batch 1.
+ let batch1: ArrayRef = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(3)]),
+ None,
+ ]);
+ let mut groups = Vec::new();
+ gv.intern(&[batch1], &mut groups).unwrap();
+ assert_eq!(
+ groups,
+ vec![0, 1, 0, 2, 1],
+ "list dedup: [1,2] -> 0, null -> 1, [3] -> 2"
+ );
+
+ // Batch 2 hits existing keys + adds one new.
+ let batch2: ArrayRef =
+ list_array(&[Some(vec![Some(3)]), Some(vec![Some(4), Some(5)])]);
+ let mut groups = Vec::new();
+ gv.intern(&[batch2], &mut groups).unwrap();
+ assert_eq!(groups, vec![2, 3]);
+
+ assert_eq!(gv.len(), 4);
+
+ let out = gv.emit(EmitTo::All).unwrap();
+ let out_list = out[0].as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(out_list.len(), 4);
+ assert!(out_list.is_null(1));
+ }
+
+ #[test]
+ fn handles_sliced_input_list_array() {
+ // Build 5 rows then slice off the first 2. append_val/equal_to must
+ // operate on logical positions of the slice, not the underlying array.
+ let full = list_array(&[
+ Some(vec![Some(99), Some(99)]), // sliced off
+ Some(vec![Some(98)]), // sliced off
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3)]),
+ ]);
+ let sliced = full.slice(2, 3);
+
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..sliced.len() {
+ builder.append_val(&sliced, i).unwrap();
+ }
+
+ let probe = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3)]),
+ Some(vec![Some(99), Some(99)]),
+ ]);
+ assert!(builder.equal_to(0, &probe, 0));
+ assert!(builder.equal_to(1, &probe, 1));
+ assert!(builder.equal_to(2, &probe, 2));
+ assert!(
+ !builder.equal_to(0, &probe, 3),
+ "sliced[0]=[1,2] != [99,99]"
+ );
+
+ // Equal_to against the SAME sliced array.
+ for i in 0..sliced.len() {
+ assert!(builder.equal_to(i, &sliced, i));
+ }
+
+ let out = Box::new(builder).build();
+ let out = out.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(out.len(), 3);
+ assert!(out.is_null(1));
+ let v0 = out.value(0);
+ let v0 = v0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v0.values(), &[1, 2]);
+ let v2 = out.value(2);
+ let v2 = v2.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v2.values(), &[3]);
+ }
+
+ #[test]
+ fn handles_sliced_input_large_list_array() {
+ let full = large_list_array(&[
+ Some(vec![Some(99)]),
+ Some(vec![Some(98)]),
+ Some(vec![Some(7), Some(8)]),
+ None,
+ ]);
+ let sliced = full.slice(2, 2);
+
+ let mut builder = ListGroupValueBuilder::<i64>::new(child_field(), make_child());
+ for i in 0..sliced.len() {
+ builder.append_val(&sliced, i).unwrap();
+ }
+
+ let probe =
+ large_list_array(&[Some(vec![Some(7), Some(8)]), None, Some(vec![Some(99)])]);
+ assert!(builder.equal_to(0, &probe, 0));
+ assert!(builder.equal_to(1, &probe, 1));
+ assert!(!builder.equal_to(0, &probe, 2), "sliced[0]=[7,8] != [99]");
+ }
+
+ #[test]
+ fn take_n_zero_and_full() {
+ let input = list_array(&[Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+
+ let none = builder.take_n(0);
+ assert_eq!(none.len(), 0);
+ assert_eq!(builder.len(), 2, "remainder unchanged after take_n(0)");
+
+ let all = builder.take_n(2);
+ let all = all.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(all.len(), 2);
+ assert_eq!(builder.len(), 0, "builder drained after take_n(len)");
+ }
+
+ #[test]
+ fn take_n_with_nulls_and_empty_rows() {
+ let input = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![]),
+ Some(vec![Some(3), Some(4), Some(5)]),
+ ]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+
+ // Take the first three (one normal + null + empty).
+ let three = builder.take_n(3);
+ let three = three.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(three.len(), 3);
+ assert!(!three.is_null(0));
+ assert!(three.is_null(1), "null carried over");
+ assert!(!three.is_null(2));
+ assert_eq!(three.value_length(2), 0, "empty list carried over");
+
+ // Remaining = row 3 only.
+ let rest = Box::new(builder).build();
+ let rest = rest.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(rest.len(), 1);
+ let v = rest.value(0);
+ let v = v.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(v.values(), &[3, 4, 5]);
+ }
+
+ #[test]
+ fn vectorized_methods_match_per_row() {
+ let input = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3)]),
+ Some(vec![Some(1), Some(2)]),
+ ]);
+
+ let mut per_row = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ per_row.append_val(&input, i).unwrap();
+ }
+
+ let mut vec_b = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ vec_b.vectorized_append(&input, &[0, 1, 2, 3]).unwrap();
+ assert_eq!(vec_b.len(), per_row.len());
+
+ let lhs = vec![0usize, 1, 2, 3];
+ let rhs = vec![0usize, 1, 2, 0];
+ let mut bb = BooleanBufferBuilder::new(4);
+ bb.append_n(4, true);
+ vec_b.vectorized_equal_to(&lhs, &input, &rhs, &mut bb);
+ for idx in 0..4 {
+ assert_eq!(
+ bb.get_bit(idx),
+ per_row.equal_to(lhs[idx], &input, rhs[idx]),
+ "row {idx}"
+ );
+ }
+
+ // Pre-set false bit must not flip back to true.
+ let mut bb = BooleanBufferBuilder::new(4);
+ bb.append_n(4, true);
+ bb.set_bit(0, false);
+ vec_b.vectorized_equal_to(&lhs, &input, &rhs, &mut bb);
+ assert!(!bb.get_bit(0), "pre-set false stays false");
+ }
+
+ #[test]
+ fn values_preserving_selects_reorders_and_repeats() {
+ let input = list_array(&[
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![]),
+ Some(vec![Some(3), Some(4), Some(5)]),
+ ]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ for i in 0..input.len() {
+ builder.append_val(&input, i).unwrap();
+ }
+
+ let selection = GroupSelection::try_from_indices(&[3, 1, 0, 3, 2], 4).unwrap();
+ let out = builder.values_preserving(selection).unwrap();
+ let out = out.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(out.len(), 5);
+ assert!(out.is_null(1), "selected null row stays null");
+ let v = |i: usize| -> Vec<i32> {
+ let v = out.value(i);
+ v.as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ };
+ assert_eq!(v(0), vec![3, 4, 5]);
+ assert_eq!(v(2), vec![1, 2]);
+ assert_eq!(v(3), vec![3, 4, 5], "repeated index materialized twice");
+ assert_eq!(out.value_length(4), 0, "empty list preserved");
+
+ // The builder is unchanged by a preserving read.
+ assert_eq!(builder.len(), 4);
+ let all = builder.values_preserving(GroupSelection::all(4)).unwrap();
+ assert_eq!(all.len(), 4);
+ let built = Box::new(builder).build();
+ assert_eq!(built.as_ref(), all.as_ref());
+ }
+
+ #[test]
+ fn values_preserving_empty_selection() {
+ let mut builder = ListGroupValueBuilder::<i64>::new(child_field(), make_child());
+ let input = large_list_array(&[Some(vec![Some(1)])]);
+ builder.append_val(&input, 0).unwrap();
+ let out = builder
+ .values_preserving(GroupSelection::try_from_indices(&[], 1).unwrap())
+ .unwrap();
+ assert_eq!(out.len(), 0);
+ }
+
+ #[test]
+ fn values_preserving_rejects_mismatched_group_count() {
+ let input = list_array(&[Some(vec![Some(1)])]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ builder.append_val(&input, 0).unwrap();
+ assert!(builder.values_preserving(GroupSelection::all(2)).is_err());
+ }
+
+ #[test]
+ fn size_grows_with_appends() {
+ let input = list_array(&[Some(vec![Some(1), Some(2), Some(3)])]);
+ let mut builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ let s0 = builder.size();
+ for _ in 0..16 {
+ builder.append_val(&input, 0).unwrap();
+ }
+ let s1 = builder.size();
+ assert!(s1 > s0, "size should grow ({s0} -> {s1})");
+ }
+
+ #[test]
+ fn build_empty_builder_returns_empty_list() {
+ let builder = ListGroupValueBuilder::<i32>::new(child_field(), make_child());
+ let out = Box::new(builder).build();
+ let out = out.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(out.len(), 0);
+ }
+
+ #[test]
+ fn nested_list_of_list_int32() {
+ // List<List<Int32>>: outer list whose elements are themselves lists of
+ // ints. Exercises a recursive child GroupColumn built via the
+ // dispatcher.
+ use crate::aggregates::group_values::new_group_values;
+ use crate::aggregates::order::GroupOrdering;
+ use arrow::array::{ListArray, builder::Int32Builder, builder::ListBuilder};
+ use arrow::datatypes::Schema;
+ use datafusion_expr::EmitTo;
+
+ let inner_field = Arc::new(Field::new("item", DataType::Int32, true));
+ let outer_field = Arc::new(Field::new(
+ "item",
+ DataType::List(Arc::clone(&inner_field)),
+ true,
+ ));
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "matrix",
+ DataType::List(Arc::clone(&outer_field)),
+ true,
+ )]));
+
+ // MatrixRow models one outer List<List<Int32>> row:
+ // None -> null outer list
+ // Some(vec![...]) -> list of sublists, each sublist itself
+ // Option<Vec<Option<i32>>>
+ type MatrixRow = Option<Vec<Option<Vec<Option<i32>>>>>;
+ let mk = |rows: &[MatrixRow]| -> ArrayRef {
+ let inner = ListBuilder::new(Int32Builder::new())
+ .with_field(Arc::clone(&inner_field));
+ let mut outer = ListBuilder::new(inner).with_field(Arc::clone(&outer_field));
+ for row in rows {
+ match row {
+ None => outer.append(false),
+ Some(sublists) => {
+ for sub in sublists {
+ match sub {
+ None => outer.values().append(false),
+ Some(items) => {
+ for v in items {
+ outer.values().values().append_option(*v);
+ }
+ outer.values().append(true);
+ }
+ }
+ }
+ outer.append(true);
+ }
+ }
+ }
+ Arc::new(outer.finish())
+ };
+
+ let mut gv = new_group_values(schema, &GroupOrdering::None).unwrap();
+
+ // Three groups: [[1,2],[3]], its duplicate, a distinct value, and a null.
+ let batch = mk(&[
+ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]),
+ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]),
+ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(4)])]),
+ None,
+ Some(vec![Some(vec![Some(1), Some(2)]), None]),
+ ]);
+ let mut groups = Vec::new();
+ gv.intern(&[batch], &mut groups).unwrap();
+ assert_eq!(
+ groups,
+ vec![0, 0, 1, 2, 3],
+ "List<List<Int32>> dedup across rows"
+ );
+
+ let out = gv.emit(EmitTo::All).unwrap();
+ let ll = out[0].as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(ll.len(), 4);
+ assert!(ll.is_null(2), "row 2 was the null outer group");
+
+ // Inspect row 0: should be [[1,2],[3]].
+ let row0 = ll.value(0);
+ let row0 = row0.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(row0.len(), 2);
+ let r0_0 = row0.value(0);
+ let r0_0 = r0_0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(r0_0.values(), &[1, 2]);
+ let r0_1 = row0.value(1);
+ let r0_1 = r0_1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(r0_1.values(), &[3]);
+
+ // Row 3 had an inner-null sublist: [[1,2], null].
+ let row3 = ll.value(3);
+ let row3 = row3.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(row3.len(), 2);
+ assert!(!row3.is_null(0));
+ assert!(row3.is_null(1), "inner null sublist preserved");
+ }
+}
diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
index 7ed0a489e2..6c1926b402 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
@@ -22,10 +22,12 @@ mod bytes;
pub mod bytes_view;
mod dictionary;
mod fixed_size_binary;
+mod list;
pub mod primitive;
pub mod row_backed;
use std::mem::{self, size_of};
+use std::sync::Arc;
use crate::aggregates::group_values::GroupValues;
use crate::aggregates::group_values::multi_group_by::{
@@ -914,15 +916,18 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
/// `$nullable`: whether the input can contains nulls
/// `$t`: the primitive type of the builder
macro_rules! instantiate_primitive {
- ($v:expr, $nullable:expr, $t:ty, $data_type:ident) => {
- if $nullable {
- let b = PrimitiveGroupValueBuilder::<$t, true>::new($data_type.to_owned());
- $v.push(Box::new(b) as _)
+ ($nullable:expr, $t:ty, $data_type:ident) => {{
+ let builder: Box<dyn GroupColumn> = if $nullable {
+ Box::new(PrimitiveGroupValueBuilder::<$t, true>::new(
+ $data_type.to_owned(),
+ ))
} else {
- let b = PrimitiveGroupValueBuilder::<$t, false>::new($data_type.to_owned());
- $v.push(Box::new(b) as _)
- }
- };
+ Box::new(PrimitiveGroupValueBuilder::<$t, false>::new(
+ $data_type.to_owned(),
+ ))
+ };
+ Some(builder)
+ }};
}
/// Returns true if the specified data type has a specialized
@@ -934,17 +939,7 @@ macro_rules! instantiate_primitive {
/// builder for. The `group_column_supported_type_matches_make_group_column`
/// test below pins this biconditional.
fn group_column_supported_type(data_type: &DataType) -> bool {
- // Nested types (Struct / List / LargeList / FixedSizeList, recursively) have
- // no type-specialized `GroupColumn`; they are handled by the generic
- // row-backed fallback in `make_group_column` whenever arrow's row format can
- // encode them. Gate the fallback to nested types so intentionally-excluded
- // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the
- // `group_column_supported_type` ⇔ `make_group_column` invariant holds.
- if data_type.is_nested() {
- return RowsGroupColumn::supports_type(data_type);
- }
- matches!(
- *data_type,
+ match data_type {
DataType::Int8
| DataType::Int16
| DataType::Int32
@@ -983,8 +978,19 @@ fn group_column_supported_type(data_type: &DataType) -> bool {
| DataType::Interval(_)
| DataType::Utf8View
| DataType::BinaryView
- | DataType::Boolean
- ) || matches!(data_type, DataType::Dictionary(_,v ) if group_column_supported_type(v))
+ | DataType::Boolean => true,
+ DataType::List(child_field) | DataType::LargeList(child_field)
+ if group_column_supported_type(child_field.data_type()) => true,
+ DataType::Dictionary(_, v) if group_column_supported_type(v) => true,
+ // Nested types (Struct / List / LargeList / FixedSizeList, recursively)
+ // that have no type-specialized `GroupColumn` as per the above are handled by the generic
+ // row-backed fallback in `make_group_column` whenever arrow's row format can
+ // encode them. Gate the fallback to nested types so intentionally-excluded
+ // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the
+ // `group_column_supported_type` ⇔ `make_group_column` invariant holds.
+ dt if dt.is_nested() => RowsGroupColumn::supports_type(dt),
+ _ => false,
+ }
}
/// Build a [`GroupColumn`] for a single schema field.
@@ -1005,136 +1011,136 @@ fn group_column_supported_type(data_type: &DataType) -> bool {
fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
let nullable = field.is_nullable();
let data_type = field.data_type();
- let mut v: Vec<Box<dyn GroupColumn>> = Vec::with_capacity(1);
- match *data_type {
- DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type),
- DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type),
- DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type),
- DataType::Int64 => instantiate_primitive!(v, nullable, Int64Type, data_type),
- DataType::UInt8 => instantiate_primitive!(v, nullable, UInt8Type, data_type),
- DataType::UInt16 => instantiate_primitive!(v, nullable, UInt16Type, data_type),
- DataType::UInt32 => instantiate_primitive!(v, nullable, UInt32Type, data_type),
- DataType::UInt64 => instantiate_primitive!(v, nullable, UInt64Type, data_type),
- DataType::Float16 => {
- instantiate_primitive!(v, nullable, Float16Type, data_type)
- }
- DataType::Float32 => {
- instantiate_primitive!(v, nullable, Float32Type, data_type)
- }
- DataType::Float64 => {
- instantiate_primitive!(v, nullable, Float64Type, data_type)
- }
- DataType::Date32 => instantiate_primitive!(v, nullable, Date32Type, data_type),
- DataType::Date64 => instantiate_primitive!(v, nullable, Date64Type, data_type),
+ let builder: Option<Box<dyn GroupColumn>> = match data_type {
+ DataType::Int8 => instantiate_primitive!(nullable, Int8Type, data_type),
+ DataType::Int16 => instantiate_primitive!(nullable, Int16Type, data_type),
+ DataType::Int32 => instantiate_primitive!(nullable, Int32Type, data_type),
+ DataType::Int64 => instantiate_primitive!(nullable, Int64Type, data_type),
+ DataType::UInt8 => instantiate_primitive!(nullable, UInt8Type, data_type),
+ DataType::UInt16 => instantiate_primitive!(nullable, UInt16Type, data_type),
+ DataType::UInt32 => instantiate_primitive!(nullable, UInt32Type, data_type),
+ DataType::UInt64 => instantiate_primitive!(nullable, UInt64Type, data_type),
+ DataType::Float16 => instantiate_primitive!(nullable, Float16Type, data_type),
+ DataType::Float32 => instantiate_primitive!(nullable, Float32Type, data_type),
+ DataType::Float64 => instantiate_primitive!(nullable, Float64Type, data_type),
+ DataType::Date32 => instantiate_primitive!(nullable, Date32Type, data_type),
+ DataType::Date64 => instantiate_primitive!(nullable, Date64Type, data_type),
DataType::Time32(t) => match t {
TimeUnit::Second => {
- instantiate_primitive!(v, nullable, Time32SecondType, data_type)
+ instantiate_primitive!(nullable, Time32SecondType, data_type)
}
TimeUnit::Millisecond => {
- instantiate_primitive!(v, nullable, Time32MillisecondType, data_type)
+ instantiate_primitive!(nullable, Time32MillisecondType, data_type)
}
// Time32 with Microsecond / Nanosecond is not a valid Arrow type
// combination; reject explicitly so group_column_supported_type
// and this dispatcher stay in lockstep (see consistency fuzz below).
- _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"),
+ _ => None,
},
DataType::Time64(t) => match t {
TimeUnit::Microsecond => {
- instantiate_primitive!(v, nullable, Time64MicrosecondType, data_type)
+ instantiate_primitive!(nullable, Time64MicrosecondType, data_type)
}
TimeUnit::Nanosecond => {
- instantiate_primitive!(v, nullable, Time64NanosecondType, data_type)
+ instantiate_primitive!(nullable, Time64NanosecondType, data_type)
}
// Time64 with Second / Millisecond is not a valid Arrow type
// combination; reject explicitly.
- _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"),
+ _ => None,
},
DataType::Timestamp(t, _) => match t {
TimeUnit::Second => {
- instantiate_primitive!(v, nullable, TimestampSecondType, data_type)
+ instantiate_primitive!(nullable, TimestampSecondType, data_type)
}
TimeUnit::Millisecond => {
- instantiate_primitive!(v, nullable, TimestampMillisecondType, data_type)
+ instantiate_primitive!(nullable, TimestampMillisecondType, data_type)
}
TimeUnit::Microsecond => {
- instantiate_primitive!(v, nullable, TimestampMicrosecondType, data_type)
+ instantiate_primitive!(nullable, TimestampMicrosecondType, data_type)
}
TimeUnit::Nanosecond => {
- instantiate_primitive!(v, nullable, TimestampNanosecondType, data_type)
+ instantiate_primitive!(nullable, TimestampNanosecondType, data_type)
}
},
DataType::Duration(t) => match t {
TimeUnit::Second => {
- instantiate_primitive!(v, nullable, DurationSecondType, data_type)
+ instantiate_primitive!(nullable, DurationSecondType, data_type)
}
TimeUnit::Millisecond => {
- instantiate_primitive!(v, nullable, DurationMillisecondType, data_type)
+ instantiate_primitive!(nullable, DurationMillisecondType, data_type)
}
TimeUnit::Microsecond => {
- instantiate_primitive!(v, nullable, DurationMicrosecondType, data_type)
+ instantiate_primitive!(nullable, DurationMicrosecondType, data_type)
}
TimeUnit::Nanosecond => {
- instantiate_primitive!(v, nullable, DurationNanosecondType, data_type)
+ instantiate_primitive!(nullable, DurationNanosecondType, data_type)
}
},
// `IntervalUnit` has exactly three variants, so this match is exhaustive
// with no fallback arm (unlike Time32 / Time64).
DataType::Interval(u) => match u {
IntervalUnit::YearMonth => {
- instantiate_primitive!(v, nullable, IntervalYearMonthType, data_type)
+ instantiate_primitive!(nullable, IntervalYearMonthType, data_type)
}
IntervalUnit::DayTime => {
- instantiate_primitive!(v, nullable, IntervalDayTimeType, data_type)
+ instantiate_primitive!(nullable, IntervalDayTimeType, data_type)
}
IntervalUnit::MonthDayNano => {
- instantiate_primitive!(v, nullable, IntervalMonthDayNanoType, data_type)
+ instantiate_primitive!(nullable, IntervalMonthDayNanoType, data_type)
}
},
DataType::Decimal128(_, _) => {
- instantiate_primitive!(v, nullable, Decimal128Type, data_type)
+ instantiate_primitive!(nullable, Decimal128Type, data_type)
}
DataType::Decimal256(_, _) => {
- instantiate_primitive!(v, nullable, Decimal256Type, data_type)
- }
- DataType::Utf8 => {
- v.push(Box::new(ByteGroupValueBuilder::<i32>::new(
- OutputType::Utf8,
- )));
- }
- DataType::LargeUtf8 => {
- v.push(Box::new(ByteGroupValueBuilder::<i64>::new(
- OutputType::Utf8,
- )));
- }
- DataType::Binary => {
- v.push(Box::new(ByteGroupValueBuilder::<i32>::new(
- OutputType::Binary,
- )));
- }
- DataType::LargeBinary => {
- v.push(Box::new(ByteGroupValueBuilder::<i64>::new(
- OutputType::Binary,
- )));
+ instantiate_primitive!(nullable, Decimal256Type, data_type)
}
+ DataType::Utf8 => Some(Box::new(ByteGroupValueBuilder::<i32>::new(
+ OutputType::Utf8,
+ ))),
+ DataType::LargeUtf8 => Some(Box::new(ByteGroupValueBuilder::<i64>::new(
+ OutputType::Utf8,
+ ))),
+ DataType::Binary => Some(Box::new(ByteGroupValueBuilder::<i32>::new(
+ OutputType::Binary,
+ ))),
+ DataType::LargeBinary => Some(Box::new(ByteGroupValueBuilder::<i64>::new(
+ OutputType::Binary,
+ ))),
// A negative width is not a valid Arrow type; it falls to the `_`
// arm below, matching `group_column_supported_type`.
DataType::FixedSizeBinary(byte_width @ 0..) => {
- v.push(Box::new(FixedSizeBinaryGroupValueBuilder::new(byte_width)));
+ Some(Box::new(FixedSizeBinaryGroupValueBuilder::new(*byte_width)))
}
DataType::Utf8View => {
- v.push(Box::new(ByteViewGroupValueBuilder::<StringViewType>::new()));
+ Some(Box::new(ByteViewGroupValueBuilder::<StringViewType>::new()))
}
DataType::BinaryView => {
- v.push(Box::new(ByteViewGroupValueBuilder::<BinaryViewType>::new()));
+ Some(Box::new(ByteViewGroupValueBuilder::<BinaryViewType>::new()))
}
DataType::Boolean => {
if nullable {
- v.push(Box::new(BooleanGroupValueBuilder::<true>::new()));
+ Some(Box::new(BooleanGroupValueBuilder::<true>::new()))
} else {
- v.push(Box::new(BooleanGroupValueBuilder::<false>::new()));
+ Some(Box::new(BooleanGroupValueBuilder::<false>::new()))
}
}
- DataType::Dictionary(ref key_dt, ref value_dt) => {
+ DataType::List(child_field) => match make_group_column(child_field.as_ref()) {
+ Ok(child) => Some(Box::new(list::ListGroupValueBuilder::<i32>::new(
+ Arc::clone(child_field),
+ child,
+ ))),
+ Err(_) => None,
+ },
+ DataType::LargeList(child_field) => match make_group_column(child_field.as_ref())
+ {
+ Ok(child) => Some(Box::new(list::ListGroupValueBuilder::<i64>::new(
+ Arc::clone(child_field),
+ child,
+ ))),
+ Err(_) => None,
+ },
+ DataType::Dictionary(key_dt, value_dt) => {
let new_field = Field::new("", *value_dt.clone(), true);
let inner = make_group_column(&new_field)?;
macro_rules! dict_col {
@@ -1159,24 +1165,22 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
);
}
};
- v.push(col)
+ Some(col)
}
+ _ => None,
+ };
+ match builder {
+ Some(b) => Ok(b),
// Generic fallback for nested types (Struct / List / LargeList /
// FixedSizeList, recursively) that lack a type-specialized builder but
// can be encoded by arrow's row format. This is what lets a mixed
// schema keep the column-wise fast path for its native columns instead
// of dropping the whole key onto `GroupValuesRows`.
- ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => {
- v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?));
+ None if data_type.is_nested() && RowsGroupColumn::supports_type(data_type) => {
+ Ok(Box::new(RowsGroupColumn::try_new(data_type.clone())?))
}
- _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"),
+ None => not_impl_err!("{data_type} not supported in GroupValuesColumn"),
}
- debug_assert_eq!(
- v.len(),
- 1,
- "make_group_column must push exactly one builder"
- );
- Ok(v.into_iter().next().unwrap())
}
impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> {
@@ -1656,6 +1660,24 @@ mod tests {
/// it should be added to unsupported_cases.
#[test]
fn group_column_supported_type_matches_make_group_column() {
+ let utf8 = || Field::new("v", DataType::Utf8, true);
+ let int32 = || Field::new("v", DataType::Int32, true);
+ let f16 = || Field::new("v", DataType::Float16, true);
+ // A leaf that neither has a specialized `GroupColumn` nor is accepted
+ // by the row-backed fallback (`RowsGroupColumn::supports_type` rejects
+ // `RunEndEncoded` anywhere in the subtree), so nesting it keeps the
+ // whole type unsupported.
+ let ree = || {
+ Field::new(
+ "v",
+ DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", DataType::Int64, true)),
+ ),
+ true,
+ )
+ };
+
let supported_cases: Vec<DataType> = vec![
DataType::Int8,
DataType::Int64,
@@ -1689,6 +1711,17 @@ mod tests {
DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth),
DataType::Interval(arrow::datatypes::IntervalUnit::DayTime),
DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano),
+ // Nested
+ DataType::List(Arc::new(int32())),
+ DataType::LargeList(Arc::new(int32())),
+ DataType::List(Arc::new(f16())),
+ DataType::LargeList(Arc::new(f16())),
+ DataType::List(Arc::new(utf8())),
+ DataType::List(Arc::new(Field::new(
+ "v",
+ DataType::List(Arc::new(int32())),
+ true,
+ ))),
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64)),
DataType::Dictionary(
@@ -1710,6 +1743,8 @@ mod tests {
group_column_supported_type(dt),
"expected group_column_supported_type=true for {dt:?}"
);
+ // Building a top-level Field and feeding it through the factory
+ // must succeed for every supported case.
let field = Field::new("col", dt.clone(), true);
make_group_column(&field).unwrap_or_else(|e| {
panic!(
@@ -1729,6 +1764,15 @@ mod tests {
DataType::Time64(arrow::datatypes::TimeUnit::Millisecond),
DataType::Time32(arrow::datatypes::TimeUnit::Microsecond),
DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond),
+ // Nested with an unsupported leaf
+ DataType::List(Arc::new(ree())),
+ DataType::LargeList(Arc::new(ree())),
+ // Deeply nested unsupported
+ DataType::List(Arc::new(Field::new(
+ "v",
+ DataType::List(Arc::new(ree())),
+ true,
+ ))),
// A negative width is representable in the DataType but is not
// a valid Arrow type; no array can be constructed for it.
DataType::FixedSizeBinary(-5),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org