andygrove commented on code in PR #4817: URL: https://github.com/apache/datafusion-comet/pull/4817#discussion_r3713808695
########## 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: Fixed, and your `-0.0` analysis holds — with one addition and one correction to the testing plan. Ordering canonicalization is now applied to the ordering column before row conversion, in both the scalar and grouped paths. I checked whether this needed the same version gate as #4782, and it does not: `SQLOrderingUtil.compareDoubles` is byte-identical on `branch-3.4`, `branch-3.5`, `branch-4.0`, `branch-4.1`, `branch-4.2` and `master`, and `PhysicalDoubleType.ordering`/`PhysicalFloatType.ordering` still route through it. So unlike `mode`, this one is stable across the matrix. **Addition: the same fix is needed for `NaN`, for the same underlying reason.** `Double.compare` goes through `doubleToLongBits`, so a sign-bit-set `NaN` is the *same value* as a positive `NaN` and still sorts above `+Infinity`. Arrow's raw-bit encoding places it below `-Infinity` instead. So `max_by(v, ord)` with `ord = -NaN` picked the wrong row too. The canonicalization folds every `NaN` to the canonical one alongside the zero fold, which makes the row bytes agree with `compareDoubles` on both counts. New test `max_by_sign_bit_nan_is_still_largest` covers it; the existing `max_by_nan_is_largest` only used a positive `NaN` and passed either way. **Correction: the SQL fixtures you asked for cannot be written meaningfully, and I have the measurement.** The divergence is only observable when the two zeros tie *with different values attached* — the canonicalization changes nothing when the tied rows share a value, because there is no ordering strictly between `-0.0` and `0.0` for a third row to separate them. But which tied row wins depends on the inter-partition merge order, and these fixture tables span 5 partitions (`spark.table(...).rdd.getNumPartitions` = 5), so that case is exactly the non-determinism `max_by` documents. I built the fixture you described anyway and measured it with the canonicalization disabled: ``` === RAW === grp=g1 v=1 ord=-0.0 bits=-9223372036854775808 grp=g1 v=2 ord=0.0 bits=0 grp=g1 v=3 ord=-5.0 bits=-4606056518893174784 grp=g2 v=4 ord=0.0 bits=0 grp=g2 v=5 ord=-0.0 bits=-9223372036854775808 grp=g2 v=6 ord=-5.0 bits=-4606056518893174784 === SPARK: g1->2, g2->4 === COMET: g1->2, g2->4 <-- canonicalization disabled ``` Both sides agree, so the fixture would have been vacuous *and* pinning unspecified behaviour — the worst combination, since it would look like coverage and could flip later on a partitioning change. So the fixtures I added give the tied rows equal values. That is deterministic under any merge order and still runs signed zeros through the canonicalization path end to end, with a comment recording why they are shaped that way and where the real guarantee lives. The order-dependent behaviour is pinned in the Rust tests instead, where row order is explicit: `max_by`/`min_by` ties in both directions, the across-batch case, the grouped case, and `Float32`. All six fail without the canonicalization and pass with it — I verified by reverting it to a no-op. One incidental find worth flagging for future float fixtures: `CAST(-0.0 AS DOUBLE)` does not produce a negative zero, because an unsuffixed `-0.0` is a `DecimalType` literal and `Decimal` has no signed zero. I read the bits back out of Parquet to confirm. `-0.0D` is needed. This is the same trap that made the `mode` fixture in #4782 vacuous. ########## 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: Done. `MaxMinByAccumulator` now holds the `RowConverter` as a field and keeps the running extremum as `OwnedRow` bytes, so the per-batch comparison is a byte compare against `best_ordering` with no `RowConverter` construction and no two-element array build — the same shape as the grouped path you pointed at. Two details that fell out of it: * `best_ordering` is an `Option<OwnedRow>` rather than a null-valued row, which also removes the `self.ordering.is_null()` special case that stood in for "nothing seen yet". * `state` needs the ordering back as a `ScalarValue`, so it reconstructs it from the row bytes via `convert_rows`. That is once per `state` call rather than once per batch, so the per-batch allocation is gone rather than moved. The value column still goes through `ScalarValue::try_from_array` at the winning index, so the scalar path needs no value converter at all. ########## 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: Reworded to state the rule rather than the conclusion: ```rust // Find the index of the extremum ordering in this batch, ignoring null orderings. `>=` // makes the later row win a tie, which is what Spark's update does: it evaluates // `If(predicate(extremumOrdering, orderingExpr), valueWithExtremumOrdering, valueExpr)` // where `predicate` is the strict `oldExpr > newExpr` for `max_by` (`<` for `min_by`), so // an equal ordering makes the predicate false and the *new* row's value is kept // (`MaxByAndMinBy.scala`). The same strictness is why the signed-zero canonicalization // above matters: without it Arrow sees a strict inequality where Spark sees a tie. ``` I verified the update expression against `MaxByAndMinBy.scala` rather than taking it from the review: `updateExpressions` is `If(predicate(extremumOrdering, orderingExpr), valueWithExtremumOrdering, valueExpr)` with `predicate` = `oldExpr > newExpr` for `MaxBy` and `oldExpr < newExpr` for `MinBy`, so last-wins is right for both, and for the reason you gave. The grouped accumulator's comment now points here rather than restating it. -- 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]
