andygrove commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r4251802656
########## native/spark-expr/src/agg_funcs/pivot_first.rs: ########## @@ -0,0 +1,497 @@ +// 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. + // `pivot_key` can fold two distinct pivot values (`0.0` and `-0.0`) onto one key, and can + // drop one entirely (NaN), so the index is not necessarily the same length as the slot + // vector - the slot count is always `pivot_values.len()`. + for (i, v) in pivot_values.iter().enumerate() { + if let Some(key) = pivot_key(v.clone()) { + pivot_index.entry(key).or_insert(i); + } + } + Self { + signature: Signature::user_defined(Immutable), + value_type, + pivot_values: Arc::new(pivot_values), + pivot_index: Arc::new(pivot_index), + } + } +} + +/// Rewrite a pivot column value into the key Spark would match it on, or `None` when Spark can +/// never match it. +/// +/// Spark's `PivotFirst` looks pivot values up in a Scala `HashMap[Any, Int]`, so matching goes +/// through `BoxesRunTime.equals` / `Statics.anyHash` on the boxed Catalyst value rather than +/// through `ScalarValue`'s own equality. The two disagree on floats in opposite directions: +/// +/// * `-0.0` and `0.0` are one key for Spark (`-0.0 == 0.0` numerically, and `doubleHash` folds +/// both onto the hash of `0L`), while `ScalarValue` keeps them apart. +/// * `NaN` matches nothing for Spark, not even another `NaN`, because Scala's `==` on `Double` +/// is IEEE. `ScalarValue` treats `NaN` as equal to itself. +/// +/// Nulls are left alone: a null pivot column value does match a null entry in the pivot list, +/// which is what Spark's `pivotIndex.getOrElse(null, -1)` does. +fn pivot_key(v: ScalarValue) -> Option<ScalarValue> { + match v { + ScalarValue::Float32(Some(f)) => { + if f.is_nan() { + None + } else if f == 0.0 { + Some(ScalarValue::Float32(Some(0.0))) + } else { + Some(ScalarValue::Float32(Some(f))) + } + } + ScalarValue::Float64(Some(f)) => { + if f.is_nan() { + None + } else if f == 0.0 { + Some(ScalarValue::Float64(Some(0.0))) + } else { + Some(ScalarValue::Float64(Some(f))) + } + } + other => Some(other), + } +} + +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), + self.pivot_values.len(), + ))) + } +} + +/// 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 { + /// `num_slots` is the pivot list's length, which is what `state_fields` declares. It can + /// exceed `pivot_index.len()` when pivot values collide under `pivot_key` or are unmatchable + /// (NaN); those slots exist in the output and stay null. + fn new( + value_type: DataType, + pivot_index: Arc<HashMap<ScalarValue, usize>>, + num_slots: usize, + ) -> Self { + let slots = vec![None; num_slots]; + 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)?; + let Some(key) = pivot_key(pivot_scalar) else { + // A NaN pivot value matches no slot in Spark, so the row is ignored. + continue; + }; + if let Some(&slot_idx) = self.pivot_index.get(&key) { Review Comment: Fixed in 1c6d38f31c by gating rather than reproducing Spark's ordering natively. The serde now admits only the pivot column types that Spark keys in a `HashMap` by value equality on every version: boolean, the integer types, float, double, decimal, UTF8_BINARY string, date, timestamp and timestamp_ntz. Arrays and structs fall back. That includes pivots on more than one column, which pivot on a struct. The native index also rejects any other pivot column type, so a gap in the gate would error out rather than return nulls. I ran Spark's `SQLQueryTestSuite` cases `pivot.sql` and `udf/udf-pivot.sql` locally against 4.1.3 with Comet on, and both pass now. Query #25 logs the `ArrayType` fallback. Against a deliberately broken native `PivotFirst` both fail at query #3, so they do exercise the native path. ########## spark/src/main/scala/org/apache/comet/serde/aggregates.scala: ########## @@ -1116,6 +1116,66 @@ object CometApproxCountDistinct extends CometAggregateExpressionSerde[HyperLogLo } } +object CometPivotFirst extends CometAggregateExpressionSerde[PivotFirst] { + + // Delegate to Spark's own PivotFirst.supportsDataType so the two lists cannot drift if a + // future Spark version adds a value type to the fast-path gate. + private def unsupportedValueTypeReason(dt: DataType): String = + s"Unsupported value data type: $dt" + + private val emptyPivotValuesReason = "Pivot values list is empty" + + override def getUnsupportedReasons(): Seq[String] = Seq( + "Value data type outside PivotFirst.supportsDataType " + + "(Boolean, Byte, Short, Int, Long, Float, Double, Decimal)", + emptyPivotValuesReason) + + override def getSupportLevel(expr: PivotFirst): SupportLevel = { + if (!PivotFirst.supportsDataType(expr.valueDataType)) { + Unsupported(Some(unsupportedValueTypeReason(expr.valueDataType))) + } else if (expr.pivotColumnValues.isEmpty) { + Unsupported(Some(emptyPivotValuesReason)) + } else { + Compatible() Review Comment: Fixed in 1c6d38f31c: binary pivot columns fall back. Spark's answer here also depends on the version. 4.2 moved binary to the `TreeMap` path and matches it by content, so falling back is the only way to be right on both. Strings with a non-UTF8_BINARY collation fall back for the same reason, since 4.2 matches those by collation. `PivotFirst.sql` checks your `X'61'` / `X'62'` case with `expect_fallback`. ########## native/spark-expr/src/agg_funcs/pivot_first.rs: ########## @@ -0,0 +1,497 @@ +// 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. + // `pivot_key` can fold two distinct pivot values (`0.0` and `-0.0`) onto one key, and can + // drop one entirely (NaN), so the index is not necessarily the same length as the slot + // vector - the slot count is always `pivot_values.len()`. + for (i, v) in pivot_values.iter().enumerate() { + if let Some(key) = pivot_key(v.clone()) { + pivot_index.entry(key).or_insert(i); + } + } + Self { + signature: Signature::user_defined(Immutable), + value_type, + pivot_values: Arc::new(pivot_values), + pivot_index: Arc::new(pivot_index), + } + } +} + +/// Rewrite a pivot column value into the key Spark would match it on, or `None` when Spark can +/// never match it. +/// +/// Spark's `PivotFirst` looks pivot values up in a Scala `HashMap[Any, Int]`, so matching goes +/// through `BoxesRunTime.equals` / `Statics.anyHash` on the boxed Catalyst value rather than +/// through `ScalarValue`'s own equality. The two disagree on floats in opposite directions: +/// +/// * `-0.0` and `0.0` are one key for Spark (`-0.0 == 0.0` numerically, and `doubleHash` folds +/// both onto the hash of `0L`), while `ScalarValue` keeps them apart. +/// * `NaN` matches nothing for Spark, not even another `NaN`, because Scala's `==` on `Double` +/// is IEEE. `ScalarValue` treats `NaN` as equal to itself. +/// +/// Nulls are left alone: a null pivot column value does match a null entry in the pivot list, +/// which is what Spark's `pivotIndex.getOrElse(null, -1)` does. +fn pivot_key(v: ScalarValue) -> Option<ScalarValue> { + match v { + ScalarValue::Float32(Some(f)) => { + if f.is_nan() { + None + } else if f == 0.0 { + Some(ScalarValue::Float32(Some(0.0))) + } else { + Some(ScalarValue::Float32(Some(f))) + } + } + ScalarValue::Float64(Some(f)) => { + if f.is_nan() { + None + } else if f == 0.0 { + Some(ScalarValue::Float64(Some(0.0))) + } else { + Some(ScalarValue::Float64(Some(f))) + } + } + other => Some(other), + } +} + +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()) Review Comment: Fixed in 1c6d38f31c. The serde compares Spark's own `indexSize` with the list length and falls back when they differ, so anything Spark treats as a duplicate, including `0.0` next to `-0.0`, never reaches native code. As a second line of defense, the native index rejects a repeated value with an internal error rather than laying out a different number of slots. `PivotFirst.sql` covers `IN (2012, 2014, 2014)` and `IN (1.5D, 0.0D, -0.0D)` with ANSI off. Neither has a row matching a repeated value, because such a row makes Spark write past the end of its own buffer. -- 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]
