From ecdd425ce58d6e3de5bcb3582e446c22de15ca36 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Tue, 23 Jun 2026 23:31:50 +0800 Subject: [PATCH 1/3] feat(physical-plan): add GroupColumn support for FixedSizeList in multi-column GROUP BY MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 2 of the EPIC #22715 sequence. Builds on the dispatcher refactor + factory framework that landed in PR 1 (#22751) to bring the first nested type into `GroupValuesColumn`, unblocking the rest of the sequence (PR 3 Struct, PR 4 List, PR 5 LargeList, PR 6 composite FSL, PR 7 Map). ## Which issue does this PR close? Part of #22682 (nested type coverage) and #22715 (EPIC). ## Rationale for this change Today a single `FixedSizeList` column in a GROUP BY drags the whole grouping onto the byte-encoded `GroupValuesRows` fallback, even when every other column would have qualified for the column-wise + cross-column short-circuit fast path in `GroupValuesColumn`. With the recursive `make_group_column` factory + `group_column_supported_type` allow-list in place from PR 1, this PR is now a self-contained addition of one builder and two dispatcher entries. ## What changes are included in this PR? - `FixedSizeListGroupValueBuilder` in a new `fixed_size_list` submodule. Storage: outer null bitmap + a child `PrimitiveGroupValueBuilder` that holds every element flat (length = `outer_len * list_len`). Element `j` of outer row `i` lives at child index `i * list_len + j`. - `group_column_supported_type` accepts `FixedSizeList` for the primitive subset wired through the dispatcher (Int8..Int64, UInt8..UInt64, Float32, Float64, Date32, Date64). Composite children (Struct / List inside FSL) are deferred to a follow-up after their respective builders land. - `make_group_column` dispatches `DataType::FixedSizeList(...)` to the appropriate `FixedSizeListGroupValueBuilder` via an `instantiate_fsl!` macro mirroring the existing `instantiate_primitive!` pattern. ## Are these changes tested? Yes. 12 new unit tests on `FixedSizeListGroupValueBuilder` plus extensions to the existing `group_column_supported_type` ⇔ `make_group_column` consistency fuzz. Per-builder: - append / build round trip with mixed outer-null and inner-null rows - `equal_to` for identical / different / outer-null / inner-null rows - `take_n` boundary cases (`n=0`, `n=len`, prefix containing null rows) - sliced input array (offset != 0) - `vectorized_equal_to` / `vectorized_append` match per-row reference - `size()` grows with appends - `build` on empty builder returns empty array - end-to-end dispatcher → `GroupValuesColumn` routing Consistency fuzz now covers four primitive FSL children (signed, unsigned, float, date) on the supported side and three non-primitive FSL children (Utf8, Decimal128, Boolean) on the rejected side, locking in the scope boundary for this PR. 127/127 aggregates tests pass, clippy clean, fmt clean. ## What follows in the EPIC - PR 3: `Struct<...>` builder + dispatcher. - PR 4: `List` builder + dispatcher (recursive child via factory). - PR 5: `LargeList` builder + dispatcher. - PR 6: Relax FSL child restriction to allow `FSL` / `FSL` once the prerequisite child builders are in. - PR 7: `Map` (`List>` Arrow representation). --- .../multi_group_by/fixed_size_list.rs | 640 ++++++++++++++++++ .../group_values/multi_group_by/mod.rs | 90 +++ 2 files changed, 730 insertions(+) create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs new file mode 100644 index 0000000000000..514bd6f596063 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs @@ -0,0 +1,640 @@ +// 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 `FixedSizeList`. +//! +//! Storage: outer null bitmap + a child [`PrimitiveGroupValueBuilder`] that +//! holds every element flat (length = outer_len * list_len). The j-th element +//! of the i-th outer row lives at child index `i * list_len + j`. + +use crate::aggregates::group_values::HashValue; +use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; +use crate::aggregates::group_values::multi_group_by::{GroupColumn, nulls_equal_to}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; + +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, BooleanBufferBuilder, FixedSizeListArray, +}; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion_common::Result; +use std::sync::Arc; + +/// A [`GroupColumn`] for `FixedSizeList` where `T` is a primitive type. +pub struct FixedSizeListGroupValueBuilder { + /// Child field, cached for `build` / `take_n`. + field: FieldRef, + /// List length per outer row. + list_len: i32, + /// Outer-level null bitmap. + outer_nulls: MaybeNullBufferBuilder, + /// Number of outer rows accumulated. + outer_len: usize, + /// Flat storage for child elements; always treated as nullable so it can + /// hold child nulls regardless of the outer row's nullability. + child: PrimitiveGroupValueBuilder, +} + +impl FixedSizeListGroupValueBuilder +where + T: ArrowPrimitiveType, + T::Native: HashValue, +{ + pub fn new(data_type: &DataType) -> Self { + let (field, list_len) = match data_type { + DataType::FixedSizeList(f, n) => (Arc::clone(f), *n), + other => unreachable!( + "FixedSizeListGroupValueBuilder built with non-FixedSizeList type {other:?}" + ), + }; + let child = PrimitiveGroupValueBuilder::::new(field.data_type().clone()); + Self { + field, + list_len, + outer_nulls: MaybeNullBufferBuilder::new(), + outer_len: 0, + child, + } + } + + fn list_len_usize(&self) -> usize { + self.list_len as usize + } +} + +impl GroupColumn for FixedSizeListGroupValueBuilder +where + T: ArrowPrimitiveType, + T::Native: HashValue, +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, 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 list_array = array + .as_any() + .downcast_ref::() + .expect("FixedSizeListGroupValueBuilder called with non-FixedSizeList array"); + let rhs_sublist: ArrayRef = list_array.value(rhs_row); + let list_len = self.list_len_usize(); + let lhs_base = lhs_row * list_len; + for j in 0..list_len { + if !self.child.equal_to(lhs_base + j, &rhs_sublist, j) { + return false; + } + } + true + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let list_array = array + .as_any() + .downcast_ref::() + .expect("FixedSizeListGroupValueBuilder called with non-FixedSizeList array"); + self.outer_nulls.append(list_array.is_null(row)); + self.outer_len += 1; + let child_array = list_array.values(); + let list_len = self.list_len_usize(); + // Use the array's own `value_offset` rather than computing `(offset + // + row) * list_len` ourselves. For sliced FixedSizeListArrays the + // `values()` slice is already advanced, so doing the arithmetic + // manually risks double-applying any future offset behavior. + let start = list_array.value_offset(row) as usize; + for j in 0..list_len { + self.child.append_val(child_array, start + j)?; + } + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + 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(lhs_row, array, rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + for &row in rows { + self.append_val(array, row)?; + } + Ok(()) + } + + fn len(&self) -> usize { + self.outer_len + } + + fn size(&self) -> usize { + self.outer_nulls.allocated_size() + self.child.size() + } + + fn build(self: Box) -> ArrayRef { + let Self { + field, + list_len, + mut outer_nulls, + outer_len: _, + child, + } = *self; + let outer_nulls = outer_nulls_take_build(&mut outer_nulls); + let child_array = Box::new(child).build(); + Arc::new(FixedSizeListArray::new( + field, + list_len, + child_array, + outer_nulls, + )) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let list_len = self.list_len_usize(); + let first_n_outer_nulls = self.outer_nulls.take_n(n); + // Drain a `PrimitiveGroupValueBuilder` cleanly through its + // `GroupColumn` impl; this also shifts the remaining elements down. + let first_n_child: ArrayRef = + as GroupColumn>::take_n( + &mut self.child, + n * list_len, + ); + self.outer_len -= n; + Arc::new(FixedSizeListArray::new( + Arc::clone(&self.field), + self.list_len, + first_n_child, + first_n_outer_nulls, + )) + } +} + +/// Helper: consume the inner builder once on `build` to release its buffer. +fn outer_nulls_take_build( + nulls: &mut MaybeNullBufferBuilder, +) -> Option { + std::mem::replace(nulls, MaybeNullBufferBuilder::new()).build() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, FixedSizeListArray, Int32Array, builder::FixedSizeListBuilder, + builder::Int32Builder, + }; + use arrow::datatypes::Int32Type; + use std::sync::Arc; + + /// Build a FixedSizeList array from a sequence of optional rows. + /// `None` outer rows are null lists. Within a non-null row, inner values come + /// from the given `[Option; 2]`. + fn fsl_array(rows: &[Option<[Option; 2]>]) -> ArrayRef { + let inner = Int32Builder::new(); + let mut builder = FixedSizeListBuilder::new(inner, 2); + for row in rows { + match row { + None => { + // Null outer row: push placeholder child values so the + // child array has the right length, then mark the outer + // row null via `append(false)`. + builder.values().append_value(0); + builder.values().append_value(0); + builder.append(false); + } + Some([a, b]) => { + match a { + Some(v) => builder.values().append_value(*v), + None => builder.values().append_null(), + } + match b { + Some(v) => builder.values().append_value(*v), + None => builder.values().append_null(), + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn data_type(child_nullable: bool) -> DataType { + DataType::FixedSizeList( + Arc::new(arrow::datatypes::Field::new( + "item", + DataType::Int32, + child_nullable, + )), + 2, + ) + } + + #[test] + fn append_and_build_round_trips_values() { + let input = fsl_array(&[ + Some([Some(1), Some(2)]), + None, + Some([Some(3), Some(4)]), + Some([Some(1), Some(2)]), + ]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + + for i in 0..input.len() { + builder.append_val(&input, i).unwrap(); + } + assert_eq!(builder.len(), 4); + + let out = Box::new(builder).build(); + let out_fsl = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out_fsl.len(), 4); + assert!(out_fsl.is_null(1)); + assert!(!out_fsl.is_null(0)); + assert!(!out_fsl.is_null(2)); + assert!(!out_fsl.is_null(3)); + + let values = out_fsl + .values() + .as_any() + .downcast_ref::() + .unwrap(); + // Row 0: [1, 2] + assert_eq!(values.value(0), 1); + assert_eq!(values.value(1), 2); + // Row 2: [3, 4] (row 1 is null, but child slot positions are 2..4) + assert_eq!(values.value(4), 3); + assert_eq!(values.value(5), 4); + // Row 3: [1, 2] + assert_eq!(values.value(6), 1); + assert_eq!(values.value(7), 2); + } + + #[test] + fn equal_to_matches_identical_rows_and_rejects_different() { + let stored = fsl_array(&[Some([Some(1), Some(2)]), Some([Some(3), Some(4)])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + builder.append_val(&stored, 0).unwrap(); + builder.append_val(&stored, 1).unwrap(); + + let probe = fsl_array(&[ + Some([Some(1), Some(2)]), // == stored row 0 + Some([Some(3), Some(5)]), // != stored row 1 (last element differs) + Some([Some(3), Some(4)]), // == stored row 1 + ]); + + assert!(builder.equal_to(0, &probe, 0), "[1,2] == [1,2]"); + assert!( + !builder.equal_to(1, &probe, 1), + "[3,4] != [3,5] (one element differs)" + ); + assert!(builder.equal_to(1, &probe, 2), "[3,4] == [3,4]"); + } + + #[test] + fn equal_to_handles_outer_nulls() { + let stored = fsl_array(&[None, Some([Some(1), Some(2)])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + builder.append_val(&stored, 0).unwrap(); + builder.append_val(&stored, 1).unwrap(); + + let probe = + fsl_array(&[None, Some([Some(1), Some(2)]), Some([Some(0), Some(0)])]); + + assert!(builder.equal_to(0, &probe, 0), "null == null"); + assert!(!builder.equal_to(0, &probe, 1), "null != [1,2]"); + assert!(!builder.equal_to(1, &probe, 0), "[1,2] != null"); + assert!(builder.equal_to(1, &probe, 1), "[1,2] == [1,2]"); + assert!(!builder.equal_to(1, &probe, 2), "[1,2] != [0,0]"); + } + + #[test] + fn equal_to_handles_inner_nulls() { + let stored = fsl_array(&[Some([Some(1), None]), Some([None, None])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + builder.append_val(&stored, 0).unwrap(); + builder.append_val(&stored, 1).unwrap(); + + let probe = fsl_array(&[ + Some([Some(1), None]), + Some([Some(1), Some(2)]), + Some([None, None]), + ]); + + assert!(builder.equal_to(0, &probe, 0), "[1, null] == [1, null]"); + assert!( + !builder.equal_to(0, &probe, 1), + "[1, null] != [1, 2] (null vs value)" + ); + assert!( + builder.equal_to(1, &probe, 2), + "[null, null] == [null, null]" + ); + } + + #[test] + fn dispatcher_routes_fixed_size_list_to_group_values_column() { + // End-to-end: feed a schema with a FixedSizeList column + // through `new_group_values` and prove that `GroupValuesColumn` (not + // `GroupValuesRows`) handles the intern. The behavioral signal is that + // dedup works correctly across batches. + use crate::aggregates::group_values::new_group_values; + use crate::aggregates::order::GroupOrdering; + use arrow::datatypes::{Field, Schema}; + use datafusion_expr::EmitTo; + + let schema = + Arc::new(Schema::new(vec![Field::new("tags", data_type(true), true)])); + + let mut gv = new_group_values(schema, &GroupOrdering::None).unwrap(); + + // Batch 1: three rows, two distinct values. + let batch1: ArrayRef = fsl_array(&[ + Some([Some(1), Some(2)]), + Some([Some(3), Some(4)]), + Some([Some(1), Some(2)]), + ]); + let mut groups = Vec::new(); + gv.intern(&[batch1], &mut groups).unwrap(); + assert_eq!( + groups, + vec![0, 1, 0], + "first batch: dedup [1,2] across rows" + ); + + // Batch 2: revisit existing groups + introduce a new one. + let batch2: ArrayRef = fsl_array(&[ + Some([Some(3), Some(4)]), // existing group 1 + Some([Some(5), Some(6)]), // new group 2 + Some([Some(1), Some(2)]), // existing group 0 + ]); + let mut groups = Vec::new(); + gv.intern(&[batch2], &mut groups).unwrap(); + assert_eq!( + groups, + vec![1, 2, 0], + "second batch: dedup hits across batches" + ); + + assert_eq!(gv.len(), 3, "exactly three distinct group keys retained"); + + // Emit and verify the materialized group keys. + let arrays = gv.emit(EmitTo::All).unwrap(); + assert_eq!(arrays.len(), 1); + let out = arrays[0] + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(out.len(), 3); + let values = out.values().as_any().downcast_ref::().unwrap(); + assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6]); + } + + #[test] + fn take_n_returns_prefix_and_shifts_remainder() { + let input = fsl_array(&[ + Some([Some(1), Some(2)]), + Some([Some(3), Some(4)]), + Some([Some(5), Some(6)]), + ]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + for i in 0..input.len() { + builder.append_val(&input, i).unwrap(); + } + assert_eq!(builder.len(), 3); + + let first_two = builder.take_n(2); + let first_two_fsl = first_two + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(first_two_fsl.len(), 2); + let v = first_two_fsl + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(v.values(), &[1, 2, 3, 4]); + + // Remaining builder should now contain only [5, 6]. + assert_eq!(builder.len(), 1); + let rest = Box::new(builder).build(); + let rest_fsl = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest_fsl.len(), 1); + let v = rest_fsl + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(v.values(), &[5, 6]); + } + + #[test] + fn handles_sliced_input_array() { + // Build a 5-row array, slice off the first 2 rows, and verify + // append_val + equal_to operate on the correct logical positions. + let full = fsl_array(&[ + Some([Some(99), Some(99)]), // sliced off + Some([Some(98), Some(98)]), // sliced off + Some([Some(1), Some(2)]), + None, + Some([Some(3), Some(4)]), + ]); + let sliced = full.slice(2, 3); + + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + for i in 0..sliced.len() { + builder.append_val(&sliced, i).unwrap(); + } + + // Probe with an UNSLICED equivalent. + let probe = fsl_array(&[ + Some([Some(1), Some(2)]), + None, + Some([Some(3), Some(4)]), + Some([Some(99), Some(99)]), + ]); + assert!(builder.equal_to(0, &probe, 0), "sliced row 0 == [1,2]"); + assert!(builder.equal_to(1, &probe, 1), "sliced row 1 == null"); + assert!(builder.equal_to(2, &probe, 2), "sliced row 2 == [3,4]"); + assert!( + !builder.equal_to(0, &probe, 3), + "sliced row 0 ([1,2]) != [99,99]" + ); + + // Probe with the SAME sliced array to make sure equal_to with sliced + // rhs also works. + for i in 0..sliced.len() { + assert!( + builder.equal_to(i, &sliced, i), + "row {i} of sliced array must equal itself" + ); + } + + let out = Box::new(builder).build(); + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 3); + assert!(out.is_null(1)); + let v = out.values().as_any().downcast_ref::().unwrap(); + // Output stores: [1,2], placeholder for null, [3,4]. Placeholder is + // whatever the source had — for our null builder we pushed actual + // values 0,0 from the test fixture, so the child layout is [1,2,0,0,3,4]. + assert_eq!(v.value(0), 1); + assert_eq!(v.value(1), 2); + assert_eq!(v.value(4), 3); + assert_eq!(v.value(5), 4); + } + + #[test] + fn take_n_zero_and_full() { + let input = fsl_array(&[Some([Some(1), Some(2)]), Some([Some(3), Some(4)])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + for i in 0..input.len() { + builder.append_val(&input, i).unwrap(); + } + + // take_n(0): empty prefix, remainder unchanged. + let none = builder.take_n(0); + let none_fsl = none.as_any().downcast_ref::().unwrap(); + assert_eq!(none_fsl.len(), 0); + assert_eq!(builder.len(), 2); + + // take_n(len): full prefix, empty remainder. + let all = builder.take_n(2); + let all_fsl = all.as_any().downcast_ref::().unwrap(); + assert_eq!(all_fsl.len(), 2); + let v = all_fsl + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(v.values(), &[1, 2, 3, 4]); + assert_eq!(builder.len(), 0, "builder drained after take_n(len)"); + } + + #[test] + fn take_n_with_null_outer_in_prefix() { + let input = + fsl_array(&[Some([Some(1), Some(2)]), None, Some([Some(3), Some(4)])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + for i in 0..input.len() { + builder.append_val(&input, i).unwrap(); + } + + let first_two = builder.take_n(2); + let first_two = first_two + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(first_two.len(), 2); + assert!(!first_two.is_null(0)); + assert!(first_two.is_null(1), "null row preserved in take_n prefix"); + + // Remaining = row 2 only. + let rest = Box::new(builder).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 1); + assert!(!rest.is_null(0)); + let v = rest.values().as_any().downcast_ref::().unwrap(); + assert_eq!(v.values(), &[3, 4]); + } + + #[test] + fn vectorized_methods_match_per_row() { + // Build two builders, one via per-row append/equal, one via vectorized + // append/equal, and prove they produce the same output and decisions. + let input = fsl_array(&[ + Some([Some(1), Some(2)]), + None, + Some([Some(3), Some(4)]), + Some([Some(1), Some(2)]), + ]); + + let mut per_row: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + for i in 0..input.len() { + per_row.append_val(&input, i).unwrap(); + } + + let mut vec_b: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + vec_b.vectorized_append(&input, &[0, 1, 2, 3]).unwrap(); + assert_eq!(vec_b.len(), per_row.len()); + + // Vectorized equal_to should produce the same per-row decisions. + 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 { + let expected = per_row.equal_to(lhs[idx], &input, rhs[idx]); + assert_eq!( + bb.get_bit(idx), + expected, + "row {idx}: vectorized={} per-row={expected}", + bb.get_bit(idx), + ); + } + + // A pre-set false entry must NOT be flipped 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 must stay false"); + } + + #[test] + fn size_grows_with_appends() { + let input = fsl_array(&[Some([Some(1), Some(2)])]); + let mut builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + let s0 = builder.size(); + for _ in 0..16 { + builder.append_val(&input, 0).unwrap(); + } + let s1 = builder.size(); + assert!(s1 > s0, "size should grow after appends ({s0} -> {s1})"); + } + + #[test] + fn build_empty_builder_returns_empty_array() { + let builder: FixedSizeListGroupValueBuilder = + FixedSizeListGroupValueBuilder::new(&data_type(true)); + let out = Box::new(builder).build(); + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 0); + } +} 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 f275d777c3279..17c62c4c34f1e 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 @@ -20,6 +20,7 @@ mod boolean; mod bytes; pub mod bytes_view; +mod fixed_size_list; pub mod primitive; use std::mem::{self, size_of}; @@ -923,6 +924,25 @@ 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 { + // FixedSizeList is the only nested type currently supported. + // List / LargeList / Struct will follow in subsequent PRs of #22715. + if let DataType::FixedSizeList(child_field, _) = data_type { + return matches!( + child_field.data_type(), + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float32 + | DataType::Float64 + | DataType::Date32 + | DataType::Date64 + ); + } matches!( *data_type, DataType::Int8 @@ -1067,6 +1087,40 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } + DataType::FixedSizeList(ref child_field, _) => { + // `group_column_supported_type` already restricts the child to + // the primitive subset supported here. Any unsupported child + // type returned `false` upstream and was routed to the + // `GroupValuesRows` fallback, so the wildcard arm below is + // only reachable from the consistency-fuzz test. + macro_rules! instantiate_fsl { + ($t:ty) => {{ + let b = fixed_size_list::FixedSizeListGroupValueBuilder::<$t>::new( + data_type, + ); + v.push(Box::new(b) as _); + }}; + } + match child_field.data_type() { + DataType::Int8 => instantiate_fsl!(Int8Type), + DataType::Int16 => instantiate_fsl!(Int16Type), + DataType::Int32 => instantiate_fsl!(Int32Type), + DataType::Int64 => instantiate_fsl!(Int64Type), + DataType::UInt8 => instantiate_fsl!(UInt8Type), + DataType::UInt16 => instantiate_fsl!(UInt16Type), + DataType::UInt32 => instantiate_fsl!(UInt32Type), + DataType::UInt64 => instantiate_fsl!(UInt64Type), + DataType::Float32 => instantiate_fsl!(Float32Type), + DataType::Float64 => instantiate_fsl!(Float64Type), + DataType::Date32 => instantiate_fsl!(Date32Type), + DataType::Date64 => instantiate_fsl!(Date64Type), + other => { + return not_impl_err!( + "FixedSizeList<{other}> not supported in GroupValuesColumn" + ); + } + } + } _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), } debug_assert_eq!( @@ -1309,6 +1363,27 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + // FixedSizeList: cover one signed integer, one + // unsigned, one float, and a date variant. The dispatcher in + // `make_group_column` enumerates the full primitive subset + // upstream of this test; the unit tests inside + // `fixed_size_list.rs` cover semantic correctness per builder. + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 3, + ), + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt64, false)), + 1, + ), + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float64, true)), + 4, + ), + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Date32, true)), + 2, + ), ]; for dt in &supported_cases { @@ -1337,6 +1412,21 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), + // FixedSizeList: only primitive children are + // covered by this PR; nested children depend on the + // List / Struct builders that follow in the EPIC sequence. + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Utf8, true)), + 2, + ), + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Decimal128(38, 10), true)), + 2, + ), + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Boolean, true)), + 2, + ), ]; for dt in &unsupported_cases { From 56ecc8d360250356ea2cf035897c99443c142ff4 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Wed, 24 Jun 2026 11:00:32 +0800 Subject: [PATCH 2/3] Address Copilot review on #23128: reject negative FixedSizeList size, drop per-compare allocation in equal_to Four Copilot inline comments handled together: * `group_column_supported_type` and `make_group_column` both now reject `FixedSizeList` with a negative `list_size`. Arrow defines that size as non-negative, but the enum lets callers construct a negative value programmatically; without this guard the `i32 as usize` cast inside the builder would wrap and risk panics / OOM. Guards are duplicated on both sides so direct callers of `make_group_column` (and the `supported_schema` planner gate) both fail safely. Consistency fuzz already covers the `make_group_column` side via the existing unsupported_cases path; an explicit dedicated test below covers both surfaces. * `FixedSizeListGroupValueBuilder::new` asserts `list_len >= 0` so any direct construction that bypasses the factory fails fast with a clear message instead of silently producing a broken builder. `list_len_usize` now uses `usize::try_from(...).expect(...)` instead of `as usize`; with the constructor's assertion the conversion is provably infallible, but going through `try_from` means a future invariant break shows up as an explicit panic rather than a silent wrap. * `equal_to` no longer calls `list_array.value(rhs_row)`, which constructs a sliced child `ArrayRef` on every comparison. The hot path now borrows `list_array.values()` and uses `value_offset(rhs_row)` to compute the child base index, mirroring the existing approach in `append_val`. Zero allocations per equality check. New test: `negative_list_size_is_rejected_by_allow_list_and_dispatcher` locks in the negative-size guard at both surfaces. 13/13 FSL tests pass, 128/128 aggregates pass, clippy clean. --- .../multi_group_by/fixed_size_list.rs | 61 ++++++++++++++++++- .../group_values/multi_group_by/mod.rs | 22 ++++++- 2 files changed, 78 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs index 514bd6f596063..4bd28bc67afc8 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs @@ -60,6 +60,15 @@ where "FixedSizeListGroupValueBuilder built with non-FixedSizeList type {other:?}" ), }; + // `group_column_supported_type` and `make_group_column` both reject + // negative `list_len` upstream so this branch should be unreachable + // for any schema that survives the allow-list. Assert defensively + // to fail fast (and with a clear message) if a direct caller ever + // bypasses the factory with an invalid Arrow type. + assert!( + list_len >= 0, + "FixedSizeListGroupValueBuilder requires non-negative list size, got {list_len}" + ); let child = PrimitiveGroupValueBuilder::::new(field.data_type().clone()); Self { field, @@ -70,8 +79,14 @@ where } } + /// Lossless widening to `usize`. The constructor asserts + /// `list_len >= 0`, so the conversion is well-defined. We still go + /// through `try_from` rather than `as usize` so any future + /// invariant break surfaces as a panic with a clear message rather + /// than a silent wrap. fn list_len_usize(&self) -> usize { - self.list_len as usize + usize::try_from(self.list_len) + .expect("list_len validated >= 0 in `new`; conversion to usize cannot fail") } } @@ -91,11 +106,16 @@ where .as_any() .downcast_ref::() .expect("FixedSizeListGroupValueBuilder called with non-FixedSizeList array"); - let rhs_sublist: ArrayRef = list_array.value(rhs_row); + // Use the borrowed child array + `value_offset` (same approach as + // `append_val` below) instead of `list_array.value(rhs_row)`, + // which would allocate a fresh sliced `ArrayRef` on every + // equality check inside grouping's hot path. + let child_array = list_array.values(); let list_len = self.list_len_usize(); let lhs_base = lhs_row * list_len; + let rhs_base = list_array.value_offset(rhs_row) as usize; for j in 0..list_len { - if !self.child.equal_to(lhs_base + j, &rhs_sublist, j) { + if !self.child.equal_to(lhs_base + j, child_array, rhs_base + j) { return false; } } @@ -417,6 +437,41 @@ mod tests { assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6]); } + /// `FixedSizeList` carries its size as `i32`; the enum lets callers + /// construct a negative size programmatically even though Arrow + /// considers it invalid. Both the supported-type allow-list and the + /// dispatcher must reject such a type so the cast inside the + /// builder cannot wrap to a huge `usize` and trigger panics / OOM. + #[test] + fn negative_list_size_is_rejected_by_allow_list_and_dispatcher() { + use crate::aggregates::group_values::multi_group_by::{ + group_column_supported_type, make_group_column, + }; + use arrow::datatypes::Field; + + let invalid = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + -1, + ); + + assert!( + !group_column_supported_type(&invalid), + "allow-list must reject FixedSizeList(_, -1)" + ); + + let field = Field::new("col", invalid, true); + let result = make_group_column(&field); + let err = match result { + Ok(_) => panic!("dispatcher must reject FixedSizeList with negative size"), + Err(e) => e, + }; + let msg = err.to_string(); + assert!( + msg.contains("negative size") && msg.contains("-1"), + "error should mention the invalid negative size, got: {msg}" + ); + } + #[test] fn take_n_returns_prefix_and_shifts_remainder() { let input = fsl_array(&[ 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 17c62c4c34f1e..77ba67d102658 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 @@ -926,7 +926,15 @@ macro_rules! instantiate_primitive { fn group_column_supported_type(data_type: &DataType) -> bool { // FixedSizeList is the only nested type currently supported. // List / LargeList / Struct will follow in subsequent PRs of #22715. - if let DataType::FixedSizeList(child_field, _) = data_type { + if let DataType::FixedSizeList(child_field, list_size) = data_type { + // Arrow defines `FixedSizeList` size as a non-negative `i32`, but + // the enum still lets callers construct a negative size. Reject + // here so the schema is routed to the `GroupValuesRows` fallback + // instead of risking a `negative_size as usize` wrap inside the + // builder. Keep this guard in lockstep with `make_group_column`. + if *list_size < 0 { + return false; + } return matches!( child_field.data_type(), DataType::Int8 @@ -1087,7 +1095,17 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } - DataType::FixedSizeList(ref child_field, _) => { + DataType::FixedSizeList(ref child_field, list_size) => { + // Defense in depth against an invalid Arrow type that the + // allow-list might have missed: negative `list_size` would + // wrap when cast to `usize` and trigger panics / OOM in the + // builder. Reject explicitly here as well so direct callers + // of `make_group_column` fail safely. + if list_size < 0 { + return not_impl_err!( + "FixedSizeList with negative size {list_size} not supported in GroupValuesColumn" + ); + } // `group_column_supported_type` already restricts the child to // the primitive subset supported here. Any unsupported child // type returned `false` upstream and was routed to the From bd4532a5d95f981e77a289eb55931b97108ecbff Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Sun, 5 Jul 2026 22:51:56 +0800 Subject: [PATCH 3/3] Address Rich-T-kid review on #23128: assert! + comment tightening * Dispatcher (`mod.rs`): change the `list_size < 0` guard from a `return not_impl_err!` to a bare `assert!`. `group_column_supported_type` already rejects negative sizes at the schema-check layer, so any reachable path with `list_size < 0` is a programmer bug rather than an unsupported case worth returning gracefully. Split the existing test into an allow-list case (checked via `bool`) and a dispatcher case (verified with `#[should_panic]`) so both surfaces stay covered. * Constructor (`fixed_size_list.rs`): fold the multi-line invariant comment above `assert!` into the assert message itself, which now reads "requires non-negative list size (allow-list / dispatcher should have rejected earlier)". Simpler to read at the panic site, and the rationale surfaces exactly where a future reader would want it (the panic backtrace). Same trim on `list_len_usize`. * Dispatcher consistency-fuzz corpus (`mod.rs`): shorten the FixedSizeList comment to a one-liner pointing at EPIC #22715, instead of the two-line prose explanation. Tests: `cargo test -p datafusion-physical-plan --lib multi_group_by` now runs 42/42 (up from 41), including the split negative-size cases and the new `#[should_panic]` guarantee. Clippy clean. --- .../multi_group_by/fixed_size_list.rs | 50 ++++++++----------- .../group_values/multi_group_by/mod.rs | 24 ++++----- 2 files changed, 33 insertions(+), 41 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs index 4bd28bc67afc8..bf59bb2c422bc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_list.rs @@ -60,14 +60,9 @@ where "FixedSizeListGroupValueBuilder built with non-FixedSizeList type {other:?}" ), }; - // `group_column_supported_type` and `make_group_column` both reject - // negative `list_len` upstream so this branch should be unreachable - // for any schema that survives the allow-list. Assert defensively - // to fail fast (and with a clear message) if a direct caller ever - // bypasses the factory with an invalid Arrow type. assert!( list_len >= 0, - "FixedSizeListGroupValueBuilder requires non-negative list size, got {list_len}" + "FixedSizeListGroupValueBuilder requires non-negative list size (allow-list / dispatcher should have rejected earlier), got {list_len}" ); let child = PrimitiveGroupValueBuilder::::new(field.data_type().clone()); Self { @@ -79,12 +74,9 @@ where } } - /// Lossless widening to `usize`. The constructor asserts - /// `list_len >= 0`, so the conversion is well-defined. We still go - /// through `try_from` rather than `as usize` so any future - /// invariant break surfaces as a panic with a clear message rather - /// than a silent wrap. fn list_len_usize(&self) -> usize { + // `try_from` (vs `as usize`) so any future invariant break + // surfaces as a panic instead of a silent wrap. usize::try_from(self.list_len) .expect("list_len validated >= 0 in `new`; conversion to usize cannot fail") } @@ -439,14 +431,11 @@ mod tests { /// `FixedSizeList` carries its size as `i32`; the enum lets callers /// construct a negative size programmatically even though Arrow - /// considers it invalid. Both the supported-type allow-list and the - /// dispatcher must reject such a type so the cast inside the - /// builder cannot wrap to a huge `usize` and trigger panics / OOM. + /// considers it invalid. The supported-type allow-list must reject + /// such a type so it never reaches the builder path. #[test] - fn negative_list_size_is_rejected_by_allow_list_and_dispatcher() { - use crate::aggregates::group_values::multi_group_by::{ - group_column_supported_type, make_group_column, - }; + fn negative_list_size_rejected_by_allow_list() { + use crate::aggregates::group_values::multi_group_by::group_column_supported_type; use arrow::datatypes::Field; let invalid = DataType::FixedSizeList( @@ -458,18 +447,23 @@ mod tests { !group_column_supported_type(&invalid), "allow-list must reject FixedSizeList(_, -1)" ); + } - let field = Field::new("col", invalid, true); - let result = make_group_column(&field); - let err = match result { - Ok(_) => panic!("dispatcher must reject FixedSizeList with negative size"), - Err(e) => e, - }; - let msg = err.to_string(); - assert!( - msg.contains("negative size") && msg.contains("-1"), - "error should mention the invalid negative size, got: {msg}" + /// The dispatcher is a belt-and-suspenders guard for direct callers + /// that bypass the allow-list. A negative `list_size` must panic + /// there rather than wrap in the builder. + #[test] + #[should_panic(expected = "FixedSizeList requires non-negative size")] + fn negative_list_size_panics_in_dispatcher() { + use crate::aggregates::group_values::multi_group_by::make_group_column; + use arrow::datatypes::Field; + + let invalid = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + -1, ); + let field = Field::new("col", invalid, true); + let _ = make_group_column(&field); } #[test] 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 77ba67d102658..5569b9e189604 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 @@ -1096,16 +1096,15 @@ fn make_group_column(field: &Field) -> Result> { } } DataType::FixedSizeList(ref child_field, list_size) => { - // Defense in depth against an invalid Arrow type that the - // allow-list might have missed: negative `list_size` would - // wrap when cast to `usize` and trigger panics / OOM in the - // builder. Reject explicitly here as well so direct callers - // of `make_group_column` fail safely. - if list_size < 0 { - return not_impl_err!( - "FixedSizeList with negative size {list_size} not supported in GroupValuesColumn" - ); - } + // `group_column_supported_type` already rejects negative + // `list_size`; this assert is a defensive belt-and-suspenders + // check so a direct caller of `make_group_column` that bypasses + // the allow-list can't wrap the `i32 as usize` cast in the + // builder into a huge value. + assert!( + list_size >= 0, + "FixedSizeList requires non-negative size, got {list_size}" + ); // `group_column_supported_type` already restricts the child to // the primitive subset supported here. Any unsupported child // type returned `false` upstream and was routed to the @@ -1430,9 +1429,8 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), - // FixedSizeList: only primitive children are - // covered by this PR; nested children depend on the - // List / Struct builders that follow in the EPIC sequence. + // FixedSizeList: primitive-only in this PR; + // nested children land later in EPIC #22715. DataType::FixedSizeList( Arc::new(Field::new("item", DataType::Utf8, true)), 2,