mbutrovich commented on code in PR #4816: URL: https://github.com/apache/datafusion-comet/pull/4816#discussion_r3615495128
########## native/spark-expr/src/agg_funcs/list_agg.rs: ########## @@ -0,0 +1,284 @@ +// 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. + +//! Spark-compatible `listagg` / `string_agg` aggregate function. +//! +//! Implements the simple form of Spark 4.0's `LISTAGG(expr, delimiter)` (no +//! `WITHIN GROUP (ORDER BY ...)`, no DISTINCT — DISTINCT is rewritten into a +//! multi-stage plan by Spark before it reaches Comet). Differences from +//! DataFusion's `string_agg`: +//! +//! * Returns `Utf8` to match Spark's `StringType` result type; DataFusion's +//! `string_agg` returns `LargeUtf8`. +//! * A `NULL` delimiter is treated as the empty string (Spark treats `NULL` as +//! the default empty delimiter; the JVM serde forwards the literal as-is). +//! * The delimiter is read once from the accumulator args (a literal is +//! enforced by Spark's analyzer). +//! +//! The intermediate state is exposed as `Binary` because Spark's `ListAgg` is +//! a `TypedImperativeAggregate` whose Catalyst buffer schema is `BinaryType`. +//! Emitting `Utf8` here would force a Comet shuffle-side cast (`Utf8` → +//! `Binary`) that the merge side then can no longer read. + +use std::hash::Hash; +use std::mem::size_of_val; + +use arrow::array::{ArrayRef, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::cast::{as_binary_array, as_string_array}; +use datafusion::common::{internal_datafusion_err, not_impl_err, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::utils::format_state_name; +use datafusion::logical_expr::{ + Accumulator, AggregateUDFImpl, Signature, TypeSignature, Volatility, +}; +use datafusion::physical_expr::expressions::Literal; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkListAgg { + signature: Signature, +} + +impl Default for SparkListAgg { + fn default() -> Self { + Self::new() + } +} + +impl SparkListAgg { + pub fn new() -> Self { + Self { + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]), + TypeSignature::Exact(vec![DataType::Utf8, DataType::Null]), + ], + Volatility::Immutable, + ), + } + } +} + +impl AggregateUDFImpl for SparkListAgg { + fn name(&self) -> &str { + "listagg" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Utf8) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + // Spark's ListAgg is a TypedImperativeAggregate — Catalyst declares its Review Comment: The module doc and `state_fields` comment justify the `Binary` intermediate state by saying a `Utf8` state would force a Comet shuffle-side `Utf8 -> Binary` cast that the merge side cannot read. I could not reconcile that with the rest of the code, so this may be worth revisiting. `CometListAgg` does not override `supportsMixedPartialFinal`, so it inherits `false` (`CometAggregateExpressionSerde.scala:93`), which as I read it means partial and final always run in the same engine (same as `CometCollectSet` and `CometPercentile`), so Spark's Catalyst buffer schema would not be involved. DataFusion's `string_agg` uses `Utf8`/`LargeUtf8` state and merges without a cast (`string_agg.rs:166`), and I did not find a `Utf8 -> Binary` cast in the shuffle layer. If that reading is right, the `Binary` state adds a per-merge `from_utf8` validation and byte round-trip (`list_agg.rs:135`, `:147`) for no benefit, and `Utf8` state would be simpler. Is there a mixed-execution or shuffle-serializati on case I am missing that the `Binary` state is guarding against? If not, consider `Utf8` state, or fold this into the `string_agg` reuse above. ########## spark/src/main/spark-4.x/org/apache/comet/serde/CometListAgg.scala: ########## @@ -0,0 +1,107 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ListAgg} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{NullType, StringType} + +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, hasNonDefaultStringCollation, serializeDataType} + +/** + * Spark 4.0+ `LISTAGG` / `STRING_AGG`. + * + * Comet only supports the simple form: a `StringType` child with a literal delimiter and no + * `WITHIN GROUP (ORDER BY ...)` clause. DISTINCT is handled by Spark's multi-stage plan rewrite + * (grouping by the child before the aggregate), so the native side never sees it. + */ +object CometListAgg extends CometAggregateExpressionSerde[ListAgg] { + + private val binaryChildReason = "`BinaryType` inputs are not supported." + private val withinGroupReason = "`WITHIN GROUP (ORDER BY ...)` is not supported." + private val nonFoldableDelimiterReason = "Non-literal delimiters are not supported." + private val collationReason = "Non-default string collations are not supported." + private val distinctReason = + "`DISTINCT` falls back to Spark because Comet rejects multi-column distinct aggregates." + + override def getUnsupportedReasons(): Seq[String] = Seq( + binaryChildReason, + withinGroupReason, + nonFoldableDelimiterReason, + collationReason, + distinctReason) + + override def getSupportLevel(expr: ListAgg): SupportLevel = { + // Spark's analyzer already enforces `delimiter.foldable`, so this only ever rejects + // non-string / non-null delimiter types. + expr.child.dataType match { + case _: StringType if hasNonDefaultStringCollation(expr.child.dataType) => + Unsupported(Some(collationReason)) + case _: StringType => + expr.delimiter.dataType match { + case _: StringType if hasNonDefaultStringCollation(expr.delimiter.dataType) => + Unsupported(Some(collationReason)) + case _: StringType | _: NullType => + if (!expr.delimiter.foldable) { + Unsupported(Some(nonFoldableDelimiterReason)) + } else if (expr.orderExpressions.nonEmpty) { + Unsupported(Some(withinGroupReason)) + } else { + Compatible() + } + case other => + Unsupported(Some(s"Unsupported delimiter data type: $other")) + } + case other => + Unsupported(Some(s"Unsupported child data type: $other")) + } + } + + override def convert( + aggExpr: AggregateExpression, + expr: ListAgg, + inputs: Seq[Attribute], + binding: Boolean, + conf: SQLConf): Option[ExprOuterClass.AggExpr] = { + val childExpr = exprToProto(expr.child, inputs, binding) + val delimiterExpr = exprToProto(expr.delimiter, inputs, binding) + val dataType = serializeDataType(expr.dataType) + + if (childExpr.isDefined && delimiterExpr.isDefined && dataType.isDefined) { + val builder = ExprOuterClass.ListAgg.newBuilder() + builder.setChild(childExpr.get) + builder.setDelimiter(delimiterExpr.get) + builder.setDatatype(dataType.get) + Some( + ExprOuterClass.AggExpr + .newBuilder() + .setListAgg(builder) + .build()) + } else if (dataType.isEmpty) { Review Comment: `convert` runs only after `getSupportLevel` returned `Compatible`, which requires the child to be `StringType`, and `serializeDataType` always returns `Some` for `StringType`. So the `dataType.isEmpty` branch is unreachable and the `datatype` proto field it sets is never read by the planner arm (`planner.rs`, which hardcodes `Utf8`). Drop the field and the branch, or replace the branch with an explicit `unreachable!`-style guard stating the child is guaranteed `StringType`, so the invariant is self-documenting rather than silent dead code. ########## spark/src/main/spark-4.x/org/apache/comet/serde/CometListAgg.scala: ########## @@ -0,0 +1,107 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ListAgg} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{NullType, StringType} + +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, hasNonDefaultStringCollation, serializeDataType} + +/** + * Spark 4.0+ `LISTAGG` / `STRING_AGG`. + * + * Comet only supports the simple form: a `StringType` child with a literal delimiter and no + * `WITHIN GROUP (ORDER BY ...)` clause. DISTINCT is handled by Spark's multi-stage plan rewrite + * (grouping by the child before the aggregate), so the native side never sees it. + */ +object CometListAgg extends CometAggregateExpressionSerde[ListAgg] { + + private val binaryChildReason = "`BinaryType` inputs are not supported." + private val withinGroupReason = "`WITHIN GROUP (ORDER BY ...)` is not supported." + private val nonFoldableDelimiterReason = "Non-literal delimiters are not supported." + private val collationReason = "Non-default string collations are not supported." + private val distinctReason = Review Comment: `distinctReason` is listed in `getUnsupportedReasons()` but is never emitted: DISTINCT falls back earlier in `aggExprToProto` (`QueryPlanSerde.scala:640`) with the generic multi-column-distinct message, which is what the SQL test at `listagg.sql:100` matches. Remove `distinctReason` or mark it informational, so the user-facing reason list reflects the actual fallback message. ########## native/spark-expr/src/agg_funcs/list_agg.rs: ########## @@ -0,0 +1,284 @@ +// 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. + +//! Spark-compatible `listagg` / `string_agg` aggregate function. +//! +//! Implements the simple form of Spark 4.0's `LISTAGG(expr, delimiter)` (no +//! `WITHIN GROUP (ORDER BY ...)`, no DISTINCT — DISTINCT is rewritten into a +//! multi-stage plan by Spark before it reaches Comet). Differences from +//! DataFusion's `string_agg`: +//! +//! * Returns `Utf8` to match Spark's `StringType` result type; DataFusion's +//! `string_agg` returns `LargeUtf8`. +//! * A `NULL` delimiter is treated as the empty string (Spark treats `NULL` as +//! the default empty delimiter; the JVM serde forwards the literal as-is). +//! * The delimiter is read once from the accumulator args (a literal is +//! enforced by Spark's analyzer). +//! +//! The intermediate state is exposed as `Binary` because Spark's `ListAgg` is +//! a `TypedImperativeAggregate` whose Catalyst buffer schema is `BinaryType`. +//! Emitting `Utf8` here would force a Comet shuffle-side cast (`Utf8` → +//! `Binary`) that the merge side then can no longer read. + +use std::hash::Hash; +use std::mem::size_of_val; + +use arrow::array::{ArrayRef, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::cast::{as_binary_array, as_string_array}; +use datafusion::common::{internal_datafusion_err, not_impl_err, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::utils::format_state_name; +use datafusion::logical_expr::{ + Accumulator, AggregateUDFImpl, Signature, TypeSignature, Volatility, +}; +use datafusion::physical_expr::expressions::Literal; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkListAgg { + signature: Signature, +} + +impl Default for SparkListAgg { + fn default() -> Self { + Self::new() + } +} + +impl SparkListAgg { + pub fn new() -> Self { + Self { + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]), + TypeSignature::Exact(vec![DataType::Utf8, DataType::Null]), + ], + Volatility::Immutable, + ), + } + } +} + +impl AggregateUDFImpl for SparkListAgg { + fn name(&self) -> &str { + "listagg" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Utf8) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + // Spark's ListAgg is a TypedImperativeAggregate — Catalyst declares its + // intermediate buffer as `BinaryType`. Match that so Comet's shuffle + // layer doesn't have to insert a Utf8 -> Binary cast that the merge + // side then can't read back. + Ok(vec![Field::new( + format_state_name(args.name, "listagg"), + DataType::Binary, + true, + ) + .into()]) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + let Some(lit) = (*acc_args.exprs[1]).downcast_ref::<Literal>() else { + return not_impl_err!( + "listagg delimiter must be a literal; got {:?}", + acc_args.exprs[1] + ); + }; + let delimiter = if lit.value().is_null() { + String::new() + } else if let Some(s) = lit.value().try_as_str() { + s.unwrap_or("").to_string() + } else { + return not_impl_err!( + "listagg delimiter literal must be Utf8; got {:?}", + lit.value() + ); + }; + Ok(Box::new(ListAggAccumulator::new(delimiter))) + } +} + +#[derive(Debug)] +struct ListAggAccumulator { Review Comment: `ListAggAccumulator` is a near-verbatim copy of DataFusion's `SimpleStringAggAccumulator` (`datafusion/functions-aggregate/src/string_agg.rs:451`): same `delimiter`/`accumulated`/`has_value` fields and the same `update_batch`/`evaluate`/`state`/`size` and delimiter-extraction logic (`string_agg.rs:122`). DataFusion's `string_agg` already implements exactly this simple form with identical null-skip and NULL-delimiter semantics, and it ships a `StringAggGroupsAccumulator` fast path (`string_agg.rs:321`) that this PR does not provide. The cited differences (`Utf8` vs `LargeUtf8` return, `Binary` vs `LargeUtf8` state) are cosmetic. Preferred fix: wire DataFusion's `string_agg` with a result cast to `Utf8` via `scalarFunctionExprToProtoWithReturnType`, which deletes this file and gains the grouped fast path. At minimum, add a `GroupsAccumulator`, since every SQL test here is `GROUP BY` and currently takes the slow per-group `Accumulator` path. ########## spark/src/main/spark-4.x/org/apache/comet/serde/CometListAgg.scala: ########## @@ -0,0 +1,107 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ListAgg} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{NullType, StringType} + +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, hasNonDefaultStringCollation, serializeDataType} + +/** + * Spark 4.0+ `LISTAGG` / `STRING_AGG`. + * + * Comet only supports the simple form: a `StringType` child with a literal delimiter and no + * `WITHIN GROUP (ORDER BY ...)` clause. DISTINCT is handled by Spark's multi-stage plan rewrite + * (grouping by the child before the aggregate), so the native side never sees it. + */ +object CometListAgg extends CometAggregateExpressionSerde[ListAgg] { + + private val binaryChildReason = "`BinaryType` inputs are not supported." + private val withinGroupReason = "`WITHIN GROUP (ORDER BY ...)` is not supported." + private val nonFoldableDelimiterReason = "Non-literal delimiters are not supported." + private val collationReason = "Non-default string collations are not supported." + private val distinctReason = + "`DISTINCT` falls back to Spark because Comet rejects multi-column distinct aggregates." + + override def getUnsupportedReasons(): Seq[String] = Seq( + binaryChildReason, + withinGroupReason, + nonFoldableDelimiterReason, + collationReason, + distinctReason) + + override def getSupportLevel(expr: ListAgg): SupportLevel = { + // Spark's analyzer already enforces `delimiter.foldable`, so this only ever rejects + // non-string / non-null delimiter types. + expr.child.dataType match { + case _: StringType if hasNonDefaultStringCollation(expr.child.dataType) => + Unsupported(Some(collationReason)) + case _: StringType => + expr.delimiter.dataType match { + case _: StringType if hasNonDefaultStringCollation(expr.delimiter.dataType) => Review Comment: The collation fallback path (`CometListAgg.scala:57`, `:61`) is reachable but untested. Add: ```sql statement CREATE TABLE la_coll(v string COLLATE UTF8_LCASE, grp string) USING parquet statement INSERT INTO la_coll VALUES ('a','g1'), ('B','g1') query expect_fallback(Non-default string collations are not supported) SELECT grp, listagg(v, ',') FROM la_coll GROUP BY grp ORDER BY grp ``` Confirm the exact fallback substring against the emitted `collationReason`. Also add a multi-partition case (repartition before the aggregate) to stress the arrival-order assumption discussed below, since every current case is single-partition. ########## spark/src/test/resources/sql-tests/expressions/aggregate/listagg.sql: ########## @@ -0,0 +1,157 @@ +-- 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. + +-- listagg / string_agg is available starting in Spark 4.0. +-- MinSparkVersion: 4.0 +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- ============================================================ +-- Setup Review Comment: Without `WITHIN GROUP`, both Spark (`collect.scala`, "without order return as is") and this accumulator concatenate in group arrival order, which Spark documents as non-deterministic after shuffle (`collect.scala:309`). This is the first order-sensitive aggregate Comet accelerates, so there is no precedent. No divergence was found because the SQL tests wrap inputs in `(SELECT ... ORDER BY grp, v)` and run single-partition, so Spark and Comet observe the same order. That protection is incidental. Add the multi-partition test above and add the non-determinism caveat to the user-guide note, since Spark itself flags it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
