diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 3023cfaa774..d34a0b61071 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -23,18 +23,18 @@ ## explode -- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). +- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`. ## 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, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map inputs fall back. +- Same `CometExplodeExec` path as `explode`. Compatible for array and map inputs; empty and NULL collections both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map entries reuse the list unnest path and expand into key/value columns. ## posexplode -- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). +- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`. ## 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, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. +- Same `CometExplodeExec` path as `posexplode`. Compatible for array and map inputs; empty and NULL collections both emit one row with null `pos` and null generated columns 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/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f9292b27343..a649bc069c1 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -337,8 +337,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci ## generator_funcs `explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via -`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back -to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via +`CometExplodeExec` for array and map inputs (operator-level, not expression-level). Enabled by default via `spark.comet.exec.explode.enabled`. | Function | Status | Implementation | Notes | diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 334d34673f2..e80f99ebc1d 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -124,12 +124,17 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx ## Generators and set operations -| Operator | Status | Notes | -| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- | -| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | ------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays and maps. `inline` / `stack` fall back. | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | + +Map generators followed by another generator through projections or filters stay on Spark's fused +generator path. Native generator chains materialize carried columns between generators for both +arrays and maps. Array chains retain their native support. A single map generator, including maps +with nested values, remains eligible for native execution. An exchange separates generator chains. ## Writes diff --git a/native/core/src/execution/expressions/map_entries.rs b/native/core/src/execution/expressions/map_entries.rs new file mode 100644 index 00000000000..8188a80730d --- /dev/null +++ b/native/core/src/execution/expressions/map_entries.rs @@ -0,0 +1,119 @@ +// 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, ListArray, MapArray, RecordBatch}; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ColumnarValue; + +/// Exposes a map's entries as a list for `ExplodeExec`, sharing the input buffers. +/// Preserve the original entry fields, including nullability and metadata; the +/// SQL `map_entries` function rebuilds those fields and loses their metadata. +#[derive(Debug, Clone)] +pub struct MapEntriesExpr { + child: Arc, +} + +impl MapEntriesExpr { + pub fn new(child: Arc) -> Self { + Self { child } + } +} + +impl Display for MapEntriesExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "map_entries({})", self.child) + } +} + +impl PartialEq for MapEntriesExpr { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) + } +} + +impl Eq for MapEntriesExpr {} + +impl Hash for MapEntriesExpr { + fn hash(&self, state: &mut H) { + self.child.hash(state); + } +} + +impl PhysicalExpr for MapEntriesExpr { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn return_field(&self, input_schema: &Schema) -> DataFusionResult { + let field = self.child.return_field(input_schema)?; + let DataType::Map(entries, _) = field.data_type() else { + return exec_err!( + "MapEntriesExpr expected Map input, got {}", + field.data_type() + ); + }; + Ok(Arc::new( + field + .as_ref() + .clone() + .with_data_type(DataType::List(Arc::clone(entries))), + )) + } + + fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult { + let array = self.child.evaluate(batch)?.into_array(batch.num_rows())?; + let Some(map) = array.as_any().downcast_ref::() else { + return exec_err!( + "MapEntriesExpr expected Map input, got {}", + array.data_type() + ); + }; + let DataType::Map(entries, _) = map.data_type() else { + unreachable!("MapArray downcast guarantees DataType::Map"); + }; + let list = ListArray::try_new( + Arc::clone(entries), + map.offsets().clone(), + Arc::new(map.entries().clone()), + map.nulls().cloned(), + )?; + Ok(ColumnarValue::Array(Arc::new(list))) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return exec_err!( + "MapEntriesExpr expects exactly 1 child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } +} diff --git a/native/core/src/execution/expressions/mod.rs b/native/core/src/execution/expressions/mod.rs index ad5c06a93b5..b180cf41bc0 100644 --- a/native/core/src/execution/expressions/mod.rs +++ b/native/core/src/execution/expressions/mod.rs @@ -21,6 +21,7 @@ pub mod arithmetic; pub mod bitwise; pub mod comparison; pub mod logical; +pub mod map_entries; pub mod nullcheck; pub mod partition; pub mod random; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a47c19d1a23..6e1fff62ce5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5008,11 +5008,23 @@ mod tests { #[tokio::test] async fn explode_evaluates_array_once_per_batch() { + check_explode_evaluates_collection_once_per_batch(false).await; + } + + #[tokio::test] + async fn explode_evaluates_map_once_per_batch() { + check_explode_evaluates_collection_once_per_batch(true).await; + } + + async fn check_explode_evaluates_collection_once_per_batch(map: bool) { + use arrow::array::{AsArray, MapArray, StructArray}; use arrow::datatypes::Int32Type; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::logical_expr::{create_udf, Volatility}; use datafusion::physical_plan::projection::ProjectionExec; - use spark_expression::data_type::{data_type_info::DatatypeStruct, DataTypeInfo, ListInfo}; + use spark_expression::data_type::{ + data_type_info::DatatypeStruct, DataTypeInfo, ListInfo, MapInfo, + }; let array_type = spark_expression::DataType { type_id: 14, @@ -5031,6 +5043,55 @@ mod tests { Some(vec![Some(20)]), ])) as ArrayRef; + let (array_type, arrays) = if map { + // Slice away a leading entry to exercise non-zero map offsets. + let values = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(99)]), + Some(vec![Some(10), None]), + Some(vec![]), + None, + Some(vec![Some(20)]), + ]) + .slice(1, 4); + let fields: Fields = vec![ + Field::new("key", DataType::Int32, false) + .with_metadata([("PARQUET:field_id".to_string(), "10".to_string())].into()), + Field::new("value", DataType::Int32, true) + .with_metadata([("PARQUET:field_id".to_string(), "11".to_string())].into()), + ] + .into(); + let entries = StructArray::new( + fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![99, 1, 2, 3])), + Arc::clone(values.values()), + ], + None, + ); + let maps = MapArray::new( + Arc::new(Field::new("entries", DataType::Struct(fields), false)), + values.offsets().clone(), + entries, + values.nulls().cloned(), + false, + ); + let map_type = spark_expression::DataType { + type_id: 15, + type_info: Some(Box::new(DataTypeInfo { + datatype_struct: Some(DatatypeStruct::Map(Box::new(MapInfo { + key_type: Some(Box::new(create_proto_datatype())), + value_type: Some(Box::new(create_proto_datatype())), + value_contains_null: true, + key_field_id: Some(10), + value_field_id: Some(11), + }))), + })), + }; + (map_type, Arc::new(maps) as ArrayRef) + } else { + (array_type, arrays) + }; + for outer in [false, true] { for position in [false, true] { for computed in [false, true] { @@ -5109,20 +5170,21 @@ mod tests { let results = collect(native_plan.execute(0, task_ctx).unwrap()) .await .unwrap(); - let context = - format!("outer={outer}, position={position}, computed={computed}"); + let context = format!( + "outer={outer}, position={position}, computed={computed}, map={map}" + ); assert_eq!( calls.load(Ordering::Relaxed), 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` + // `pos` and `value` references, so a computed child or map wrapper 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 && computed), + 1 + usize::from(position && (computed || map)), "{context}" ); let expected_values = if outer { @@ -5142,6 +5204,38 @@ mod tests { }) .collect(); assert_eq!(values, expected_values.repeat(2), "{context}"); + if map { + let schema = native_plan.schema(); + for (index, id) in [ + (schema.fields().len() - 2, "10"), + (schema.fields().len() - 1, "11"), + ] { + assert_eq!( + schema + .field(index) + .metadata() + .get("PARQUET:field_id") + .map(String::as_str), + Some(id) + ); + } + assert_eq!(schema.field(schema.fields().len() - 2).is_nullable(), outer); + let expected_keys = if outer { + vec![Some(1), Some(2), None, None, Some(3)] + } else { + vec![Some(1), Some(2), Some(3)] + }; + let keys: Vec<_> = results + .iter() + .flat_map(|batch| { + batch + .column(batch.num_columns() - 2) + .as_primitive::() + .iter() + }) + .collect(); + assert_eq!(keys, expected_keys.repeat(2), "{context}"); + } if position { let expected_positions = if outer { vec![Some(0), Some(1), None, None, Some(0)] diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 1159f242a53..0f25bc0a8d4 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -33,7 +33,9 @@ use jni::objects::{Global, JObject}; use super::{operator_registry::OperatorBuilder, PhysicalPlanner, PlanCreationResult}; use crate::{ - execution::{operators::ExecutionError, spark_plan::SparkPlan}, + execution::{ + expressions::map_entries::MapEntriesExpr, operators::ExecutionError, spark_plan::SparkPlan, + }, extract_op, }; @@ -55,7 +57,7 @@ impl OperatorBuilder for ExplodeBuilder { let (scans, shuffle_scans, child) = planner.create_plan(&children[0], inputs, partition_count)?; - // Create the expression for the array to explode + // Create the expression for the collection to explode let child_expr = if let Some(child_expr) = &explode.child { planner.create_expr(child_expr, child.schema())? } else { @@ -71,6 +73,26 @@ impl OperatorBuilder for ExplodeBuilder { .name() .to_string(); + // Expose maps as List>, sharing their entries buffers. + // Reuse the list path for outer rows, positions and bounded output batches, + // then flatten the entry struct into Spark's two output columns. + let map_fields = match child_expr.data_type(&child_schema)? { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) => Some(fields.clone()), + _ => { + return Err(ExecutionError::GeneralError( + "Map entries must be a struct".to_string(), + )); + } + }, + _ => None, + }; + let child_expr: Arc = if map_fields.is_some() { + Arc::new(MapEntriesExpr::new(child_expr)) + } else { + 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. @@ -162,11 +184,22 @@ impl OperatorBuilder for ExplodeBuilder { } }; - output_fields.push(Field::new( - array_field.name(), - element_type, - true, // Element is nullable after unnesting - )); + let struct_unnests = if let Some(fields) = map_fields { + output_fields.extend(fields.iter().map(|field| { + field + .as_ref() + .clone() + .with_nullable(explode.outer || field.is_nullable()) + })); + vec![array_input_index] + } else { + output_fields.push(Field::new( + array_field.name(), + element_type, + true, // Element is nullable after unnesting + )); + vec![] + }; let output_schema = Arc::new(Schema::new(output_fields)); @@ -196,7 +229,7 @@ impl OperatorBuilder for ExplodeBuilder { let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_unnests, output_schema, unnest_options, )?); diff --git a/native/operators/benches/explode.rs b/native/operators/benches/explode.rs index ef14e722548..829b6152325 100644 --- a/native/operators/benches/explode.rs +++ b/native/operators/benches/explode.rs @@ -28,7 +28,9 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, Int32Array, Int64Array, ListArray, StringArray, StructArray}; +use arrow::array::{ + Array, ArrayRef, AsArray, Int32Array, Int64Array, ListArray, StringArray, StructArray, +}; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -40,6 +42,7 @@ use datafusion::physical_plan::unnest::ListUnnest; use datafusion::physical_plan::{common::collect, ExecutionPlan}; use datafusion::prelude::SessionConfig; use datafusion_comet_operators::ExplodeExec; +use futures::StreamExt; use tokio::runtime::Runtime; /// Input rows per batch, and batches per run. 8192 is DataFusion's default `batch_size`, so the @@ -264,6 +267,61 @@ fn run(runtime: &Runtime, plan: &Arc, ctx: &Arc) assert!(!batches.is_empty()); } +fn wide_carried_plan() -> Arc { + let values = Arc::new(Int64Array::from_iter_values( + 0..(ROWS_PER_BATCH * 300) as i64, + )); + let list = ListArray::new( + Arc::new(Field::new("item", DataType::Int64, false)), + OffsetBuffer::from_lengths(std::iter::repeat_n(300, ROWS_PER_BATCH)), + values, + None, + ); + let payload = "x".repeat(1024); + let batch = RecordBatch::try_from_iter(vec![ + ("l", Arc::new(list) as ArrayRef), + ( + "payload", + Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + payload.as_str(), + ROWS_PER_BATCH, + ))) as ArrayRef, + ), + ]) + .unwrap(); + let source = + MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap(); + Arc::new( + ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + Arc::new(Schema::new(vec![ + Field::new("l", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])), + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(), + ) +} + +fn run_wide(runtime: &Runtime, plan: &Arc, ctx: &Arc) -> usize { + runtime.block_on(async { + let mut stream = plan.execute(0, Arc::clone(ctx)).unwrap(); + let mut bytes = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + bytes += batch.column(1).as_string::().value_data().len(); + std::hint::black_box(batch); + } + bytes + }) +} + fn criterion_benchmark(c: &mut Criterion) { let runtime = Runtime::new().unwrap(); let ctx = Arc::new( @@ -320,6 +378,12 @@ fn criterion_benchmark(c: &mut Criterion) { } } group.finish(); + + let plan = wide_carried_plan(); + assert_eq!(run_wide(&runtime, &plan, &ctx), ROWS_PER_BATCH * 300 * 1024); + c.bench_function("explode_1kib_carried_fan_out_300", |b| { + b.iter(|| run_wide(&runtime, &plan, &ctx)) + }); } criterion_group!(benches, criterion_benchmark); diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 6aae0e2d349..5134705faee 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! A fork of DataFusion's `UnnestExec`, kept for two unnesting kernels Comet has specialized. +//! A fork of DataFusion's `UnnestExec`, kept for Comet's specialized unnesting kernels. //! //! # Why this fork exists //! @@ -23,16 +23,20 @@ //! 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 fork. What is left are performance improvements that have not been upstreamed: //! -//! * `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; +//! * `list_output_lens`, which computes per-row output lengths in one pass for `List` columns +//! sharing offsets and validity, instead of chaining 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`. +//! * splitting individual List rows before repeating carried values, including a repetition +//! bound that prevents copied variable-width buffers and child offsets from overflowing, and +//! trimming list children so retained slices do not inflate Arrow take's preallocation. +//! * skipping gather indices and placeholders when every column is being unnested. //! //! # Deleting this file //! -//! Upstream those two paths, then delete this module and go back to +//! Upstream these improvements, 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/operators/benches/explode.rs` first. //! @@ -51,15 +55,16 @@ //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; //! * `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. +//! * the performance improvements described above. //! //! `ExplodeExec` and `ExplodeStream` were always Comet's own. use arrow::array::{ - new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, - LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, - StructArray, + new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, + GenericListArray, Int64Array, LargeListArray, LargeListViewArray, ListArray, ListViewArray, + MapArray, OffsetSizeTrait, PrimitiveArray, Scalar, StructArray, }; +use arrow::buffer::OffsetBuffer; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; @@ -276,10 +281,8 @@ impl ExecutionPlan for ExplodeExec { pending_input: None, }); - // Chunking the input bounds each build to roughly `batch_size` rows, but two cases can - // still produce an oversized batch (see `predict_output_lens`), so the output goes - // through DataFusion's splitter to make the bound unconditional. Note this only bounds - // the emitted row count; bounding peak memory is the job of the input chunking. + // Depth-1 List inputs are split before materialization, including within one row. + // Keep the output splitter for the generic shapes described in `predict_output_lens`. Ok(Box::pin(BatchSplitStream::new( stream, batch_size, @@ -302,6 +305,10 @@ struct PendingInput { batch: RecordBatch, /// Index of the next input row to unnest. row_offset: usize, + /// Elements already emitted from the current input row. + element_offset: usize, + /// Repetition safe for every slice of this input. None for unsupported list shapes. + safe_repetitions: Option, /// How many output rows each input row expands into, indexed by input row. /// /// `None` when the expansion factor cannot be predicted from the input alone, in which @@ -318,10 +325,7 @@ impl PendingInput { /// How many input rows to unnest next so the resulting batch holds at most `batch_size` /// rows. /// - /// Always returns at least 1 while rows remain, so the stream always makes progress: a - /// single input row is never split across output batches, so one row whose array is - /// longer than `batch_size` still produces one oversized build, which `BatchSplitStream` - /// slices down on the way out. + /// Oversized List rows are handled separately by `next_chunk`. fn next_chunk_rows(&self, batch_size: usize) -> usize { let Some(output_lens) = &self.output_lens else { return self.remaining_rows(); @@ -346,6 +350,136 @@ impl PendingInput { .as_ref() .map(|lens| lens.slice(self.row_offset, rows)) } + + fn next_chunk( + &mut self, + batch_size: usize, + list_columns: &[ListUnnest], + ) -> Result<(RecordBatch, Option)> { + let mut rows = self.next_chunk_rows(batch_size); + let mut chunk = self.batch.slice(self.row_offset, rows); + let mut lengths = self.chunk_lengths(rows); + if let Some(safe_repetitions) = self.safe_repetitions { + // Check the selected rows, not buffers retained from the rest of the input. + // A chunk's largest repetition times its offset span bounds every copied offset. + // The input-wide bound only skips checks; it never shrinks a chunk. + let limit = if safe_repetitions >= batch_size { + batch_size + } else { + loop { + let limit = batch_size.min(repetition_limit(&chunk, list_columns)); + let max_repeat = *lengths.as_ref().unwrap().values().iter().max().unwrap(); + if rows == 1 || max_repeat as usize <= limit { + break limit; + } + rows = rows.div_ceil(2); + chunk = self.batch.slice(self.row_offset, rows); + lengths = self.chunk_lengths(rows); + } + }; + let row_length = self.output_lens.as_ref().unwrap().value(self.row_offset) as usize; + if row_length > limit { + let length = limit.min(row_length - self.element_offset); + let mut columns = chunk.columns().to_vec(); + for unnest in list_columns { + let list = columns[unnest.index_in_input_schema].as_list::(); + let DataType::List(field) = list.data_type() else { + unreachable!(); + }; + let available = if list.is_null(0) { + 0 + } else { + list.value_length(0) as usize + }; + let skip = self.element_offset.min(available); + let start = list.value_offsets()[0] + skip as i32; + let end = start + length.min(available - skip) as i32; + // Narrow the entries and parallel positions together, preserving their + // original values and sharing their buffers. Shorter lists still pad. + columns[unnest.index_in_input_schema] = Arc::new(ListArray::new( + Arc::clone(field), + OffsetBuffer::new(vec![start, end].into()), + Arc::clone(list.values()), + list.nulls().cloned(), + )); + } + self.element_offset += length; + if self.element_offset == row_length { + self.row_offset += 1; + self.element_offset = 0; + } + return Ok(( + RecordBatch::try_new(chunk.schema(), columns)?, + Some(Int64Array::from(vec![length as i64])), + )); + } + } + self.row_offset += rows; + Ok((chunk, lengths)) + } +} + +/// Bound repetition by the visible spans addressed by 32-bit offsets. Fixed-width +/// buffers do not constrain these offsets; a sliced list still needs its selected child count +/// even when that child has no buffers, as with List. +fn repetition_limit(batch: &RecordBatch, list_columns: &[ListUnnest]) -> usize { + fn max_offset_span(array: &dyn Array) -> usize { + match array.data_type() { + DataType::Utf8 => { + let offsets = array.as_string::().value_offsets(); + (offsets[array.len()] - offsets[0]) as usize + } + DataType::Binary => { + let offsets = array.as_binary::().value_offsets(); + (offsets[array.len()] - offsets[0]) as usize + } + DataType::List(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let length = offsets[list.len()] as usize - start; + length.max(max_offset_span(list.values().slice(start, length).as_ref())) + } + DataType::Map(_, _) => { + let map = array.as_map(); + let offsets = map.value_offsets(); + let start = offsets[0] as usize; + let length = offsets[map.len()] as usize - start; + length.max(max_offset_span(&map.entries().slice(start, length))) + } + DataType::Struct(_) => array + .as_struct() + .columns() + .iter() + .map(|column| max_offset_span(column.as_ref())) + .max() + .unwrap_or(0), + DataType::LargeList(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0] as usize; + let length = offsets[list.len()] as usize - start; + max_offset_span(list.values().slice(start, length).as_ref()) + } + DataType::FixedSizeList(_, _) => { + max_offset_span(array.as_fixed_size_list().values().as_ref()) + } + _ => 0, + } + } + let span = batch + .columns() + .iter() + .enumerate() + .filter(|(index, _)| { + !list_columns + .iter() + .any(|list| list.index_in_input_schema == *index) + }) + .map(|(_, column)| max_offset_span(column.as_ref())) + .max() + .unwrap_or(0); + (i32::MAX as usize / span.max(1)).max(1) } /// A stream that unnests its input, bounding output batches to `batch_size` rows. @@ -364,10 +498,8 @@ struct ExplodeStream { /// batch can produce arbitrarily many output rows, so the input is consumed in chunks /// small enough that each chunk's output stays near `batch_size`. /// - /// Note the scope of the memory bound this buys: chunking removes the input batch size - /// from the peak, but not the length of an individual list. A single row whose list is - /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be - /// predicted up front), both still materialize their full expansion in one build. + /// Comet's depth-1 List rows are also split within a row before carried values repeat. + /// Generic recursive unnesting still materializes its expansion before output splitting. pending_input: Option, } @@ -390,14 +522,11 @@ impl ExplodeStream { loop { // Unnest the next chunk of the input batch already in hand. if let Some(pending) = self.pending_input.as_mut() { - // `PendingInput` is only built from a non-empty batch and `next_chunk_rows` - // always consumes at least one row, so it is dropped the moment it drains. + // Pending input is non-empty; each chunk advances a row or its element cursor. debug_assert!(pending.remaining_rows() > 0); - let rows = pending.next_chunk_rows(self.batch_size); - let chunk = pending.batch.slice(pending.row_offset, rows); - let chunk_lengths = pending.chunk_lengths(rows); - pending.row_offset += rows; + let (chunk, chunk_lengths) = + pending.next_chunk(self.batch_size, &self.list_type_columns)?; let drained = pending.remaining_rows() == 0; let timer = self.baseline_metrics.elapsed_compute().timer(); @@ -438,10 +567,21 @@ impl ExplodeStream { if batch.num_rows() > 0 { let timer = self.baseline_metrics.elapsed_compute().timer(); let output_lens = self.predict_output_lens(&batch); + let safe_repetitions = (!self.list_type_columns.is_empty() + && self.list_type_columns.iter().all(|unnest| { + unnest.depth == 1 + && matches!( + batch.column(unnest.index_in_input_schema).data_type(), + DataType::List(_) + ) + })) + .then(|| repetition_limit(&batch, &self.list_type_columns)); timer.done(); self.pending_input = Some(PendingInput { batch, row_offset: 0, + element_offset: 0, + safe_repetitions, output_lens: output_lens?, }); } @@ -496,8 +636,20 @@ 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::() { + if let Some(list) = list_arrays[0].as_any().downcast_ref::() { + // Shared offset and validity views guarantee identical per-row lengths. + // Check validity's bit offset too; bitmap bytes alone do not suffice. + if list_arrays.iter().skip(1).all(|other| { + let Some(other) = other.as_any().downcast_ref::() else { + return false; + }; + list.offsets().ptr_eq(other.offsets()) + && match (list.nulls(), other.nulls()) { + (None, None) => true, + (Some(a), Some(b)) => a.inner().ptr_eq(b.inner()), + _ => false, + } + }) { return Ok(Some(list_output_lens(list, &self.options))); } } @@ -513,8 +665,8 @@ impl ExplodeStream { /// 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. +/// and parallel lists sharing offsets and validity have the same per-row length, so one +/// pass also covers `posexplode` when its position and value lists share those views. fn list_output_lens(list: &ListArray, options: &UnnestOptions) -> PrimitiveArray { // 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, @@ -683,8 +835,6 @@ fn list_unnest_at_level( let unnested_temp_arrays = unnest_list_arrays(arrs_to_unnest.as_ref(), unnested_length, total_length)?; - // Create the take indices array for other columns - let take_indices = create_take_indices(unnested_length, total_length); unnested_temp_arrays .into_iter() .zip(list_unnest_specs.iter()) @@ -713,7 +863,13 @@ fn list_unnest_at_level( // Dimension of arrays in batch is untouched, but the values are repeated // as the side effect of unnesting - let ret = repeat_arrs_from_indices(batch, &take_indices, &repeat_mask)?; + let ret = if repeat_mask.iter().any(|&repeat| repeat) { + let take_indices = create_take_indices(unnested_length, total_length); + repeat_arrs_from_indices(batch, &take_indices, &repeat_mask)? + } else { + // Every column is replaced by an unnested result before the output is built. + batch.to_vec() + }; Ok(Some(ret)) } @@ -1242,6 +1398,79 @@ fn create_take_indices( PrimitiveArray::::from(indices) } +/// Arrow take estimates list child capacity from its full length, even for a sliced parent. +/// Restrict that estimate to the referenced children before repeating carried columns. +/// See https://github.com/apache/arrow-rs/pull/11360. +fn trim_list_children(array: &ArrayRef) -> Result> { + fn trim_list(list: &GenericListArray) -> Result> { + let start = list.value_offsets()[0].as_usize(); + let end = list.value_offsets()[list.len()].as_usize(); + if start == 0 && end == list.values().len() { + return Ok(None); + } + let (field, offsets, values, nulls) = list.clone().into_parts(); + let base = offsets[0]; + Ok(Some(Arc::new(GenericListArray::::try_new( + field, + offsets.subtract(base), + values.slice(start, end - start), + nulls, + )?))) + } + + let trimmed: ArrayRef = match array.data_type() { + DataType::List(_) => return trim_list(array.as_list::()), + DataType::LargeList(_) => return trim_list(array.as_list::()), + DataType::Map(field, ordered) => { + let map = array.as_map(); + let start = map.value_offsets()[0] as usize; + let end = map.value_offsets()[map.len()] as usize; + if start == 0 && end == map.entries().len() { + return Ok(None); + } + Arc::new(MapArray::try_new( + Arc::clone(field), + map.offsets().clone().subtract(start as i32), + map.entries().slice(start, end - start), + map.nulls().cloned(), + *ordered, + )?) + } + DataType::Struct(fields) => { + let array = array.as_struct(); + let mut columns = None; + for (index, column) in array.columns().iter().enumerate() { + if let Some(trimmed) = trim_list_children(column)? { + columns.get_or_insert_with(|| array.columns().to_vec())[index] = trimmed; + } + } + let Some(columns) = columns else { + return Ok(None); + }; + Arc::new(StructArray::try_new( + fields.clone(), + columns, + array.nulls().cloned(), + )?) + } + DataType::FixedSizeList(field, size) => { + let array = array.as_fixed_size_list(); + let Some(values) = trim_list_children(array.values())? else { + return Ok(None); + }; + Arc::new(FixedSizeListArray::try_new_with_length( + Arc::clone(field), + *size, + values, + array.nulls().cloned(), + array.len(), + )?) + } + _ => return Ok(None), + }; + Ok(Some(trimmed)) +} + /// Create a batch of arrays based on an input `batch` and a `indices` array. /// The `indices` array is used by the take kernel to repeat values in the arrays /// that are marked with `true` in the `repeat_mask`. Arrays marked with `false` @@ -1298,7 +1527,12 @@ fn repeat_arrs_from_indices( .zip(repeat_mask.iter()) .map(|(arr, &repeat)| { if repeat { - Ok(kernels::take::take(arr, indices, None)?) + let trimmed = trim_list_children(arr)?; + Ok(kernels::take::take( + trimmed.as_ref().unwrap_or(arr), + indices, + None, + )?) } else { Ok(new_null_array(arr.data_type(), arr.len())) } @@ -1309,7 +1543,7 @@ fn repeat_arrs_from_indices( #[cfg(test)] mod tests { use super::*; - use arrow::array::Int32Array; + use arrow::array::{Int32Array, StringArray}; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{Field, Int32Type, Schema}; use datafusion::common::NullHandling; @@ -1439,7 +1673,6 @@ 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, NullHandling::Preserve) .await .unwrap(); @@ -1447,6 +1680,536 @@ mod tests { assert_eq!(values(&batches), seq(25)); } + #[tokio::test] + async fn ordinary_carried_strings_keep_multi_row_chunks() { + let lists = list_batch(&vec![Some(300); 8192]); + let payload = "x".repeat(1024); + let strings = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + payload.as_str(), + 8192, + ))) as ArrayRef; + let schema = Arc::new(Schema::new(vec![ + lists.schema().field(0).clone(), + Field::new("payload", DataType::Utf8, false), + ])); + let input = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::clone(lists.column(0)), strings], + ) + .unwrap(); + let source = MemorySourceConfig::try_new_exec(&[vec![input]], schema, None).unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("l", DataType::Int32, true), + Field::new("payload", DataType::Utf8, false), + ])); + let exec = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + let context = Arc::new( + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(8192)), + ); + let mut stream = exec.execute(0, context).unwrap(); + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 27 * 300); + assert_eq!( + first.column(1).as_string::().value_data().len(), + 8100 * 1024 + ); + } + + #[test] + fn repetition_limit_ignores_fixed_width_buffers() { + let batch = RecordBatch::try_from_iter(vec![( + "id", + Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef, + )]) + .unwrap(); + assert_eq!(repetition_limit(&batch, &[]), i32::MAX as usize); + } + + #[tokio::test] + async fn carried_nested_columns_do_not_preallocate_for_unselected_children() { + use arrow::array::MapArray; + + let lists = list_batch(&vec![Some(2); 1024]); + let list = lists.column(0).as_list::(); + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::clone(list.values()), + ), + ( + Arc::new(Field::new("value", DataType::Int32, true)), + Arc::clone(list.values()), + ), + ]); + let map = Arc::new(MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + list.offsets().clone(), + entries, + None, + false, + )) as ArrayRef; + let nested_field = Arc::new(Field::new("item", map.data_type().clone(), true)); + let carried: Vec = vec![ + Arc::clone(&map), + Arc::clone(lists.column(0)), + Arc::new(LargeListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::new(list.value_offsets().iter().map(|&o| i64::from(o)).collect()), + Arc::clone(list.values()), + None, + )), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("m", map.data_type().clone(), true)), + Arc::clone(&map), + )])), + Arc::new(FixedSizeListArray::new( + Arc::clone(&nested_field), + 1, + Arc::clone(&map), + None, + )), + Arc::new(ListArray::new( + nested_field, + OffsetBuffer::new((0..=1024).collect()), + map, + None, + )), + ]; + + fn assert_small_buffers(array: &ArrayRef) { + match array.data_type() { + DataType::Map(_, _) => { + for column in array.as_map().entries().columns() { + assert_small_buffers(column); + } + } + DataType::List(_) => assert_small_buffers(array.as_list::().values()), + DataType::LargeList(_) => assert_small_buffers(array.as_list::().values()), + DataType::FixedSizeList(_, _) => { + assert_small_buffers(array.as_fixed_size_list().values()); + } + DataType::Struct(_) => { + for column in array.as_struct().columns() { + assert_small_buffers(column); + } + } + _ => assert!( + array.get_buffer_memory_size() <= 1024, + "eight output rows must not allocate for the 1024-row source: {} bytes", + array.get_buffer_memory_size() + ), + } + } + + for carried in carried { + let input = RecordBatch::try_from_iter(vec![ + ("l", Arc::clone(lists.column(0))), + ("carried", Arc::clone(&carried)), + ]) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("l", DataType::Int32, true), + input.schema().field(1).clone(), + ])); + let input_schema = input.schema(); + let source = + MemorySourceConfig::try_new_exec(&[vec![input]], input_schema, None).unwrap(); + let exec = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + Arc::clone(&schema), + UnnestOptions::new(), + ) + .unwrap(); + let context = Arc::new( + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(8)), + ); + let mut stream = exec.execute(0, context).unwrap(); + for start in [0, 8] { + let actual = stream.next().await.unwrap().unwrap(); + let expected = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(start..start + 8)), + kernels::take::take( + &carried, + &Int64Array::from_iter_values( + (start / 2..start / 2 + 4).flat_map(|row| [i64::from(row); 2]), + ), + None, + ) + .unwrap(), + ], + ) + .unwrap(); + assert_eq!(actual, expected); + assert_small_buffers(actual.column(1)); + } + } + } + + #[test] + fn repetition_limit_uses_selected_nested_offsets() { + use arrow::array::NullArray; + + let length = i32::MAX as usize / 2 + 1; + let list = ListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![0, length as i32 - 1, length as i32].into()), + Arc::new(NullArray::new(length)), + None, + ); + let batch = + RecordBatch::try_from_iter(vec![("carried", Arc::new(list) as ArrayRef)]).unwrap(); + assert_eq!(repetition_limit(&batch.slice(1, 1), &[]), i32::MAX as usize); + } + + #[test] + fn repetition_limit_follows_nested_variable_width_slices() { + use arrow::array::{BinaryArray, MapArray}; + + let strings = Arc::new(StringArray::from(vec!["ignored prefix", "a", "bc"])) as ArrayRef; + let field = Arc::new(Field::new("item", DataType::Utf8, false)); + let list = ListArray::new( + Arc::clone(&field), + OffsetBuffer::new(vec![0, 1, 3].into()), + Arc::clone(&strings), + None, + ); + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef, + ), + ( + Arc::new(Field::new("value", DataType::Utf8, false)), + Arc::clone(&strings), + ), + ]); + let map = MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + list.offsets().clone(), + entries, + None, + false, + ); + let arrays: Vec = vec![ + strings.slice(1, 2), + Arc::new(BinaryArray::from(vec![b"prefix".as_slice(), b"a", b"bc"]).slice(1, 2)), + Arc::new(list.slice(1, 1)), + Arc::new(map.slice(1, 1)), + Arc::new( + StructArray::from(vec![(Arc::clone(&field), Arc::clone(&strings))]).slice(1, 2), + ), + Arc::new( + LargeListArray::new( + Arc::clone(&field), + OffsetBuffer::new(vec![0i64, 1, 3].into()), + Arc::clone(&strings), + None, + ) + .slice(1, 1), + ), + Arc::new(FixedSizeListArray::new(field, 1, strings, None).slice(1, 2)), + ]; + for array in arrays { + let batch = RecordBatch::try_from_iter(vec![("carried", array)]).unwrap(); + assert_eq!(repetition_limit(&batch, &[]), i32::MAX as usize / 3); + } + } + + #[test] + fn chunk_offset_bound_counts_all_repeated_rows() { + use arrow::array::NullArray; + + let child_length = i32::MAX / 3; + let carried = ListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![0, child_length, 2 * child_length].into()), + Arc::new(NullArray::new(2 * child_length as usize)), + None, + ); + let lists = list_batch(&[Some(2), Some(2)]); + let batch = RecordBatch::try_from_iter(vec![ + ("l", Arc::clone(lists.column(0))), + ("carried", Arc::new(carried) as ArrayRef), + ]) + .unwrap(); + let list_columns = [ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + let mut pending = PendingInput { + batch, + row_offset: 0, + element_offset: 0, + safe_repetitions: Some(1), + output_lens: Some(Int64Array::from(vec![2, 2])), + }; + // Each row can repeat twice, but their combined child offsets would overflow. + for _ in 0..2 { + let (chunk, lengths) = pending.next_chunk(4, &list_columns).unwrap(); + assert_eq!(chunk.num_rows(), 1); + assert_eq!(lengths.unwrap().values(), &[2]); + } + assert_eq!(pending.remaining_rows(), 0); + } + + #[tokio::test] + async fn splits_one_list_before_repeating_carried_values() { + let list = list_batch(&[Some(25)]); + let schema = Arc::new(Schema::new(vec![ + list.schema().field(0).clone(), + Field::new("payload", DataType::Utf8, false), + ])); + let payload = "x".repeat(1024); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::clone(list.column(0)), + Arc::new(StringArray::from(vec![payload])), + ], + ) + .unwrap(); + let source = MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("l", DataType::Int32, true), + Field::new("payload", DataType::Utf8, false), + ])); + let exec = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + let context = Arc::new( + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(10)), + ); + let mut stream = exec.execute(0, context).unwrap(); + let mut actual = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let strings = batch.column(1).as_string::(); + assert_eq!( + strings.value_data().len(), + batch.num_rows() * 1024, + "carried values must be materialized only for this output chunk" + ); + actual.extend(batch.column(0).as_primitive::().iter()); + } + assert_eq!(actual, seq(25)); + } + + #[tokio::test] + async fn wide_carried_string_respects_offset_limit_at_large_batch_size() { + let list = list_batch(&[Some(2049)]); + let schema = Arc::new(Schema::new(vec![ + list.schema().field(0).clone(), + Field::new("payload", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::clone(list.column(0)), + Arc::new(StringArray::from(vec!["x".repeat(1024 * 1024)])), + ], + ) + .unwrap(); + let source = MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("l", DataType::Int32, true), + Field::new("payload", DataType::Utf8, false), + ])); + let exec = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + let context = Arc::new( + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(8192)), + ); + let mut stream = exec.execute(0, context).unwrap(); + let mut next = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + assert!(batch.num_rows() < 2049); + let strings = batch.column(1).as_string::(); + for (i, value) in batch + .column(0) + .as_primitive::() + .iter() + .enumerate() + { + assert_eq!(value, Some(next)); + assert_eq!(strings.value(i).len(), 1024 * 1024); + next += 1; + } + } + assert_eq!(next, 2049); + } + + #[tokio::test] + async fn splitting_preserves_map_entries_positions_and_nested_carried_columns() { + use arrow::array::{BinaryArray, BooleanArray}; + let lists = list_batch(&[Some(1), None, Some(11), Some(0), Some(3)]); + let list = lists.column(0).as_list::(); + let fields = vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Boolean, true), + ]; + let entries = StructArray::new( + fields.clone().into(), + vec![ + Arc::clone(list.values()), + Arc::new(BooleanArray::from_iter( + (0..list.values().len()).map(|i| (!i.is_multiple_of(3)).then_some(i % 5 < 2)), + )), + ], + None, + ); + let entries = ListArray::new( + Arc::new(Field::new("entry", entries.data_type().clone(), false)), + list.offsets().clone(), + Arc::new(entries), + list.nulls().cloned(), + ); + let positions = ListArray::new( + Arc::new(Field::new("pos", DataType::Int32, true)), + list.offsets().clone(), + Arc::new(Int32Array::from_iter_values( + list.offsets().windows(2).flat_map(|w| 0..w[1] - w[0]), + )), + list.nulls().cloned(), + ); + let carried = StructArray::from(vec![ + ( + Arc::new(Field::new("s", DataType::Utf8, true)), + Arc::new(StringArray::from(vec![ + Some("skip"), + Some("null map"), + Some("populated"), + None, + Some("tail"), + ])) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Binary, true)), + Arc::new(BinaryArray::from(vec![ + Some(b"skip".as_slice()), + None, + Some(b"binary"), + Some(b"empty"), + Some(b"tail"), + ])) as ArrayRef, + ), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("pos", positions.data_type().clone(), true), + Field::new("entries", entries.data_type().clone(), true), + Field::new("carried", carried.data_type().clone(), true), + ])); + let input = RecordBatch::try_new( + schema, + vec![Arc::new(positions), Arc::new(entries), Arc::new(carried)], + ) + .unwrap() + .slice(1, 4); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("pos", DataType::Int32, true), + fields[0].clone().with_nullable(true), + fields[1].clone(), + input.schema().field(2).clone(), + ])); + let list_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + for handling in [NullHandling::Drop, NullHandling::PreserveAndExpandEmpty] { + let options = UnnestOptions::new().with_null_handling(handling); + let expected = build_batch( + &input, + &output_schema, + &list_columns, + &HashSet::from([1]), + &options, + None, + ) + .unwrap() + .unwrap(); + let source = + MemorySourceConfig::try_new_exec(&[vec![input.clone()]], input.schema(), None) + .unwrap(); + let exec = ExplodeExec::new( + source, + list_columns.clone(), + vec![1], + Arc::clone(&output_schema), + options, + ) + .unwrap(); + let context = Arc::new( + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(4)), + ); + let output = + datafusion::physical_plan::common::collect(exec.execute(0, context).unwrap()) + .await + .unwrap(); + assert!(output.iter().all(|batch| batch.num_rows() <= 4)); + let actual = arrow::compute::concat_batches(&output_schema, &output).unwrap(); + assert_eq!(actual, expected); + } + } + + #[test] + fn repetition_limit_counts_zero_width_children() { + use arrow::array::NullArray; + let length = i32::MAX as usize / 2 + 1; + let list = ListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![0, length as i32].into()), + Arc::new(NullArray::new(length)), + None, + ); + let schema = Arc::new(Schema::new(vec![Field::new( + "carried", + list.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap(); + assert_eq!(repetition_limit(&batch, &[]), 1); + } + /// Chunked output must match unchunked output exactly, and hold the invariants that /// apply regardless of null handling: bounded, non-empty, and totalling `expected_rows`. async fn assert_chunking_matches_whole( @@ -1833,6 +2596,128 @@ mod tests { } } + #[tokio::test] + async fn parallel_lists_match_upstream_with_shared_and_distinct_views() { + use datafusion::physical_plan::{common::collect, unnest::UnnestExec}; + + let offsets = OffsetBuffer::new(vec![0, 2, 2, 5, 10, 11].into()); + let values = Arc::new(Int32Array::from_iter_values(0..11)) as ArrayRef; + let other_values = Arc::new(Int32Array::from_iter_values(100..111)) as ArrayRef; + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let validity = NullBuffer::from(vec![true, false, true, false, true]); + let nulls = validity.slice(0, 4); + // These views share the same bitmap bytes and null count, but different valid rows. + let shifted_nulls = validity.slice(1, 4); + let validity_pairs = [ + (None, None), + (Some(nulls.clone()), Some(nulls.clone())), + (Some(nulls.clone()), Some(shifted_nulls)), + ( + Some(nulls.clone()), + Some(NullBuffer::from(vec![true, false, true, false])), + ), + (Some(nulls.clone()), None), + (None, Some(nulls)), + ]; + + for (left_nulls, right_nulls) in validity_pairs { + for offset in [0, 1] { + let left = ListArray::new( + Arc::clone(&field), + offsets.slice(0, 4), + Arc::clone(&values), + left_nulls.clone(), + ); + let right = ListArray::new( + Arc::clone(&field), + offsets.slice(offset, 4), + Arc::clone(&other_values), + right_nulls.clone(), + ); + let large_right = LargeListArray::new( + Arc::clone(&field), + OffsetBuffer::new( + right + .value_offsets() + .iter() + .map(|&o| i64::from(o)) + .collect(), + ), + Arc::clone(&other_values), + right_nulls.clone(), + ); + for right in [Arc::new(right) as ArrayRef, Arc::new(large_right)] { + let batch = RecordBatch::try_from_iter(vec![ + ("left", Arc::new(left.clone()) as ArrayRef), + ("right", right), + ("id", Arc::new(Int32Array::from_iter_values(0..4))), + ]) + .unwrap(); + for batch in [batch.clone(), batch.slice(1, 3)] { + let source: Arc = MemorySourceConfig::try_new_exec( + &[vec![batch.clone()]], + batch.schema(), + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + batch.schema().field(2).clone(), + ])); + let specs: Vec<_> = (0..2) + .map(|index_in_input_schema| ListUnnest { + index_in_input_schema, + depth: 1, + }) + .collect(); + for mode in [ + NullHandling::Drop, + NullHandling::Preserve, + NullHandling::PreserveAndExpandEmpty, + ] { + let options = UnnestOptions::new().with_null_handling(mode); + let comet = ExplodeExec::new( + Arc::clone(&source), + specs.clone(), + vec![], + Arc::clone(&schema), + options.clone(), + ) + .unwrap(); + let upstream = UnnestExec::new( + Arc::clone(&source), + specs.clone(), + vec![], + Arc::clone(&schema), + options, + ) + .unwrap(); + for size in [1, 3, 8] { + let context = Arc::new(TaskContext::default().with_session_config( + SessionConfig::new().with_batch_size(size), + )); + let actual = + collect(comet.execute(0, Arc::clone(&context)).unwrap()) + .await + .unwrap(); + let expected = collect(upstream.execute(0, context).unwrap()) + .await + .unwrap(); + assert!(actual.iter().all(|b| b.num_rows() <= size)); + assert_eq!( + arrow::compute::concat_batches(&schema, &actual).unwrap(), + arrow::compute::concat_batches(&schema, &expected).unwrap(), + "offset={offset}, {mode:?}, batch_size={size}, {batch:?}" + ); + } + } + } + } + } + } + } + #[test] fn output_lens_substitute_per_null_handling() { // Row 2 is empty with no validity buffer at all, which is the case diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 8e720a8897c..3b6373cd196 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1106,9 +1106,9 @@ message MergeOutputRow { } message Explode { - // The array expression to explode into multiple rows + // The array or map expression to explode into multiple rows spark.spark_expression.Expr child = 1; - // Whether this is explode_outer (produces null row for empty/null arrays) + // Whether this is explode_outer (produces null row for empty/null collections) bool outer = 2; // Expressions for other columns to project alongside the exploded values repeated spark.spark_expression.Expr project_list = 3; diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d79927f0c04..47721935c10 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,7 +25,7 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, ExplodeBase, Expression, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.rules.Rule @@ -190,6 +190,9 @@ object CometExecRule { private val TYPED_DATASET_PARTIAL_READER: TreeNodeTag[String] = TreeNodeTag[String]("comet.typedDatasetPartialReader") + private val CHAINED_MAP_GENERATOR: TreeNodeTag[String] = + TreeNodeTag[String]("comet.chainedMapGenerator") + /** Why Comet does not convert rows to Arrow in a plan that reads `InputFileBlockHolder`. */ private val INPUT_FILE_BLOCK_FALLBACK_REASON: String = "Spark to Arrow conversion is not compatible with input_file_name, " + @@ -439,12 +442,19 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) private def transform(plan: SparkPlan): SparkPlan = { // Walks the whole plan, so it is lazy: only consulted once a node could be converted. lazy val readsInputFileBlock = CometScanRule.readsInputFileBlock(plan) + tagChainedMapGenerators(plan) if (CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf)) { tagPartiallyReadTypedDatasetOutputs(plan) } def convertNode(op: SparkPlan): SparkPlan = op match { + case generate: GenerateExec + if generate.getTagValue(CometExecRule.CHAINED_MAP_GENERATOR).isDefined => + withFallbackReason( + generate, + generate.getTagValue(CometExecRule.CHAINED_MAP_GENERATOR).get) + // Scan marker produced by an optional, out-of-tree scan contrib (e.g. contrib/delta). // Matched by trait (no compile-time dependency on the contrib) and present only when that // contrib is on the classpath. The marker carries its own serde handler and typically wraps @@ -1444,6 +1454,29 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) visit(plan, if (producesObjects(plan)) Some("code reading Dataset.rdd") else None) } + // Native generator chains materialize carried columns for both arrays and maps. Preserve + // Spark's fused path for newly supported map chains and the existing eligibility of arrays. + // Tags survive child copies made by bottom-up conversion; exchanges end the shared pipeline. + private def tagChainedMapGenerators(plan: SparkPlan): Unit = { + val enabled = CometConf.COMET_EXEC_EXPLODE_ENABLED.get(conf) + def visit(op: SparkPlan, generatorAbove: Boolean): Unit = op match { + case generate: GenerateExec => + generate.unsetTagValue(CometExecRule.CHAINED_MAP_GENERATOR) + generate.generator match { + case explode: ExplodeBase + if enabled && generatorAbove && explode.child.dataType.isInstanceOf[MapType] => + generate.setTagValue( + CometExecRule.CHAINED_MAP_GENERATOR, + "Chained map generators stay on Spark to avoid carried-column materialization") + case _ => + } + visit(generate.child, generatorAbove = true) + case _: ProjectExec | _: FilterExec => visit(op.children.head, generatorAbove) + case _ => op.children.foreach(visit(_, generatorAbove = false)) + } + visit(plan, generatorAbove = false) + } + /** * Converts the rows a typed Dataset operation produces to Arrow, so the operators above it can * run natively. See [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index da26c91f9c7..dda2dea267a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1928,12 +1928,8 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return Unsupported(Some(s"Unsupported generator: ${op.generator.nodeName}")) } op.generator.children.head.dataType match { - case _: ArrayType => + case _: ArrayType | _: MapType => Compatible() - case _: MapType => - // TODO add support for map types - // https://github.com/apache/datafusion-comet/issues/2837 - Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) case other => Unsupported(Some(s"Unsupported data type: $other")) } 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 b1bdb8b7493..a958d9a3907 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/explode.sql @@ -328,20 +328,6 @@ CREATE TABLE test_explode_empty(id int, arr array) USING parquet query SELECT id, explode_outer(arr) FROM test_explode_empty --- ===== Map falls back to Spark: outer form must also fall back ===== - -statement -CREATE TABLE test_explode_map(id int, m map) USING parquet - -statement -INSERT INTO test_explode_map VALUES - (1, map('k1', 1, 'k2', 2)), - (2, map()), - (3, NULL) - -query expect_fallback(Comet only supports explode/explode_outer for arrays, not maps) -SELECT id, explode_outer(m) FROM test_explode_map - -- ===== TIMESTAMP_NTZ (Spark 3.4+): distinct Arrow layout (no tz) from LTZ ===== statement 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 eb537586f0e..f8028b88565 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql @@ -85,18 +85,6 @@ INSERT INTO test_posexplode_struct VALUES query SELECT id, pos, value.v1 AS v1, value.v2 AS v2 FROM test_posexplode_struct LATERAL VIEW posexplode(arr) p AS pos, value -statement -CREATE TABLE test_posexplode_map(id int, m map) USING parquet - -statement -INSERT INTO test_posexplode_map VALUES - (1, map('a', 1, 'b', 2)), - (2, map('c', 3)) - --- posexplode over a map falls back to Spark (Comet only supports array inputs, not maps) -query expect_fallback(Comet only supports explode/explode_outer for arrays, not maps) -SELECT id, posexplode(m) FROM test_posexplode_map - -- ===== posexplode_outer across non-int element types ===== -- posexplode_outer with nullable int elements (in-array NULLs plus outer null row) diff --git a/spark/src/test/resources/sql-tests/expressions/map/explode.sql b/spark/src/test/resources/sql-tests/expressions/map/explode.sql new file mode 100644 index 00000000000..0bb29116910 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/explode.sql @@ -0,0 +1,156 @@ +-- 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.comet.exec.explode.enabled=true +-- Config: spark.comet.batchSize=4 + +statement +CREATE TABLE explode_maps(id int, m map) USING parquet + +statement +INSERT INTO explode_maps VALUES + (1, map('b', 20, 'a', 10)), (2, map()), (3, NULL), (4, map('null', NULL)) + +query +SELECT id, explode(m) FROM explode_maps + +query +SELECT id, m, explode_outer(m) FROM explode_maps + +query +SELECT id, posexplode(m) FROM explode_maps + +query +SELECT id, m, posexplode_outer(m) FROM explode_maps + +query +SELECT posexplode_outer(m) FROM explode_maps + +query +SELECT id, k, v FROM explode_maps LATERAL VIEW OUTER explode(m) e AS k, v +WHERE k IS NULL OR v IS NULL + +query +SELECT id, p, k, v FROM explode_maps LATERAL VIEW OUTER posexplode(m) e AS p, k, v +WHERE p = 0 OR p IS NULL + +statement +CREATE TABLE explode_empty_maps(id int, m map) USING parquet + +statement +INSERT INTO explode_empty_maps SELECT id, CASE WHEN id % 2 = 0 THEN NULL + ELSE cast(map() AS map) END FROM range(16) + +query +SELECT id, m, explode_outer(m) FROM explode_empty_maps + +query +SELECT id, m, posexplode_outer(m) FROM explode_empty_maps + +-- Native map producers feed the generator directly, without a Parquet map column. +statement +CREATE TABLE explode_map_inputs( + id int, k array, v array, + entries array>, text string) USING parquet + +statement +INSERT INTO explode_map_inputs VALUES + (1, array(2, 1), array('b', 'a'), array(struct(2, 'b'), struct(1, 'a')), 'b:2,a:1'), + (2, array(), array(), array(), ''), + (3, NULL, NULL, NULL, NULL), + (4, array(3), array(NULL), array(struct(3, NULL)), 'missing') + +query expect_native(map_from_arrays) +SELECT id, explode(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, explode_outer(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, posexplode(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, posexplode_outer(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, explode(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, explode_outer(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, posexplode(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, posexplode_outer(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, explode(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, explode_outer(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, posexplode(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, posexplode_outer(str_to_map(text)) FROM explode_map_inputs + +-- Computed and scalar maps, including non-nullable values and nested fields. +query +SELECT id, posexplode(map(id, named_struct('v', coalesce(id, 0)))) FROM explode_maps + +query +SELECT id, explode(map(1, 10, 2, 20)) FROM explode_maps + +query +SELECT id, posexplode_outer(cast(NULL AS map)) FROM explode_maps + +-- The untyped map() constructor still falls back before the cast. +query expect_fallback(MapType(NullType,NullType,false)) +SELECT id, explode_outer(cast(map() AS map)) FROM explode_maps + +statement +CREATE TABLE explode_nested_maps( + id int, m map, struct, m: map>>) USING parquet + +statement +INSERT INTO explode_nested_maps VALUES + (1, map(array(2, 1), named_struct('a', array(10, NULL), 'm', map('x', unhex('FF00'))), + array(3), NULL)), + (2, map()), (3, NULL) + +query +SELECT id, explode(m) FROM explode_nested_maps + +query +SELECT id, explode_outer(m) FROM explode_nested_maps + +query +SELECT id, posexplode(m) FROM explode_nested_maps + +query +SELECT id, posexplode_outer(m) FROM explode_nested_maps + +-- Project the output fields and compose map expansion with array expansion. +query +SELECT id, k, v.a, v.m FROM explode_nested_maps LATERAL VIEW OUTER explode(m) e AS k, v + +query expect_fallback(Chained map generators stay on Spark) +SELECT id, p, k, item FROM explode_nested_maps +LATERAL VIEW OUTER posexplode(m) e AS p, k, v +LATERAL VIEW OUTER explode(v.a) a AS item diff --git a/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql b/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql new file mode 100644 index 00000000000..58b04c51f98 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql @@ -0,0 +1,47 @@ +-- 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.comet.exec.explode.enabled=true +-- Config: spark.sql.legacy.disableMapKeyNormalization=true + +-- Preserve -0.0 when writing the fixture on Spark 4.0+. The generator must return +-- the stored keys unchanged. Native map construction's normalization is tracked +-- separately in https://github.com/apache/datafusion-comet/issues/6549. +statement +CREATE TABLE explode_float_maps(id int, m map) USING parquet + +statement +INSERT INTO explode_float_maps VALUES + (1, map(cast('-0.0' AS double), 10, cast('NaN' AS double), 20)), + (2, map(cast('NaN' AS double), NULL)), (3, map()), (4, NULL) + +-- Render keys as strings too, so a comparison cannot silently equate the two zero signs. +query +SELECT id, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW explode(m) e AS key, value + +query +SELECT id, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW OUTER explode(m) e AS key, value + +query +SELECT id, pos, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW posexplode(m) e AS pos, key, value + +query +SELECT id, pos, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW OUTER posexplode(m) e AS pos, key, value 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 1ccfacdffe5..7c6be5f15fc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -20,7 +20,7 @@ package org.apache.comet.exec import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.comet.CometExplodeExec +import org.apache.spark.sql.comet.{CometExplodeExec, CometNativeScanExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.col @@ -224,16 +224,148 @@ class CometGenerateExecSuite extends CometTestBase { } } - test("explode with map input falls back") { - withSQLConf( - CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", - CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { - val df = Seq((1, Map("a" -> 1, "b" -> 2)), (2, Map("c" -> 3))) - .toDF("id", "map") - .selectExpr("id", "explode(map) as (key, value)") - checkSparkAnswerAndFallbackReason( - df, - "Comet only supports explode/explode_outer for arrays, not maps") + for (generator <- Seq("explode", "explode_outer", "posexplode", "posexplode_outer")) { + test(s"$generator with JVM-fed map input") { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + val input = Seq( + (1, Map("b" -> Integer.valueOf(20), "a" -> Integer.valueOf(10))), + (2, Map.empty[String, Integer]), + (3, null.asInstanceOf[Map[String, Integer]]), + (4, Map("null" -> null.asInstanceOf[Integer]))) + .toDF("id", "m") + checkMapGenerator(input, "m", generator) + } + } + + test(s"$generator with map input across batch boundaries") { + withSQLConf( + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "4") { + val rows = (0 until 16).map { i => + val m = i % 4 match { + case 0 => null.asInstanceOf[Map[Int, java.lang.Integer]] + case 1 => Map.empty[Int, java.lang.Integer] + case _ => + (0 until 13).map { j => + j -> (if (j % 3 == 0) null else java.lang.Integer.valueOf(i * 100 + j)) + }.toMap + } + val booleans = Option(m) + .map(_.map { case (key, value) => + // A period of five cannot hide a wrong bitmap offset at four-row batch boundaries. + key -> (if (value == null) null else java.lang.Boolean.valueOf((i + key) % 5 < 2)) + }) + .orNull + (i, m, booleans) + } + withParquetDataFrame(rows) { input => + // One map exceeds the output batch size. Carry the map through too, so + // outer padding cannot accidentally replace the original empty map. + val maps = input.toDF("id", "ints", "booleans") + Seq("ints", "booleans").foreach { name => + withClue(s"$name: ") { checkMapGenerator(maps, name, generator) } + } + } + } + } + + test(s"$generator splits a map before repeating a wide carried column") { + withSQLConf( + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "256") { + val payload = "x" * (1024 * 1024) + val entries = (0 until 2049).map(i => i -> i).toMap + withParquetDataFrame(Seq((payload, entries))) { input => + val generated = input.toDF("payload", "m").selectExpr("payload", s"$generator(m)") + val aggregates = Seq( + "sum(CAST(length(concat(payload, CAST(key AS STRING))) AS BIGINT))", + "sum(CAST(value AS BIGINT))") ++ + (if (generator.startsWith("pos")) Seq("sum(CAST(pos AS BIGINT))") else Nil) + val query = generated.selectExpr(aggregates: _*) + checkSparkAnswerAndOperator(query, Seq(classOf[CometExplodeExec])) + } + } + } + } + + private def checkMapGenerator(input: DataFrame, name: String, generator: String): Unit = { + val query = input.selectExpr("id", name, s"$generator($name)") + checkSparkAnswerAndOperator(query, Seq(classOf[CometExplodeExec])) + checkSparkSchema(query) + } + + for (generator <- Seq("explode", "explode_outer", "posexplode", "posexplode_outer"); + adaptive <- Seq(false, true)) { + test(s"chained map generators stay on Spark: $generator, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + val rows = Seq( + (1, Map("keep" -> Seq[Integer](10, null, 20), "drop" -> Seq[Integer](30))), + (2, Map.empty[String, Seq[Integer]]), + (3, null.asInstanceOf[Map[String, Seq[Integer]]])) + withParquetDataFrame(rows) { input => + val mapAliases = if (generator.startsWith("pos")) "(map_pos, k, a)" else "(k, a)" + val arrayAliases = if (generator.startsWith("pos")) "(array_pos, v)" else "v" + val generated = input + .toDF("id", "m") + .selectExpr("id", s"$generator(m) AS $mapAliases") + .filter("k IS NULL OR k <> 'drop'") + .selectExpr("*", s"$generator(a) AS $arrayAliases") + .repartition(2, col("id")) + val (_, plan) = + checkSparkAnswerAndFallbackReason(generated, "Chained map generators stay on Spark") + assert(collect(plan) { case e: CometExplodeExec => e }.isEmpty) + assert(collect(plan) { case e: CometNativeScanExec => e }.nonEmpty) + } + } + } + } + + test("single map generator with array values stays native") { + withSQLConf(CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + withParquetDataFrame(Seq((1, Map("a" -> Seq(10, 20), "b" -> Seq.empty[Int])))) { input => + checkMapGenerator(input.toDF("id", "m"), "m", "explode") + } + } + } + + test("array-only generator chain stays native") { + withSQLConf(CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + withParquetDataFrame(Seq((1, Seq(Seq(10, 20), Seq(30))))) { input => + val generated = input + .toDF("id", "a") + .selectExpr("id", "explode(a) AS inner_array") + .selectExpr("id", "explode(inner_array) AS value") + checkSparkAnswerAndOperator(generated, Seq(classOf[CometExplodeExec])) + assert(collect(generated.queryExecution.executedPlan) { case e: CometExplodeExec => + e + }.size == 2) + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"exchange separates map and array generator chains: AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withParquetDataFrame(Seq((1, Map("a" -> Seq(10, 20), "b" -> Seq(30))))) { input => + val generated = input + .toDF("id", "m") + .selectExpr("id", "explode(m) AS (k, a)") + .repartition(2, col("id")) + .selectExpr("id", "k", "explode(a) AS value") + checkSparkAnswerAndOperator(generated, Seq(classOf[CometExplodeExec])) + assert(collect(generated.queryExecution.executedPlan) { case e: CometExplodeExec => + e + }.size == 2) + } + } } } @@ -390,16 +522,16 @@ class CometGenerateExecSuite extends CometTestBase { } } - test("posexplode with map input falls back") { + test("posexplode with map input falls back when disabled") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", - CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "false") { val df = Seq((1, Map("a" -> 1, "b" -> 2)), (2, Map("c" -> 3))) .toDF("id", "map") .selectExpr("id", "posexplode(map) as (pos, key, value)") checkSparkAnswerAndFallbackReason( df, - "Comet only supports explode/explode_outer for arrays, not maps") + "Native support for operator GenerateExec is disabled") } } @@ -698,8 +830,7 @@ class CometGenerateExecSuite extends CometTestBase { } test("explode of structs keeps boolean fields when one row exceeds the batch size") { - // A row longer than the batch size is unnested in one build, which is then sliced into - // batch-size pieces on the way out, here at elements 100 and 200. + // A single row spans batches at elements 100 and 200, including non-byte-aligned bitmaps. withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "100") { withBooleanArrays(numRows = 1, arrayLength = 250) { checkSparkAnswerAndOperator(sql("SELECT id, explode(structs) FROM t")) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala index 86909f839c5..b217e42e4a4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala @@ -22,6 +22,7 @@ package org.apache.spark.sql.benchmark import java.io.File import org.apache.spark.sql.catalyst.optimizer.InferFiltersFromGenerate +import org.apache.spark.sql.comet.{CometExplodeExec, CometNativeScanExec} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -59,18 +60,16 @@ import org.apache.comet.CometConf * processes. Worth knowing when reading these numbers: real queries do get that filter, so a * plain `explode` in production usually sees an array column with no nulls and no empty rows. * - * Only array inputs are covered. Comet declines to convert a generator over a map - * (https://github.com/apache/datafusion-comet/issues/2837), so the Comet arm of such a case would - * be Spark's `GenerateExec` behind a columnar-to-row transition and its timing would say nothing - * about `CometExplodeExec`. The nesting group below therefore reaches its event list through an - * `array>` where a map would be the more natural modeling choice. + * Map cases consume both keys and values. The nesting group compares maps keyed by platform with + * equivalent arrays of structs, including a second generator over each platform's event list. The + * chained map case retains Spark generators above the native scan. */ object CometExplodeBenchmark extends CometBenchmarkBase { private val numRows = 256 * 1024 /** One benchmark case: the query to time, and the single row it must produce. */ - private case class Case(query: String, expected: Seq[Long]) + private case class Case(query: String, expected: Seq[Long], nativeGenerators: Int = 1) /** A temp view the benchmark reads: its row count, and the expressions that build it. */ private case class TempView(name: String, rows: Int, columns: Seq[String]) @@ -109,6 +108,11 @@ object CometExplodeBenchmark extends CometBenchmarkBase { |END""".stripMargin } + /** Maps with the same NULL/empty distribution as the array cases and ordered platform keys. */ + private def mapColumn(elementExpr: String, len: Int): String = + s"map_from_arrays(${arrayColumn("concat('p_', CAST(x AS STRING))", len)}, " + + s"${arrayColumn(elementExpr, len)})" + /** * Types `expr` as nullable without ever evaluating to null, `guard` being a predicate that is * never true. @@ -153,20 +157,23 @@ object CometExplodeBenchmark extends CometBenchmarkBase { | 'page', concat('page_', CAST(id + y AS STRING)), | 'source', concat('src_', CAST(id % 4 AS STRING)))""".stripMargin - /** One platform's event list, the element of the array the nesting group explodes. */ + /** One platform's event list, the value of the map the nesting group explodes. */ private val eventStruct = s"""named_struct( - | 'platform', concat('p_', CAST(x AS STRING)), | 'entries', ${fullArray(entryStruct, entriesLen, "y")})""".stripMargin - private val eventsArray = arrayColumn(eventStruct, eventsLen) + private val eventsMap = mapColumn(eventStruct, eventsLen) + private val eventsArray = arrayColumn( + s"named_struct('platform', concat('p_', CAST(x AS STRING)), " + + s"'entries', ${fullArray(entryStruct, entriesLen, "y")})", + eventsLen) - /** [[eventsArray]] wrapped in the struct chain that puts it eight accessors down. */ - private val profileColumn = { + /** The collection wrapped in the struct chain that puts it eight accessors down. */ + private def profileColumn(events: String): String = { val sessions = s"""named_struct( | 'device', named_struct('type', 'mobile', 'os', 'iOS'), - | 'events', $eventsArray)""".stripMargin + | 'events', $events)""".stripMargin val account = accountPath.foldRight(s"named_struct('sessions', $sessions)") { (field, inner) => s"named_struct('$field', $inner)" } @@ -179,16 +186,19 @@ object CometExplodeBenchmark extends CometBenchmarkBase { /** * The temp views the benchmark reads. * - * Each array column gets its own view rather than sharing one wide table, so that a case is - * never charged for scanning an array column it does not read. The nesting group is the - * exception: `events` and `profile` hold the same array at two different depths, and keeping - * them in one view is what makes them the same array. Nested schema pruning stops either case - * from reading the other's copy. + * Each collection column gets its own view rather than sharing one wide table, so that a case + * is never charged for scanning a column it does not read. The nesting group is the exception: + * `events` and `profile` hold the same collection at two different depths, and keeping them in + * one view is what makes them the same collection. Nested schema pruning stops either case from + * reading the other's copy. */ private val views: Seq[TempView] = Seq( TempView("arr_len2", numRows, Seq(s"${arrayColumn("id + x", 2)} AS arr")), TempView("arr_len10", numRows, Seq(s"${arrayColumn("id + x", 10)} AS arr")), TempView("arr_len100", numRows, Seq(s"${arrayColumn("id + x", 100)} AS arr")), + TempView("map_len2", numRows, Seq(s"${mapColumn("id + x", 2)} AS m")), + TempView("map_len10", numRows, Seq(s"${mapColumn("id + x", 10)} AS m")), + TempView("map_len100", numRows, Seq(s"${mapColumn("id + x", 100)} AS m")), TempView("arr_str10", numRows, Seq(s"${arrayColumn(stringElement, 10)} AS arr")), TempView( "arr_struct10", @@ -202,14 +212,18 @@ object CometExplodeBenchmark extends CometBenchmarkBase { s"${nullableExpr("id", "id < 0")} AS k", s"${nullableExpr("CAST(id AS STRING)", "id < 0")} AS s", s"${nullableExpr("id * 2", "id < 0")} AS v")), + nestedView("nested_map", eventsMap), + nestedView("nested_array", eventsArray)) + + private def nestedView(name: String, events: String): TempView = TempView( - "nested", + name, nestedRows, Seq( s"${nullableExpr("id", "id < 0")} AS k", s"${nullableExpr("concat('r_', CAST(id % 8 AS STRING))", "id < 0")} AS region", - s"$eventsArray AS events", - s"$profileColumn AS profile"))) + s"$events AS events", + s"${profileColumn(events)} AS profile")) /** Writes a view's rows to Parquet and registers it. */ private def createView(dir: File, view: TempView): Unit = { @@ -219,8 +233,8 @@ object CometExplodeBenchmark extends CometBenchmarkBase { } /** - * A case that applies `generator` to `arrayExpr` over `view`, carries `carried` through the - * generator alongside it, and aggregates every column that comes out. + * A case that applies `generator` to `collectionExpr` over `view`, carries `carried` through + * the generator alongside it, and aggregates every column that comes out. * * The position column of the `posexplode` variants is summed rather than counted, because * counting it does not read a position. `pos` is declared non-null, so `NullPropagation` @@ -230,24 +244,26 @@ object CometExplodeBenchmark extends CometBenchmarkBase { * `ListPositionsExpr` and unnests it alongside the values — so a sink that ignores them is not * measuring the generator it names. `sum` over the same column is value-dependent. * - * Everything else is counted rather than summed: the element type varies across these cases and - * most of the element types cannot be summed, and unlike `pos` the generated column is nullable - * everywhere, so the count is not rewritten away. + * Map keys are also non-nullable. Sum their lengths to keep the key column in the plan and + * consume its bytes. Values and carried columns are nullable and counted, since the value types + * vary and most cannot be summed. */ private def generatorCase( generator: String, view: String, len: Int, rows: Int = numRows, - arrayExpr: String = "arr", + collectionExpr: String = "arr", + mapInput: Boolean = false, carried: Seq[String] = Nil, where: Option[String] = None): Case = { val position = generator.startsWith("posexplode") - val generated = if (position) Seq("pos", "col") else Seq("col") + val values = if (mapInput) Seq("key", "col") else Seq("col") + val generated = if (position) "pos" +: values else values val alias = if (generated.length == 1) s"AS ${generated.head}" else generated.mkString("AS (", ", ", ")") - val projectList = (carried :+ s"$generator($arrayExpr) $alias").mkString(", ") + val projectList = (carried :+ s"$generator($collectionExpr) $alias").mkString(", ") val filter = where.map(w => s" WHERE $w").getOrElse("") val aggregates = (carried ++ generated).map(aggregate) val query = @@ -262,14 +278,20 @@ object CometExplodeBenchmark extends CometBenchmarkBase { // Positions run 0 until len on every row that emits, and are null on the rows only an // outer variant emits, which `sum` skips. case "pos" => elements * (len - 1) / 2 + case "key" => nonEmptyRows(rows) * (1 to len).map(i => s"p_$i".length.toLong).sum case _ => elements } Case(query, expected) } - /** See [[generatorCase]] for why `pos` alone is summed. */ - private def aggregate(column: String): String = - if (column == "pos") s"sum($column)" else s"count($column)" + /** + * Map keys are non-nullable, so consume their bytes rather than letting count become count(1). + */ + private def aggregate(column: String): String = column match { + case "pos" => s"sum($column)" + case "key" | "platform" => s"sum(length($column))" + case _ => s"count($column)" + } /** * Runs `query` under both engines, untimed, and fails unless each produces exactly `expected`. @@ -284,28 +306,42 @@ object CometExplodeBenchmark extends CometBenchmarkBase { * work per-row; it cannot change a result, and these queries hold no constant subexpression for * it to fold in any case. */ - private def verifySink(name: String, query: String, expected: Seq[Long]): Unit = + private def verifySink(name: String, benchmarkCase: Case): Unit = Seq(false, true).foreach { cometEnabled => val engine = if (cometEnabled) "Comet" else "Spark" withSQLConf( CometConf.COMET_ENABLED.key -> cometEnabled.toString, CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString) { - val actual = spark.sql(query).collect().head.toSeq.map { + val df = spark.sql(benchmarkCase.query) + val actual = df.collect().head.toSeq.map { case null => null case value: Number => value.longValue() case other => other } - if (actual != expected) { + if (actual != benchmarkCase.expected) { throw new AssertionError( - s"$name: $engine produced $actual, expected $expected. The case is not aggregating " + - s"what it names.\n$query") + s"$name: $engine produced $actual, expected ${benchmarkCase.expected}. " + + s"The case is not aggregating what it names.\n${benchmarkCase.query}") + } + if (cometEnabled) { + val expectedGenerators = benchmarkCase.nativeGenerators + val nativeGenerators = collect(df.queryExecution.executedPlan) { + case e: CometExplodeExec => e + }.size + assert( + nativeGenerators == expectedGenerators, + s"$name executed $nativeGenerators of $expectedGenerators generators natively:\n" + + df.queryExecution.executedPlan) + assert(collect(df.queryExecution.executedPlan) { case scan: CometNativeScanExec => + scan + }.nonEmpty) } } } /** Verifies a case's aggregate under both engines, then times it. */ private def runCase(name: String, benchmarkCase: Case, rows: Int = numRows): Unit = { - verifySink(name, benchmarkCase.query, benchmarkCase.expected) + verifySink(name, benchmarkCase) runExpressionBenchmark(name, rows.toLong, benchmarkCase.query) } @@ -338,6 +374,24 @@ object CometExplodeBenchmark extends CometBenchmarkBase { } } + runBenchmark("Explode - maps") { + Seq(2, 10, 100).foreach { len => + runCase( + s"explode map[$len]", + generatorCase( + "explode", + s"map_len$len", + len, + collectionExpr = "m", + mapInput = true)) + } + Seq("posexplode", "explode_outer", "posexplode_outer").foreach { generator => + runCase( + s"$generator map[10]", + generatorCase(generator, "map_len10", 10, collectionExpr = "m", mapInput = true)) + } + } + runBenchmark("Explode - element type") { Seq("bigint" -> "arr_len10", "string" -> "arr_str10", "struct" -> "arr_struct10") .foreach { case (elementType, view) => @@ -382,32 +436,37 @@ object CometExplodeBenchmark extends CometBenchmarkBase { * * The shape is the one asked for in review: a customer profile whose event list sits eight * struct accessors down, holding a second array of four-field structs inside each element. The - * requested outer container was a map keyed by platform, which is where a real schema would put - * it; that is an `array>` here because Comet has no native generator - * over maps yet (#2837) and the Comet arm would silently be Spark. + * outer container is a `map>>>` keyed by platform, or + * the equivalent `array>`. * * The whole event struct is counted rather than one of its fields, so nested schema pruning * cannot narrow the exploded element and leave the case measuring a two-column gather. It does * prune the siblings no case reads: on the executed plan the `depth 8` scan's `profile` column - * is `struct>>>` with `billing`, - * `addresses` and `device` gone, and the `depth 1` scan does not project `profile` at all. + * retains only the event collection, with `billing`, `addresses` and `device` gone, and the + * `depth 1` scan does not project `profile` at all. * - * `depth 1` and `depth 8` explode the same array, written twice into the same file, so the pair - * measures what the struct chain costs: extra definition levels in the Parquet column and a - * chain of `GetStructField` above the scan. `then its inner array` chains a second generator + * `depth 1` and `depth 8` explode the same collection, written twice into the same file, so the + * pair measures what the struct chain costs: extra definition levels in the Parquet column and + * a chain of `GetStructField` above the scan. `then its inner array` chains a second generator * onto the first, which is the shape the review asked about; its fan-out is 20 against the * other two's 4, so read it on its own rather than against them. */ - private def nestedInputCases: Seq[(String, Case)] = { + private def nestedInputCases: Seq[(String, Case)] = + Seq(true, false).flatMap(mapInput => nestedInputCases(mapInput)) + + private def nestedInputCases(mapInput: Boolean): Seq[(String, Case)] = { + val view = if (mapInput) "nested_map" else "nested_array" + val element = if (mapInput) "map" else "array" val carried = Seq("k", "region") val single = Seq("depth 1" -> "events", "depth 8" -> deepEvents).map { case (label, path) => - s"explode array[$eventsLen] at $label" -> + s"explode $element[$eventsLen] at $label" -> generatorCase( "explode", - "nested", + view, eventsLen, rows = nestedRows, - arrayExpr = path, + collectionExpr = path, + mapInput = mapInput, carried = carried) } @@ -415,14 +474,22 @@ object CometExplodeBenchmark extends CometBenchmarkBase { // fan-out is the product of the two lengths. val events = nonEmptyRows(nestedRows) * eventsLen val entries = events * entriesLen - val outer = + val outer = if (mapInput) { + "SELECT k, region, platform, ev.entries AS entries " + + s"FROM (SELECT k, region, explode($deepEvents) AS (platform, ev) FROM $view)" + } else { "SELECT k, region, ev.platform AS platform, ev.entries AS entries " + - s"FROM (SELECT k, region, explode($deepEvents) AS ev FROM nested)" + s"FROM (SELECT k, region, explode($deepEvents) AS ev FROM $view)" + } val chained = s"SELECT ${Seq("k", "region", "platform", "entry").map(aggregate).mkString(", ")} " + s"FROM (SELECT k, region, platform, explode(entries) AS entry FROM ($outer))" - single :+ (s"explode array[$eventsLen] at depth 8, then its inner array" -> - Case(chained, Seq.fill(4)(entries))) + val execution = if (mapInput) "Spark generators" else "native generators" + single :+ (s"$element at depth 8, then its inner array ($execution)" -> + Case( + chained, + Seq(entries, entries, entries * 3, entries), + nativeGenerators = if (mapInput) 0 else 2)) } }