diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 133daf681a..549a07e8f6 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -260,7 +260,7 @@ index cf40e944c09..bdd5be4f462 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 1cc09c3d7fc..9e1e883d450 100644 +index 1cc09c3d7fc..e25dc3d1993 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -26,8 +26,9 @@ import org.scalatest.matchers.must.Matchers.the @@ -285,7 +285,17 @@ index 1cc09c3d7fc..9e1e883d450 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -750,12 +753,12 @@ class DataFrameAggregateSuite extends QueryTest +@@ -700,7 +703,8 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggOrSortAggPlan = + stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) + if (useObjectHashAgg) { +- assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec]) ++ assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] || ++ objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec]) + } else { + assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) + } +@@ -750,12 +754,12 @@ class DataFrameAggregateSuite extends QueryTest assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 428e964bad..179381342b 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -241,7 +241,7 @@ index e5494726695..00937f025c2 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 6f3090d8908..4774aad5019 100644 +index 6f3090d8908..52523429c87 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -27,8 +27,9 @@ import org.apache.spark.{SparkException, SparkThrowable} @@ -266,7 +266,17 @@ index 6f3090d8908..4774aad5019 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -788,12 +791,12 @@ class DataFrameAggregateSuite extends QueryTest +@@ -738,7 +741,8 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggOrSortAggPlan = + stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) + if (useObjectHashAgg) { +- assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec]) ++ assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] || ++ objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec]) + } else { + assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) + } +@@ -788,12 +792,12 @@ class DataFrameAggregateSuite extends QueryTest assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff index 8ee55fea1d..510b9b510d 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -378,7 +378,7 @@ index 0f42502f1d9..e9ff802141f 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 9db406ff12f..b3d55394d25 100644 +index 9db406ff12f..e59a931c06a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -29,8 +29,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -403,7 +403,17 @@ index 9db406ff12f..b3d55394d25 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -850,12 +853,12 @@ class DataFrameAggregateSuite extends QueryTest +@@ -800,7 +803,8 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggOrSortAggPlan = + stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) + if (useObjectHashAgg) { +- assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec]) ++ assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] || ++ objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec]) + } else { + assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) + } +@@ -850,12 +854,12 @@ class DataFrameAggregateSuite extends QueryTest assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff index 599e2cd077..b934fd568b 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -392,7 +392,7 @@ index 0d807aeae4d..6d7744e771b 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index bfe15b33768..13aeb3f6610 100644 +index bfe15b33768..df8fb4c0a35 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -417,7 +417,17 @@ index bfe15b33768..13aeb3f6610 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -851,12 +854,12 @@ class DataFrameAggregateSuite extends QueryTest +@@ -801,7 +804,8 @@ class DataFrameAggregateSuite extends QueryTest + val objHashAggOrSortAggPlan = + stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) + if (useObjectHashAgg) { +- assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec]) ++ assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] || ++ objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec]) + } else { + assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) + } +@@ -851,12 +855,12 @@ class DataFrameAggregateSuite extends QueryTest assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md index 63f4ca84b7..cb1061180b 100644 --- a/docs/source/contributor-guide/expression-audits/agg_funcs.md +++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md @@ -57,7 +57,7 @@ - Spark 3.5.8 (audited 2026-06-24): identical to 3.4.3. - Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class (no behavior change). - Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1. -- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (same pattern already in use for `collect_set`). +- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array. Comet rewrites intermediate `Partial`/`PartialMerge` object-aggregate buffer attributes from Spark's serialized `BinaryType` to that native list state. - Spark 4.2 (preview): `CollectList` and `CollectSet` gain an `ignoreNulls` field (default `true`); `RESPECT NULLS` sets it to `false` and keeps null elements. The native path always drops nulls, so `CometCollectShim` reads the field per Spark version (always `true` on 3.4-4.1) and `CometCollectList` / `CometCollectSet` report `Unsupported` when it is `false`, falling back to Spark. ## median diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs index 5ea26115bf..a0c1753145 100644 --- a/native/core/src/execution/merge_as_partial.rs +++ b/native/core/src/execution/merge_as_partial.rs @@ -30,7 +30,7 @@ use std::fmt::Debug; use std::hash::{Hash, Hasher}; use arrow::array::{ArrayRef, BooleanArray}; -use arrow::datatypes::{DataType, FieldRef}; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion::common::Result; use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::function::StateFieldsArgs; @@ -41,6 +41,8 @@ use datafusion::logical_expr::{ use datafusion::physical_expr::aggregate::AggregateFunctionExpr; use datafusion::scalar::ScalarValue; +use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; + /// An AggregateUDF wrapper that gives merge semantics in Partial mode. /// /// When DataFusion runs an AggregateExec in Partial mode, it calls `update_batch` @@ -60,6 +62,10 @@ pub struct MergeAsPartialUDF { cached_state_fields: Vec, /// Cached signature that accepts state field types. signature: Signature, + /// Override argument fields used only when constructing the inner accumulator. + accumulator_expr_fields: Option>, + /// Decoder for Spark JVM state, or pass-through for native-compatible state. + state_decoder: PartialMergeStateDecoder, /// Name for this wrapper. name: String, } @@ -83,6 +89,9 @@ impl MergeAsPartialUDF { let name = format!("merge_as_partial_{}", inner_expr.name()); let return_type = inner_expr.field().data_type().clone(); let cached_state_fields = inner_expr.state_fields()?; + let accumulator_expr_fields = + Self::accumulator_expr_fields(inner_expr, &cached_state_fields); + let state_decoder = PartialMergeStateDecoder::new(inner_expr, &cached_state_fields); // Use a permissive signature since we accept state field types which // vary per aggregate function. @@ -93,9 +102,44 @@ impl MergeAsPartialUDF { return_type, cached_state_fields, signature, + accumulator_expr_fields, + state_decoder, name, }) } + + fn accumulator_expr_fields( + inner_expr: &AggregateFunctionExpr, + state_fields: &[FieldRef], + ) -> Option> { + // Spark collect_list/collect_set use a single list state field but their accumulator + // constructors expect the original scalar input type. MergeAsPartial receives state-typed + // input columns, so passing those fields through would make the collect accumulator think + // it is aggregating arrays and create nested list state (List(List(T))). + match (inner_expr.fun().name(), state_fields) { + ("collect_list" | "collect_set", [state_field]) => match state_field.data_type() { + DataType::List(item_field) => Some(vec![Field::new( + item_field.name(), + item_field.data_type().clone(), + item_field.is_nullable(), + ) + .into()]), + _ => None, + }, + _ => None, + } + } + + fn inner_accumulator_args<'a>(&'a self, args: AccumulatorArgs<'a>) -> AccumulatorArgs<'a> { + let expr_fields = self + .accumulator_expr_fields + .as_deref() + .unwrap_or(args.expr_fields); + AccumulatorArgs { + expr_fields, + ..args + } + } } impl AggregateUDFImpl for MergeAsPartialUDF { @@ -120,27 +164,30 @@ impl AggregateUDFImpl for MergeAsPartialUDF { } fn accumulator(&self, args: AccumulatorArgs) -> Result> { - // args.exprs are state-typed (match this wrapper's signature), not the - // inner aggregate's original inputs. Safe for built-ins (SUM/COUNT/ - // MIN/MAX/AVG) which build accumulators from return_type; aggregates - // that inspect args.exprs types would need reconsideration. - let inner_acc = self.inner_udf.accumulator(args)?; - Ok(Box::new(MergeAsPartialAccumulator { inner: inner_acc })) + let inner_acc = self + .inner_udf + .accumulator(self.inner_accumulator_args(args))?; + Ok(Box::new(MergeAsPartialAccumulator { + inner: inner_acc, + state_decoder: self.state_decoder.clone(), + })) } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - // See `accumulator`: args.exprs are state-typed. - self.inner_udf.groups_accumulator_supported(args) + self.inner_udf + .groups_accumulator_supported(self.inner_accumulator_args(args)) } fn create_groups_accumulator( &self, args: AccumulatorArgs, ) -> Result> { - // See `accumulator`: args.exprs are state-typed. - let inner_acc = self.inner_udf.create_groups_accumulator(args)?; + let inner_acc = self + .inner_udf + .create_groups_accumulator(self.inner_accumulator_args(args))?; Ok(Box::new(MergeAsPartialGroupsAccumulator { inner: inner_acc, + state_decoder: self.state_decoder.clone(), })) } @@ -160,6 +207,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF { /// Accumulator wrapper that redirects update_batch to merge_batch. struct MergeAsPartialAccumulator { inner: Box, + state_decoder: PartialMergeStateDecoder, } impl Debug for MergeAsPartialAccumulator { @@ -171,11 +219,13 @@ impl Debug for MergeAsPartialAccumulator { impl Accumulator for MergeAsPartialAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { // Redirect update to merge — this is the key trick. - self.inner.merge_batch(values) + let decoded = self.state_decoder.decode(values)?; + self.inner.merge_batch(decoded.as_ref()) } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.inner.merge_batch(states) + let decoded = self.state_decoder.decode(states)?; + self.inner.merge_batch(decoded.as_ref()) } fn evaluate(&mut self) -> Result { @@ -194,6 +244,7 @@ impl Accumulator for MergeAsPartialAccumulator { /// GroupsAccumulator wrapper that redirects update_batch to merge_batch. struct MergeAsPartialGroupsAccumulator { inner: Box, + state_decoder: PartialMergeStateDecoder, } impl Debug for MergeAsPartialGroupsAccumulator { @@ -211,8 +262,13 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator { total_num_groups: usize, ) -> Result<()> { // Redirect update to merge — this is the key trick. - self.inner - .merge_batch(values, group_indices, opt_filter, total_num_groups) + let decoded = self.state_decoder.decode(values)?; + self.inner.merge_batch( + decoded.as_ref(), + group_indices, + opt_filter, + total_num_groups, + ) } fn merge_batch( @@ -222,8 +278,13 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator { opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { - self.inner - .merge_batch(values, group_indices, opt_filter, total_num_groups) + let decoded = self.state_decoder.decode(values)?; + self.inner.merge_batch( + decoded.as_ref(), + group_indices, + opt_filter, + total_num_groups, + ) } fn evaluate(&mut self, emit_to: EmitTo) -> Result { diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index ec247f72b7..5145843970 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -26,6 +26,7 @@ pub(crate) mod planner; pub mod serde; pub use datafusion_comet_shuffle as shuffle; pub(crate) mod sort; +pub(crate) mod spark_aggregate_state; pub(crate) mod spark_plan; pub use datafusion_comet_spark_expr::timezone; mod memory_pools; diff --git a/native/core/src/execution/spark_aggregate_state.rs b/native/core/src/execution/spark_aggregate_state.rs new file mode 100644 index 0000000000..222581ca2b --- /dev/null +++ b/native/core/src/execution/spark_aggregate_state.rs @@ -0,0 +1,379 @@ +// 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. + +//! Decoders for Spark JVM aggregate state consumed by native PartialMerge. + +use std::{borrow::Cow, sync::Arc}; + +use arrow::array::{ + builder::{make_builder, ArrayBuilder, ListBuilder}, + Array, ArrayRef, BinaryArray, GenericByteArray, LargeBinaryArray, OffsetSizeTrait, +}; +use arrow::datatypes::{DataType, FieldRef, GenericBinaryType}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::physical_expr::aggregate::AggregateFunctionExpr; +use datafusion_comet_shuffle::spark_unsafe::list::{append_to_builder, SparkUnsafeArray}; + +/// Decoder used by `MergeAsPartial` before forwarding state to the inner accumulator. +#[derive(Clone, Debug)] +pub(crate) enum PartialMergeStateDecoder { + PassThrough, + SparkCollect(SparkCollectStateDecoder), +} + +impl PartialMergeStateDecoder { + pub(crate) fn new(inner_expr: &AggregateFunctionExpr, state_fields: &[FieldRef]) -> Self { + SparkCollectStateDecoder::try_new(inner_expr, state_fields) + .map(Self::SparkCollect) + .unwrap_or(Self::PassThrough) + } + + pub(crate) fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result> { + match self { + Self::PassThrough => Ok(Cow::Borrowed(values)), + Self::SparkCollect(decoder) => decoder.decode(values), + } + } +} + +/// Decodes Spark JVM collect aggregate buffers into DataFusion collect state. +/// +/// Spark's `CollectList` / `CollectSet` are `TypedImperativeAggregate`s. When Spark runs the +/// lower Partial aggregate, each buffer is serialized as a `BinaryType` value containing a +/// single-field `UnsafeRow`; field 0 is the `UnsafeArrayData` with the collected elements. +/// DataFusion's collect accumulators expect the merge input to be a list-typed state column, so +/// mixed Spark-Partial -> Comet-PartialMerge plans must materialize those unsafe bytes into an +/// Arrow `ListArray` before calling the inner accumulator's `merge_batch`. +#[derive(Clone, Debug)] +pub(crate) struct SparkCollectStateDecoder { + item_field: FieldRef, +} + +impl SparkCollectStateDecoder { + fn try_new(inner_expr: &AggregateFunctionExpr, state_fields: &[FieldRef]) -> Option { + match (inner_expr.fun().name(), state_fields) { + ("collect_list" | "collect_set", [state_field]) => match state_field.data_type() { + DataType::List(item_field) => Some(Self { + item_field: Arc::clone(item_field), + }), + _ => None, + }, + _ => None, + } + } + + fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result> { + if values.len() != 1 { + return Err(DataFusionError::Internal(format!( + "Spark collect state decoder expected one state column, got {}", + values.len() + ))); + } + + match values[0].data_type() { + DataType::Binary => { + let array = values[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Self::decode_error("expected BinaryArray for Binary collect state") + })?; + Ok(Cow::Owned(vec![self.decode_binary_array(array)?])) + } + DataType::LargeBinary => { + let array = values[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Self::decode_error( + "expected LargeBinaryArray for LargeBinary collect state", + ) + })?; + Ok(Cow::Owned(vec![self.decode_binary_array(array)?])) + } + _ => Ok(Cow::Borrowed(values)), + } + } + + fn decode_binary_array( + &self, + array: &GenericByteArray>, + ) -> Result { + let mut builder = self.new_list_builder(array.len()); + + for row_idx in 0..array.len() { + if array.is_null(row_idx) { + builder.append_null(); + } else { + self.append_unsafe_row_array(array.value(row_idx), &mut builder)?; + } + } + + Ok(Arc::new(builder.finish())) + } + + fn new_list_builder(&self, capacity: usize) -> ListBuilder> { + let value_builder = make_builder(self.item_field.data_type(), capacity); + ListBuilder::with_capacity(value_builder, capacity).with_field(Arc::clone(&self.item_field)) + } + + fn append_unsafe_row_array( + &self, + row_bytes: &[u8], + builder: &mut ListBuilder>, + ) -> Result<()> { + match Self::spark_array_from_single_field_unsafe_row(row_bytes)? { + Some(array) => { + append_to_builder::(self.item_field.data_type(), builder.values(), &array) + .map_err(|e| Self::decode_error(e.to_string()))?; + builder.append(true); + } + None => builder.append_null(), + } + Ok(()) + } + + fn spark_array_from_single_field_unsafe_row( + row_bytes: &[u8], + ) -> Result> { + const BITSET_WIDTH: usize = 8; + const FIXED_FIELD_WIDTH: usize = 8; + const ARRAY_FIELD_INDEX: usize = 0; + const MIN_ROW_WIDTH: usize = BITSET_WIDTH + FIXED_FIELD_WIDTH; + + if row_bytes.len() < MIN_ROW_WIDTH { + return Err(Self::decode_error(format!( + "UnsafeRow collect buffer is too small: {} bytes", + row_bytes.len() + ))); + } + + let null_bits = i64::from_le_bytes( + row_bytes[0..BITSET_WIDTH] + .try_into() + .expect("slice length checked"), + ); + if (null_bits & (1_i64 << ARRAY_FIELD_INDEX)) != 0 { + return Ok(None); + } + + let offset_and_size = i64::from_le_bytes( + row_bytes[BITSET_WIDTH..MIN_ROW_WIDTH] + .try_into() + .expect("slice length checked"), + ); + let offset = (offset_and_size >> 32) as i32; + let size = offset_and_size as i32; + + if offset < MIN_ROW_WIDTH as i32 || size < 0 { + return Err(Self::decode_error(format!( + "Invalid UnsafeRow array field offset/size: offset={offset}, size={size}, row_size={}", + row_bytes.len() + ))); + } + + let offset = offset as usize; + let size = size as usize; + let end = offset.checked_add(size).ok_or_else(|| { + Self::decode_error(format!( + "UnsafeRow array field range overflows: offset={offset}, size={size}" + )) + })?; + if end > row_bytes.len() { + return Err(Self::decode_error(format!( + "UnsafeRow array field range is out of bounds: offset={offset}, size={size}, row_size={}", + row_bytes.len() + ))); + } + if size < 8 { + return Err(Self::decode_error(format!( + "UnsafeArrayData collect buffer is too small: {size} bytes" + ))); + } + + let array_addr = unsafe { row_bytes.as_ptr().add(offset) } as i64; + Ok(Some(SparkUnsafeArray::new(array_addr))) + } + + fn decode_error(message: impl Into) -> DataFusionError { + DataFusionError::Execution(format!( + "Failed to decode Spark UnsafeRow collect aggregate buffer: {}", + message.into() + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, ListArray, StringArray}; + use arrow::datatypes::Field; + + fn collect_state_decoder(element_type: DataType) -> SparkCollectStateDecoder { + SparkCollectStateDecoder { + item_field: Arc::new(Field::new_list_field(element_type, true)), + } + } + + fn unsafe_row_with_array(array: Vec) -> Vec { + const ARRAY_OFFSET: usize = 16; + let mut row = vec![0_u8; ARRAY_OFFSET]; + let offset_and_size = ((ARRAY_OFFSET as i64) << 32) | array.len() as i64; + row[8..16].copy_from_slice(&offset_and_size.to_le_bytes()); + row.extend_from_slice(&array); + row + } + + fn unsafe_array_i32(values: &[Option]) -> Vec { + let num_elements = values.len(); + let bitset_words = num_elements.div_ceil(64); + let header_len = 8 + bitset_words * 8; + let mut bytes = vec![0_u8; header_len + num_elements * std::mem::size_of::()]; + bytes[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + + let mut null_bits = 0_i64; + for (idx, value) in values.iter().enumerate() { + let value_offset = header_len + idx * std::mem::size_of::(); + match value { + Some(value) => { + bytes[value_offset..value_offset + 4].copy_from_slice(&value.to_le_bytes()) + } + None => null_bits |= 1_i64 << idx, + } + } + if bitset_words > 0 { + bytes[8..16].copy_from_slice(&null_bits.to_le_bytes()); + } + + bytes + } + + fn unsafe_array_utf8(values: &[Option<&str>]) -> Vec { + let num_elements = values.len(); + let bitset_words = num_elements.div_ceil(64); + let header_len = 8 + bitset_words * 8; + let fixed_len = num_elements * 8; + let mut bytes = vec![0_u8; header_len + fixed_len]; + bytes[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + + let mut null_bits = 0_i64; + for (idx, value) in values.iter().enumerate() { + let slot_offset = header_len + idx * 8; + match value { + Some(value) => { + let value_offset = bytes.len(); + let value_bytes = value.as_bytes(); + bytes.extend_from_slice(value_bytes); + let padded_len = value_bytes.len().next_multiple_of(8); + bytes.resize(value_offset + padded_len, 0); + + let offset_and_size = ((value_offset as i64) << 32) | value_bytes.len() as i64; + bytes[slot_offset..slot_offset + 8] + .copy_from_slice(&offset_and_size.to_le_bytes()); + } + None => null_bits |= 1_i64 << idx, + } + } + if bitset_words > 0 { + bytes[8..16].copy_from_slice(&null_bits.to_le_bytes()); + } + + bytes + } + + #[test] + fn decodes_spark_collect_binary_int_state() { + let first = unsafe_row_with_array(unsafe_array_i32(&[Some(1)])); + // The first row is intentionally 36 bytes, so the second row starts at an unaligned + // address in the Arrow Binary values buffer. The decoder must not rely on UnsafeRow + // alignment once Spark bytes have been materialized into Arrow BinaryArray storage. + assert_eq!(first.len(), 36); + let second = unsafe_row_with_array(unsafe_array_i32(&[Some(2), None, Some(4)])); + let third = unsafe_row_with_array(unsafe_array_i32(&[])); + + let binary = BinaryArray::from_iter([ + Some(first.as_slice()), + Some(second.as_slice()), + None, + Some(third.as_slice()), + ]); + let input = [Arc::new(binary) as ArrayRef]; + let decoder = collect_state_decoder(DataType::Int32); + let decoded = decoder.decode(&input).unwrap(); + let list = decoded.as_ref()[0] + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(list.len(), 4); + assert!(!list.is_null(0)); + assert!(!list.is_null(1)); + assert!(list.is_null(2)); + assert!(!list.is_null(3)); + + let first_values = list + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(first_values, vec![1]); + + let second_values = list + .value(1) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(second_values.values(), &[2, 0, 4]); + assert!(!second_values.is_null(0)); + assert!(second_values.is_null(1)); + assert!(!second_values.is_null(2)); + + assert_eq!(list.value(3).len(), 0); + } + + #[test] + fn decodes_spark_collect_binary_string_state() { + let row = unsafe_row_with_array(unsafe_array_utf8(&[ + Some("alpha"), + Some("βeta"), + None, + Some("spark"), + ])); + let binary = BinaryArray::from_iter([Some(row.as_slice())]); + let input = [Arc::new(binary) as ArrayRef]; + let decoder = collect_state_decoder(DataType::Utf8); + let decoded = decoder.decode(&input).unwrap(); + let list = decoded.as_ref()[0] + .as_any() + .downcast_ref::() + .unwrap(); + let values = list + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + + assert_eq!(values.value(0), "alpha"); + assert_eq!(values.value(1), "βeta"); + assert!(values.is_null(2)); + assert_eq!(values.value(3), "spark"); + } +} 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 bd3b92f2cc..5d4252d6b2 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -876,29 +876,6 @@ case class CometExecRule(session: SparkSession) } } } - - // CollectList/CollectSet round-trip an ArrayType buffer that Spark declares as BinaryType. - // In a multi-stage aggregate with a PartialMerge stage (e.g. Spark's distinct-aggregate - // rewrite), Comet cannot represent that buffer consistently across the intermediate stages - // (issue #4724), so a fully-native pipeline crashes. Force the whole chain to fall back to - // Spark by tagging the feeding pure-Partial; the PartialMerge/Final stages then fall back - // via the buffer-source check in doConvert. - // - // This block is intentionally separate from the tagging block just above: that one only - // fires when the Final itself cannot be converted, but `canAggregateBeConverted` skips - // the child-native check, so an all-native distinct `collect_list` chain converts its - // Final and slips past the earlier tagging pass. This block catches that case. - if (agg.aggregateExpressions.exists(_.mode == PartialMerge) && - QueryPlanSerde.hasNativeArrayBufferAgg(agg.aggregateExpressions)) { - findPartialAggInPlan(agg.child).foreach { partial => - if (canAggregateBeConverted(partial, Partial)) { - partial.setTagValue( - CometExecRule.COMET_UNSAFE_PARTIAL, - "Partial aggregate disabled: part of a multi-stage CollectList/CollectSet " + - "aggregate whose intermediate buffer cannot round-trip in Comet (issue #4724)") - } - } - } case _ => } } diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 4892d6f966..d7362096c4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -443,25 +443,6 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } } - /** - * Returns true if any aggregate is CollectList/CollectSet. These produce a native ArrayType - * intermediate buffer while Spark declares BinaryType for its serialized - * TypedImperativeAggregate buffer, so Comet cannot interpret Spark's Binary buffer, and Comet - * cannot yet represent this buffer consistently across the intermediate PartialMerge stages of - * a multi-stage aggregate (issue #4724). These aggregates are therefore only safe to run - * natively when every stage runs in Comet and there are at most two stages (Partial + Final). - * - * Percentile has a similar Array-shaped intermediate buffer (see `adjustOutputForNativeState`) - * but is not matched here: it already passes through the general mixed-execution guard, so this - * check is scoped narrowly to the collect functions. - */ - def hasNativeArrayBufferAgg(aggExprs: Seq[AggregateExpression]): Boolean = { - aggExprs.exists(_.aggregateFunction match { - case _: CollectList | _: CollectSet => true - case _ => false - }) - } - // A unique id for each expression. ~used to look up QueryContext during error creation. private val exprIdCounter = new AtomicLong(0) 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 fcd48d69b1..0867a6ba21 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 @@ -1610,17 +1610,18 @@ trait CometBaseAggregate { val modes = aggregate.aggregateExpressions.map(_.mode).distinct val modeSet = modes.toSet val hasPartialMerge = modeSet.contains(PartialMerge) + val cometPartialAgg = findCometPartialAgg(aggregate.child) // In distinct aggregates there can be a combination of modes. // We support {Partial, PartialMerge} mix; other combinations are rejected. val multiMode = modes.size > 1 && modeSet != Set(Partial, PartialMerge) - // An aggregate that consumes intermediate buffers (Final, or the PartialMerge stages of a - // distinct-aggregate rewrite) must have a Comet aggregate producing those buffers below it. - // Otherwise Comet would try to read a Spark partial's buffer, which is only safe when every - // aggregate has a buffer format compatible between Spark and Comet. This guards the - // Spark-Partial to Comet-Merge direction; the Comet-Partial to Spark-Final direction is - // guarded by the COMET_UNSAFE_PARTIAL tagging pass in CometExecRule. See issues #1389, #4813. - val consumesBuffers = modes.contains(Final) || modes.contains(PartialMerge) - val missingCometProducer = consumesBuffers && findCometPartialAgg(aggregate.child).isEmpty + // For a final mode HashAggregate, we only need to transform the HashAggregate + // if there is Comet partial aggregation, unless all aggregates have compatible + // intermediate buffer formats (safe for mixed Spark/Comet execution). + val sparkFinalMode = modes.contains(Final) && cometPartialAgg.isEmpty + // For PartialMerge, the child aggregate may still be Spark/JVM even when this node's direct + // child is a Comet shuffle. In that case Spark's intermediate buffers must be compatible with + // the native merge accumulator state expected by Comet. + val sparkPartialMergeMode = hasPartialMerge && cometPartialAgg.isEmpty if (multiMode) { withFallbackReason( @@ -1629,7 +1630,7 @@ trait CometBaseAggregate { return None } - if (missingCometProducer) { + if (sparkFinalMode) { val incompatibleAggs = QueryPlanSerde.aggsNotSupportingMixedExecution(aggregate.aggregateExpressions) if (incompatibleAggs.nonEmpty) { @@ -1643,6 +1644,22 @@ trait CometBaseAggregate { } } + if (sparkPartialMergeMode) { + val hasUnsupportedAgg = aggregate.aggregateExpressions.exists { aggExpr => + aggExpr.mode == PartialMerge && + !QueryPlanSerde.allAggsSupportMixedExecution(Seq(aggExpr)) && + !aggExpr.aggregateFunction.isInstanceOf[CollectList] && + !aggExpr.aggregateFunction.isInstanceOf[CollectSet] + } + if (hasUnsupportedAgg) { + withFallbackReason( + aggregate, + "Spark PartialMerge aggregate without Comet Partial requires compatible " + + "intermediate buffer formats") + return None + } + } + // Check if this aggregate has been tagged as unsafe for mixed execution // (Comet partial + Spark final with incompatible intermediate buffers) val unsafeReason = aggregate.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL) @@ -2030,6 +2047,48 @@ object CometObjectHashAggregateExec op.child, SerializedPlan(None)) } + + /** + * For intermediate aggregates containing TypedImperativeAggregate functions (like CollectSet or + * CollectList), Spark declares buffer columns as BinaryType because it serializes the JVM + * state. Native Comet keeps the actual state type instead: ArrayType(elementType) with + * containsNull true for CollectSet/CollectList. Rewrite the Spark-side output attributes for + * Partial, PartialMerge, and mixed {Partial, PartialMerge} stages so shuffle and downstream + * native aggregate stages see the schema that native execution really produces. + * + * Final aggregates output user-visible values rather than intermediate state, so their Spark + * result schema is left unchanged. + */ + private def adjustOutputForNativeState(op: ObjectHashAggregateExec): Seq[Attribute] = { + val modes = op.aggregateExpressions.map(_.mode).distinct + if (modes.exists(mode => mode != Partial && mode != PartialMerge)) { + return op.output + } + + val numGrouping = op.groupingExpressions.length + val output = op.output.toArray + + var bufferIdx = numGrouping + for (aggExpr <- op.aggregateExpressions) { + val aggFunc = aggExpr.aggregateFunction + val bufferAttrs = aggFunc.aggBufferAttributes + aggFunc match { + case _: CollectSet | _: CollectList => + val elementType = aggFunc.children.head.dataType + val nativeStateType = ArrayType(elementType, containsNull = true) + output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) + case _: Percentile => + // DataFusion's percentile_cont keeps all values in a List partial state. + // Comet casts the child to double, so the native state is ArrayType(DoubleType). + val nativeStateType = ArrayType(DoubleType, containsNull = true) + output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) + case _ => + } + bufferIdx += bufferAttrs.length + } + + output.toSeq + } } case class CometHashAggregateExec( diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql index e00d2b1969..b88d0c8c1b 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -435,3 +435,18 @@ INSERT INTO cl_src_map VALUES query spark_answer_only SELECT size(collect_list(m)) FROM cl_src_map + +-- ============================================================ +-- PartialMerge: combined with distinct and collect-style aggregates +-- ============================================================ + +query +SELECT grp, count(DISTINCT a), sort_array(collect_list(b)) +FROM cl_src_multi GROUP BY grp ORDER BY grp + +query +SELECT grp, + count(DISTINCT a), + sort_array(collect_list(b)), + sort_array(collect_set(b)) +FROM cl_src_multi GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql index cd528272af..e1b30d0d3f 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql @@ -291,6 +291,15 @@ FROM cs_src_multi GROUP BY grp ORDER BY grp query SELECT grp, sort_array(collect_set(DISTINCT i)) FROM cs_src_int GROUP BY grp ORDER BY grp +-- ============================================================ +-- PartialMerge: collect_set combined with a distinct aggregate +-- exercises native intermediate ArrayType state through shuffle +-- ============================================================ + +query +SELECT grp, count(DISTINCT a), sort_array(collect_set(b)) +FROM cs_src_multi GROUP BY grp ORDER BY grp + -- ============================================================ -- HAVING clause with collect_set -- ============================================================ diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 3248101a2b..79b15fdfd6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -23,12 +23,14 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.SparkConf -import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.comet.CometHashAggregateExec +import org.apache.spark.sql.execution.LocalTableScanExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.functions.{avg, col, count_distinct, sum} +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec +import org.apache.spark.sql.functions.{avg, col, collect_list, collect_set, count_distinct, sort_array, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -49,27 +51,18 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") test("collect_list over struct with non-nullable fields") { - // Building a struct from non-nullable columns yields non-nullable struct fields. The native - // collect_list accumulator emits a list whose element fields are all nullable, so the produced - // array must be reconciled with the declared aggregate output type, otherwise the grouped - // native aggregate fails with "column types must match schema types". - import org.apache.spark.sql.functions.{collect_list, expr} - import org.apache.spark.sql.execution.LocalTableScanExec - // One row per (a, b) group keeps each collected list deterministic for the answer comparison. - // repartition inserts a Comet exchange so the aggregate runs natively over a Comet child; the - // in-memory source stays a (non-Comet) LocalTableScan, which we allow via excludedClasses. + import org.apache.spark.sql.functions.expr + val df = Seq(("1", "2", 1), ("2", "3", 3)) .toDF("a", "b", "c") .repartition(col("a")) .withColumn("d", expr("named_struct('a', a, 'b', b, 'c', c)")) val include = Seq(classOf[CometHashAggregateExec]) - // Single grouped collect_list of a struct with a non-nullable field. + checkSparkAnswerAndOperator( df.groupBy("a", "b").agg(collect_list("d")), include, classOf[LocalTableScanExec]) - // Multi-stage: the array result flows through a projection into a second collect_list - // (the SPARK-22223 plan shape), so the nested struct field nullability must round-trip. checkSparkAnswerAndOperator( df.groupBy("a", "b") .agg(collect_list("d").as("e")) @@ -80,29 +73,58 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { classOf[LocalTableScanExec]) } - test("collect_list/collect_set combined with distinct aggregate falls back safely") { - // SPARK-17616: a distinct aggregate combined with collect_list/collect_set produces a - // multi-stage plan where the buffer-producing Partial may run in Spark (e.g. over a - // non-native LocalTableScan). Comet cannot read Spark's serialized Binary buffer, so the - // dependent PartialMerge/Final stages must also fall back rather than crash. See issue #4724 - // for enabling the fully-native distinct path. - import org.apache.spark.sql.functions.{collect_list, collect_set, sort_array} - // Non-native source (LocalTableScan): the buffer-producing Partial runs in Spark. - val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")) - .toDF("x", "y", "z") - checkSparkAnswer( - df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_list(col("z"))))) - checkSparkAnswer( - df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_set(col("z"))))) - - // Native source (Parquet): the whole multi-stage distinct chain must still fall back to - // Spark consistently (issue #4724), rather than running a fully-native pipeline that crashes. + test("collect_list/collect_set combined with distinct aggregate runs fully native") { + // SPARK-17616: combining a distinct aggregate with collect_list/collect_set creates a + // multi-stage aggregate plan with a PartialMerge collect stage. These collect functions declare + // BinaryType JVM buffers in Spark but produce ArrayType native state in Comet; the intermediate + // Comet aggregate outputs must advertise the native ArrayType so that the state can round-trip + // through shuffle and downstream native PartialMerge/Final stages. See issue #4724. withParquetTable( Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")), - "t17616") { + "t_collect_distinct", + withDictionary = false) { for (fn <- Seq("collect_list", "collect_set")) { - checkSparkAnswer( - sql(s"SELECT _1, count(distinct _2), sort_array($fn(_3)) FROM t17616 GROUP BY _1")) + val query = + s"SELECT _1, count(DISTINCT _2), sort_array($fn(_3)) " + + "FROM t_collect_distinct GROUP BY _1" + val (_, cometPlan) = + checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometHashAggregateExec])) + assert( + cometPlan.toString.contains(s"merge_$fn"), + s"Expected $fn PartialMerge stage to run as CometHashAggregateExec; plan:\n$cometPlan") + } + } + } + + test("collect_list/collect_set with distinct aggregate decodes Spark PartialMerge state") { + // This mirrors Spark's DataFrameAggregateSuite SPARK-17616 test. With LocalTableScan disabled, + // the lower collect partial aggregate stays on Spark and produces a BinaryType JVM buffer. + // Native PartialMerge decodes that UnsafeRow/UnsafeArray buffer into DataFusion's list-typed + // state before merging. Run the same hash-map configurations used by the inherited Spark + // aggregate suites that failed in CI. + for ((twoLevelAggMap, vectorizedHashMap) <- Seq( + (false, false), + (true, false), + (true, true))) { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "false", + SQLConf.CODEGEN_FALLBACK.key -> "false", + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelAggMap.toString, + SQLConf.ENABLE_VECTORIZED_HASH_MAP.key -> vectorizedHashMap.toString) { + val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")) + .toDF("x", "y", "z") + for ((name, collect) <- Seq[(String, Column => Column)]( + "collect_list" -> (col => collect_list(col)), + "collect_set" -> (col => collect_set(col)))) { + val (_, cometPlan) = checkSparkAnswerAndOperator( + df.groupBy($"x").agg(count_distinct($"y"), sort_array(collect($"z"))), + Seq(classOf[CometHashAggregateExec]), + classOf[ObjectHashAggregateExec], + classOf[LocalTableScanExec]) + assert( + cometPlan.toString.contains(s"merge_$name"), + s"Expected $name PartialMerge stage to run natively; plan:\n$cometPlan") + } } } }