From 9601e81ba8fa2e8b1e7b3d92fbdc976e91f7a380 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Mon, 21 Sep 2026 17:56:22 +0800 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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/5] 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]