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-6132-10741ae97220be5d055c561bd2c4273ebc9f1c97 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 5d59317630e19d22cc9be5dbdae59ae1b87b4f41 Author: Oleks V <[email protected]> AuthorDate: Thu Sep 24 01:53:37 2026 +0000 feat: use DataFusion `unnest_outer` instead of Comet's `ListEmptyToNullExpr` (#6132) DataFusion 55.1.0, which Comet already pins, carries `unnest_outer` (apache/datafusion#22100) as `NullHandling::PreserveAndExpandEmpty`. That is exactly Spark's `explode_outer` semantics, so the planner can ask for it directly instead of rewriting empty lists to NULL first. Closes #5210. --- .../expression-audits/generator_funcs.md | 4 +- native/core/benches/explode.rs | 176 +++++++++--- .../execution/expressions/list_empty_to_null.rs | 311 --------------------- native/core/src/execution/expressions/mod.rs | 1 - native/core/src/execution/operators/explode.rs | 311 ++++++++++++++------- native/core/src/execution/planner.rs | 40 ++- .../sql-tests/expressions/array/explode.sql | 10 +- .../sql-tests/expressions/array/posexplode.sql | 15 +- .../apache/comet/exec/CometGenerateExecSuite.scala | 15 +- 9 files changed, 382 insertions(+), 501 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 4d160c6ac3..3023cfaa77 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -27,7 +27,7 @@ ## explode_outer -- Same `CometExplodeExec` path as `explode`. Compatible for array inputs; empty and NULL arrays both emit one null-valued row per Spark's `outer` semantics via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)). Map inputs fall back. +- Same `CometExplodeExec` path as `explode`. Compatible for array inputs; empty and NULL arrays both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map inputs fall back. ## posexplode @@ -35,6 +35,6 @@ ## posexplode_outer -- Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics via the `ListEmptyToNullExpr` planner bridge (works around [datafusion#19053](https://github.com/apache/datafusion/issues/19053)). +- Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/native/core/benches/explode.rs b/native/core/benches/explode.rs index b72aa524e5..54778bb193 100644 --- a/native/core/benches/explode.rs +++ b/native/core/benches/explode.rs @@ -22,18 +22,19 @@ //! 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. +//! being unnested, how many columns are replicated alongside the generated one, whether the input +//! holds the NULL and empty rows that outer semantics pad, and whether a parallel positions column +//! is unnested alongside the array as `posexplode` does. use std::sync::Arc; -use arrow::array::{Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray}; +use arrow::array::{Array, ArrayRef, Int32Array, 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::common::{NullHandling, UnnestOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::TaskContext; use datafusion::physical_plan::unnest::ListUnnest; @@ -99,27 +100,103 @@ impl Element { } } -/// One input batch: a `List` column of `fan_out`-element rows, plus `carried` passthrough -/// columns that unnesting has to replicate. +/// Which rows the input holds. /// -/// 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<i32> = (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))); +/// `Dense` is the plain `explode` shape, where every row fans out and the unnested column can be +/// sliced straight out of the child. `NullsAndEmpties` is the shape `explode_outer` exists for: +/// both a NULL row and an empty row are padded to one NULL, which breaks the contiguous run and +/// forces the gather. Spark treats the two identically, so a benchmark of the outer path that +/// holds only NULL rows leaves the empty-row substitution unmeasured. +#[derive(Clone, Copy, PartialEq)] +enum RowMix { + Dense, + NullsAndEmpties, +} + +impl RowMix { + fn name(self) -> &'static str { + match self { + RowMix::Dense => "dense", + RowMix::NullsAndEmpties => "nulls_and_empties", + } + } + + /// The per-row element count. Every tenth row is NULL and every tenth is empty, so a batch + /// holds a fifth padded rows. + fn row_len(self, fan_out: usize, row: usize) -> usize { + match self { + RowMix::Dense => fan_out, + RowMix::NullsAndEmpties if self.is_null(row) || row % 10 == 5 => 0, + RowMix::NullsAndEmpties => fan_out, + } + } + + fn is_null(self, row: usize) -> bool { + self == RowMix::NullsAndEmpties && row.is_multiple_of(10) + } + + /// The options the planner builds for this shape. `Dense` stands in for plain `explode`, + /// which drops NULL and empty rows alike. + fn unnest_options(self) -> UnnestOptions { + UnnestOptions::new().with_null_handling(match self { + RowMix::Dense => NullHandling::Drop, + RowMix::NullsAndEmpties => NullHandling::PreserveAndExpandEmpty, + }) + } +} + +/// One input batch: a `List` column of `mix`-shaped rows, optionally a parallel `List<Int32>` of +/// positions, plus `carried` passthrough columns that unnesting has to replicate. +/// +/// The positions column is what `ListPositionsExpr` builds for `posexplode`: the same offsets and +/// the same validity as the array, with values `0..len`. Unnesting the two together is the only +/// shape that exercises the multi-array row-wise maximum in `find_longest_length`. +fn input_batch( + element: Element, + fan_out: usize, + carried: usize, + mix: RowMix, + positions: bool, +) -> RecordBatch { + let offsets: Vec<i32> = std::iter::once(0) + .chain((0..ROWS_PER_BATCH).scan(0i32, |end, row| { + *end += mix.row_len(fan_out, row) as i32; + Some(*end) + })) + .collect(); + let total = *offsets.last().unwrap() as usize; + let offsets = OffsetBuffer::new(offsets.into()); + + let nulls = (mix != RowMix::Dense) + .then(|| NullBuffer::from_iter((0..ROWS_PER_BATCH).map(|row| !mix.is_null(row)))); let list = ListArray::new( Arc::new(Field::new("item", element.data_type(), true)), - OffsetBuffer::new(offsets.into()), + offsets.clone(), element.values(total), - null_buffer, + nulls.clone(), ); - let mut fields = vec![Field::new("arr", list.data_type().clone(), true)]; - let mut columns: Vec<ArrayRef> = vec![Arc::new(list)]; + let mut fields = Vec::new(); + let mut columns: Vec<ArrayRef> = Vec::new(); + + if positions { + // Per-row lengths come back off the offsets rather than a second pass over `row_len`. + let pos_values = + Int32Array::from_iter_values(offsets.windows(2).flat_map(|w| 0..w[1] - w[0])); + let pos = ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + offsets, + Arc::new(pos_values), + nulls, + ); + fields.push(Field::new("pos", pos.data_type().clone(), true)); + columns.push(Arc::new(pos)); + } + + fields.push(Field::new("arr", list.data_type().clone(), true)); + columns.push(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( @@ -130,12 +207,17 @@ fn input_batch(element: Element, fan_out: usize, carried: usize, nulls: bool) -> RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() } -/// The operator's output schema: the unnested element column, then the passthrough columns. +/// The operator's output schema: the unnested position column when positional, 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)]; +fn output_schema(element: Element, carried: usize, positions: bool) -> SchemaRef { + let mut fields = Vec::new(); + if positions { + fields.push(Field::new("pos", DataType::Int32, true)); + } + fields.push(Field::new("arr", element.data_type(), true)); for c in 0..carried { fields.push(Field::new(format!("k{c}"), DataType::Int64, true)); } @@ -146,24 +228,31 @@ fn explode_plan( element: Element, fan_out: usize, carried: usize, - outer: bool, + mix: RowMix, + positions: bool, ) -> Arc<dyn ExecutionPlan> { let batches: Vec<RecordBatch> = (0..BATCHES) - .map(|_| input_batch(element, fan_out, carried, outer)) + .map(|_| input_batch(element, fan_out, carried, mix, positions)) .collect(); let schema = batches[0].schema(); let source = MemorySourceConfig::try_new_exec(&[batches], schema, None).unwrap(); + // Positional unnesting targets the positions column and the array together, in the order the + // planner projects them: `0..=1` when positional, just the array at 0 otherwise. + let list_unnests = (0..=usize::from(positions)) + .map(|index_in_input_schema| ListUnnest { + index_in_input_schema, + depth: 1, + }) + .collect(); + Arc::new( ExplodeExec::new( source, - vec![ListUnnest { - index_in_input_schema: 0, - depth: 1, - }], + list_unnests, vec![], - output_schema(element, carried), - UnnestOptions::new().with_preserve_nulls(outer), + output_schema(element, carried, positions), + mix.unnest_options(), ) .unwrap(), ) @@ -184,7 +273,7 @@ fn criterion_benchmark(c: &mut Criterion) { 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); + let plan = explode_plan(Element::Int64, fan_out, 0, RowMix::Dense, false); group.bench_with_input(BenchmarkId::from_parameter(fan_out), &fan_out, |b, _| { b.iter(|| run(&runtime, &plan, &ctx)) }); @@ -193,29 +282,44 @@ fn criterion_benchmark(c: &mut Criterion) { 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); + let plan = explode_plan(element, 10, 0, RowMix::Dense, 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); + let plan = explode_plan(Element::Int64, 10, carried, RowMix::Dense, 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. + // NULL and empty 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); + let plan = explode_plan(element, 10, 0, RowMix::NullsAndEmpties, false); group.bench_function(element.name(), |b| b.iter(|| run(&runtime, &plan, &ctx))); } group.finish(); + + // `posexplode` unnests the positions column alongside the array, which is the only shape that + // reaches the multi-array row-wise maximum in `find_longest_length`. Short arrays are the + // interesting case: the per-batch length work is fixed, so the shorter the rows the larger + // its share of the total. + let mut group = c.benchmark_group("posexplode_fan_out"); + for mix in [RowMix::Dense, RowMix::NullsAndEmpties] { + for fan_out in [2usize, 10] { + let plan = explode_plan(Element::Int64, fan_out, 0, mix, true); + group.bench_with_input(BenchmarkId::new(mix.name(), fan_out), &fan_out, |b, _| { + b.iter(|| run(&runtime, &plan, &ctx)) + }); + } + } + group.finish(); } criterion_group!(benches, criterion_benchmark); diff --git a/native/core/src/execution/expressions/list_empty_to_null.rs b/native/core/src/execution/expressions/list_empty_to_null.rs deleted file mode 100644 index ad6d9452ac..0000000000 --- a/native/core/src/execution/expressions/list_empty_to_null.rs +++ /dev/null @@ -1,311 +0,0 @@ -// 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. - -use std::fmt::{Display, Formatter}; -use std::hash::{Hash, Hasher}; -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, ListArray, RecordBatch}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::datatypes::{DataType, Field, FieldRef, Schema}; -use datafusion::common::{exec_err, Result as DataFusionResult}; -use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_plan::ColumnarValue; - -/// A `PhysicalExpr` that marks every empty row of a `List<T>` input as null. -/// Bridges DataFusion's `UnnestExec` (which drops empty rows under -/// `preserve_nulls=true`) to Spark's `explode_outer`/`posexplode_outer` -/// semantics. See <https://github.com/apache/datafusion/issues/19053>. -#[derive(Debug, Clone)] -pub struct ListEmptyToNullExpr { - child: Arc<dyn PhysicalExpr>, -} - -impl ListEmptyToNullExpr { - pub fn new(child: Arc<dyn PhysicalExpr>) -> Self { - Self { child } - } -} - -impl Display for ListEmptyToNullExpr { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "list_empty_to_null({})", self.child) - } -} - -impl PartialEq for ListEmptyToNullExpr { - fn eq(&self, other: &Self) -> bool { - self.child.eq(&other.child) - } -} - -impl Eq for ListEmptyToNullExpr {} - -impl Hash for ListEmptyToNullExpr { - fn hash<H: Hasher>(&self, state: &mut H) { - self.child.hash(state); - } -} - -impl PhysicalExpr for ListEmptyToNullExpr { - fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - Display::fmt(self, f) - } - - fn return_field(&self, input_schema: &Schema) -> DataFusionResult<FieldRef> { - // Preserve the child field's name and element type; force the outer - // list to nullable because we mark empty rows as null. - let child_field = self.child.return_field(input_schema)?; - Ok(Arc::new(Field::new( - child_field.name(), - child_field.data_type().clone(), - true, - ))) - } - - fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> { - let value = self.child.evaluate(batch)?; - let array = value.into_array(batch.num_rows())?; - - let Some(list) = array.as_any().downcast_ref::<ListArray>() else { - return exec_err!( - "ListEmptyToNullExpr expected List input, got {}", - array.data_type() - ); - }; - - let offsets = list.offsets(); - let len = list.len(); - let existing_nulls = list.nulls(); - - // Fast path: no currently-valid row is empty, so the input already - // satisfies outer semantics. `is_valid` returns true when `nulls` is - // `None`, so this single scan short-circuits on the first empty - // valid row without allocating. - let has_valid_empty = (0..len) - .any(|i| offsets[i + 1] == offsets[i] && existing_nulls.is_none_or(|n| n.is_valid(i))); - if !has_valid_empty { - return Ok(ColumnarValue::Array(Arc::clone(&array))); - } - - let combined = BooleanBuffer::collect_bool(len, |i| { - offsets[i + 1] > offsets[i] && existing_nulls.is_none_or(|n| n.is_valid(i)) - }); - let new_nulls = NullBuffer::new(combined); - - let DataType::List(element_field) = list.data_type() else { - unreachable!("ListArray downcast guarantees DataType::List"); - }; - - let result = ListArray::try_new( - Arc::clone(element_field), - offsets.clone(), - Arc::clone(list.values()), - Some(new_nulls), - )?; - - Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) - } - - fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { - vec![&self.child] - } - - fn with_new_children( - self: Arc<Self>, - children: Vec<Arc<dyn PhysicalExpr>>, - ) -> DataFusionResult<Arc<dyn PhysicalExpr>> { - if children.len() != 1 { - return exec_err!( - "ListEmptyToNullExpr expects exactly 1 child, got {}", - children.len() - ); - } - Ok(Arc::new(ListEmptyToNullExpr::new(Arc::clone(&children[0])))) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{Int32Array, ListArray}; - use arrow::buffer::{NullBuffer, OffsetBuffer}; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion::physical_expr::expressions::Column; - - fn element_field() -> Arc<Field> { - Arc::new(Field::new("item", DataType::Int32, true)) - } - - fn list_field(nullable: bool) -> Arc<Field> { - Arc::new(Field::new("arr", DataType::List(element_field()), nullable)) - } - - fn build_batch(list: ArrayRef) -> RecordBatch { - let schema = Schema::new(vec![Field::new("arr", list.data_type().clone(), true)]); - RecordBatch::try_new(Arc::new(schema), vec![list]).unwrap() - } - - fn evaluate_ref(list: ArrayRef) -> ArrayRef { - let batch = build_batch(list); - let expr = ListEmptyToNullExpr::new(Arc::new(Column::new("arr", 0))); - let result = expr.evaluate(&batch).unwrap(); - result.into_array(batch.num_rows()).unwrap() - } - - fn evaluate(list: ListArray) -> ListArray { - let array: ArrayRef = Arc::new(list); - let out = evaluate_ref(array); - out.as_any().downcast_ref::<ListArray>().unwrap().clone() - } - - #[test] - fn fast_path_no_empty_rows_returns_input_untouched() { - // Rows: [1,2,3], [4], [5,6] -- no empty rows, so the fast path should - // return the original array pointer without allocating a new bitmap. - let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); - let offsets = OffsetBuffer::new(vec![0, 3, 4, 6].into()); - let input: ArrayRef = Arc::new(ListArray::new( - element_field(), - offsets, - Arc::new(values), - None, - )); - - let out = evaluate_ref(Arc::clone(&input)); - assert!( - Arc::ptr_eq(&out, &input), - "fast path should return the same ArrayRef" - ); - let list = out.as_any().downcast_ref::<ListArray>().unwrap(); - assert!(list.nulls().is_none()); - assert_eq!(list.len(), 3); - } - - #[test] - fn fast_path_when_only_empty_row_is_already_null() { - // Rows: [1,2], NULL (offsets 2..2 -- looks empty), [3]. The middle row - // is already null so the fast path applies without materializing a new - // bitmap. - let values = Int32Array::from(vec![1, 2, 3]); - let offsets = OffsetBuffer::new(vec![0, 2, 2, 3].into()); - let nulls = NullBuffer::from(vec![true, false, true]); - let input: ArrayRef = Arc::new(ListArray::new( - element_field(), - offsets, - Arc::new(values), - Some(nulls), - )); - - let out = evaluate_ref(Arc::clone(&input)); - assert!( - Arc::ptr_eq(&out, &input), - "fast path should return the same ArrayRef when only empty rows are already null" - ); - } - - #[test] - fn mixed_empty_null_and_non_empty_rows() { - // Rows: [10, 20], [], NULL, [30]. The empty row (index 1) must become - // null; the already-null row (index 2) must stay null; the two data - // rows must survive intact. - let values = Int32Array::from(vec![10, 20, 30]); - let offsets = OffsetBuffer::new(vec![0, 2, 2, 2, 3].into()); - let nulls = NullBuffer::from(vec![true, true, false, true]); - let input = ListArray::new(element_field(), offsets, Arc::new(values), Some(nulls)); - let output = evaluate(input); - let out_nulls = output.nulls().expect("nulls buffer must be present"); - assert!(out_nulls.is_valid(0)); - assert!(!out_nulls.is_valid(1), "empty row must be marked null"); - assert!(!out_nulls.is_valid(2), "already-null row must stay null"); - assert!(out_nulls.is_valid(3)); - // Offsets and values must be preserved so downstream unnest still - // reads the original element slices for the valid rows. - assert_eq!(output.value(0).len(), 2); - assert_eq!(output.value(3).len(), 1); - } - - #[test] - fn empty_row_without_prior_null_bitmap() { - // Fresh (no nulls) input containing one empty row must materialize a - // null bitmap with only that row cleared. - let values = Int32Array::from(vec![1, 2, 3]); - let offsets = OffsetBuffer::new(vec![0, 2, 2, 3].into()); - let input = ListArray::new(element_field(), offsets, Arc::new(values), None); - let output = evaluate(input); - let out_nulls = output.nulls().expect("nulls buffer must be materialized"); - assert!(out_nulls.is_valid(0)); - assert!(!out_nulls.is_valid(1)); - assert!(out_nulls.is_valid(2)); - } - - #[test] - fn zero_row_batch_takes_fast_path() { - // A zero-row batch has no rows to inspect, so the fast path returns - // the input untouched. - let values = Int32Array::from(Vec::<i32>::new()); - let offsets = OffsetBuffer::new(vec![0].into()); - let input: ArrayRef = Arc::new(ListArray::new( - element_field(), - offsets, - Arc::new(values), - None, - )); - - let out = evaluate_ref(Arc::clone(&input)); - assert_eq!(out.len(), 0); - assert!( - Arc::ptr_eq(&out, &input), - "zero-row batch should take the fast path" - ); - } - - #[test] - fn sliced_input_with_non_zero_offset() { - // Build a 5-row list, then slice out rows 1..4 -- exposes non-zero - // logical offset while `values()` stays unsliced. The empty row that - // now lives at logical index 1 (was index 2 in the underlying array) - // must be flipped to null. - let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); - let offsets = OffsetBuffer::new(vec![0, 2, 3, 3, 5, 6].into()); - let input = ListArray::new(element_field(), offsets, Arc::new(values), None); - let sliced = input.slice(1, 3); - let output = evaluate(sliced); - assert_eq!(output.len(), 3); - let out_nulls = output.nulls().expect("nulls buffer must be materialized"); - // Logical row 0: was [3] -- non-empty - assert!(out_nulls.is_valid(0)); - // Logical row 1: was [] -- must be null - assert!(!out_nulls.is_valid(1)); - // Logical row 2: was [4, 5] -- non-empty - assert!(out_nulls.is_valid(2)); - // The non-empty rows must still expose their original elements. - assert_eq!(output.value(0).len(), 1); - assert_eq!(output.value(2).len(), 2); - } - - #[test] - fn return_field_forces_nullable() { - // The output field must be nullable regardless of the child's - // nullability, because empty rows are marked null downstream. - let schema = Schema::new(vec![list_field(false).as_ref().clone()]); - let expr = ListEmptyToNullExpr::new(Arc::new(Column::new("arr", 0))); - let field = expr.return_field(&schema).unwrap(); - assert!(field.is_nullable()); - assert_eq!(field.name(), "arr"); - } -} diff --git a/native/core/src/execution/expressions/mod.rs b/native/core/src/execution/expressions/mod.rs index 4c25109af6..e174bd3747 100644 --- a/native/core/src/execution/expressions/mod.rs +++ b/native/core/src/execution/expressions/mod.rs @@ -20,7 +20,6 @@ pub mod arithmetic; pub mod bitwise; pub mod comparison; -pub mod list_empty_to_null; pub mod list_positions; pub mod logical; pub mod nullcheck; diff --git a/native/core/src/execution/operators/explode.rs b/native/core/src/execution/operators/explode.rs index e3a5e3bb29..763a45eb36 100644 --- a/native/core/src/execution/operators/explode.rs +++ b/native/core/src/execution/operators/explode.rs @@ -15,31 +15,26 @@ // specific language governing permissions and limitations // under the License. -//! A temporary fork of DataFusion's `UnnestExec` that respects -//! `datafusion.execution.batch_size`. +//! A fork of DataFusion's `UnnestExec`, kept for two unnesting kernels Comet has specialized. //! //! # Why this fork exists //! -//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many -//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means -//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and -//! peak memory scales with input batch size times array length rather than with -//! `batch_size`. +//! It was created because `UnnestExec` emitted exactly one output batch per input batch, +//! however many rows the unnesting produced, and never consulted +//! `datafusion.execution.batch_size`. That fix is now upstream +//! (apache/datafusion#24384, in DataFusion 55.1.0), so it is no longer the reason to keep +//! the fork. What is left are two performance paths that have not been upstreamed: //! -//! The fix has been submitted upstream: -//! -//! * <https://github.com/apache/datafusion/issues/24383> -//! * <https://github.com/apache/datafusion/pull/24384> +//! * `list_output_lens`, which computes the per-row output lengths of a single `List` column +//! in one pass over the offsets instead of chaining six arrow kernels; +//! * 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`. //! //! # Deleting this file //! -//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this -//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner. -//! -//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different* -//! unnest cleanup — it tracks adopting upstream `unnest_outer` -//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can -//! land in different releases, so closing 5210 is not a signal to delete this fork. +//! Upstream those two paths, then delete this module and go back to +//! `datafusion::physical_plan::unnest::UnnestExec` in the planner. Deleting it before that +//! would regress `explode`, so measure with `native/core/benches/explode.rs` first. //! //! # What was forked //! @@ -51,21 +46,14 @@ //! 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; +//! * the `eq` and `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 -//! 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`. -//! -//! 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. +//! * `find_longest_length` applies the empty-list bump once after the row-wise maximum rather +//! than once per array, which is equivalent because `max` is associative; +//! * `list_output_lens` and the contiguous-run fast path described above. //! -//! 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 -//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics. +//! `ExplodeExec` and `ExplodeStream` were always Comet's own. use arrow::array::{ new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, @@ -77,9 +65,9 @@ use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; use arrow::datatypes::{DataType, Int64Type, SchemaRef}; use arrow::record_batch::RecordBatch; -// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`, +// Upstream imports these as `arrow_ord::cmp::{eq, lt}`; Comet reaches them through `arrow`, // which does not have `arrow_ord` as a direct dependency. -use arrow::compute::kernels::cmp::lt; +use arrow::compute::kernels::cmp::{eq, lt}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{ exec_datafusion_err, exec_err, internal_err, Constraints, HashMap, HashSet, Result, @@ -430,8 +418,10 @@ impl ExplodeStream { } // A chunk can legitimately produce no rows at all (for example rows whose - // arrays are all empty and `preserve_nulls` is false); `build_batch` signals - // that with `None` rather than an empty batch, so move on to the next chunk. + // arrays are all empty under `NullHandling::Drop`, which is plain `explode`); + // `build_batch` signals that with `None` rather than an empty batch, so move on + // to the next chunk. Under the outer handling every row yields at least one row, + // so this cannot happen there. if let Some(batch) = result? { debug_assert!(batch.num_rows() > 0); (&batch).record_output(&self.baseline_metrics); @@ -508,7 +498,7 @@ impl ExplodeStream { // 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::<ListArray>() { - return Ok(Some(list_output_lens(list, self.options.preserve_nulls()))); + return Ok(Some(list_output_lens(list, &self.options))); } } let longest_length = find_longest_length(&list_arrays, &self.options)?; @@ -516,27 +506,36 @@ impl ExplodeStream { } } -/// The per-row unnested length of a single `List` column: the row's list length, or `null_length` -/// for a NULL row. +/// The per-row unnested length of a single `List` column: the row's list length, with NULL and +/// empty rows substituted according to [`datafusion::common::NullHandling`]. /// /// 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<Int64Type> { - let null_length = if preserve_nulls { 1 } else { 0 }; +/// rather than the 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, and a second `eq`/`zip` pair to bump the empty rows. Comet only ever plans `List`, +/// and `explode` plans exactly one of them, so this is the path every non-positional explode +/// takes; `posexplode` unnests two arrays in parallel and uses the general version. +fn list_output_lens(list: &ListArray, options: &UnnestOptions) -> PrimitiveArray<Int64Type> { + // The floor every row's length is raised to. Under `PreserveAndExpandEmpty` an empty row + // expands to one NULL just like a NULL row, including in an array with no validity buffer, + // which can still hold empty rows. + let min_length = if options.expand_empty_as_null() { 1 } else { 0 }; + // `find_longest_length` applies that floor after substituting the NULL length, so it reaches + // substituted NULL rows too. Fold it in rather than leaning on `expand_empty_as_null` + // implying `preserve_nulls`, so the two agree for any future `NullHandling`. + let null_length = cmp::max(if options.preserve_nulls() { 1 } else { 0 }, min_length); + let row_length = |w: &[i32]| cmp::max((w[1] - w[0]) as i64, min_length); 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<i64> = match list.nulls() { - None => offsets.windows(2).map(|w| (w[1] - w[0]) as i64).collect(), + None => offsets.windows(2).map(row_length).collect(), Some(nulls) => offsets .windows(2) .enumerate() .map(|(row, w)| { if nulls.is_valid(row) { - (w[1] - w[0]) as i64 + row_length(w) } else { null_length } @@ -917,24 +916,25 @@ fn build_batch( /// Find the longest list length among the given list arrays for each row. /// -/// For example if we have the following two list arrays: +/// The per-row length of one array is its list length, with NULL and empty rows substituted +/// according to [`datafusion::common::NullHandling`]. For a single array: /// /// ```ignore -/// l1: [1, 2, 3], null, [], [3] -/// l2: [4,5], [], null, [6, 7] -/// ``` +/// l1: [1, 2, 3], null, [], [3] /// -/// If `preserve_nulls` is false, the longest length array will be: -/// -/// ```ignore -/// longest_length: [3, 0, 0, 2] +/// Drop: 3, 0, 0, 1 +/// Preserve: 3, 1, 0, 1 +/// PreserveAndExpandEmpty: 3, 1, 1, 1 /// ``` /// -/// whereas if `preserve_nulls` is true, the longest length array will be: -/// +/// The substitution happens per array, before the row-wise maximum, so a row that is empty in +/// one array still takes the longer length from another: /// /// ```ignore -/// longest_length: [3, 1, 1, 2] +/// l2: [4, 5], [], null, [6, 7] +/// +/// PreserveAndExpandEmpty: 2, 1, 1, 2 +/// longest_length(l1, l2): 3, 1, 1, 2 /// ``` fn find_longest_length(list_arrays: &[ArrayRef], options: &UnnestOptions) -> Result<ArrayRef> { // The length of a NULL list @@ -961,7 +961,18 @@ fn find_longest_length(list_arrays: &[ArrayRef], options: &UnnestOptions) -> Res zip(&is_lt, ¤t, &longest) }, )?; - Ok(longest_length) + + if !options.expand_empty_as_null() { + return Ok(longest_length); + } + // Bump empty lists to length 1 so they produce a single NULL-padded output row. Upstream + // does this per array, inside the map above. `max` is associative and commutative, so + // `max(max(a, 1), max(b, 1)) == max(max(a, b), 1)` and one pass here is equivalent to one + // pass per array. The NULL substitution above has already set NULL rows to the preserved + // length, which is at least 1 whenever this mode is set, so they are not matched. + let zero = Scalar::new(Int64Array::from_value(0, 1)); + let one = Scalar::new(Int64Array::from_value(1, 1)); + Ok(zip(&eq(&longest_length, &zero)?, &one, &longest_length)?) } /// Trait defining common methods used for unnesting, implemented by list array types. @@ -1250,7 +1261,7 @@ fn create_take_indices( /// ``` /// /// then the `unnested_list_arrays` contains the unnest column that will replace `c1` in -/// the final batch if `preserve_nulls` is true: +/// the final batch under [`datafusion::common::NullHandling::Preserve`]: /// /// ```ignore /// c1: 1, null, 2, 3, 4, null, 5, 6 @@ -1306,6 +1317,7 @@ mod tests { use arrow::array::Int32Array; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{Field, Int32Type, Schema}; + use datafusion::common::NullHandling; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; @@ -1343,18 +1355,18 @@ mod tests { async fn explode( input: Vec<RecordBatch>, batch_size: usize, - preserve_nulls: bool, + null_handling: NullHandling, ) -> Result<Vec<RecordBatch>> { let input_schema = input[0].schema(); let source = MemorySourceConfig::try_new_exec(&[input], input_schema, None)?; - explode_child(source, batch_size, preserve_nulls).await + explode_child(source, batch_size, null_handling).await } /// As [`explode`], but over an arbitrary child plan. async fn explode_child( child: Arc<dyn ExecutionPlan>, batch_size: usize, - preserve_nulls: bool, + null_handling: NullHandling, ) -> Result<Vec<RecordBatch>> { let output_schema = Arc::new(Schema::new(vec![Field::new("l", DataType::Int32, true)])); let explode = ExplodeExec::new( @@ -1365,7 +1377,7 @@ mod tests { }], vec![], output_schema, - UnnestOptions::new().with_preserve_nulls(preserve_nulls), + UnnestOptions::new().with_null_handling(null_handling), )?; let task_ctx = Arc::new( TaskContext::default() @@ -1403,7 +1415,7 @@ mod tests { // *several* input rows (2 rows -> 6 rows out; a third would overshoot 8). Using // arrays longer than batch_size would send every row down the oversized-build path // instead, which `single_row_exceeding_batch_size_is_sliced` already covers. - let batches = explode(vec![list_batch(&[Some(3); 10])], 8, true) + let batches = explode(vec![list_batch(&[Some(3); 10])], 8, NullHandling::Preserve) .await .unwrap(); assert_eq!(sizes(&batches), vec![6, 6, 6, 6, 6]); @@ -1415,9 +1427,13 @@ mod tests { // Pins *how* the limit is met, which is what bounds peak memory. 3 rows of 3 // elements at batch_size=4 gives [3, 3, 3] when the input is chunked per row; // building all 9 first and slicing would give [4, 4, 1]. - let batches = explode(vec![list_batch(&[Some(3), Some(3), Some(3)])], 4, true) - .await - .unwrap(); + let batches = explode( + vec![list_batch(&[Some(3), Some(3), Some(3)])], + 4, + NullHandling::Preserve, + ) + .await + .unwrap(); assert_eq!( sizes(&batches), vec![3, 3, 3], @@ -1429,7 +1445,7 @@ mod tests { #[tokio::test] async fn single_row_exceeding_batch_size_is_sliced() { // One row cannot be chunked on the input side, so the oversized build is sliced. - let batches = explode(vec![list_batch(&[Some(25)])], 10, true) + let batches = explode(vec![list_batch(&[Some(25)])], 10, NullHandling::Preserve) .await .unwrap(); assert_eq!(sizes(&batches), vec![10, 10, 5]); @@ -1440,13 +1456,13 @@ mod tests { /// apply regardless of null handling: bounded, non-empty, and totalling `expected_rows`. async fn assert_chunking_matches_whole( lens: &[Option<usize>], - preserve_nulls: bool, + null_handling: NullHandling, expected_rows: usize, ) { - let chunked = explode(vec![list_batch(lens)], 2, preserve_nulls) + let chunked = explode(vec![list_batch(lens)], 2, null_handling) .await .unwrap(); - let whole = explode(vec![list_batch(lens)], 1024, preserve_nulls) + let whole = explode(vec![list_batch(lens)], 1024, null_handling) .await .unwrap(); @@ -1461,17 +1477,27 @@ mod tests { #[tokio::test] async fn chunking_preserves_outer_semantics() { - // With preserve_nulls (Spark's explode_outer, after the planner has rewritten empty - // arrays to NULL), a NULL array yields one NULL row. The per-row counts driving - // chunking must agree, or boundaries drift out of step with the unnesting. - assert_chunking_matches_whole(&[Some(3), None, Some(2), None], true, 7).await; + // Spark's explode_outer: a NULL array and an empty one each yield one NULL row. The + // per-row counts driving chunking must agree, or boundaries drift out of step with the + // unnesting. + assert_chunking_matches_whole( + &[Some(3), None, Some(0), Some(2), None], + NullHandling::PreserveAndExpandEmpty, + 8, + ) + .await; } #[tokio::test] async fn chunking_preserves_non_outer_semantics() { - // Without preserve_nulls (plain explode), NULL arrays produce nothing. Chunks made - // up entirely of such rows must not stall the stream or emit an empty batch. - assert_chunking_matches_whole(&[None, Some(4), None, Some(1)], false, 5).await; + // Plain explode: NULL and empty arrays both produce nothing. Chunks made up entirely of + // such rows must not stall the stream or emit an empty batch. + assert_chunking_matches_whole( + &[None, Some(4), Some(0), None, Some(1)], + NullHandling::Drop, + 5, + ) + .await; } #[tokio::test] @@ -1483,7 +1509,7 @@ mod tests { list_batch(&[Some(1)]), list_batch(&[Some(7), Some(2)]), ]; - let batches = explode(input, 4, true).await.unwrap(); + let batches = explode(input, 4, NullHandling::Preserve).await.unwrap(); let sizes = sizes(&batches); assert!( sizes.iter().all(|s| *s <= 4), @@ -1542,7 +1568,7 @@ mod tests { depth: 1, }], struct_column_indices: HashSet::new(), - options: UnnestOptions::new().with_preserve_nulls(true), + options: UnnestOptions::new(), baseline_metrics: BaselineMetrics::new(&metrics, 0), input_batches: MetricBuilder::new(&metrics).counter("input_batches", 0), input_rows: MetricBuilder::new(&metrics).counter("input_rows", 0), @@ -1605,7 +1631,7 @@ mod tests { }], vec![], output_schema, - UnnestOptions::new().with_preserve_nulls(true), + UnnestOptions::new(), ) .unwrap(); @@ -1679,16 +1705,26 @@ mod tests { #[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. + // Row 1 is NULL and row 2 is empty, and neither holds elements. let list = int_list( vec![0, 2, 2, 2, 5], vec![1, 2, 3, 4, 5], Some(vec![true, false, true, true]), ); + + // Plain `explode` drops both, so nothing is padded and the run stays unbroken. 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::<Vec<_>>()); + + // `explode_outer` pads each of them to one NULL, which no slice of the child can + // produce, so the same array must now take the gather. + let (values, fast) = unnest(&list, vec![2, 1, 1, 3]); + assert!(!fast, "padded rows must break the run"); + assert_eq!( + values, + vec![Some(1), Some(2), None, None, Some(3), Some(4), Some(5)] + ); } #[test] @@ -1745,41 +1781,100 @@ mod tests { // --------------------------------------------------------------------------------------- /// `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::new().with_preserve_nulls(preserve_nulls); - let arrays = vec![Arc::new(list.clone()) as ArrayRef]; - let expected = find_longest_length(&arrays, &options).unwrap(); - let expected = expected.as_primitive::<Int64Type>(); - 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"); + /// chunking in `ExplodeStream` and the unnesting in `build_batch` both consume it. Both are + /// also pinned against `expected`, given per [`NullHandling`] in the order below, so that a + /// mistake made in both at once cannot pass. + fn assert_output_lens(list: ListArray, expected_per_mode: [&[i64]; 3]) { + let modes = [ + NullHandling::Drop, + NullHandling::Preserve, + NullHandling::PreserveAndExpandEmpty, + ]; + for (null_handling, expected) in modes.into_iter().zip(expected_per_mode) { + let options = UnnestOptions::new().with_null_handling(null_handling); + let arrays = vec![Arc::new(list.clone()) as ArrayRef]; + let general = find_longest_length(&arrays, &options).unwrap(); + let general = general.as_primitive::<Int64Type>(); + let fused = list_output_lens(&list, &options); + + assert_eq!( + &general.values()[..], + expected, + "general, {null_handling:?}" + ); + assert_eq!(&fused.values()[..], expected, "fused, {null_handling:?}"); + // `create_take_indices` reads the values buffer directly and relies on this. + assert_eq!(general.null_count(), 0, "lengths must never be NULL"); + assert_eq!(fused.null_count(), 0, "lengths must never be NULL"); + } } + /// `find_longest_length` bumps empty rows once after the row-wise maximum rather than once + /// per array. That is only observably different with more than one array, which is the + /// `posexplode` shape, so pin the combined result directly. #[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); + fn longest_length_combines_arrays_before_the_empty_bump() { + // Row 0 is non-empty in both. Row 1 is empty in l1 but 2 long in l2, so the bump must + // not inflate it past the other array. Row 2 is empty in both, the only row where + // `PreserveAndExpandEmpty` differs from `Preserve`. Row 3 is empty in l1 and NULL in l2. + let l1 = int_list(vec![0, 2, 2, 2, 2], vec![1, 2], None); + let l2 = int_list( + vec![0, 1, 3, 3, 3], + vec![4, 5, 6], + Some(vec![true, true, true, false]), + ); + let arrays = vec![Arc::new(l1) as ArrayRef, Arc::new(l2) as ArrayRef]; + + for (null_handling, expected) in [ + (NullHandling::Drop, [2, 2, 0, 0]), + (NullHandling::Preserve, [2, 2, 0, 1]), + (NullHandling::PreserveAndExpandEmpty, [2, 2, 1, 1]), + ] { + let options = UnnestOptions::new().with_null_handling(null_handling); + let longest = find_longest_length(&arrays, &options).unwrap(); + let longest = longest.as_primitive::<Int64Type>(); + assert_eq!(&longest.values()[..], &expected, "{null_handling:?}"); + assert_eq!(longest.null_count(), 0, "lengths must never be NULL"); + } + } - let with_nulls = int_list( - vec![0, 2, 2, 2, 5], - vec![1, 2, 3, 4, 5], - Some(vec![true, false, true, true]), + #[test] + fn output_lens_substitute_per_null_handling() { + // Row 2 is empty with no validity buffer at all, which is the case + // `PreserveAndExpandEmpty` must still bump. + assert_output_lens( + int_list(vec![0, 3, 4, 4, 6], vec![1, 2, 3, 4, 5, 6], None), + [&[3, 1, 0, 2], &[3, 1, 0, 2], &[3, 1, 1, 2]], + ); + + // Row 1 is NULL and row 2 is an empty non-null row, so the two substitutions are + // distinguishable: only `PreserveAndExpandEmpty` bumps both. + assert_output_lens( + int_list( + vec![0, 2, 2, 2, 5], + vec![1, 2, 3, 4, 5], + Some(vec![true, false, true, true]), + ), + [&[2, 0, 0, 3], &[2, 1, 0, 3], &[2, 1, 1, 3]], ); - 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); + assert_output_lens(int_list(vec![0], vec![], None), [&[], &[], &[]]); } #[test] - fn fused_lengths_handle_a_sliced_input() { + fn output_lens_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); + assert_output_lens(list.slice(1, 3), [&[1, 0, 3], &[1, 0, 3], &[1, 1, 3]]); + + // `ListArray::slice` slices the validity buffer alongside the offsets, so the row index + // the fused pass hands to `is_valid` has to be in the sliced row space, not the + // original. Rows 1..4 of this array are NULL, empty, non-empty. + let with_nulls = int_list( + vec![0, 2, 4, 4, 7], + vec![1, 2, 3, 4, 5, 6, 7], + Some(vec![true, false, true, true]), + ); + assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5406df4f04..52d5cc5266 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -39,7 +39,6 @@ use crate::execution::operators::IcebergScanExec; use crate::execution::operators::IcebergWriteExec; use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind}; use crate::execution::{ - expressions::list_empty_to_null::ListEmptyToNullExpr, expressions::list_positions::ListPositionsExpr, expressions::subquery::Subquery, operators::{ @@ -127,7 +126,7 @@ use arrow::array::{ use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; use arrow::row::{OwnedRow, RowConverter, SortField}; use datafusion::common::utils::SingleRowListArrayBuilder; -use datafusion::common::UnnestOptions; +use datafusion::common::{NullHandling, UnnestOptions}; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::joins::NestedLoopJoinExec; use datafusion::physical_plan::limit::GlobalLimitExec; @@ -2058,7 +2057,7 @@ impl PhysicalPlanner { self.create_plan(&children[0], inputs, partition_count)?; // Create the expression for the array to explode - let raw_child_expr = if let Some(child_expr) = &explode.child { + let child_expr = if let Some(child_expr) = &explode.child { self.create_expr(child_expr, child.schema())? } else { return Err(ExecutionError::GeneralError( @@ -2067,27 +2066,12 @@ impl PhysicalPlanner { }; let child_schema = child.schema(); - let child_field_name = raw_child_expr + let child_field_name = child_expr .return_field(&child_schema) .expect("Failed to get field from child expression") .name() .to_string(); - // Bridge Spark's outer semantics: DataFusion's `UnnestExec` with - // `preserve_nulls = true` emits one null row for a NULL list but drops rows - // whose list is empty. Spark's `explode_outer`/`posexplode_outer` must emit - // exactly one null row in both cases, so we mark empty rows as null before - // unnesting. See https://github.com/apache/datafusion/issues/19053. Once - // Comet moves to a DataFusion release carrying - // https://github.com/apache/datafusion/pull/22100, `ListEmptyToNullExpr` - // can be removed in favor of `NullHandling::PreserveAndExpandEmpty`. See - // https://github.com/apache/datafusion-comet/issues/5210. - let child_expr: Arc<dyn PhysicalExpr> = if explode.outer { - Arc::new(ListEmptyToNullExpr::new(raw_child_expr)) - } else { - raw_child_expr - }; - // Both posexplode variants reference the array twice: once for positions // and once for values. Materialize computed arrays so both references // share one evaluation. A plain Column is already materialized. @@ -2202,9 +2186,17 @@ impl PhysicalPlanner { depth: 1, }); - let unnest_options = UnnestOptions::new().with_preserve_nulls(explode.outer); + // Spark's `explode_outer`/`posexplode_outer` emit exactly one null row for both + // a NULL array and an empty one, which is `PreserveAndExpandEmpty`. Plain + // `explode` drops both. + let null_handling = if explode.outer { + NullHandling::PreserveAndExpandEmpty + } else { + NullHandling::Drop + }; + let unnest_options = UnnestOptions::new().with_null_handling(null_handling); - // Comet's batch-size-respecting fork of `UnnestExec`; see `operators::explode`. + // Comet's specialized fork of `UnnestExec`; see `operators::explode`. let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, @@ -6385,9 +6377,13 @@ mod tests { if computed { 2 } else { 0 }, "{context}" ); + // The array is pre-projected only to share one evaluation between the + // `pos` and `value` references, so only a computed child needs it. `outer` + // does not, since it is now a `UnnestOptions` mode rather than a wrapper + // expression around the child. assert_eq!( projections, - 1 + usize::from(position && (outer || computed)), + 1 + usize::from(position && computed), "{context}" ); let expected_values = if outer { diff --git a/spark/src/test/resources/sql-tests/expressions/array/explode.sql b/spark/src/test/resources/sql-tests/expressions/array/explode.sql index 68a2bcc943..b1bdb8b749 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/explode.sql @@ -304,16 +304,16 @@ INSERT INTO test_explode_multi VALUES query SELECT id, name, extra, explode_outer(arr) AS v FROM test_explode_multi --- ===== Pre-projection wiring: carry the array column through alongside its +-- ===== Projection wiring: carry the array column through alongside its -- explosion. The passthrough `arr` shows the original array (empty rows stay --- []) while the exploded value is NULL for empty rows, so this is the only --- shape where the difference between the original array and the null-marked --- copy is observable at the query level. +-- []) while the exploded value is NULL for those same rows, so this is the +-- shape that shows the outer substitution applying to the generated column +-- only. query SELECT id, arr, explode_outer(arr) FROM test_explode_int --- ===== Pre-projection wiring: no passthrough columns. This drives the +-- ===== Projection wiring: no passthrough columns. This drives the -- planner's `project_list` to empty (no columns carried through) and covers -- the codepath where the second projection contains only the exploded array. diff --git a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql index cca5650b05..eb537586f0 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql @@ -15,8 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- posexplode_outer is now supported natively; see DataFusion #19053 handling --- via ListEmptyToNullExpr in the planner. +-- posexplode_outer runs natively: the planner asks DataFusion's unnest for +-- NullHandling::PreserveAndExpandEmpty, which matches Spark's outer semantics. statement CREATE TABLE test_posexplode_int(id int, arr array<int>) USING parquet @@ -168,13 +168,12 @@ query SELECT id, pos, value FROM test_posexplode_int LATERAL VIEW OUTER posexplode(arr) p AS pos, value --- ===== Pre-projection wiring for posexplode_outer ===== +-- ===== Projection wiring for posexplode_outer ===== -- Carry the array column through alongside posexplode_outer. The passthrough -- `arr` shows the original array (empty rows stay []) while the exploded pos --- and value are NULL for empty rows, so this is the only shape where the --- difference between the original array and the null-marked copy is --- observable at the query level. +-- and value are NULL for those same rows, so this is the shape that shows the +-- outer substitution applying to the generated columns only. query SELECT id, arr, posexplode_outer(arr) FROM test_posexplode_int @@ -184,8 +183,8 @@ SELECT id, arr, posexplode_outer(arr) FROM test_posexplode_int query SELECT posexplode_outer(arr) FROM test_posexplode_int --- posexplode_outer batch of only-empty arrays exercises the slow path with an --- all-zeros non-empty bitmap; only-null exercises the fast-path passthrough. +-- A batch of only-empty arrays and a batch of only-null arrays: every row is +-- substituted, so the whole output batch comes from padding. statement CREATE TABLE test_posexplode_all_empty(id int, arr array<int>) USING parquet diff --git a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index bf35b8d3fd..4cb6a7f7ef 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -463,9 +463,9 @@ class CometGenerateExecSuite extends CometTestBase { } test("explode_outer across batch boundary with mixed empty/null rows") { - // Mix null, empty, and non-empty rows and force multiple small batches so that - // `ListEmptyToNullExpr` runs on each batch and its fast/slow path split is exercised - // more than once with different offset patterns. + // Mix null, empty, and non-empty rows and force multiple small batches so that the + // per-row output lengths are recomputed on each batch with a different offset pattern, + // and so that some chunk boundaries fall on a substituted row. withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", @@ -487,8 +487,8 @@ class CometGenerateExecSuite extends CometTestBase { test("posexplode_outer across batch boundary with mixed empty/null rows") { // Same shape as the explode_outer counterpart but exercises the parallel positions - // branch. With the pre-projection introduced for outer, `ListEmptyToNullExpr` runs once - // per batch and both branches share the same materialized array. + // branch, where the `pos` and `value` arrays are unnested together and must be padded + // to the same per-row length. withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", @@ -578,9 +578,8 @@ class CometGenerateExecSuite extends CometTestBase { } test("explode_outer over limit with offset") { - // Exercises `ListEmptyToNullExpr` on a sliced input with a non-zero offset base. The - // helper preserves the base offset and passes it through unchanged, so the fix in - // `ListPositionsExpr` is what actually keeps the parallel `pos` branch safe. This test + // Exercises the outer path on a sliced input with a non-zero offset base, which is what + // the fix in `ListPositionsExpr` keeps the parallel `pos` branch safe against. This test // covers the `explode_outer` shape without the `pos` branch. withSQLConf( "spark.sql.adaptive.enabled" -> "false", --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
