From b5d2dd331f0a8bbf78a54179177645c9bbd7de35 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 16:34:02 -0600 Subject: [PATCH 1/2] bench: add a native micro-benchmark for ExplodeExec Measures the operator over in-memory batches, so a change to the unnesting kernels is not diluted by the Parquet scan and the aggregate that CometExplodeBenchmark necessarily includes. --- native/core/Cargo.toml | 4 + native/core/benches/explode.rs | 225 +++++++++++++++++++++++++++++++++ 2 files changed, 229 insertions(+) create mode 100644 native/core/benches/explode.rs diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 8dc8d73273f..65414ddb8b5 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -118,3 +118,7 @@ crate-type = ["cdylib", "rlib"] [[bench]] name = "array_element_append" harness = false + +[[bench]] +name = "explode" +harness = false diff --git a/native/core/benches/explode.rs b/native/core/benches/explode.rs new file mode 100644 index 00000000000..6c05e9cc818 --- /dev/null +++ b/native/core/benches/explode.rs @@ -0,0 +1,225 @@ +// 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. + +//! Micro-benchmarks for `ExplodeExec`, the operator behind Spark's `explode` and `posexplode`. +//! +//! `CometExplodeBenchmark` on the JVM side measures the same operator end to end, where the +//! Parquet scan and the counting aggregate are a large share of the total. This one runs the +//! operator over in-memory batches so a change to the unnesting kernels shows up undiluted. +//! +//! The dimensions are the ones that drive its cost: how far each row fans out, the element type +//! being unnested, how many columns are replicated alongside the generated one, and whether the +//! input has the NULL rows that force outer semantics. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use comet::execution::operators::ExplodeExec; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::UnnestOptions; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::unnest::ListUnnest; +use datafusion::physical_plan::{common::collect, ExecutionPlan}; +use datafusion::prelude::SessionConfig; +use tokio::runtime::Runtime; + +/// Input rows per batch, and batches per run. 8192 is DataFusion's default `batch_size`, so the +/// operator chunks the input rather than seeing it whole. +const ROWS_PER_BATCH: usize = 8192; +const BATCHES: usize = 8; + +/// The element types worth distinguishing: a fixed-width primitive, a variable-width type whose +/// gather has to rebuild offsets and copy bytes, and a nested type that gathers per field. +#[derive(Clone, Copy, PartialEq)] +enum Element { + Int64, + Utf8, + Struct, +} + +impl Element { + fn name(self) -> &'static str { + match self { + Element::Int64 => "bigint", + Element::Utf8 => "string", + Element::Struct => "struct", + } + } + + fn data_type(self) -> DataType { + match self { + Element::Int64 => DataType::Int64, + Element::Utf8 => DataType::Utf8, + Element::Struct => DataType::Struct(self.struct_fields()), + } + } + + fn struct_fields(self) -> Fields { + Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ]) + } + + /// A flat child array of `count` elements, which the list offsets then carve into rows. + fn values(self, count: usize) -> ArrayRef { + let ints = Int64Array::from_iter_values((0..count).map(|i| i as i64)); + match self { + Element::Int64 => Arc::new(ints), + Element::Utf8 => Arc::new(StringArray::from_iter_values( + (0..count).map(|i| format!("str_{i}")), + )), + Element::Struct => { + let strings = StringArray::from_iter_values((0..count).map(|i| format!("str_{i}"))); + Arc::new(StructArray::new( + self.struct_fields(), + vec![Arc::new(ints), Arc::new(strings)], + None, + )) + } + } + } +} + +/// One input batch: a `List` column of `fan_out`-element rows, plus `carried` passthrough +/// columns that unnesting has to replicate. +/// +/// With `nulls`, every tenth row is a NULL list. That is the shape `explode_outer` sees, and it +/// is also what decides whether the unnested column can be sliced out of the child or has to be +/// gathered: a NULL row under outer semantics is padded, which breaks the run. +fn input_batch(element: Element, fan_out: usize, carried: usize, nulls: bool) -> RecordBatch { + let total = ROWS_PER_BATCH * fan_out; + let offsets: Vec = (0..=ROWS_PER_BATCH).map(|r| (r * fan_out) as i32).collect(); + let null_buffer = + nulls.then(|| NullBuffer::from_iter((0..ROWS_PER_BATCH).map(|row| row % 10 != 0))); + + let list = ListArray::new( + Arc::new(Field::new("item", element.data_type(), true)), + OffsetBuffer::new(offsets.into()), + element.values(total), + null_buffer, + ); + + let mut fields = vec![Field::new("arr", list.data_type().clone(), true)]; + let mut columns: Vec = vec![Arc::new(list)]; + for c in 0..carried { + fields.push(Field::new(format!("k{c}"), DataType::Int64, true)); + columns.push(Arc::new(Int64Array::from_iter_values( + (0..ROWS_PER_BATCH).map(|r| (r + c) as i64), + ))); + } + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +/// The operator's output schema: the unnested element column, then the passthrough columns. +/// +/// This mirrors what the planner builds, except that the planner puts the passthrough columns +/// first; the order does not change the work, only which index the unnest targets. +fn output_schema(element: Element, carried: usize) -> SchemaRef { + let mut fields = vec![Field::new("arr", element.data_type(), true)]; + for c in 0..carried { + fields.push(Field::new(format!("k{c}"), DataType::Int64, true)); + } + Arc::new(Schema::new(fields)) +} + +fn explode_plan( + element: Element, + fan_out: usize, + carried: usize, + outer: bool, +) -> Arc { + let batches: Vec = (0..BATCHES) + .map(|_| input_batch(element, fan_out, carried, outer)) + .collect(); + let schema = batches[0].schema(); + let source = MemorySourceConfig::try_new_exec(&[batches], schema, None).unwrap(); + + Arc::new( + ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + output_schema(element, carried), + UnnestOptions { + preserve_nulls: outer, + recursions: vec![], + }, + ) + .unwrap(), + ) +} + +fn run(runtime: &Runtime, plan: &Arc, ctx: &Arc) { + let stream = plan.execute(0, Arc::clone(ctx)).unwrap(); + let batches = runtime.block_on(collect(stream)).unwrap(); + assert!(!batches.is_empty()); +} + +fn criterion_benchmark(c: &mut Criterion) { + let runtime = Runtime::new().unwrap(); + let ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(ROWS_PER_BATCH)), + ); + + let mut group = c.benchmark_group("explode_fan_out"); + for fan_out in [2usize, 10, 100] { + let plan = explode_plan(Element::Int64, fan_out, 0, false); + group.bench_with_input(BenchmarkId::from_parameter(fan_out), &fan_out, |b, _| { + b.iter(|| run(&runtime, &plan, &ctx)) + }); + } + group.finish(); + + let mut group = c.benchmark_group("explode_element_type"); + for element in [Element::Int64, Element::Utf8, Element::Struct] { + let plan = explode_plan(element, 10, 0, false); + group.bench_function(element.name(), |b| b.iter(|| run(&runtime, &plan, &ctx))); + } + group.finish(); + + let mut group = c.benchmark_group("explode_carried_columns"); + for carried in [0usize, 3] { + let plan = explode_plan(Element::Int64, 10, carried, false); + group.bench_with_input(BenchmarkId::from_parameter(carried), &carried, |b, _| { + b.iter(|| run(&runtime, &plan, &ctx)) + }); + } + group.finish(); + + // NULL rows under outer semantics are padded, so this is the shape that cannot be served by + // slicing the child and has to gather instead. Kept as its own group so the two paths are + // not averaged together. + let mut group = c.benchmark_group("explode_outer_with_nulls"); + for element in [Element::Int64, Element::Utf8] { + let plan = explode_plan(element, 10, 0, true); + group.bench_function(element.name(), |b| b.iter(|| run(&runtime, &plan, &ctx))); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 53ef92990f7da2a433f348df5f6ffbe3e709ca7a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 16:48:27 -0600 Subject: [PATCH 2/2] perf: slice the child instead of gathering it when unnesting Unnesting a single list column pads nothing, so `unnest_list_array` was building an index array holding one i64 per output element and then gathering the child through it, when the indices it built were the contiguous run the elements already sit in. Return a slice of the child instead. Comet takes this path for plain `explode` and for both columns of `posexplode`, whose position array is built with the same per-row lengths; `explode_outer` still gathers, since a NULL or empty row is padded and breaks the run. Guarding on the offset span alone is not enough. Arrow permits a NULL list slot to span elements, which the gather skips and a slice would not, and such a row can cancel out padding elsewhere and leave the totals agreeing with a run that is not the one to take. So the check also rejects a NULL row holding elements, and the view types, whose per-row offsets are independent. Two smaller things on the way past. `predict_output_lens` derived its per-row lengths through `find_longest_length`, which chains `length`, `cast`, `is_not_null` and `zip` to stay generic over list types: four allocating passes to subtract adjacent offsets, once per input batch. Comet only plans `List`, so take that case in one pass and keep the general version as the fallback. And `create_take_indices` appended the repeat indices one element at a time through a builder that has no validity to track; fill the buffer per run instead. Measured with the new native benchmark, on an Apple M3 Max: explode_fan_out/2 -74% explode_fan_out/10 -82% explode_fan_out/100 -90% explode_element_type/bigint -82% explode_element_type/string -91% explode_element_type/struct -91% explode_carried_columns/0 -82% explode_carried_columns/3 -60% explode_outer_with_nulls/bigint -13% explode_outer_with_nulls/string -8% The outer cases keep gathering, so they get only the length and index changes. --- .../core/src/execution/operators/explode.rs | 346 ++++++++++++++++-- 1 file changed, 321 insertions(+), 25 deletions(-) diff --git a/native/core/src/execution/operators/explode.rs b/native/core/src/execution/operators/explode.rs index d4da87e5737..b5484593f95 100644 --- a/native/core/src/execution/operators/explode.rs +++ b/native/core/src/execution/operators/explode.rs @@ -43,27 +43,25 @@ //! //! # What was forked //! -//! The unnesting kernels below (`build_batch` and everything it calls) are copied from +//! The unnesting kernels below (`build_batch` and everything it calls) started as a copy of //! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision //! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to //! `datafusion-physical-plan`, so they cannot be called from here without copying them. -//! Leave them semantically unmodified so the eventual deletion is mechanical; the -//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`. //! -//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and -//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt` -//! reflows some signatures. To audit for real changes, reformat this region at -//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference -//! to a single cosmetic line wrap in `flatten_struct_cols`. The deliberate edits are: +//! They have since been specialized for the shapes Comet actually plans, so this is no longer a +//! copy that can be diffed against upstream line by line. The deliberate divergences are: //! //! * the `lt` import path noted below, since Comet does not depend on `arrow_ord` directly; //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; -//! * the `precomputed_lengths` parameter on `build_batch` and `list_unnest_at_level`, which -//! is itself part of apache/datafusion#24384 and so disappears with the rest of the fork. +//! * the `precomputed_lengths` parameter on `build_batch` and `list_unnest_at_level`, which is +//! part of apache/datafusion#24384; +//! * the contiguous-run fast path in `unnest_list_array`, which returns a slice of the child +//! values instead of gathering them, and the buffer fills in `create_take_indices`. //! -//! Everything else this fork does — chunking, plan properties, EOF handling — mirrors that -//! same upstream PR, so keep the two in step: a change made here that is not in #24384 -//! either belongs upstream or does not belong at all. +//! The performance work is Comet-specific and is not held to upstream's shape. When the fork is +//! eventually retired in favor of `UnnestExec`, these paths are what would have to be measured +//! again — or upstreamed first — rather than simply deleted. `ExplodeExec` and `ExplodeStream` +//! were always Comet's own. //! //! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses //! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with @@ -498,15 +496,50 @@ impl ExplodeStream { // This is exactly the per-row length that `list_unnest_at_level` derives when it // actually unnests, so the chunk boundaries are exact rather than estimated, and each // chunk's slice of it is handed back to `build_batch` instead of recomputed there. + if let [single] = list_arrays.as_slice() { + if let Some(list) = single.as_any().downcast_ref::() { + return Ok(Some(list_output_lens(list, self.options.preserve_nulls))); + } + } let longest_length = find_longest_length(&list_arrays, &self.options)?; Ok(Some(longest_length.as_primitive::().clone())) } } +/// The per-row unnested length of a single `List` column: the row's list length, or `null_length` +/// for a NULL row. +/// +/// What [`find_longest_length`] computes when handed one array, in one pass over the offsets +/// rather than the four allocating kernels it chains to stay generic over list types — `length` +/// (which returns `Int32` for `List`), `cast` to widen it, `is_not_null`, and `zip` to substitute +/// the NULL length. Comet only ever plans `List`, and only ever one or two of them, so this is +/// the path every explode takes; anything else still falls back to the general version. +fn list_output_lens(list: &ListArray, preserve_nulls: bool) -> PrimitiveArray { + let null_length = if preserve_nulls { 1 } else { 0 }; + let offsets = list.offsets(); + // Like `find_longest_length`, the result is non-null throughout: a NULL row reports + // `null_length` rather than a NULL length, which `create_take_indices` relies on. + let lens: Vec = match list.nulls() { + None => offsets.windows(2).map(|w| (w[1] - w[0]) as i64).collect(), + Some(nulls) => offsets + .windows(2) + .enumerate() + .map(|(row, w)| { + if nulls.is_valid(row) { + (w[1] - w[0]) as i64 + } else { + null_length + } + }) + .collect(), + }; + PrimitiveArray::::from(lens) +} + // --------------------------------------------------------------------------------------- -// Everything below is copied from DataFusion 54.1.0 `physical-plan/src/unnest.rs` -// (revision cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13). See the module docs for why, and -// for how to audit it against upstream. Do not change it semantically. +// Everything below started as a copy of DataFusion 54.1.0 `physical-plan/src/unnest.rs` +// (revision cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13), since specialized for Comet. See the +// module docs for what diverges and why. // --------------------------------------------------------------------------------------- /// Given a set of struct column indices to flatten @@ -928,6 +961,14 @@ trait ListArrayType: Array { /// Returns the start and end offset of the values for the given row. fn value_offsets(&self, row: usize) -> (i64, i64); + + /// Whether consecutive rows occupy consecutive ranges of [`values`](Self::values), so that a + /// run of rows is one slice of it. + /// + /// True for the offset-based list types, where row `i` is `[offsets[i], offsets[i + 1])` and + /// so ends exactly where row `i + 1` begins. False for the view types, whose per-row offsets + /// are independent and may overlap, repeat, or leave gaps. + fn is_contiguous(&self) -> bool; } impl ListArrayType for ListArray { @@ -939,6 +980,10 @@ impl ListArrayType for ListArray { let offsets = self.value_offsets(); (offsets[row].into(), offsets[row + 1].into()) } + + fn is_contiguous(&self) -> bool { + true + } } impl ListArrayType for LargeListArray { @@ -950,6 +995,10 @@ impl ListArrayType for LargeListArray { let offsets = self.value_offsets(); (offsets[row], offsets[row + 1]) } + + fn is_contiguous(&self) -> bool { + true + } } impl ListArrayType for FixedSizeListArray { @@ -961,6 +1010,10 @@ impl ListArrayType for FixedSizeListArray { let start = self.value_offset(row) as i64; (start, start + self.value_length() as i64) } + + fn is_contiguous(&self) -> bool { + true + } } impl ListArrayType for ListViewArray { @@ -973,6 +1026,10 @@ impl ListArrayType for ListViewArray { let size = self.value_sizes()[row] as i64; (offset, offset + size) } + + fn is_contiguous(&self) -> bool { + false + } } impl ListArrayType for LargeListViewArray { @@ -985,6 +1042,10 @@ impl ListArrayType for LargeListViewArray { let size = self.value_sizes()[row]; (offset, offset + size) } + + fn is_contiguous(&self) -> bool { + false + } } /// Unnest multiple list arrays according to the length array. @@ -1015,6 +1076,52 @@ fn unnest_list_arrays( .collect::>() } +/// Whether unnesting `list_array` against `length_array` would gather exactly the contiguous run +/// `values[offsets.first()..offsets.last()]`, in which case [`unnest_list_array`] can slice the +/// child instead of building an index array and gathering through it. +/// +/// The run is the right answer only if the loop in [`unnest_list_array`] would emit each row's +/// values in order with nothing added and nothing skipped, which needs three things: +/// +/// * The rows are laid out consecutively in the child, so a run of them is one slice. The view +/// types are excluded here rather than checked, since their offsets are unordered. +/// * No row is padded. `target >= value` holds per row, so it is enough that the totals agree: +/// `capacity` is the sum of the targets and the offset span is the sum of the values. +/// * No NULL row holds elements. Arrow permits a NULL list slot to span a non-empty range, and +/// the loop skips those elements while the slice would include them. Builders and the Parquet +/// reader emit an empty range, so this scan almost always confirms rather than rejects, and it +/// only runs when the array has nulls at all. +/// +/// The last two conditions are independent: a NULL row spanning elements can cancel out padding +/// elsewhere and leave the totals matching a run that is not the one to take. +fn is_contiguous_unnest(list_array: &dyn ListArrayType, capacity: usize) -> bool { + let len = list_array.len(); + if len == 0 || !list_array.is_contiguous() { + return false; + } + + let (first, _) = list_array.value_offsets(0); + let (_, last) = list_array.value_offsets(len - 1); + if last - first != capacity as i64 { + return false; + } + + if list_array.null_count() > 0 { + let has_populated_null = (0..len).any(|row| { + if !list_array.is_null(row) { + return false; + } + let (start, end) = list_array.value_offsets(row); + end > start + }); + if has_populated_null { + return false; + } + } + + true +} + /// Unnest a list array according the target length array. /// /// Consider a list array like this: @@ -1041,6 +1148,26 @@ fn unnest_list_array( capacity: usize, ) -> Result { let values = list_array.values(); + + // Unnesting a single list column pads nothing, so the elements come out in the order they + // are already stored and the gather below would read straight through them. Hand back a + // slice of the child instead: no index buffer, no copy of the element data, which for a + // string or nested element type is the bulk of the operator's work. Comet reaches this for + // plain `explode`, and for both columns of `posexplode`, whose position array is built with + // the same per-row lengths. `explode_outer` falls through as soon as a row is NULL or empty, + // because those rows are padded. + // + // The result aliases the child rather than owning a compacted copy, and `ListArray::slice` + // leaves `values` whole, so this is the child of the whole input batch and not of the chunk. + // Every chunk of one input batch therefore pins the same buffer, which between them they + // fill; a downstream operator that keeps only some of those chunks pins all of it. That is + // bounded by one input batch's expansion, which `pending_input` already holds materialized, + // and slicing the input is what `BatchSplitStream` above does too. + if is_contiguous_unnest(list_array, capacity) { + let (first, _) = list_array.value_offsets(0); + return Ok(values.slice(first as usize, capacity)); + } + let mut take_indices_builder = PrimitiveArray::::builder(capacity); for row in 0..list_array.len() { let mut value_length = 0; @@ -1057,9 +1184,7 @@ fn unnest_list_array( "value length is beyond the longest length" ); // Pad with NULL values - for _ in value_length..target_length { - take_indices_builder.append_null(); - } + take_indices_builder.append_nulls((target_length - value_length) as usize); } Ok(kernels::take::take( &values, @@ -1091,13 +1216,14 @@ fn create_take_indices( length_array.null_count() == 0, "length array should not contain nulls" ); - let mut builder = PrimitiveArray::::builder(capacity); - for (index, repeat) in length_array.iter().enumerate() { - // The length array should not contain nulls, so unwrap is safe - let repeat = repeat.unwrap(); - (0..repeat).for_each(|_| builder.append_value(index as i64)); + // A run of one index at a time, so fill the buffer directly rather than appending element by + // element through a builder: there is no validity to track, and each row becomes one fill of + // `repeat` slots rather than `repeat` calls. + let mut indices: Vec = Vec::with_capacity(capacity); + for (index, repeat) in length_array.values().iter().enumerate() { + indices.resize(indices.len() + *repeat as usize, index as i64); } - builder.finish() + PrimitiveArray::::from(indices) } /// Create a batch of arrays based on an input `batch` and a `indices` array. @@ -1168,6 +1294,7 @@ fn repeat_arrs_from_indices( mod tests { use super::*; use arrow::array::Int32Array; + use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{Field, Int32Type, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -1488,4 +1615,173 @@ mod tests { assert_eq!(ordering.len(), 1); assert_eq!(ordering[0].expr.as_ref(), key.as_ref()); } + + // --------------------------------------------------------------------------------------- + // Contiguous-run fast path in `unnest_list_array` + // --------------------------------------------------------------------------------------- + + fn int_list(offsets: Vec, values: Vec, nulls: Option>) -> ListArray { + ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(Int32Array::from(values)), + nulls.map(NullBuffer::from), + ) + } + + /// Unnest `list` against `lens`, returning the elements and whether the fast path was taken. + fn unnest(list: &ListArray, lens: Vec) -> (Vec>, bool) { + let capacity = lens.iter().sum::() as usize; + let length_array = PrimitiveArray::::from(lens); + let took_fast_path = is_contiguous_unnest(list as &dyn ListArrayType, capacity); + let out = unnest_list_array(list as &dyn ListArrayType, &length_array, capacity).unwrap(); + let out = out.as_primitive::().iter().collect(); + (out, took_fast_path) + } + + /// The pointer the array's first buffer starts at, to tell a shared child from a copy. + fn buffer_ptr(array: &dyn Array) -> *const u8 { + array.to_data().buffers()[0].as_ptr() + } + + #[test] + fn contiguous_unnest_returns_a_slice_of_the_child() { + // Rows [1,2,3], [4], [5,6] unnest to the child verbatim, so the result must share the + // child's buffer rather than gather into a fresh one. This is the plain-`explode` path. + let list = int_list(vec![0, 3, 4, 6], vec![1, 2, 3, 4, 5, 6], None); + let (values, fast) = unnest(&list, vec![3, 1, 2]); + assert!( + fast, + "a single unpadded list column must take the fast path" + ); + assert_eq!(values, (1..=6).map(Some).collect::>()); + + let length_array = PrimitiveArray::::from(vec![3i64, 1, 2]); + let out = unnest_list_array(&list as &dyn ListArrayType, &length_array, 6).unwrap(); + assert_eq!( + buffer_ptr(out.as_ref()), + buffer_ptr(ListArrayType::values(&list).as_ref()), + "the fast path must alias the child, not copy it" + ); + } + + #[test] + fn contiguous_unnest_honors_a_non_zero_offset_base() { + // Slicing leaves `offsets.first() > 0` while `values` stays whole, so a fast path that + // sliced from 0 would silently return the wrong elements. + let list = int_list(vec![0, 2, 3, 5, 6], vec![1, 2, 3, 4, 5, 6], None); + let sliced = list.slice(1, 2); + let (values, fast) = unnest(&sliced, vec![1, 2]); + assert!(fast); + assert_eq!(values, vec![Some(3), Some(4), Some(5)]); + } + + #[test] + fn contiguous_unnest_covers_empty_rows_and_dropped_nulls() { + // Plain `explode`: a NULL row and an empty row both contribute no elements, and with + // `preserve_nulls` false neither is padded, so the run stays unbroken across them. + let list = int_list( + vec![0, 2, 2, 2, 5], + vec![1, 2, 3, 4, 5], + Some(vec![true, false, true, true]), + ); + let (values, fast) = unnest(&list, vec![2, 0, 0, 3]); + assert!(fast, "rows contributing nothing must not break the run"); + assert_eq!(values, (1..=5).map(Some).collect::>()); + } + + #[test] + fn padded_rows_fall_back_to_the_gather() { + // `explode_outer`: the NULL row is padded to one element, so the output interleaves a + // NULL that no slice of the child contains. + let list = int_list( + vec![0, 2, 2, 4], + vec![1, 2, 3, 4], + Some(vec![true, false, true]), + ); + let (values, fast) = unnest(&list, vec![2, 1, 2]); + assert!(!fast, "a padded row cannot be served by a slice"); + assert_eq!(values, vec![Some(1), Some(2), None, Some(3), Some(4)]); + } + + #[test] + fn populated_null_row_falls_back_even_when_the_totals_agree() { + // The arithmetic check alone is not enough. Arrow allows a NULL slot to span elements; + // here row 0 is NULL over two of them and row 1 is padded by two, so the offset span + // (5) equals the capacity (5) while the correct output skips the NULL row's elements. + // Slicing would return [1,2,3,4,5] instead. + let list = int_list(vec![0, 2, 5], vec![1, 2, 3, 4, 5], Some(vec![false, true])); + let (values, fast) = unnest(&list, vec![0, 5]); + assert!( + !fast, + "a NULL row holding elements breaks the run even when the totals match" + ); + assert_eq!(values, vec![Some(3), Some(4), Some(5), None, None]); + } + + #[test] + fn list_view_input_falls_back() { + // View offsets are independent per row, so consecutive rows need not be adjacent and a + // run cannot be assumed. Here they are deliberately out of order. + let view = ListViewArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + ScalarBuffer::from(vec![3, 0]), + ScalarBuffer::from(vec![2, 3]), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + None, + ); + let capacity = 5; + assert!(!is_contiguous_unnest(&view as &dyn ListArrayType, capacity)); + + let length_array = PrimitiveArray::::from(vec![2i64, 3]); + let out = unnest_list_array(&view as &dyn ListArrayType, &length_array, capacity).unwrap(); + let values: Vec> = out.as_primitive::().iter().collect(); + assert_eq!(values, vec![Some(4), Some(5), Some(1), Some(2), Some(3)]); + } + + // --------------------------------------------------------------------------------------- + // Fused per-row length computation + // --------------------------------------------------------------------------------------- + + /// `list_output_lens` must agree with `find_longest_length` element for element, since the + /// chunking in `ExplodeStream` and the unnesting in `build_batch` both consume it. + fn assert_lens_match_general(list: ListArray, preserve_nulls: bool) { + let options = UnnestOptions { + preserve_nulls, + recursions: vec![], + }; + let arrays = vec![Arc::new(list.clone()) as ArrayRef]; + let expected = find_longest_length(&arrays, &options).unwrap(); + let expected = expected.as_primitive::(); + let actual = list_output_lens(&list, preserve_nulls); + assert_eq!(&actual, expected, "preserve_nulls = {preserve_nulls}"); + assert_eq!(actual.null_count(), 0, "lengths must never be NULL"); + } + + #[test] + fn fused_lengths_match_the_general_kernel() { + let plain = int_list(vec![0, 3, 4, 4, 6], vec![1, 2, 3, 4, 5, 6], None); + assert_lens_match_general(plain.clone(), true); + assert_lens_match_general(plain, false); + + let with_nulls = int_list( + vec![0, 2, 2, 2, 5], + vec![1, 2, 3, 4, 5], + Some(vec![true, false, true, true]), + ); + assert_lens_match_general(with_nulls.clone(), true); + assert_lens_match_general(with_nulls, false); + + let empty = int_list(vec![0], vec![], None); + assert_lens_match_general(empty.clone(), true); + assert_lens_match_general(empty, false); + } + + #[test] + fn fused_lengths_handle_a_sliced_input() { + // Sliced offsets start away from zero; the length is still the per-row difference. + let list = int_list(vec![0, 2, 3, 3, 6], vec![1, 2, 3, 4, 5, 6], None); + let sliced = list.slice(1, 3); + assert_lens_match_general(sliced, true); + } }