andygrove commented on code in PR #4816: URL: https://github.com/apache/datafusion-comet/pull/4816#discussion_r3635631091
########## 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 `Binary` state is actually required. I switched `state_fields` to `Utf8` to try it and it fails at runtime: ``` could not cast array of type Binary to arrow_array::...::GenericStringType<i32> ``` `CometHashAggregateExec.output` for a Partial aggregate is Spark's `groupingAttributes ++ aggBufferAttributes`, and Spark's `ListAgg` is a `TypedImperativeAggregate` whose `aggBufferAttributes` is a single `BinaryType` column. The shuffle Exchange between the partial and final aggregate is serialized with that Spark-declared schema, so the native partial state must be `Binary`. This holds even though `supportsMixedPartialFinal=false`: partial and final both run in Comet, but the Exchange schema between them is Spark's plan schema, not DataFusion's. That is why `string_agg` using `Utf8`/`LargeUtf8` state does not carry over here. You were right that the old comment was hard to reconcile: the "`Utf8` -> `Binary` cast the merge side cannot read" wording was imprecise. I rewrote the module doc and the `state_fields` comment to explain the real mechanism. ########## 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: Added the multi-partition test (`REPARTITION(4)` over identical per-group values, so the result stays deterministic while the group is split across partitions and merged in the final aggregate). The collation fallback turns out not to be testable at the `listagg` level. Comet's Parquet scan rejects a collated-string schema before the aggregate is considered: ``` Unsupported schema StructType(StructField(v,StringType(UTF8_LCASE),true),...) ``` so the whole plan falls back to Spark and `CometListAgg`'s collation guard is never reached. (Also, the delimiter must share the child's collation or Spark's analyzer rejects the call outright with `DATATYPE_MISMATCH`.) I kept the guard as defense-in-depth (it keeps `listagg` correct if Comet later gains collated-scan support) and documented this in both the serde and the SQL file rather than adding a test that cannot exercise it. ########## 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: Added the multi-partition merge test and added the non-determinism caveat to the user-guide note ("Without `WITHIN GROUP`, concatenation order is the (non-deterministic) group arrival order, matching Spark"). -- 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]
