mbutrovich commented on code in PR #4816: URL: https://github.com/apache/datafusion-comet/pull/4816#discussion_r4066317996
########## 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: You're right that `scalarFunctionExprToProtoWithReturnType` doesn't apply to aggregates, and `StringAggGroupsAccumulator` is private in DataFusion 55 ([`string_agg.rs:323`](https://github.com/apache/datafusion/blob/55.0.0/datafusion/functions-aggregate/src/string_agg.rs#L323)). The `Binary` state requirement from the other thread rules out a straight reuse anyway, because `string_agg` keeps `LargeUtf8` state. I'm fine with the Comet-owned UDAF. -- 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]
