mbutrovich commented on code in PR #4817: URL: https://github.com/apache/datafusion-comet/pull/4817#discussion_r3659340300
########## native/spark-expr/src/agg_funcs/max_min_by.rs: ########## @@ -0,0 +1,632 @@ +// 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::SortOptions; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion::common::{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::mem::size_of_val; +use std::sync::Arc; + +/// Spark-compatible `max_by(value, ordering)` / `min_by(value, ordering)` aggregate. +/// +/// Returns the `value` associated with the maximum (`max_by`) or minimum (`min_by`) +/// non-null `ordering`. Rows with a null `ordering` are ignored. The returned value +/// may itself be null when it is the value paired with the extremum ordering. If every +/// `ordering` in the group is null, the result is null. +/// +/// Spark's `MaxBy`/`MinBy` are `DeclarativeAggregate`s that keep a `(value, ordering)` +/// buffer and, on a tie in the ordering, the later row wins. Because ties across +/// partitions are processed in an unspecified order, Spark documents the function as +/// non-deterministic when several rows share the extremum ordering. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MaxMinBy { + name: String, + signature: Signature, + /// `true` for `max_by`, `false` for `min_by`. + is_max: bool, +} + +impl std::hash::Hash for MaxMinBy { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.is_max.hash(state); + } +} + +impl MaxMinBy { + /// Create a `max_by` aggregate. + pub fn new_max_by() -> Self { + Self { + name: "max_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: true, + } + } + + /// Create a `min_by` aggregate. + pub fn new_min_by() -> Self { + Self { + name: "min_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: false, + } + } +} + +impl AggregateUDFImpl for MaxMinBy { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + // The result has the same type as the `value` argument. + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + let value_type = acc_args.exprs[0].data_type(acc_args.schema)?; + let ordering_type = acc_args.exprs[1].data_type(acc_args.schema)?; + Ok(Box::new(MaxMinByAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + let value_type = args.input_fields[0].data_type().clone(); + let ordering_type = args.input_fields[1].data_type().clone(); + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "value"), + value_type, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "ordering"), + ordering_type, + true, + )), + ]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + let value_type = args.exprs[0].data_type(args.schema)?; + let ordering_type = args.exprs[1].data_type(args.schema)?; + Ok(Box::new(MaxMinByGroupsAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } +} + +/// Sort options that make the wanted extremum encode to the largest row bytes: ascending for +/// `max_by` (largest ordering wins), descending for `min_by` (smallest ordering wins). Nulls +/// sort first (smallest) so they are never selected as the extremum; null orderings are also +/// skipped explicitly. +fn extremum_sort_options(is_max: bool) -> SortOptions { Review Comment: The Arrow row encoding does not reproduce Spark's ordering for `-0.0` and `0.0`, so `max_by`/`min_by` can pick a different row than Spark when the ordering column is a float containing signed zeros. Spark orders doubles with `SQLOrderingUtil.compareDoubles`, which is `if (x == y) 0 else java.lang.Double.compare(x, y)` and documents "3. `-0.0 == 0.0`" (`SQLOrderingUtil.scala:27-31`, wired in via `PhysicalDoubleType.ordering`, `PhysicalDataType.scala:183-184`). So Spark treats `-0.0` and `0.0` as a tie, and the tie is resolved by row order: the update chain keeps the existing value when `predicate(old, new)` holds, so for `max_by` an equal ordering leaves the earlier row's value in place. Arrow's row format instead encodes floats by flipping bits off the sign (`arrow-row/src/fixed.rs:149-157`), which is a total order that places `-0.0` strictly below `0.0`. So `rows.row(i) >= rows.row(b)` at `:196` and `candidate >= self.best_ordering[...]` at `:340` see a strict inequality where Spark sees equality. Concretely, for `max_by(v, ord)` over `(1, 0.0), (2, -0.0)`: Spark ties, keeps the first, returns 1. Comet ranks `0.0 > -0.0`, so the second row does not displace the first and it also returns 1. But reverse the input to `(1, -0.0), (2, 0.0)` and Spark still ties and returns 1, while Comet sees `0.0 > -0.0`, takes the second, and returns 2. `min_by` diverges symmetrically. Neither fixture covers this. `max_by.sql` has NaN (`:122`) and both infinities (`:198`) but no signed zero, and the Rust tests only cover NaN (`max_min_by.rs:474`). Please add a signed-zero case to both SQL fixtures and a Rust test, and normalize `-0.0` to `0.0` on the ordering column before row conversion so ties fall through to the row-order tie-break. Note this is the mirror image of the `mode` finding in #4782: there the fix is to stop collapsing `-0.0`, because `mode` keys on `OpenHashSet.equals`, which distinguishes them. Here the fix is to start collapsing, because `max_by` compares with `SQLOrderingUtil`, which does not. The two aggregates genuinely differ, so it is worth a comment at each site recording which Spark comparison path applies. ########## native/spark-expr/src/agg_funcs/max_min_by.rs: ########## @@ -0,0 +1,632 @@ +// 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::SortOptions; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion::common::{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::mem::size_of_val; +use std::sync::Arc; + +/// Spark-compatible `max_by(value, ordering)` / `min_by(value, ordering)` aggregate. +/// +/// Returns the `value` associated with the maximum (`max_by`) or minimum (`min_by`) +/// non-null `ordering`. Rows with a null `ordering` are ignored. The returned value +/// may itself be null when it is the value paired with the extremum ordering. If every +/// `ordering` in the group is null, the result is null. +/// +/// Spark's `MaxBy`/`MinBy` are `DeclarativeAggregate`s that keep a `(value, ordering)` +/// buffer and, on a tie in the ordering, the later row wins. Because ties across +/// partitions are processed in an unspecified order, Spark documents the function as +/// non-deterministic when several rows share the extremum ordering. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MaxMinBy { + name: String, + signature: Signature, + /// `true` for `max_by`, `false` for `min_by`. + is_max: bool, +} + +impl std::hash::Hash for MaxMinBy { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.is_max.hash(state); + } +} + +impl MaxMinBy { + /// Create a `max_by` aggregate. + pub fn new_max_by() -> Self { + Self { + name: "max_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: true, + } + } + + /// Create a `min_by` aggregate. + pub fn new_min_by() -> Self { + Self { + name: "min_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: false, + } + } +} + +impl AggregateUDFImpl for MaxMinBy { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + // The result has the same type as the `value` argument. + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + let value_type = acc_args.exprs[0].data_type(acc_args.schema)?; + let ordering_type = acc_args.exprs[1].data_type(acc_args.schema)?; + Ok(Box::new(MaxMinByAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + let value_type = args.input_fields[0].data_type().clone(); + let ordering_type = args.input_fields[1].data_type().clone(); + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "value"), + value_type, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "ordering"), + ordering_type, + true, + )), + ]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + let value_type = args.exprs[0].data_type(args.schema)?; + let ordering_type = args.exprs[1].data_type(args.schema)?; + Ok(Box::new(MaxMinByGroupsAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } +} + +/// Sort options that make the wanted extremum encode to the largest row bytes: ascending for +/// `max_by` (largest ordering wins), descending for `min_by` (smallest ordering wins). Nulls +/// sort first (smallest) so they are never selected as the extremum; null orderings are also +/// skipped explicitly. +fn extremum_sort_options(is_max: bool) -> SortOptions { + SortOptions { + descending: !is_max, + nulls_first: true, + } +} + +/// Accumulator that tracks the running `(value, ordering)` pair for the extremum ordering. +#[derive(Debug)] +struct MaxMinByAccumulator { + /// The value paired with the current extremum ordering. May be null. + value: ScalarValue, + /// The current extremum ordering. Null means no non-null ordering has been seen yet. + ordering: ScalarValue, + /// `true` for `max_by`, `false` for `min_by`. + is_max: bool, +} + +impl MaxMinByAccumulator { + fn try_new(value_type: DataType, ordering_type: DataType, is_max: bool) -> Result<Self> { + Ok(Self { + value: ScalarValue::try_from(&value_type)?, + ordering: ScalarValue::try_from(&ordering_type)?, + is_max, + }) + } + + fn sort_options(&self) -> SortOptions { + // Encode the ordering column into arrow's row format so that the extremum can be + // found for any orderable type with a single comparison. + extremum_sort_options(self.is_max) + } + + /// Apply a batch of `(value, ordering)` columns, keeping the value paired with the + /// extremum ordering. Rows with a null ordering are ignored. + fn update_from(&mut self, value_arr: &ArrayRef, ordering_arr: &ArrayRef) -> Result<()> { + if ordering_arr.is_empty() { + return Ok(()); + } + + let converter = RowConverter::new(vec![SortField::new_with_options( + ordering_arr.data_type().clone(), + self.sort_options(), + )])?; + let rows = converter.convert_columns(&[Arc::clone(ordering_arr)])?; + + // Find the index of the extremum ordering in this batch (last one wins on a tie, Review Comment: The comment says "last one wins on a tie, matching Spark's sequential row processing", which is backwards for the value that gets kept. Spark's update keeps the existing value on a tie: `If(predicate(extremumOrdering, orderingExpr), valueWithExtremumOrdering, valueExpr)` where `predicate` for `max_by` is `oldExpr > newExpr` (`MaxByAndMinBy.scala:69-77`, `:113-114`). On a tie the predicate is false, so the expression evaluates to `valueExpr`, the new row. So last-wins is correct for `max_by` here, but the reasoning is not "sequential processing", it is the specific strictness of the predicate, and the same strictness is what makes the `-0.0` case above diverge. Please state the actual rule so the next reader can check it against `MaxByAndMinBy.scala` rather than re-deriving it. ########## native/spark-expr/src/agg_funcs/max_min_by.rs: ########## @@ -0,0 +1,632 @@ +// 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::SortOptions; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion::common::{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::mem::size_of_val; +use std::sync::Arc; + +/// Spark-compatible `max_by(value, ordering)` / `min_by(value, ordering)` aggregate. +/// +/// Returns the `value` associated with the maximum (`max_by`) or minimum (`min_by`) +/// non-null `ordering`. Rows with a null `ordering` are ignored. The returned value +/// may itself be null when it is the value paired with the extremum ordering. If every +/// `ordering` in the group is null, the result is null. +/// +/// Spark's `MaxBy`/`MinBy` are `DeclarativeAggregate`s that keep a `(value, ordering)` +/// buffer and, on a tie in the ordering, the later row wins. Because ties across +/// partitions are processed in an unspecified order, Spark documents the function as +/// non-deterministic when several rows share the extremum ordering. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MaxMinBy { + name: String, + signature: Signature, + /// `true` for `max_by`, `false` for `min_by`. + is_max: bool, +} + +impl std::hash::Hash for MaxMinBy { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.is_max.hash(state); + } +} + +impl MaxMinBy { + /// Create a `max_by` aggregate. + pub fn new_max_by() -> Self { + Self { + name: "max_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: true, + } + } + + /// Create a `min_by` aggregate. + pub fn new_min_by() -> Self { + Self { + name: "min_by".to_string(), + signature: Signature::any(2, Volatility::Immutable), + is_max: false, + } + } +} + +impl AggregateUDFImpl for MaxMinBy { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + // The result has the same type as the `value` argument. + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + let value_type = acc_args.exprs[0].data_type(acc_args.schema)?; + let ordering_type = acc_args.exprs[1].data_type(acc_args.schema)?; + Ok(Box::new(MaxMinByAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + let value_type = args.input_fields[0].data_type().clone(); + let ordering_type = args.input_fields[1].data_type().clone(); + Ok(vec![ + Arc::new(Field::new( + format_state_name(&self.name, "value"), + value_type, + true, + )), + Arc::new(Field::new( + format_state_name(&self.name, "ordering"), + ordering_type, + true, + )), + ]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + let value_type = args.exprs[0].data_type(args.schema)?; + let ordering_type = args.exprs[1].data_type(args.schema)?; + Ok(Box::new(MaxMinByGroupsAccumulator::try_new( + value_type, + ordering_type, + self.is_max, + )?)) + } +} + +/// Sort options that make the wanted extremum encode to the largest row bytes: ascending for +/// `max_by` (largest ordering wins), descending for `min_by` (smallest ordering wins). Nulls +/// sort first (smallest) so they are never selected as the extremum; null orderings are also +/// skipped explicitly. +fn extremum_sort_options(is_max: bool) -> SortOptions { + SortOptions { + descending: !is_max, + nulls_first: true, + } +} + +/// Accumulator that tracks the running `(value, ordering)` pair for the extremum ordering. +#[derive(Debug)] +struct MaxMinByAccumulator { + /// The value paired with the current extremum ordering. May be null. + value: ScalarValue, + /// The current extremum ordering. Null means no non-null ordering has been seen yet. + ordering: ScalarValue, + /// `true` for `max_by`, `false` for `min_by`. + is_max: bool, +} + +impl MaxMinByAccumulator { + fn try_new(value_type: DataType, ordering_type: DataType, is_max: bool) -> Result<Self> { + Ok(Self { + value: ScalarValue::try_from(&value_type)?, + ordering: ScalarValue::try_from(&ordering_type)?, + is_max, + }) + } + + fn sort_options(&self) -> SortOptions { + // Encode the ordering column into arrow's row format so that the extremum can be + // found for any orderable type with a single comparison. + extremum_sort_options(self.is_max) + } + + /// Apply a batch of `(value, ordering)` columns, keeping the value paired with the + /// extremum ordering. Rows with a null ordering are ignored. + fn update_from(&mut self, value_arr: &ArrayRef, ordering_arr: &ArrayRef) -> Result<()> { + if ordering_arr.is_empty() { + return Ok(()); + } + + let converter = RowConverter::new(vec![SortField::new_with_options( + ordering_arr.data_type().clone(), + self.sort_options(), + )])?; + let rows = converter.convert_columns(&[Arc::clone(ordering_arr)])?; + + // Find the index of the extremum ordering in this batch (last one wins on a tie, + // matching Spark's sequential row processing), ignoring null orderings. + let mut best: Option<usize> = None; + for i in 0..ordering_arr.len() { + if ordering_arr.is_null(i) { + continue; + } + best = match best { + None => Some(i), + Some(b) if rows.row(i) >= rows.row(b) => Some(i), + Some(b) => Some(b), + }; + } + + let Some(b) = best else { + return Ok(()); + }; + + let candidate_ordering = ScalarValue::try_from_array(ordering_arr, b)?; + let take = if self.ordering.is_null() { + true + } else { + // Compare the batch's extremum ordering against the running extremum using the + // same row encoding. Build a two-row array [running, candidate] and compare. + let pair = ScalarValue::iter_to_array(vec![ Review Comment: The scalar accumulator allocates a `RowConverter` per call to `update_from` (`:181`) and then, for every batch that has a non-null ordering, builds a two-element array and converts it just to compare the batch extremum against the running one (`:211-216`). The grouped path already avoids this by holding `ordering_converter` as a field (`:260`) and keeping the running best as `OwnedRow` bytes, so comparison is a byte compare with no allocation (`:340`). The scalar accumulator can do the same: hold the converter and store `ordering` as an `OwnedRow` alongside the `ScalarValue` needed by `state()`. As written, the global (non-grouped) path pays two `RowConverter` constructions and an extra array build per batch, which is the path the benchmark table shows at 24ms and therefore the one where fixed per-batch overhead is most visible. -- 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]
