andygrove commented on code in PR #4816: URL: https://github.com/apache/datafusion-comet/pull/4816#discussion_r3635629757
########## 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: Added a `GroupsAccumulator` (`ListAggGroupsAccumulator`) so grouped queries no longer take the per-group `Accumulator` path. I kept the Comet-owned UDAF rather than reusing `string_agg` directly. `string_agg` returns `LargeUtf8` and its `SimpleStringAggAccumulator` / `StringAggGroupsAccumulator` are `pub(crate)`, so reuse would mean wiring the whole UDAF plus a `LargeUtf8` -> `Utf8` cast on the aggregate result. The cited `scalarFunctionExprToProtoWithReturnType` is a scalar-function path and does not apply to aggregates. Emitting `Utf8` at the JVM FFI boundary is the reason the custom UDAF exists. ########## 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: Removed the `datatype` proto field and the unreachable `dataType.isEmpty` branch. `getSupportLevel` guarantees a `StringType` child, so `serializeDataType` always returns `Some`, and the planner hardcodes `Utf8`. ########## 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: Removed `distinctReason`. Added a comment noting DISTINCT falls back earlier in `aggExprToProto` with the generic multi-column-distinct message, which is what the SQL test matches. -- 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]
