andygrove commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r3969478008
########## native/spark-expr/src/agg_funcs/pivot_first.rs: ########## @@ -0,0 +1,349 @@ +// 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's `PivotFirst` aggregate. Used only by the second phase of the optimized pivot plan +//! generated by `PivotTransformer`. For each group, `PivotFirst` maintains an array of +//! `pivot_values.len()` slots; on each input row it evaluates the pivot column, looks up its +//! index in `pivot_values`, and writes the value column into that slot when a match is found +//! and the value is non-null. Rows with unmatched pivot values are ignored; matched rows with +//! a null value column leave the slot unchanged (matches Spark). +//! +//! State layout is one column per pivot slot, matching Spark's `aggBufferAttributes` (which +//! declares `indexSize` `AttributeReference`s, one per pivot value). This keeps the shuffle +//! schema between Partial and Final consistent with what Spark catalyst declared; otherwise +//! the shuffle exchange rejects the batch. `evaluate()` reassembles the slots into a +//! `ListArray` matching `PivotFirst.dataType = ArrayType(value_type)`. + +use arrow::array::{Array, ArrayRef}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::utils::SingleRowListArrayBuilder; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature}; +use datafusion::physical_expr::expressions::format_state_name; +use std::collections::HashMap; +use std::sync::Arc; + +/// UDAF implementation of Spark's `PivotFirst`. +/// +/// `pivot_values` is a fixed, plan-time list of the pivot column values that occupy each +/// output slot; `pivot_index[v] = i` means an input row whose pivot column equals `v` writes +/// into slot `i`. Both the vector and the map are wrapped in `Arc` because `accumulator()` +/// fires once per group in a grouped aggregate and we want that path to bump a refcount +/// rather than deep-clone. +#[derive(Debug)] +pub struct SparkPivotFirst { + signature: Signature, + value_type: DataType, + // Kept for `PartialEq`/`Hash` (identity of the aggregate for plan comparison) and for the + // deterministic slot ordering `state_fields` needs. `HashMap` alone would give us the map + // but not a stable order or a `Hash` impl. + pivot_values: Arc<Vec<ScalarValue>>, + pivot_index: Arc<HashMap<ScalarValue, usize>>, +} + +impl PartialEq for SparkPivotFirst { + fn eq(&self, other: &Self) -> bool { + self.value_type == other.value_type && self.pivot_values == other.pivot_values + } +} + +impl Eq for SparkPivotFirst {} + +impl std::hash::Hash for SparkPivotFirst { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.value_type.hash(state); + self.pivot_values.hash(state); + } +} + +impl SparkPivotFirst { + pub fn new(value_type: DataType, pivot_values: Vec<ScalarValue>) -> Self { + let mut pivot_index = HashMap::with_capacity(pivot_values.len()); + // Spark's PivotFirst uses the FIRST occurrence's index (HashMap/TreeMap semantics), so + // when duplicates are somehow present we mirror that by only inserting the first one. + for (i, v) in pivot_values.iter().enumerate() { + pivot_index.entry(v.clone()).or_insert(i); + } + Self { + signature: Signature::user_defined(Immutable), + value_type, + pivot_values: Arc::new(pivot_values), + pivot_index: Arc::new(pivot_index), + } + } +} + +impl AggregateUDFImpl for SparkPivotFirst { + fn name(&self) -> &str { + "pivot_first" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> { + Ok(DataType::List(Arc::new(Field::new_list_field( + self.value_type.clone(), + true, + )))) + } + + fn state_fields(&self, args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> { + // One field per pivot slot, matching Spark's aggBufferAttributes so the shuffle + // exchange sees the same schema catalyst declared. `format_state_name` is the same + // helper other aggregates in this crate use (see `avg.rs`, `stddev.rs`). + Ok((0..self.pivot_values.len()) + .map(|i| { + Arc::new(Field::new( + format_state_name(args.name, &i.to_string()), + self.value_type.clone(), + true, + )) + }) + .collect()) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> DFResult<Box<dyn Accumulator>> { + Ok(Box::new(PivotFirstAccumulator::new( + self.value_type.clone(), + Arc::clone(&self.pivot_index), + ))) + } +} + +/// Per-group state: `slots[i]` holds the latest non-null value assigned to pivot slot `i`, or +/// `None` when nothing has written to that slot yet. +#[derive(Debug)] +struct PivotFirstAccumulator { + value_type: DataType, + pivot_index: Arc<HashMap<ScalarValue, usize>>, + slots: Vec<Option<ScalarValue>>, +} + +impl PivotFirstAccumulator { + fn new(value_type: DataType, pivot_index: Arc<HashMap<ScalarValue, usize>>) -> Self { + let slots = vec![None; pivot_index.len()]; + Self { + value_type, + pivot_index, + slots, + } + } + + /// Turn slot `i` into a `ScalarValue`, substituting a typed null when the slot is empty. + fn slot_or_null(&self, i: usize) -> DFResult<ScalarValue> { + Ok(match &self.slots[i] { + Some(v) => v.clone(), + None => ScalarValue::try_from(&self.value_type)?, + }) + } +} + +impl Accumulator for PivotFirstAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> { + if values.len() != 2 { + return Err(DataFusionError::Internal(format!( + "PivotFirst expects 2 inputs (pivot, value); got {}", + values.len() + ))); + } + let pivot_arr = &values[0]; + let value_arr = &values[1]; + if pivot_arr.len() != value_arr.len() { + return Err(DataFusionError::Internal( + "PivotFirst pivot and value arrays have different lengths".into(), + )); + } + for row in 0..pivot_arr.len() { + // Spark ignores the row entirely if either the pivot value is unmatched (index<0) + // or the value is null. Matching Spark exactly here is important because + // `PivotFirst.update` never writes for a null value, so a mid-batch null does not + // clobber an earlier non-null. + let pivot_scalar = ScalarValue::try_from_array(pivot_arr, row)?; + if let Some(&slot_idx) = self.pivot_index.get(&pivot_scalar) { Review Comment: Fixed in bae922556. Both halves were real, and they were wrong in opposite directions, which is what made this worth tracing rather than patching. Spark looks pivot values up in a Scala `HashMap[Any, Int]`, so matching runs through `BoxesRunTime.equals` / `Statics.anyHash` on the boxed Catalyst value — not through IEEE, and not through `ScalarValue`: - `-0.0` and `0.0` are **one** key (`-0.0 == 0.0` numerically, and `doubleHash` folds both onto the hash of `0L`). `ScalarValue` keeps them apart, so the column came back NULL. - `NaN` matches **nothing**, not even another `NaN`, because Scala's `==` on `Double` is IEEE. `ScalarValue` treats `NaN` as equal to itself, so the column got populated. The fix is a `pivot_key` function applied on both sides — building the index and looking each row up — returning `None` for a value Spark can never match. One thing that fell out of it and is worth flagging, because it would have been a separate bug: dropping NaN from the index means `pivot_index.len()` no longer equals the pivot list length, and the accumulator was sizing its slot vector from the index. That would have made `state()` return fewer columns than `state_fields` declares and the shuffle exchange would reject the batch. The slot count now comes from `pivot_values.len()`, so an unmatchable pivot value still gets its (always-null) output column. `unmatchable_pivot_value_keeps_its_state_slot` covers it. **On the tests.** Four Rust unit tests, and I checked each one fails against the old key handling rather than only passing against the new — the first version of the signed-zero test passed either way, because a later `-0.0` row won the slot regardless of whether the `0.0` row had matched. The input side now carries only positive zero, plus a mirror test for a `0.0` entry catching a `-0.0` row. For native execution specifically, `CometAggregateSuite` has a new *PivotFirst matches Spark's float key semantics natively*, which asserts a `CometHashAggregateExec` containing `PivotFirst` is in the plan and then compares to Spark over both `-0.0D` and `0.0D` pivot lists. Against the pre-fix native library it fails with `Results do not match`, so it is not passing by falling back. `PivotFirst.sql` gained double and float pivot-column sections too. Both fixtures write `-0.0D` / `-0.0F` rather than `CAST(-0.0 AS DOUBLE)` — the cast constant-folds to `+0.0` and would have made the whole case vacuous. The pivot benchmark you asked for in the review comment is not in this push; I'm treating it as a separate piece of work and will follow up on 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]
