From 9601e81ba8fa2e8b1e7b3d92fbdc976e91f7a380 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Mon, 21 Sep 2026 17:56:22 +0800 Subject: [PATCH 1/9] feat: support map inputs for explode --- .../expression-audits/generator_funcs.md | 8 +- docs/source/user-guide/latest/expressions.md | 3 +- docs/source/user-guide/latest/operators.md | 12 +- .../src/execution/expressions/map_entries.rs | 119 ++++++++++++++++ native/core/src/execution/expressions/mod.rs | 1 + native/core/src/execution/planner.rs | 131 ++++++++++++++++-- native/proto/src/proto/operator.proto | 4 +- .../apache/spark/sql/comet/operators.scala | 6 +- .../sql-tests/expressions/array/explode.sql | 4 +- .../expressions/array/posexplode.sql | 4 +- .../sql-tests/expressions/map/explode.sql | 93 +++++++++++++ .../comet/exec/CometGenerateExecSuite.scala | 41 ++++-- .../sql/benchmark/CometExplodeBenchmark.scala | 11 +- 13 files changed, 381 insertions(+), 56 deletions(-) create mode 100644 native/core/src/execution/expressions/map_entries.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/map/explode.sql 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 cbc0b6e6406..efda85c89db 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -332,8 +332,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 c9ba5e491e1..0acb2b616ac 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -108,12 +108,12 @@ omitted from the tables below and may be reconsidered based on demand: ## 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` | ✅ | | ## 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 e174bd37475..b1da89aacc2 100644 --- a/native/core/src/execution/expressions/mod.rs +++ b/native/core/src/execution/expressions/mod.rs @@ -22,6 +22,7 @@ pub mod bitwise; pub mod comparison; pub mod list_positions; 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 80f8d890a34..20b383fe2e5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -40,6 +40,7 @@ use crate::execution::operators::IcebergWriteExec; use crate::execution::operators::{PartitionedRankLimitExec, WindowFnKind}; use crate::execution::{ expressions::list_positions::ListPositionsExpr, + expressions::map_entries::MapEntriesExpr, expressions::subquery::Subquery, operators::{ CometFilterExec, ExecutionError, ExpandExec, ExplodeExec, ParquetCompression, @@ -2076,7 +2077,7 @@ impl PhysicalPlanner { let (scans, shuffle_scans, child) = self.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 { self.create_expr(child_expr, child.schema())? } else { @@ -2092,6 +2093,22 @@ impl PhysicalPlanner { .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()), + _ => unreachable!("Map entries must be a struct"), + }, + _ => 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. @@ -2186,11 +2203,22 @@ impl PhysicalPlanner { } }; - 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)); @@ -2220,7 +2248,7 @@ impl PhysicalPlanner { let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_unnests, output_schema, unnest_options, )?); @@ -6342,11 +6370,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, @@ -6365,6 +6405,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] { @@ -6443,20 +6532,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 { @@ -6476,6 +6566,23 @@ mod tests { }) .collect(); assert_eq!(values, expected_values.repeat(2), "{context}"); + if map { + 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/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 7b3a7aec0d2..2619cf0be71 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -945,9 +945,9 @@ message Expand { } 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/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 952b53def8e..a465cebba16 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 @@ -1711,12 +1711,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..881186e264e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/explode.sql @@ -328,7 +328,7 @@ 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 ===== +-- ===== Map inputs: outer form preserves empty and null maps ===== statement CREATE TABLE test_explode_map(id int, m map) USING parquet @@ -339,7 +339,7 @@ INSERT INTO test_explode_map VALUES (2, map()), (3, NULL) -query expect_fallback(Comet only supports explode/explode_outer for arrays, not maps) +query SELECT id, explode_outer(m) FROM test_explode_map -- ===== TIMESTAMP_NTZ (Spark 3.4+): distinct Arrow layout (no tz) from LTZ ===== 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..d31c00619c4 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql @@ -93,8 +93,8 @@ 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) +-- Map entries expand to position, key and value columns. +query SELECT id, posexplode(m) FROM test_posexplode_map -- ===== posexplode_outer across non-int element types ===== 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..5e21151cfd2 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/explode.sql @@ -0,0 +1,93 @@ +-- 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 + +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 + +-- 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 +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/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index 4cb6a7f7efd..a81e73dbecd 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -224,16 +224,31 @@ 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 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 + } + (i, m) + } + 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 query = input.toDF("id", "m").selectExpr("id", "m", s"$generator(m)") + val (_, plan) = checkSparkAnswerAndOperator(query) + assert(collect(plan) { case e: CometExplodeExec => e }.nonEmpty) + checkSparkSchema(query) + } + } } } @@ -390,16 +405,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") } } 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 aa16417ca2c..06d5008c0b1 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 @@ -59,11 +59,8 @@ 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. + * Only array inputs are covered. The nesting group below reaches its event list through an + * `array>`. */ object CometExplodeBenchmark extends CometBenchmarkBase { @@ -382,9 +379,7 @@ 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 an `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 From d940195338997ba6879e5720fa42b89afa643991 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Fri, 2 Oct 2026 18:05:36 +0800 Subject: [PATCH 2/9] test: recognize missing Parquet field IDs through nested causes --- .../comet/parquet/ParquetReadSuite.scala | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index c5b1d04f540..107dab57360 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -2456,12 +2456,24 @@ abstract class ParquetReadSuite extends CometTestBase { } } - // Spark's own assertion from `ParquetFieldIdIOSuite`: the `SparkException` a read raises has - // the `RuntimeException` from `ParquetReadSupport` as its cause. + test("missing field ids are recognized through nested read exceptions") { + val message = "Parquet file schema doesn't contain any field Ids" + val missingIds = new RuntimeException(message) + val readError = new SparkException("Encountered error while reading file", missingIds) + val taskError = new SparkException("Job aborted", readError) + Seq(taskError, readError, missingIds, new RuntimeException(null: String, missingIds)) + .foreach(error => assert(isMissingFieldIdsError(error))) + + Seq(new RuntimeException("unrelated failure"), new RuntimeException(), new Exception(message)) + .foreach(error => assert(!isMissingFieldIdsError(error))) + } + + // Spark may wrap the `RuntimeException` from `ParquetReadSupport` in multiple read/task errors. private def isMissingFieldIdsError(error: Throwable): Boolean = { - val cause = error.getCause - cause.isInstanceOf[RuntimeException] && - cause.getMessage.contains("Parquet file schema doesn't contain any field Ids") + causeChain(error).exists { cause => + cause.isInstanceOf[RuntimeException] && + Option(cause.getMessage).exists(_.contains("Parquet file schema doesn't contain any field Ids")) + } } // Spark and Comet both raise the missing field ids error for `df`, and Comet plans the read From 6a23a4f83de77edd9bc342ec2df8fc71929020e8 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 3 Oct 2026 04:26:00 +0800 Subject: [PATCH 3/9] test: strengthen native map explode coverage and benchmarks Cover JVM-fed maps, nullable Boolean batch boundaries, all-NULL/empty outer batches, native map producers, and stored floating-point keys. Require native generators and matching Spark results and schemas. Benchmark map variants and nested map/array inputs with sinks that consume keys and positions. Remove the unrelated Parquet field-ID assertion change so it can be reviewed independently. --- .../sql-tests/expressions/map/explode.sql | 49 ++++++++ .../expressions/map/explode_float_keys.sql | 47 +++++++ .../comet/exec/CometGenerateExecSuite.scala | 52 +++++++- .../comet/parquet/ParquetReadSuite.scala | 23 +--- .../sql/benchmark/CometExplodeBenchmark.scala | 119 ++++++++++++------ 5 files changed, 232 insertions(+), 58 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql diff --git a/spark/src/test/resources/sql-tests/expressions/map/explode.sql b/spark/src/test/resources/sql-tests/expressions/map/explode.sql index 5e21151cfd2..6a47e40aa96 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/explode.sql @@ -47,6 +47,55 @@ 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 +-- 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 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 da82bbee659..7feefee3b26 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -225,6 +225,20 @@ class CometGenerateExecSuite extends CometTestBase { } 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", @@ -238,20 +252,48 @@ class CometGenerateExecSuite extends CometTestBase { j -> (if (j % 3 == 0) null else java.lang.Integer.valueOf(i * 100 + j)) }.toMap } - (i, m) + 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 query = input.toDF("id", "m").selectExpr("id", "m", s"$generator(m)") - val (_, plan) = checkSparkAnswerAndOperator(query) - assert(collect(plan) { case e: CometExplodeExec => e }.nonEmpty) - checkSparkSchema(query) + val maps = input.toDF("id", "ints", "booleans") + Seq("ints", "booleans").foreach { name => + withClue(s"$name: ") { checkMapGenerator(maps, name, generator) } + } } } } } + for (generator <- Seq("explode_outer", "posexplode_outer")) { + test(s"$generator with batches containing only null and empty maps") { + withSQLConf( + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "4") { + val rows = (0 until 16).map { i => + val m = if (i % 2 == 0) null else Map.empty[Int, java.lang.Boolean] + (i, m) + } + withParquetDataFrame(rows) { input => + checkMapGenerator(input.toDF("id", "m"), "m", generator) + } + } + } + } + + 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) + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 51d28862718..2a595398367 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -2634,25 +2634,12 @@ abstract class ParquetReadSuite extends CometTestBase { } } - test("missing field ids are recognized through nested read exceptions") { - val message = "Parquet file schema doesn't contain any field Ids" - val missingIds = new RuntimeException(message) - val readError = new SparkException("Encountered error while reading file", missingIds) - val taskError = new SparkException("Job aborted", readError) - Seq(taskError, readError, missingIds, new RuntimeException(null: String, missingIds)) - .foreach(error => assert(isMissingFieldIdsError(error))) - - Seq(new RuntimeException("unrelated failure"), new RuntimeException(), new Exception(message)) - .foreach(error => assert(!isMissingFieldIdsError(error))) - } - - // Spark may wrap the `RuntimeException` from `ParquetReadSupport` in multiple read/task errors. + // Spark's own assertion from `ParquetFieldIdIOSuite`: the `SparkException` a read raises has + // the `RuntimeException` from `ParquetReadSupport` as its cause. private def isMissingFieldIdsError(error: Throwable): Boolean = { - causeChain(error).exists { cause => - cause.isInstanceOf[RuntimeException] && - Option(cause.getMessage).exists( - _.contains("Parquet file schema doesn't contain any field Ids")) - } + val cause = error.getCause + cause.isInstanceOf[RuntimeException] && + cause.getMessage.contains("Parquet file schema doesn't contain any field Ids") } // Spark and Comet both raise the missing field ids error for `df`, and Comet plans the read 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 360e3885e81..faafcb4d98a 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,8 @@ package org.apache.spark.sql.benchmark import java.io.File import org.apache.spark.sql.catalyst.optimizer.InferFiltersFromGenerate +import org.apache.spark.sql.catalyst.plans.logical.Generate +import org.apache.spark.sql.comet.CometExplodeExec import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -59,8 +61,8 @@ 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. The nesting group below reaches its event list through an - * `array>`. + * Map cases consume both keys and values. The nesting group models event lists as a map keyed by + * platform, including a second generator over the entries in each map value. */ object CometExplodeBenchmark extends CometBenchmarkBase { @@ -106,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. @@ -150,20 +157,19 @@ 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) - /** [[eventsArray]] wrapped in the struct chain that puts it eight accessors down. */ + /** [[eventsMap]] wrapped in the struct chain that puts it eight accessors down. */ private val profileColumn = { val sessions = s"""named_struct( | 'device', named_struct('type', 'mobile', 'os', 'iOS'), - | 'events', $eventsArray)""".stripMargin + | 'events', $eventsMap)""".stripMargin val account = accountPath.foldRight(s"named_struct('sessions', $sessions)") { (field, inner) => s"named_struct('$field', $inner)" } @@ -176,16 +182,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 map at two different depths, and keeping them in one + * view is what makes them the same map. 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", @@ -205,7 +214,7 @@ object CometExplodeBenchmark extends CometBenchmarkBase { 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"$eventsMap AS events", s"$profileColumn AS profile"))) /** Writes a view's rows to Parquet and registers it. */ @@ -216,8 +225,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` @@ -227,24 +236,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 = @@ -259,14 +270,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`. @@ -287,7 +304,8 @@ object CometExplodeBenchmark extends CometBenchmarkBase { 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(query) + val actual = df.collect().head.toSeq.map { case null => null case value: Number => value.longValue() case other => other @@ -297,6 +315,18 @@ object CometExplodeBenchmark extends CometBenchmarkBase { s"$name: $engine produced $actual, expected $expected. The case is not aggregating " + s"what it names.\n$query") } + if (cometEnabled) { + val expectedGenerators = df.queryExecution.optimizedPlan.collect { case _: Generate => + 1 + }.size + val nativeGenerators = collect(df.queryExecution.executedPlan) { + case e: CometExplodeExec => e + }.size + assert( + expectedGenerators > 0 && nativeGenerators == expectedGenerators, + s"$name executed $nativeGenerators of $expectedGenerators generators natively:\n" + + df.queryExecution.executedPlan) + } } } @@ -335,6 +365,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) => @@ -379,15 +427,15 @@ 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 - * outer container is an `array>`. + * outer container is a `map>>>` keyed by platform. * * 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. + * is `struct>>>` 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 + * `depth 1` and `depth 8` explode the same map, 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 @@ -396,13 +444,14 @@ object CometExplodeBenchmark extends CometBenchmarkBase { private def nestedInputCases: Seq[(String, Case)] = { 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 map[$eventsLen] at $label" -> generatorCase( "explode", "nested", eventsLen, rows = nestedRows, - arrayExpr = path, + collectionExpr = path, + mapInput = true, carried = carried) } @@ -411,13 +460,13 @@ object CometExplodeBenchmark extends CometBenchmarkBase { val events = nonEmptyRows(nestedRows) * eventsLen val entries = events * entriesLen val outer = - "SELECT k, region, ev.platform AS platform, ev.entries AS entries " + - s"FROM (SELECT k, region, explode($deepEvents) AS ev FROM nested)" + "SELECT k, region, platform, ev.entries AS entries " + + s"FROM (SELECT k, region, explode($deepEvents) AS (platform, ev) FROM nested)" 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))) + single :+ (s"explode map[$eventsLen] at depth 8, then its inner array" -> + Case(chained, Seq(entries, entries, entries * 3, entries))) } } From e8156c26692d2cf23908850fd9e3b3d2e940ee0d Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 3 Oct 2026 04:26:00 +0800 Subject: [PATCH 4/9] perf: avoid unused allocations when unnesting collections Keep input arrays as temporary placeholders until generated columns replace them. Build take indices only when a carried column consumes them. Cover recursive unnest with and without carried columns and both drop and outer null handling. Preserve the existing bounded map and array execution path. --- native/operators/src/explode.rs | 130 ++++++++++++++++++++++++++------ 1 file changed, 106 insertions(+), 24 deletions(-) diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 5af1e052d7b..9cd69d42cba 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -23,16 +23,18 @@ //! 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; //! * 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`. +//! * retaining input arrays as unnest placeholders and skipping take indices when no column +//! consumes them, avoiding allocations whose contents never reach the output. //! //! # 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,14 +53,13 @@ //! * 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, + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, LargeListArray, + LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray, }; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; @@ -685,8 +686,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()) @@ -715,7 +714,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 a placeholder for an unnest result. No column consumes take indices. + batch.to_vec() + }; Ok(Some(ret)) } @@ -1250,8 +1255,8 @@ fn create_take_indices( /// 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` -/// in the `repeat_mask` will be replaced with arrays filled with nulls of the -/// appropriate length. +/// are placeholders: they are replaced by their unnested versions before the +/// output batch is built. Keep those arrays instead of allocating unused null arrays. /// /// For example if we have the following batch: /// @@ -1281,18 +1286,6 @@ fn create_take_indices( /// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' /// ``` /// -/// The `repeat_mask` determines whether an array's values are repeated or replaced with nulls. -/// For example, if the `repeat_mask` is: -/// -/// ```ignore -/// [true, false] -/// ``` -/// -/// The final batch will look like: -/// -/// ```ignore -/// c1: 1, null, 2, 3, 4, null, 5, 6 // Repeated using `indices` -/// c2: null, null, null, null, null, null, null, null // Replaced with nulls fn repeat_arrs_from_indices( batch: &[ArrayRef], indices: &PrimitiveArray, @@ -1305,7 +1298,7 @@ fn repeat_arrs_from_indices( if repeat { Ok(kernels::take::take(arr, indices, None)?) } else { - Ok(new_null_array(arr.data_type(), arr.len())) + Ok(Arc::clone(arr)) } }) .collect() @@ -1325,6 +1318,95 @@ mod tests { use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; + #[test] + fn unnest_placeholders_share_the_input() { + let list = Arc::clone(list_batch(&[Some(2), Some(1)]).column(0)); + let carried: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let indices = Int64Array::from(vec![0, 0, 1]); + let result = + repeat_arrs_from_indices(&[Arc::clone(&list), carried], &indices, &[false, true]) + .unwrap(); + assert!(Arc::ptr_eq(&result[0], &list)); + assert_eq!( + result[1].as_primitive::().values(), + &[10, 10, 20] + ); + } + + #[test] + fn recursive_unnest_replaces_placeholders_with_and_without_carried_columns() { + let inner = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(10), None]), + None, + Some(vec![]), + Some(vec![Some(20)]), + Some(vec![Some(30), Some(40)]), + ]); + let outer: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 5, 5, 5])), + Arc::new(inner), + Some(NullBuffer::from(vec![true, true, false, true])), + )); + for carried in [false, true] { + for handling in [NullHandling::Drop, NullHandling::PreserveAndExpandEmpty] { + let mut columns = vec![Arc::clone(&outer)]; + let mut fields = vec![Field::new("lists", outer.data_type().clone(), true)]; + let mut output_fields = vec![Field::new("value", DataType::Int32, true)]; + if carried { + columns.push(Arc::new(Int32Array::from(vec![1, 2, 3, 4]))); + fields.push(Field::new("id", DataType::Int32, false)); + output_fields.push(Field::new("id", DataType::Int32, false)); + } + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let output = build_batch( + &input, + &Arc::new(Schema::new(output_fields)), + &[ListUnnest { + index_in_input_schema: 0, + depth: 2, + }], + &HashSet::new(), + &UnnestOptions::new().with_null_handling(handling), + None, + ) + .unwrap() + .unwrap(); + let (expected, ids) = if handling == NullHandling::Drop { + ( + vec![Some(10), None, Some(20), Some(30), Some(40)], + vec![1, 1, 2, 2, 2], + ) + } else { + ( + vec![ + Some(10), + None, + None, + None, + Some(20), + Some(30), + Some(40), + None, + None, + ], + vec![1, 1, 1, 1, 2, 2, 2, 3, 4], + ) + }; + assert_eq!( + output.column(0).as_primitive::(), + &Int32Array::from(expected) + ); + if carried { + assert_eq!( + output.column(1).as_primitive::(), + &Int32Array::from(ids) + ); + } + } + } + } + /// Build a single-column `List` batch where row `i` holds `lens[i]` elements, /// numbered consecutively across the whole batch. `None` is a NULL list. fn list_batch(lens: &[Option]) -> RecordBatch { From ff55b13cdd88d983465876697dda19482e38f198 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 3 Oct 2026 21:58:28 +0800 Subject: [PATCH 5/9] perf: repeat carried short strings by input row in explode --- native/operators/benches/explode.rs | 193 +++++++++++++++++++- native/operators/src/explode.rs | 273 ++++++++++++++++++++++++++-- 2 files changed, 443 insertions(+), 23 deletions(-) diff --git a/native/operators/benches/explode.rs b/native/operators/benches/explode.rs index ef14e722548..c5e624fa641 100644 --- a/native/operators/benches/explode.rs +++ b/native/operators/benches/explode.rs @@ -28,11 +28,13 @@ 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::datatypes::{DataType, Field, Fields, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use datafusion::common::{NullHandling, UnnestOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::TaskContext; @@ -264,6 +266,173 @@ fn run(runtime: &Runtime, plan: &Arc, ctx: &Arc) assert!(!batches.is_empty()); } +// The inner generator in CometExplodeBenchmark's chained map case carries k, region and +// platform through five entries per map value. The existing carried-column group only has +// Int64 columns, so it cannot expose the string gathers in this shape. +const CHAINED_FAN_OUT: usize = 5; +const K: usize = 0; +const REGION: usize = 1; +const PLATFORM: usize = 2; + +fn chained_entries(start: usize) -> ListArray { + let entry_fields = Fields::from(vec![ + Field::new("type", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("page", DataType::Utf8, true), + Field::new("source", DataType::Utf8, true), + ]); + let entry_start = start * CHAINED_FAN_OUT; + let entry_count = ROWS_PER_BATCH * CHAINED_FAN_OUT; + let entries = StructArray::new( + entry_fields, + vec![ + Arc::new(StringArray::from_iter_values( + (0..entry_count).map(|i| format!("t_{}", i % CHAINED_FAN_OUT)), + )), + Arc::new(Int64Array::from_iter_values( + (entry_start..entry_start + entry_count).map(|i| i as i64), + )), + Arc::new(StringArray::from_iter_values( + (entry_start..entry_start + entry_count).map(|i| format!("page_{i}")), + )), + Arc::new(StringArray::from_iter_values( + (0..entry_count).map(|i| format!("src_{}", i % 4)), + )), + ], + None, + ); + ListArray::new( + Arc::new(Field::new("item", entries.data_type().clone(), true)), + OffsetBuffer::from_lengths(std::iter::repeat_n(CHAINED_FAN_OUT, ROWS_PER_BATCH)), + Arc::new(entries), + None, + ) +} + +fn chained_plan(carried: &[usize], nullable: bool) -> (Arc, Vec) { + let batches: Vec<_> = (0..BATCHES) + .map(|batch| { + let start = batch * ROWS_PER_BATCH; + let list = chained_entries(start); + let k = Int64Array::from_iter( + (start..start + ROWS_PER_BATCH) + .map(|i| (!nullable || !i.is_multiple_of(10)).then_some((i / 4) as i64)), + ); + let region = + StringArray::from_iter((start..start + ROWS_PER_BATCH).map(|i| { + (!nullable || !i.is_multiple_of(10)).then(|| format!("r_{}", i / 4 % 8)) + })); + let platform = StringArray::from_iter_values( + (start..start + ROWS_PER_BATCH).map(|i| format!("p_{}", i % 4 + 1)), + ); + let payload: [ArrayRef; 3] = [Arc::new(k), Arc::new(region), Arc::new(platform)]; + let names = ["k", "region", "platform"]; + let mut fields = vec![Field::new("entries", list.data_type().clone(), true)]; + let mut columns: Vec = vec![Arc::new(list)]; + for &i in carried { + fields.push(Field::new( + names[i], + payload[i].data_type().clone(), + i != PLATFORM, + )); + columns.push(Arc::clone(&payload[i])); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + }) + .collect(); + let schema = batches[0].schema(); + let output_fields = std::iter::once(Arc::new(Field::new( + "entry", + batches[0].column(0).as_list::().value_type(), + true, + ))) + .chain(schema.fields().iter().skip(1).cloned()) + .collect::>(); + let source = + MemorySourceConfig::try_new_exec(std::slice::from_ref(&batches), schema, None).unwrap(); + let plan = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + Arc::new(Schema::new(output_fields)), + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + (Arc::new(plan), batches) +} + +fn verify_chained( + runtime: &Runtime, + plan: &Arc, + ctx: &Arc, + input: &[RecordBatch], + carried: &[usize], + nullable: bool, +) { + let output = runtime + .block_on(collect(plan.execute(0, Arc::clone(ctx)).unwrap())) + .unwrap(); + let mut row = 0; + let mut null_counts = vec![0; carried.len()]; + for batch in output { + assert!(batch.num_rows() <= ROWS_PER_BATCH); + let entries = batch.column(0).as_struct(); + let source_batch = row / (ROWS_PER_BATCH * CHAINED_FAN_OUT); + let source_list = input[source_batch].column(0).as_list::(); + let source_entries = source_list.values().as_struct(); + // String slicing retains the full values buffer. A gather would allocate a new one. + let page = entries.column(2).as_string::(); + let source_page = source_entries.column(2).as_string::(); + assert_eq!( + page.value_data().as_ptr(), + source_page.value_data().as_ptr() + ); + let ts = entries.column(1).as_primitive::(); + for (position, &column) in carried.iter().enumerate() { + let array = batch.column(position + 1); + null_counts[position] += array.null_count(); + for j in 0..batch.num_rows() { + let input_row = (row + j) / CHAINED_FAN_OUT; + let is_null = nullable && column != PLATFORM && input_row.is_multiple_of(10); + assert_eq!(array.is_null(j), is_null); + if is_null { + continue; + } + if column == K { + assert_eq!( + array.as_primitive::().value(j), + (input_row / 4) as i64 + ); + } else { + let expected = if column == REGION { + format!("r_{}", input_row / 4 % 8) + } else { + format!("p_{}", input_row % 4 + 1) + }; + let value = array.as_string::().value(j); + assert_eq!(value, expected); + } + } + } + for j in 0..batch.num_rows() { + assert_eq!(ts.value(j), (row + j) as i64); + } + row += batch.num_rows(); + } + assert_eq!(row, ROWS_PER_BATCH * BATCHES * CHAINED_FAN_OUT); + for (position, &column) in carried.iter().enumerate() { + let expected = if nullable && column != PLATFORM { + (ROWS_PER_BATCH * BATCHES).div_ceil(10) * CHAINED_FAN_OUT + } else { + 0 + }; + assert_eq!(null_counts[position], expected); + } +} + fn criterion_benchmark(c: &mut Criterion) { let runtime = Runtime::new().unwrap(); let ctx = Arc::new( @@ -320,6 +489,24 @@ fn criterion_benchmark(c: &mut Criterion) { } } group.finish(); + + let mut group = c.benchmark_group("explode_chained_carry"); + group.throughput(Throughput::Elements( + (ROWS_PER_BATCH * BATCHES * CHAINED_FAN_OUT) as u64, + )); + for (name, carried, nullable) in [ + ("none", &[][..], false), + ("k", &[K][..], false), + ("region", &[REGION][..], false), + ("platform", &[PLATFORM][..], false), + ("full", &[K, REGION, PLATFORM][..], false), + ("full_nulls", &[K, REGION, PLATFORM][..], true), + ] { + let (plan, input) = chained_plan(carried, nullable); + verify_chained(&runtime, &plan, &ctx, &input, carried, nullable); + group.bench_function(name, |b| b.iter(|| run(&runtime, &plan, &ctx))); + } + group.finish(); } criterion_group!(benches, criterion_benchmark); diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 9cd69d42cba..cc2159db142 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 //! @@ -31,6 +31,8 @@ //! values instead of gathering them, and the buffer fills in `create_take_indices`. //! * retaining input arrays as unnest placeholders and skipping take indices when no column //! consumes them, avoiding allocations whose contents never reach the output. +//! * repeating carried short strings by input row, with fixed-size byte copies instead of +//! a variable-size copy for every output row. //! //! # Deleting this file //! @@ -59,8 +61,10 @@ use arrow::array::{ Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, LargeListArray, - LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray, + LargeListViewArray, ListArray, ListViewArray, NullBufferBuilder, PrimitiveArray, Scalar, + StringArray, 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}; @@ -715,8 +719,7 @@ 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 = 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)? + repeat_arrs(batch, unnested_length, total_length, &repeat_mask)? } else { // Every column is a placeholder for an unnest result. No column consumes take indices. batch.to_vec() @@ -1252,8 +1255,8 @@ fn create_take_indices( PrimitiveArray::::from(indices) } -/// 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 +/// Create a batch of arrays based on an input `batch` and per-row repetition lengths. +/// Take indices are created only for columns that need the take kernel to repeat values /// that are marked with `true` in the `repeat_mask`. Arrays marked with `false` /// are placeholders: they are replaced by their unnested versions before the /// output batch is built. Keep those arrays instead of allocating unused null arrays. @@ -1272,11 +1275,10 @@ fn create_take_indices( /// c1: 1, null, 2, 3, 4, null, 5, 6 /// ``` /// -/// And the `indices` array contains the indices that are used by `take` kernel to -/// repeat the values in `c2`: +/// Each value in `c2` is repeated by its row's length: /// /// ```ignore -/// 0, 1, 2, 2, 2, 3, 4, 4 +/// 1, 1, 3, 1, 2 /// ``` /// /// so that the final batch will look like: @@ -1286,16 +1288,25 @@ fn create_take_indices( /// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' /// ``` /// -fn repeat_arrs_from_indices( +fn repeat_arrs( batch: &[ArrayRef], - indices: &PrimitiveArray, + lengths: &PrimitiveArray, + total_length: usize, repeat_mask: &[bool], ) -> Result>> { + let mut indices = None; batch .iter() .zip(repeat_mask.iter()) .map(|(arr, &repeat)| { if repeat { + if let Some(strings) = arr.as_string_opt::() { + if let Some(strings) = repeat_short_strings(strings, lengths, total_length) { + return Ok(Arc::new(strings) as ArrayRef); + } + } + let indices = + indices.get_or_insert_with(|| create_take_indices(lengths, total_length)); Ok(kernels::take::take(arr, indices, None)?) } else { Ok(Arc::clone(arr)) @@ -1304,6 +1315,115 @@ fn repeat_arrs_from_indices( .collect() } +/// Repeat short Utf8 values using the input runs rather than gathering every output value. +/// Batches dominated by long strings or with less than twofold expansion use Arrow's take instead. +fn repeat_short_strings( + array: &StringArray, + lengths: &PrimitiveArray, + total_length: usize, +) -> Option { + const SHORT_STRING_BYTES: usize = 16; + + debug_assert_eq!(array.len(), lengths.len()); + debug_assert_eq!(lengths.null_count(), 0); + debug_assert!(lengths.values().iter().all(|&n| n >= 0)); + if total_length < array.len().saturating_mul(2) { + return None; + } + if array.null_count() == array.len() { + return Some(StringArray::new_null(total_length)); + } + + // Reject wide input batches in O(1), before scanning their repetition lengths. This + // also avoids paying a guard pass followed by take for mixed short/long batches. + let input_offsets = array.value_offsets(); + let input_bytes = (input_offsets[array.len()] - input_offsets[0]) as usize; + if input_bytes > array.len().checked_mul(SHORT_STRING_BYTES)? { + return None; + } + let mut capacity = 0usize; + let mut valid_rows = 0usize; + for (i, (range, &repeat)) in input_offsets.windows(2).zip(lengths.values()).enumerate() { + if repeat == 0 || array.is_null(i) { + continue; + } + let len = (range[1] - range[0]) as usize; + capacity = capacity.checked_add(len.checked_mul(repeat as usize)?)?; + valid_rows = valid_rows.checked_add(repeat as usize)?; + // Stop early on long-string batches. Occasional longer values among short strings + // can use exact-size copies below, avoiding a full guard pass followed by take. + if capacity > valid_rows.checked_mul(SHORT_STRING_BYTES)? { + return None; + } + } + // Leave offset overflow to the existing take kernel, including its error reporting. + i32::try_from(capacity).ok()?; + let mut values = vec![0u8; capacity.checked_add(SHORT_STRING_BYTES)?]; + let mut offsets = Vec::with_capacity(total_length + 1); + offsets.push(0i32); + let mut nulls = (array.null_count() > 0).then(|| NullBufferBuilder::new(total_length)); + let data = array.value_data(); + let mut position = 0usize; + for (i, (range, &repeat)) in input_offsets.windows(2).zip(lengths.values()).enumerate() { + let repeat = repeat as usize; + if repeat == 0 { + continue; + } + if array.is_null(i) { + offsets.resize(offsets.len() + repeat, position as i32); + nulls.as_mut().unwrap().append_n_nulls(repeat); + continue; + } + let start = range[0] as usize; + let len = (range[1] - range[0]) as usize; + // Fixed-size copies avoid a tiny variable-size memcpy per output value. The source + // must contain the whole copy, even for a sliced array's final value. Extra bytes + // overwrite only future values or the output padding; later copies replace + // them before we truncate to the logical byte length. + if len == 0 { + offsets.resize(offsets.len() + repeat, position as i32); + } else if len <= 8 && data.len() - start >= 8 { + let value = &data[start..start + 8]; + for _ in 0..repeat { + values[position..position + 8].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } else if len <= SHORT_STRING_BYTES && data.len() - start >= SHORT_STRING_BYTES { + let value = &data[start..start + SHORT_STRING_BYTES]; + for _ in 0..repeat { + values[position..position + SHORT_STRING_BYTES].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } else { + let value = &data[start..start + len]; + for _ in 0..repeat { + values[position..position + len].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } + if let Some(nulls) = &mut nulls { + nulls.append_n_non_nulls(repeat); + } + } + debug_assert_eq!(position, capacity); + debug_assert_eq!(offsets.len(), total_length + 1); + values.truncate(capacity); + let nulls = nulls + .and_then(|mut nulls| nulls.finish()) + .filter(|nulls| nulls.null_count() > 0); + // SAFETY: offsets are monotonic, fit i32, and end at the checked byte capacity. Each + // logical range copies a complete input Utf8 value, relying on the input array's Utf8 + // invariant; null ranges are empty. Safe slices bound every copy, values is fully + // initialized, and there is one validity entry per output value, by construction. + // Avoid validating the same input Utf8 again for every repetition. + Some(unsafe { + StringArray::new_unchecked(OffsetBuffer::new(offsets.into()), values.into(), nulls) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -1318,19 +1438,132 @@ mod tests { use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; + #[test] + fn short_string_repetition_matches_take_buffers() { + let words = [ + "", + "a", + "abc", + "four", + "eight888", + "nine99999", + "sixteen_12345678", + "中文", + "😀", + "z", + ]; + for null_stride in [0usize, 2, 7, 1] { + let source = StringArray::from_iter((0..60).map(|i| { + (null_stride == 0 || i % null_stride != 0).then_some(words[i % words.len()]) + })); + for start in [0usize, 3, 59] { + let source = source.slice(start, 60 - start); + for fanout in [1i64, 5, 100] { + let lengths = Int64Array::from_iter_values((0..source.len()).map(|i| { + if i % 4 == 0 { + 0 + } else { + fanout + } + })); + let total = lengths.values().iter().sum::() as usize; + if total > source.len() { + assert_string_repeat_buffers(&source, &lengths); + } + } + } + } + // The last value has no following bytes to support a wider source copy. + for word in words { + assert_string_repeat_buffers( + &StringArray::from(vec![word]), + &Int64Array::from(vec![5]), + ); + } + // A null can own nonempty bytes; take must discard them. The middle slice also + // starts inside both the offset buffer and the validity bitmap. + let source = StringArray::new( + OffsetBuffer::new(vec![0, 3, 8, 9, 12].into()), + arrow::buffer::Buffer::from_slice_ref(b"abcHIDDENxyz"), + Some(NullBuffer::from(vec![true, false, false, true])), + ); + for lengths in [vec![5, 5, 5, 5], vec![5, 0, 0, 5], vec![0, 5, 5, 0]] { + assert_string_repeat_buffers(&source, &Int64Array::from(lengths)); + } + assert_string_repeat_buffers(&source.slice(1, 3), &Int64Array::from(vec![1, 0, 5])); + let valid = StringArray::new( + source.offsets().clone(), + source.values().clone(), + Some(NullBuffer::new_valid(source.len())), + ); + assert_string_repeat_buffers(&valid, &Int64Array::from(vec![5, 5, 5, 5])); + } + + fn assert_string_repeat_buffers(source: &StringArray, lengths: &Int64Array) { + let total = lengths.values().iter().sum::() as usize; + let actual = repeat_short_strings(source, lengths, total).expect("short string fast path"); + let indices = create_take_indices(lengths, total); + let expected = kernels::take::take(source, &indices, None).unwrap(); + let expected = expected.as_string::(); + assert_eq!(actual.value_offsets(), expected.value_offsets()); + assert_eq!(actual.value_data(), expected.value_data()); + assert_eq!(actual.nulls(), expected.nulls()); + actual.to_data().validate_full().unwrap(); + } + + #[test] + fn short_string_repetition_defers_long_unexpanded_and_overflowing_values() { + let lengths = Int64Array::from(vec![5]); + let long = StringArray::from(vec!["this string is longer than sixteen bytes"]); + assert!(repeat_short_strings(&long, &lengths, 5).is_none()); + let short = StringArray::from(vec!["abc"]); + assert!(repeat_short_strings(&short, &Int64Array::from(vec![1]), 1).is_none()); + assert!(repeat_short_strings(&short, &Int64Array::from(vec![0]), 0).is_none()); + let overflow = i32::MAX as i64; + assert!( + repeat_short_strings(&short, &Int64Array::from(vec![overflow]), overflow as usize) + .is_none() + ); + // Long values that are not selected do not disqualify the short values. + let mixed = StringArray::from(vec!["a very long unselected value", "abc"]); + assert_string_repeat_buffers(&mixed, &Int64Array::from(vec![0, 5])); + // A late long value among short strings must not trigger a complete guard pass + // followed by take. It uses an exact-size copy while the short values use wider ones. + let mixed = StringArray::from_iter_values((0..32).map(|i| { + if i == 31 { + "x".repeat(64) + } else { + "abc".to_owned() + } + })); + assert_string_repeat_buffers(&mixed, &Int64Array::from(vec![5; 32])); + let almost_flat = StringArray::from(vec!["a", "b", "c", "d"]); + assert!( + repeat_short_strings(&almost_flat, &Int64Array::from(vec![2, 1, 1, 1]), 5).is_none() + ); + } + #[test] fn unnest_placeholders_share_the_input() { let list = Arc::clone(list_batch(&[Some(2), Some(1)]).column(0)); - let carried: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); - let indices = Int64Array::from(vec![0, 0, 1]); - let result = - repeat_arrs_from_indices(&[Arc::clone(&list), carried], &indices, &[false, true]) - .unwrap(); + let batch: Vec = vec![ + Arc::clone(&list), + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(StringArray::from(vec!["abc", "中文"])), + Arc::new(StringArray::from(vec![ + "a long value that uses Arrow take", + "another value longer than sixteen bytes", + ])), + Arc::new(Int64Array::from(vec![Some(1), None])), + ]; + let lengths = Int64Array::from(vec![2, 2]); + let result = repeat_arrs(&batch, &lengths, 4, &[false, true, true, true, true]).unwrap(); assert!(Arc::ptr_eq(&result[0], &list)); - assert_eq!( - result[1].as_primitive::().values(), - &[10, 10, 20] - ); + let indices = create_take_indices(&lengths, 4); + for (input, output) in batch.iter().zip(&result).skip(1) { + let expected = kernels::take::take(input, &indices, None).unwrap(); + assert_eq!(output.to_data(), expected.to_data()); + } } #[test] From 6415ed3a89aab5837964c669001c10ee945f51ad Mon Sep 17 00:00:00 2001 From: rich7420 Date: Wed, 28 Oct 2026 17:24:22 +0800 Subject: [PATCH 6/9] fix: bound map expansion before repeating carried values Split individual List rows before materialization and cap repetition so carried buffers and child offsets stay within Arrow's offset limits. Preserve shared entry buffers, positions and outer semantics. --- native/operators/src/explode.rs | 372 +++++++++++++++++- .../comet/exec/CometGenerateExecSuite.scala | 18 + 2 files changed, 369 insertions(+), 21 deletions(-) diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 1b494d12ba5..f3999778c00 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -33,6 +33,8 @@ //! consumes them, avoiding allocations whose contents never reach the output. //! * repeating carried short strings by input row, with fixed-size byte copies instead of //! a variable-size copy for every output row. +//! * splitting individual List rows before repeating carried values, including a repetition +//! bound that prevents copied variable-width buffers and child offsets from overflowing. //! //! # Deleting this file //! @@ -60,9 +62,9 @@ //! `ExplodeExec` and `ExplodeStream` were always Comet's own. use arrow::array::{ - Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, LargeListArray, - LargeListViewArray, ListArray, ListViewArray, NullBufferBuilder, PrimitiveArray, Scalar, - StringArray, StructArray, + Array, ArrayData, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, + LargeListArray, LargeListViewArray, ListArray, ListViewArray, NullBufferBuilder, + PrimitiveArray, Scalar, StringArray, StructArray, }; use arrow::buffer::OffsetBuffer; use arrow::compute::kernels::length::length; @@ -281,10 +283,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, @@ -307,6 +307,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, + /// Maximum expansion of one row per build, for Comet's depth-1 List inputs. + split_limit: 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 @@ -323,10 +327,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(); @@ -336,7 +337,10 @@ impl PendingInput { let batch_size = batch_size as i64; let mut output_rows = 0i64; for (rows, len) in lens.iter().enumerate() { - if rows > 0 && output_rows + len > batch_size { + if rows > 0 + && (output_rows + len > batch_size + || self.split_limit.is_some_and(|limit| *len as usize > limit)) + { return rows; } output_rows += len; @@ -351,6 +355,82 @@ 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)> { + if let Some(limit) = self.split_limit { + 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 row = self.batch.slice(self.row_offset, 1); + let mut columns = row.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(row.schema(), columns)?, + Some(Int64Array::from(vec![length as i64])), + )); + } + } + let rows = self.next_chunk_rows(batch_size); + let chunk = self.batch.slice(self.row_offset, rows); + let lengths = self.chunk_lengths(rows); + self.row_offset += rows; + Ok((chunk, lengths)) + } +} + +/// If every row repeats at most R times, no copied buffer or child offset span exceeds R +/// times its entire source span. Count elements too: List has no child value buffer. +fn repetition_limit(batch: &RecordBatch, list_columns: &[ListUnnest]) -> usize { + fn max_span(data: &ArrayData) -> usize { + std::iter::once(data.len()) + .chain(data.buffers().iter().map(|buffer| buffer.len())) + .chain(data.child_data().iter().map(max_span)) + .max() + .unwrap_or(0) + } + let span = batch + .columns() + .iter() + .enumerate() + .filter(|(index, _)| { + !list_columns + .iter() + .any(|list| list.index_in_input_schema == *index) + }) + .map(|(_, column)| max_span(&column.to_data())) + .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. @@ -369,10 +449,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, } @@ -399,10 +477,8 @@ impl ExplodeStream { // always consumes at least one row, so it is dropped the moment it drains. 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(); @@ -443,10 +519,24 @@ impl ExplodeStream { if batch.num_rows() > 0 { let timer = self.baseline_metrics.elapsed_compute().timer(); let output_lens = self.predict_output_lens(&batch); + let split_limit = (!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(|| { + self.batch_size + .min(repetition_limit(&batch, &self.list_type_columns)) + }); timer.done(); self.pending_input = Some(PendingInput { batch, row_offset: 0, + element_offset: 0, + split_limit, output_lens: output_lens?, }); } @@ -1754,7 +1844,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(); @@ -1762,6 +1851,247 @@ mod tests { assert_eq!(values(&batches), seq(25)); } + #[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( 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 7feefee3b26..682f8dbac43 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -270,6 +270,24 @@ class CometGenerateExecSuite extends CometTestBase { } } } + + 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])) + } + } + } } for (generator <- Seq("explode_outer", "posexplode_outer")) { From c4525cb85f0e7768a0fc40e4973896ff85b4e5c1 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Wed, 28 Oct 2026 18:27:14 +0800 Subject: [PATCH 7/9] perf: repeat carried values by input runs in explode Fill non-null primitive runs directly for expansions of at least five and build short-string offsets with one Vec extension per run. Retain Arrow take for smaller expansions and nullable primitives. --- native/operators/benches/explode.rs | 6 ++ native/operators/src/explode.rs | 85 ++++++++++++++++++++++++++++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/native/operators/benches/explode.rs b/native/operators/benches/explode.rs index c5e624fa641..6d37bf4ecce 100644 --- a/native/operators/benches/explode.rs +++ b/native/operators/benches/explode.rs @@ -463,6 +463,12 @@ fn criterion_benchmark(c: &mut Criterion) { b.iter(|| run(&runtime, &plan, &ctx)) }); } + for fan_out in [1usize, 2, 100] { + let plan = explode_plan(Element::Int64, fan_out, 3, RowMix::Dense, false); + group.bench_function(format!("fan_out_{fan_out}"), |b| { + b.iter(|| run(&runtime, &plan, &ctx)) + }); + } group.finish(); // NULL and empty rows under outer semantics are padded, so this is the shape that cannot be diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index f3999778c00..034b17259c6 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -33,6 +33,7 @@ //! consumes them, avoiding allocations whose contents never reach the output. //! * repeating carried short strings by input row, with fixed-size byte copies instead of //! a variable-size copy for every output row. +//! * filling non-null primitive runs directly instead of building and gathering 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. //! @@ -1390,6 +1391,15 @@ fn repeat_arrs( return Ok(Arc::new(strings) as ArrayRef); } } + if arr.null_count() == 0 && total_length >= arr.len().saturating_mul(5) { + let repeated = arrow::array::downcast_primitive_array! { + arr => { Some(Arc::new(repeat_primitive(arr, lengths, total_length)) as ArrayRef) } + _ => None, + }; + if let Some(repeated) = repeated { + return Ok(repeated); + } + } let indices = indices.get_or_insert_with(|| create_take_indices(lengths, total_length)); Ok(kernels::take::take(arr, indices, None)?) @@ -1400,6 +1410,19 @@ fn repeat_arrs( .collect() } +/// Repeat non-null primitives by input runs without allocating gather indices. +fn repeat_primitive( + array: &PrimitiveArray, + lengths: &Int64Array, + total_length: usize, +) -> PrimitiveArray { + let mut values = Vec::with_capacity(total_length); + for (&value, &repeat) in array.values().iter().zip(lengths.values()) { + values.resize(values.len() + repeat as usize, value); + } + PrimitiveArray::new(values.into(), None).with_data_type(array.data_type().clone()) +} + /// Repeat short Utf8 values using the input runs rather than gathering every output value. /// Batches dominated by long strings or with less than twofold expansion use Arrow's take instead. fn repeat_short_strings( @@ -1461,6 +1484,9 @@ fn repeat_short_strings( } let start = range[0] as usize; let len = (range[1] - range[0]) as usize; + if len > 0 { + offsets.extend((1..=repeat).map(|j| (position + len * j) as i32)); + } // Fixed-size copies avoid a tiny variable-size memcpy per output value. The source // must contain the whole copy, even for a sliced array's final value. Extra bytes // overwrite only future values or the output padding; later copies replace @@ -1472,21 +1498,18 @@ fn repeat_short_strings( for _ in 0..repeat { values[position..position + 8].copy_from_slice(value); position += len; - offsets.push(position as i32); } } else if len <= SHORT_STRING_BYTES && data.len() - start >= SHORT_STRING_BYTES { let value = &data[start..start + SHORT_STRING_BYTES]; for _ in 0..repeat { values[position..position + SHORT_STRING_BYTES].copy_from_slice(value); position += len; - offsets.push(position as i32); } } else { let value = &data[start..start + len]; for _ in 0..repeat { values[position..position + len].copy_from_slice(value); position += len; - offsets.push(position as i32); } } if let Some(nulls) = &mut nulls { @@ -1523,6 +1546,62 @@ mod tests { use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; + #[test] + fn primitive_repetition_preserves_buffers_and_type_metadata() { + use arrow::array::{Decimal128Array, Float64Array, TimestampMicrosecondArray}; + + let arrays: Vec = vec![ + Arc::new(Int64Array::from(vec![1, -2, 3, 4, 5])), + Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + ])), + Arc::new(Int64Array::from(vec![None; 5])), + Arc::new(Float64Array::from(vec![ + f64::from_bits(0x7ff0000000000001), + -0.0, + f64::NEG_INFINITY, + f64::from_bits(0x7ff8000000000002), + 1.0, + ])), + Arc::new( + TimestampMicrosecondArray::from(vec![0, -1, 1, 100, 200]) + .with_timezone("Asia/Taipei"), + ), + Arc::new( + Decimal128Array::from(vec![0, -12345, 12345, 100, 200]) + .with_precision_and_scale(20, 4) + .unwrap(), + ), + ]; + for array in arrays { + for start in [0, 1] { + let array = array.slice(start, array.len() - start); + for fan_out in [0, 1, 2, 5, 100] { + let lengths = Int64Array::from_iter_values((0..array.len()).map(|i| { + if i % 3 == 1 { + 0 + } else { + fan_out + } + })); + let total = lengths.values().iter().sum::() as usize; + let indices = create_take_indices(&lengths, total); + let expected = kernels::take::take(&array, &indices, None).unwrap(); + let actual = + repeat_arrs(std::slice::from_ref(&array), &lengths, total, &[true]) + .unwrap(); + assert_eq!(actual[0].data_type(), expected.data_type()); + assert_eq!(actual[0].to_data().buffers(), expected.to_data().buffers()); + assert_eq!(actual[0].nulls(), expected.nulls()); + } + } + } + } + #[test] fn short_string_repetition_matches_take_buffers() { let words = [ From 4fe5f30642905202a52dee151a63e6014a27818a Mon Sep 17 00:00:00 2001 From: rich7420 Date: Wed, 28 Oct 2026 18:27:14 +0800 Subject: [PATCH 8/9] fix: retain Spark execution for chained map generators Native generators materialize carried columns between expansions and were slower than Spark's pipeline in the nested benchmark. Keep maps followed by another generator on Spark across projections and filters, while retaining native single maps and chains separated by an exchange. --- docs/source/user-guide/latest/operators.md | 4 + .../apache/comet/rules/CometExecRule.scala | 35 ++++++++- .../sql-tests/expressions/map/explode.sql | 2 +- .../comet/exec/CometGenerateExecSuite.scala | 78 ++++++++++++++++++- .../sql/benchmark/CometExplodeBenchmark.scala | 33 ++++---- 5 files changed, 131 insertions(+), 21 deletions(-) diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index ffc565ea911..6504175233f 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -131,6 +131,10 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx | `UnionExec` | ✅ | | | `CoalesceExec` | ✅ | | +Map generators followed by another generator through projections or filters stay on Spark to +avoid repeated carried-column materialization. A single map generator, including maps with nested +values, remains eligible for native execution. An exchange separates generator chains. + ## Writes | Operator | Status | Notes | 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..8778b424c23 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) } + // Spark fuses a generator chain without materializing its carried columns between generators. + // Keep that path until native map expansion can avoid the extra repetition too. Tags survive + // the 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/test/resources/sql-tests/expressions/map/explode.sql b/spark/src/test/resources/sql-tests/expressions/map/explode.sql index 6a47e40aa96..a0c18f9062e 100644 --- a/spark/src/test/resources/sql-tests/expressions/map/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/map/explode.sql @@ -136,7 +136,7 @@ SELECT id, posexplode_outer(m) FROM explode_nested_maps query SELECT id, k, v.a, v.m FROM explode_nested_maps LATERAL VIEW OUTER explode(m) e AS k, v -query +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/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index 682f8dbac43..741e4d19069 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 @@ -312,6 +312,79 @@ class CometGenerateExecSuite extends CometTestBase { 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) + } + } + } + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", @@ -773,8 +846,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 faafcb4d98a..6c1818fceff 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,8 +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.catalyst.plans.logical.Generate -import org.apache.spark.sql.comet.CometExplodeExec +import org.apache.spark.sql.comet.{CometExplodeExec, CometNativeScanExec} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -62,14 +61,15 @@ import org.apache.comet.CometConf * plain `explode` in production usually sees an array column with no nulls and no empty rows. * * Map cases consume both keys and values. The nesting group models event lists as a map keyed by - * platform, including a second generator over the entries in each map value. + * platform, including a second generator over the entries in each map value. That chained 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]) @@ -298,41 +298,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 df = spark.sql(query) + 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 = df.queryExecution.optimizedPlan.collect { case _: Generate => - 1 - }.size + val expectedGenerators = benchmarkCase.nativeGenerators val nativeGenerators = collect(df.queryExecution.executedPlan) { case e: CometExplodeExec => e }.size assert( - expectedGenerators > 0 && nativeGenerators == expectedGenerators, + 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) } @@ -466,7 +467,7 @@ object CometExplodeBenchmark extends CometBenchmarkBase { 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 map[$eventsLen] at depth 8, then its inner array" -> - Case(chained, Seq(entries, entries, entries * 3, entries))) + single :+ (s"map at depth 8, then its inner array (Spark generators)" -> + Case(chained, Seq(entries, entries, entries * 3, entries), nativeGenerators = 0)) } } From 259c891c2c18d85b62e738ec3f25eafc894131d3 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Thu, 29 Oct 2026 19:41:27 +0800 Subject: [PATCH 9/9] fix: remove redundant interpolation in explode benchmark --- .../org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 6c1818fceff..b05709417c9 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 @@ -467,7 +467,7 @@ object CometExplodeBenchmark extends CometBenchmarkBase { s"SELECT ${Seq("k", "region", "platform", "entry").map(aggregate).mkString(", ")} " + s"FROM (SELECT k, region, platform, explode(entries) AS entry FROM ($outer))" - single :+ (s"map at depth 8, then its inner array (Spark generators)" -> + single :+ ("map at depth 8, then its inner array (Spark generators)" -> Case(chained, Seq(entries, entries, entries * 3, entries), nativeGenerators = 0)) } }