rich7420 commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r3951496223
########## 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: For rows `(0.0, 10)` and `(-0.0, 20)`, pivoting `sum(v)` on `-0.0` returns `NULL` natively, while Spark returns `30`. NaN keys also match natively where Spark returns `NULL`. Please preserve Spark's key matching and add Float/Double tests that assert native execution. -- 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]
