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/10/15 20:53:47 UTC
(datafusion-comet) 01/01: fix: compact oversized nested list_extract buffers (#6498)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6498-7c58073982a2ad68a17bb0eb6ab06b06871d6a50
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 1a8dc4064b3f4775e5f574c78bf9841b4656a40c
Author: KUAN-HAO HUANG <10...@users.noreply.github.com>
AuthorDate: Thu Oct 15 18:02:13 2026 +0000
fix: compact oversized nested list_extract buffers (#6498)
* fix: compact oversized nested list_extract buffers
* test: exercise nested list extraction compaction through JVM
* test: deduplicate list_extract reference gathers
---
native/spark-expr/benches/list_extract.rs | 158 ++++++++-
native/spark-expr/src/array_funcs/list_extract.rs | 385 ++++++++++++++++++++-
.../expressions/array/list_extract_nested.sql | 77 +++++
3 files changed, 616 insertions(+), 4 deletions(-)
diff --git a/native/spark-expr/benches/list_extract.rs b/native/spark-expr/benches/list_extract.rs
index 3b9c62de9a..318a853f22 100644
--- a/native/spark-expr/benches/list_extract.rs
+++ b/native/spark-expr/benches/list_extract.rs
@@ -15,7 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{Array, ArrayRef, Int32Array, ListArray, RecordBatch, StringArray};
+use arrow::array::{
+ Array, ArrayRef, Int32Array, ListArray, MapArray, RecordBatch, StringArray, StructArray,
+};
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{DataType, Field, Schema};
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
@@ -218,5 +220,157 @@ fn bench_defaults(c: &mut Criterion) {
);
}
-criterion_group!(benches, criterion_benchmark, bench_defaults);
+/// Skewed nested lengths expose `take`'s reservation based on the entire input,
+/// rather than the selected children. Include the opposite (long selection), too.
+fn bench_nested(c: &mut Criterion) {
+ let mut group = c.benchmark_group("list_extract_nested");
+ for selected_len in [0, 1, 128] {
+ let lengths = (0..ROWS).flat_map(|_| [selected_len, 128]);
+ let offsets = OffsetBuffer::<i32>::from_lengths(lengths);
+ let count = *offsets.last().unwrap() as usize;
+ let ints: ArrayRef = Arc::new(Int32Array::from_iter_values(0..count as i32));
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let lists = Arc::new(ListArray::new(field, offsets.clone(), ints.clone(), None));
+ let fields = vec![
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Field::new("value", DataType::Int32, true)),
+ ];
+ let entries = StructArray::new(fields.into(), vec![ints.clone(), ints], None);
+ let maps = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+ offsets,
+ entries,
+ None,
+ false,
+ ));
+ let structs = Arc::new(StructArray::new(
+ vec![Arc::new(Field::new(
+ "nested",
+ lists.data_type().clone(),
+ true,
+ ))]
+ .into(),
+ vec![lists.clone()],
+ None,
+ ));
+
+ bench_nested_values(
+ &mut group,
+ selected_len,
+ [
+ ("list", lists as ArrayRef),
+ ("map", maps as ArrayRef),
+ ("struct-list", structs as ArrayRef),
+ ],
+ );
+ }
+ group.finish();
+}
+
+fn bench_nested_values(
+ group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
+ selected_len: usize,
+ values: impl IntoIterator<Item = (&'static str, ArrayRef)>,
+) {
+ for (ty, values) in values {
+ for null_percent in [0, 25, 75] {
+ let nulls = (null_percent > 0).then(|| {
+ (0..ROWS)
+ .map(|row| row % 4 >= null_percent / 25)
+ .collect::<NullBuffer>()
+ });
+ let outer: ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new("item", values.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(2, ROWS)),
+ Arc::clone(&values),
+ nulls,
+ ));
+ let batch = single_col_batch(outer);
+ let expr = ListExtract::new(
+ Arc::new(Column::new("c0", 0)),
+ Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+ None,
+ false,
+ false,
+ None,
+ create_query_context_map(),
+ );
+ group.bench_function(
+ format!("{ty}/{selected_len}-selected/{null_percent}%-nulls"),
+ |b| b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())),
+ );
+ }
+ }
+}
+
+/// Middle lists have equal lengths. Only the deeper selected lists are short,
+/// exposing reservations propagated into their primitive buffers.
+fn bench_deep_nested(c: &mut Criterion) {
+ const WIDTH: usize = 32;
+ let mut group = c.benchmark_group("list_extract_deep_nested");
+ for selected_len in [0, 1, 16] {
+ let offsets = OffsetBuffer::from_lengths((0..ROWS).flat_map(|_| {
+ std::iter::repeat_n(selected_len, WIDTH).chain(std::iter::repeat_n(16, WIDTH))
+ }));
+ let count = *offsets.last().unwrap();
+ let inner: ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new("item", DataType::Int32, true)),
+ offsets,
+ Arc::new(Int32Array::from_iter_values(0..count)),
+ None,
+ ));
+ let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
+ let lists: ArrayRef = Arc::new(ListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(WIDTH, 2 * ROWS)),
+ Arc::clone(&inner),
+ None,
+ ));
+ let fixed_lists: ArrayRef = Arc::new(arrow::array::FixedSizeListArray::new(
+ field,
+ WIDTH as i32,
+ Arc::clone(&inner),
+ None,
+ ));
+ let entries = StructArray::new(
+ vec![
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Field::new("value", inner.data_type().clone(), true)),
+ ]
+ .into(),
+ vec![
+ Arc::new(Int32Array::from_iter_values(
+ (0..2 * ROWS).flat_map(|_| 0..WIDTH as i32),
+ )),
+ inner,
+ ],
+ None,
+ );
+ let maps: ArrayRef = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(WIDTH, 2 * ROWS)),
+ entries,
+ None,
+ false,
+ ));
+ bench_nested_values(
+ &mut group,
+ selected_len,
+ [
+ ("list-list", lists),
+ ("map-list", maps),
+ ("fixed-list", fixed_lists),
+ ],
+ );
+ }
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ criterion_benchmark,
+ bench_defaults,
+ bench_nested,
+ bench_deep_nested
+);
criterion_main!(benches);
diff --git a/native/spark-expr/src/array_funcs/list_extract.rs b/native/spark-expr/src/array_funcs/list_extract.rs
index 061760c01b..e31175b39e 100644
--- a/native/spark-expr/src/array_funcs/list_extract.rs
+++ b/native/spark-expr/src/array_funcs/list_extract.rs
@@ -16,7 +16,8 @@
// under the License.
use arrow::array::{
- Array, GenericListArray, Int32Array, MutableArrayData, OffsetSizeTrait, UInt64Builder,
+ Array, ArrayData, AsArray, GenericListArray, Int32Array, MutableArrayData, OffsetSizeTrait,
+ UInt64Builder,
};
use arrow::compute::take;
use arrow::datatypes::{DataType, FieldRef, Schema};
@@ -359,6 +360,113 @@ fn list_extract<O: OffsetSizeTrait>(
)))
}
+// Inspect owned output capacity directly: nested builders can propagate a reservation
+// into deeper children even when the immediate child count was estimated correctly.
+// Keep ordinary growth headroom and small aligned buffers to avoid unnecessary copies.
+fn has_excess_nested_capacity(array: &dyn Array) -> bool {
+ let oversized = has_excess_capacity;
+ if array
+ .nulls()
+ .is_some_and(|nulls| oversized(nulls.buffer().capacity(), nulls.len().div_ceil(8)))
+ {
+ return true;
+ }
+ match array.data_type() {
+ DataType::List(_) => {
+ let list = array.as_list::<i32>();
+ has_excess_nested_capacity(list.values().as_ref())
+ || oversized(
+ list.offsets().inner().inner().capacity(),
+ list.offsets().len() * 4,
+ )
+ }
+ DataType::LargeList(_) => {
+ let list = array.as_list::<i64>();
+ has_excess_nested_capacity(list.values().as_ref())
+ || oversized(
+ list.offsets().inner().inner().capacity(),
+ list.offsets().len() * 8,
+ )
+ }
+ DataType::Map(_, _) => {
+ let map = array.as_map();
+ has_excess_nested_capacity(map.entries())
+ || oversized(
+ map.offsets().inner().inner().capacity(),
+ map.offsets().len() * 4,
+ )
+ }
+ DataType::Struct(_) => array
+ .as_struct()
+ .columns()
+ .iter()
+ .any(|column| has_excess_nested_capacity(column.as_ref())),
+ DataType::FixedSizeList(_, _) => {
+ has_excess_nested_capacity(array.as_fixed_size_list().values().as_ref())
+ }
+ data_type => {
+ let used = match data_type {
+ DataType::Boolean => array.len().div_ceil(8),
+ DataType::Utf8 => {
+ array.as_string::<i32>().value_data().len() + (array.len() + 1) * 4
+ }
+ DataType::LargeUtf8 => {
+ array.as_string::<i64>().value_data().len() + (array.len() + 1) * 8
+ }
+ DataType::Binary => {
+ array.as_binary::<i32>().value_data().len() + (array.len() + 1) * 4
+ }
+ DataType::LargeBinary => {
+ array.as_binary::<i64>().value_data().len() + (array.len() + 1) * 8
+ }
+ DataType::FixedSizeBinary(width) => array.len().saturating_mul(*width as usize),
+ _ => match data_type.primitive_width() {
+ Some(width) => array.len().saturating_mul(width),
+ None => return false,
+ },
+ };
+ let null_bytes = array.nulls().map_or(0, |nulls| nulls.len().div_ceil(8));
+ oversized(
+ array.get_buffer_memory_size(),
+ used.saturating_add(null_bytes),
+ )
+ }
+ }
+}
+
+fn has_excess_capacity(capacity: usize, used: usize) -> bool {
+ // Allow growth headroom plus one alignment block. Without the padding allowance,
+ // an offset buffer just above four times its used size can trigger compaction.
+ capacity > used.saturating_mul(4).saturating_add(64)
+}
+
+fn compact_nested_data(data: ArrayData) -> ArrayData {
+ let (data_type, len, mut nulls, offset, mut buffers, children) = data.into_parts();
+ for buffer in &mut buffers {
+ if has_excess_capacity(buffer.capacity(), buffer.len()) {
+ buffer.shrink_to_fit();
+ }
+ }
+ if let Some(nulls) = &mut nulls {
+ if has_excess_capacity(nulls.buffer().capacity(), nulls.buffer().len()) {
+ nulls.shrink_to_fit();
+ }
+ }
+ let children = children.into_iter().map(compact_nested_data).collect();
+ // SAFETY: this data came from a valid Arrow take result. Only buffer capacities
+ // changed: bytes, lengths, offsets, types, and validity are preserved. Rechecking
+ // every nested offset or string would add a scan without changing these invariants.
+ unsafe {
+ ArrayData::builder(data_type)
+ .len(len)
+ .offset(offset)
+ .nulls(nulls)
+ .buffers(buffers)
+ .child_data(children)
+ .build_unchecked()
+ }
+}
+
fn list_extract_without_default<O: OffsetSizeTrait>(
list_array: &GenericListArray<O>,
index_array: &Int32Array,
@@ -389,7 +497,13 @@ fn list_extract_without_default<O: OffsetSizeTrait>(
let indices = indices.finish();
// Every gathered index was validated by `resolve_row` against its row's list bounds,
// and `take` masks null index slots, so the default unchecked `TakeOptions` is safe.
- Ok(ColumnarValue::Array(take(values.as_ref(), &indices, None)?))
+ let mut result = take(values.as_ref(), &indices, None)?;
+ if values.data_type().is_nested() && has_excess_nested_capacity(result.as_ref()) {
+ // Convert only oversized results, releasing the array before shrinking its
+ // owned buffers. Preserve headroom in other buffers to avoid needless copies.
+ result = arrow::array::make_array(compact_nested_data(result.into_data()));
+ }
+ Ok(ColumnarValue::Array(result))
}
impl Display for ListExtract {
@@ -411,6 +525,273 @@ mod test {
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::ColumnarValue;
+ fn reference_gather(
+ values: &dyn Array,
+ source_rows: &[Option<usize>],
+ ) -> Result<arrow::array::ArrayRef> {
+ let data = values.to_data();
+ let mut result = MutableArrayData::new(vec![&data], true, source_rows.len());
+ for &row in source_rows {
+ match row {
+ Some(index) => result.try_extend(0, index, index + 1)?,
+ None => result.try_extend_nulls(1)?,
+ }
+ }
+ Ok(arrow::array::make_array(result.freeze()))
+ }
+
+ fn skewed_nested_values(selected_len: usize) -> Vec<arrow::array::ArrayRef> {
+ use arrow::array::{LargeListArray, MapArray, StructArray};
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ const ROWS: usize = 8192;
+ let offsets =
+ OffsetBuffer::<i32>::from_lengths((0..ROWS).flat_map(|_| [selected_len, 128]));
+ let count = *offsets.last().unwrap() as usize;
+ let ints: arrow::array::ArrayRef = Arc::new(Int32Array::from_iter_values(0..count as i32));
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let lists = Arc::new(ListArray::new(
+ Arc::clone(&field),
+ offsets.clone(),
+ Arc::clone(&ints),
+ None,
+ ));
+ let large_lists = Arc::new(LargeListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::<i64>::from_lengths((0..ROWS).flat_map(|_| [selected_len, 128])),
+ Arc::clone(&ints),
+ None,
+ ));
+ let entries = StructArray::new(
+ vec![
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Field::new("value", DataType::Int32, true)),
+ ]
+ .into(),
+ vec![Arc::clone(&ints), ints],
+ None,
+ );
+ let maps = Arc::new(MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+ offsets,
+ entries,
+ None,
+ false,
+ ));
+ let structs = Arc::new(StructArray::new(
+ vec![Arc::new(Field::new(
+ "nested",
+ lists.data_type().clone(),
+ true,
+ ))]
+ .into(),
+ vec![Arc::<ListArray>::clone(&lists)],
+ None,
+ ));
+ let nullable_lists = Arc::new(ListArray::new(
+ field,
+ OffsetBuffer::from_lengths((0..ROWS).flat_map(|_| [selected_len, 128])),
+ Arc::new(Int32Array::from_iter(
+ (0..count as i32).map(|i| (i % 11 != 0).then_some(i)),
+ )),
+ Some((0..2 * ROWS).map(|i| i % 7 != 0).collect::<NullBuffer>()),
+ ));
+ vec![lists, large_lists, maps, structs, nullable_lists]
+ }
+
+ #[test]
+ fn test_list_extract_nested_retained_capacity() -> Result<()> {
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ const ROWS: usize = 8192;
+ let error_wrapper = |error: SparkError| DataFusionError::from(error);
+ for selected_len in [0, 1] {
+ for values in skewed_nested_values(selected_len) {
+ for nullable in [false, true] {
+ let outer = ListArray::new(
+ Arc::new(Field::new("item", values.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(2, ROWS)),
+ Arc::clone(&values),
+ nullable.then(|| (0..ROWS).map(|row| row % 4 != 0).collect::<NullBuffer>()),
+ );
+ let source_rows = (0..ROWS)
+ .map(|row| outer.is_valid(row).then_some(2 * row))
+ .collect::<Vec<_>>();
+ let expected = reference_gather(values.as_ref(), &source_rows)?;
+ let null_default = ScalarValue::try_from(values.data_type())?;
+ for default in [None, Some(&null_default)] {
+ let ColumnarValue::Array(result) = list_extract(
+ &outer,
+ &Int32Array::from(vec![0; ROWS]),
+ default,
+ false,
+ false,
+ |idx, len| zero_based_index(idx, len, &error_wrapper),
+ &error_wrapper,
+ )?
+ else {
+ unreachable!()
+ };
+ assert_eq!(result.to_data(), expected.to_data());
+ // Allow builder alignment/headroom without accepting capacity sized
+ // from the unselected 128-element children (at least 2 MiB).
+ assert!(
+ result.get_buffer_memory_size() <= 2 * expected.get_buffer_memory_size(),
+ "{:?}, selected_len={selected_len}, nullable={nullable}: retained {} bytes vs {} for the row-wise gather",
+ values.data_type(), result.get_buffer_memory_size(), expected.get_buffer_memory_size(),
+ );
+ }
+ }
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn test_list_extract_deep_nested_retained_capacity() -> Result<()> {
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ const ROWS: usize = 256;
+ const WIDTH: usize = 32;
+ let error_wrapper = |error: SparkError| DataFusionError::from(error);
+ for selected_len in [0, 1] {
+ let offsets = OffsetBuffer::from_lengths((0..ROWS).flat_map(|_| {
+ std::iter::repeat_n(selected_len, WIDTH).chain(std::iter::repeat_n(16, WIDTH))
+ }));
+ let count = *offsets.last().unwrap();
+ let inner: arrow::array::ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new("item", DataType::Int32, true)),
+ offsets,
+ Arc::new(Int32Array::from_iter_values(0..count)),
+ None,
+ ));
+ let values: arrow::array::ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new("item", inner.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(WIDTH, 2 * ROWS)),
+ inner,
+ None,
+ ));
+ for null_percent in [0, 25, 75] {
+ let outer = ListArray::new(
+ Arc::new(Field::new("item", values.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(2, ROWS)),
+ Arc::clone(&values),
+ (null_percent > 0).then(|| {
+ (0..ROWS)
+ .map(|row| row % 4 >= null_percent / 25)
+ .collect::<NullBuffer>()
+ }),
+ );
+ let source_rows = (0..ROWS)
+ .map(|row| outer.is_valid(row).then_some(2 * row))
+ .collect::<Vec<_>>();
+ let expected = reference_gather(values.as_ref(), &source_rows)?;
+ let ColumnarValue::Array(result) = list_extract(
+ &outer,
+ &Int32Array::from(vec![0; ROWS]),
+ None,
+ false,
+ false,
+ |idx, len| zero_based_index(idx, len, &error_wrapper),
+ &error_wrapper,
+ )?
+ else {
+ unreachable!()
+ };
+ assert_eq!(result.to_data(), expected.to_data());
+ if selected_len == 1 && null_percent == 75 {
+ let indices = arrow::array::UInt64Array::from_iter(
+ (0..ROWS).map(|row| outer.is_valid(row).then_some((2 * row) as u64)),
+ );
+ let gathered = take(values.as_ref(), &indices, None)?;
+ // Preserve normal 4x headroom, including alignment padding. Compacting
+ // this result used to copy offsets even though the leaf is not oversized.
+ assert_eq!(
+ result.get_buffer_memory_size(),
+ gathered.get_buffer_memory_size()
+ );
+ }
+ let ints = result.as_list::<i32>().values().as_list::<i32>().values();
+ // An accurate immediate child reservation must not leave an oversized
+ // primitive buffer when propagated into empty or short inner lists.
+ assert!(
+ ints.get_buffer_memory_size() <= ints.len() * 16 + 64,
+ "selected_len={selected_len}, null_percent={null_percent}: {} integers retain {} bytes",
+ ints.len(),
+ ints.get_buffer_memory_size(),
+ );
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn test_list_extract_nested_sliced_nulls_and_indices() -> Result<()> {
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ let error_wrapper = |error: SparkError| DataFusionError::from(error);
+ for values in skewed_nested_values(1) {
+ let outer = ListArray::new(
+ Arc::new(Field::new("item", values.data_type().clone(), true)),
+ OffsetBuffer::from_lengths(std::iter::repeat_n(2, 8192)),
+ Arc::clone(&values),
+ Some((0..8192).map(|row| row % 4 != 0).collect::<NullBuffer>()),
+ )
+ .slice(1, 8);
+ // The sliced rows are: first child, last child, null index, null list,
+ // out-of-bounds, first child, last child, null list.
+ let expected = reference_gather(
+ values.as_ref(),
+ &[Some(2), Some(5), None, None, None, Some(12), Some(15), None],
+ )?;
+ for (one_based, indices) in [
+ (
+ false,
+ vec![
+ Some(0),
+ Some(1),
+ None,
+ Some(0),
+ Some(2),
+ Some(0),
+ Some(1),
+ Some(-1),
+ ],
+ ),
+ (
+ true,
+ vec![
+ Some(1),
+ Some(-1),
+ None,
+ Some(1),
+ Some(3),
+ Some(1),
+ Some(-1),
+ Some(0),
+ ],
+ ),
+ ] {
+ let ColumnarValue::Array(result) = list_extract(
+ &outer,
+ &Int32Array::from(indices),
+ None,
+ false,
+ one_based,
+ |idx, len| {
+ if one_based {
+ one_based_index(idx, len, &error_wrapper)
+ } else {
+ zero_based_index(idx, len, &error_wrapper)
+ }
+ },
+ &error_wrapper,
+ )?
+ else {
+ unreachable!()
+ };
+ assert_eq!(result.to_data(), expected.to_data());
+ }
+ }
+ Ok(())
+ }
+
#[test]
fn test_nullable_when_array_is_nullable() -> Result<()> {
// Regression test for SPARK-55747: GetArrayItem over a nullable array (e.g. the
diff --git a/spark/src/test/resources/sql-tests/expressions/array/list_extract_nested.sql b/spark/src/test/resources/sql-tests/expressions/array/list_extract_nested.sql
new file mode 100644
index 0000000000..8517686da3
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/array/list_extract_nested.sql
@@ -0,0 +1,77 @@
+-- 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.
+
+
+-- Config: spark.sql.ansi.enabled=false
+-- Config: spark.comet.batchSize=3
+
+statement
+CREATE TABLE test_list_extract_nested(a array<array<int>>, m array<map<int,int>>, s array<struct<n:array<int>>>, idx int) USING parquet
+
+-- Long unselected children force compaction when the empty first element is selected.
+statement
+INSERT INTO test_list_extract_nested VALUES
+ (array(array(), sequence(1, 64)), array(map(), map_from_arrays(sequence(1, 64), sequence(1, 64))), array(named_struct('n', array()), named_struct('n', sequence(1, 64))), 0),
+ (array(array(), array(1, NULL, 3)), array(map(), map(1, 2)), array(named_struct('n', array()), named_struct('n', array(1, NULL, 3))), 0),
+ (array(array(10), array(20, 21)), array(map(10, 11), map(20, 21)), array(named_struct('n', array(10)), named_struct('n', array(20, 21))), 1),
+ (array(NULL, array(1)), array(NULL, map(1, 2)), array(NULL, named_struct('n', array(1))), 0),
+ (NULL, NULL, NULL, 0),
+ (array(array(), array(2)), array(map(), map(2, 3)), array(named_struct('n', array()), named_struct('n', array(2))), NULL),
+ (array(array(1)), array(map(1, 2)), array(named_struct('n', array(1))), 5),
+ (array(array(1)), array(map(1, 2)), array(named_struct('n', array(1))), -1)
+
+-- Column ordinals, selected NULL children, NULL lists and ordinals, and invalid indices.
+query expect_native(getarrayitem)
+SELECT a[idx], m[idx], s[idx] FROM test_list_extract_nested
+
+-- Literal ordinals select heterogeneous empty and short children.
+query expect_native(getarrayitem)
+SELECT a[0], m[0], s[0] FROM test_list_extract_nested
+
+-- One-based indexing and negative indexing share the native ListExtract kernel.
+query expect_native(element_at)
+SELECT element_at(a, 1), element_at(m, 1), element_at(s, 1) FROM test_list_extract_nested
+
+query expect_native(element_at)
+SELECT element_at(a, -1), element_at(m, -1), element_at(s, -1) FROM test_list_extract_nested
+
+query expect_native(element_at)
+SELECT element_at(a, idx + 1), element_at(m, idx + 1), element_at(s, idx + 1) FROM test_list_extract_nested WHERE idx >= 0 OR idx IS NULL
+
+
+-- Equal middle lengths can still propagate oversized reservations into deeper children.
+statement
+CREATE TABLE test_list_extract_deep(a array<array<array<int>>>, m array<map<int,array<int>>>, idx int) USING parquet
+
+-- Equal middle lengths with long unselected descendants exercise recursive compaction.
+statement
+INSERT INTO test_list_extract_deep VALUES
+ (array(array_repeat(array(), 32), array_repeat(sequence(1, 4), 32)), array(map_from_arrays(sequence(1, 32), array_repeat(array(), 32)), map_from_arrays(sequence(1, 32), array_repeat(sequence(1, 4), 32))), 0),
+ (array(array(array(), array()), array(array(1, 2), array(3, 4))), array(map(1, array(), 2, array()), map(1, array(1, 2), 2, array(3, 4))), 0),
+ (array(array(array(10), array(20)), array(array(30, 31), array(40, 41))), array(map(1, array(10), 2, array(20)), map(1, array(30, 31), 2, array(40, 41))), 1),
+ (array(array(NULL, array()), array(array(1))), array(map(1, NULL, 2, array()), map(1, array(1))), 0),
+ (NULL, NULL, 0),
+ (array(array(array(1))), array(map(1, array(1))), NULL)
+
+query expect_native(getarrayitem)
+SELECT a[idx], m[idx] FROM test_list_extract_deep
+
+query expect_native(element_at)
+SELECT element_at(a, 1), element_at(m, 1) FROM test_list_extract_deep
+
+query expect_native(element_at)
+SELECT element_at(a, -1), element_at(m, -1) FROM test_list_extract_deep
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org