sunchao commented on code in PR #7191: URL: https://github.com/apache/datafusion-comet/pull/7191#discussion_r4407341369
########## native/spark-expr/src/agg_funcs/min_max_struct.rs: ########## @@ -0,0 +1,589 @@ +// 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. + +use arrow::array::{new_null_array, Array, ArrayRef, BooleanArray}; +use arrow::compute::{not, nullif, prep_null_mask_filter, SortOptions}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::row::{OwnedRow, RowConverter, Rows, SortField}; +use datafusion::common::{internal_err, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::{ + Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature, Volatility, +}; +use datafusion::physical_expr::expressions::format_state_name; +use std::fmt::Debug; +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; + +/// Spark's `max` or `min` over a struct. +/// +/// Spark orders structs field by field, with a null field below every value. Each input batch is +/// encoded once with Arrow's row format, ascending with nulls first, whose bytes sort in exactly +/// that order. Each group keeps the bytes of its current winner, and memory is accounted by the +/// size of those buffers. Float and double fields are not supported: the row format orders `-0.0` +/// below `0.0`, while Spark treats them as equal. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStructMinMax { + signature: Signature, + is_max: bool, +} + +impl SparkStructMinMax { + pub fn new(is_max: bool) -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + is_max, + } + } +} + +impl AggregateUDFImpl for SparkStructMinMax { + fn name(&self) -> &str { + if self.is_max { + "max" + } else { + "min" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(StructMinMaxAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(args.name, self.name()), + args.return_field.data_type().clone(), + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + Ok(Box::new(StructMinMaxGroupsAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } +} + +/// Encodes struct values to Spark-ordered row bytes and decodes them back. +#[derive(Debug)] +struct StructRows { + converter: RowConverter, + /// The encoding of a null struct, emitted for a group that saw no value. + null_row: OwnedRow, + is_max: bool, +} + +impl StructRows { + fn try_new(data_type: &DataType, is_max: bool) -> Result<Self> { + if !matches!(data_type, DataType::Struct(_)) { + return internal_err!("Spark struct min/max expects a struct, got {data_type}"); + } + let options = SortOptions { + descending: false, + nulls_first: true, + }; + let converter = RowConverter::new(vec![SortField::new_with_options( + data_type.clone(), + options, + )])?; + let null_row = converter + .convert_columns(&[new_null_array(data_type, 1)])? + .row(0) + .owned(); + Ok(Self { + converter, + null_row, + is_max, + }) + } + + fn encode(&self, values: &ArrayRef) -> Result<Rows> { + Ok(self.converter.convert_columns(&[Arc::clone(values)])?) Review Comment: [P2] Avoid the wide-payload regression before enabling struct min/max by default Could we address the cost of encoding every complete struct before choosing the winner, and add a benchmark for wide payloads? For the motivating `max(struct(order_key, payload))` shape, this serializes the payload even when the first key already decides the comparison. A matched current-head control shows the newly enabled path is slower than keeping the aggregate on Spark. I wrote these 200,000 rows to one Parquet file and read the same file for both variants: ```sql SELECT id % 1000 AS g, 200000 - id AS k, CASE WHEN id % 10 = 0 THEN NULL ELSE repeat(sha2(CAST(id AS STRING), 256), 64) END AS payload FROM range(200000) ``` Then I collected `SELECT max(struct(k, payload)) FROM t` and `SELECT g, max(struct(k, payload)) FROM t GROUP BY g`. With Spark 4.1.3/JDK 17, `local[1]`, AQE off, one shuffle partition, Comet enabled, on-heap execution and JVM columnar shuffle, I changed only `spark.comet.expression.Max.enabled`. Plans confirmed the same Comet scan and, for grouped input, the same native sorts and columnar shuffle. The control uses Spark partial/final aggregates and the enabled variant uses Comet partial/final aggregates. After two warmup rounds, five alternating measured rounds gave median collection times of **607 ms fallback vs 798 ms native** ungrouped, and **2,075 ms fallback vs 4,059 ms native** grouped. Every native sample exceeded the slowest fallback sample for that query, and all answers matched. Planning and result comparison were outside the timed region. These are local pipeline measurements using this head's compiled JVM sources and CI native library; they establish the regression for this shape without attributing all elapsed time to this function. The 64-byte payload case was faster than plain Spark, so a benchmark should retain both narrow and wide cases. -- 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]
