rich7420 commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r4255138412
########## native/spark-expr/src/agg_funcs/pivot_first.rs: ########## @@ -0,0 +1,1179 @@ +// 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 one slot per pivot +//! value; on each input row it looks the pivot column up in the pivot list, 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. Spark sizes that buffer from its index of the pivot +//! values, so the serde falls back when the list holds duplicates, and the native slot count is +//! always the list length. The result reassembles the slots into a `ListArray` matching +//! `PivotFirst.dataType = ArrayType(value_type)`. +//! +//! The optimized pivot plan groups by every column of the pivoted relation other than the pivot +//! and aggregated columns, so the aggregate runs grouped, through [`PivotFirstGroupsAccumulator`], +//! unless the relation has no other column. That case uses [`PivotFirstAccumulator`]. + +use crate::float_semantics::ieee_equality_key; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, BooleanBufferBuilder, ListArray, + PrimitiveArray, +}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::cast; +use arrow::datatypes::{ + DataType, Date32Type, Decimal128Type, Field, FieldRef, Float32Type, Float64Type, Int16Type, + Int32Type, Int64Type, Int8Type, TimeUnit, TimestampMicrosecondType, +}; +use datafusion::common::utils::SingleRowListArrayBuilder; +use datafusion::common::{internal_err, not_impl_err, HashMap, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::Volatility::Immutable; +use datafusion::logical_expr::{ + Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature, +}; +use datafusion::physical_expr::expressions::format_state_name; +use std::mem::size_of_val; +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, and `pivot_index` maps an input row's pivot column value to its slot. The index +/// is wrapped in `Arc` so that every accumulator shares it rather than rebuilding it. +#[derive(Debug)] +pub struct SparkPivotFirst { + signature: Signature, + value_type: DataType, + // Kept for `PartialEq`/`Hash` (identity of the aggregate for plan comparison) and for the + // slot count `state_fields` needs. + pivot_values: Vec<ScalarValue>, + pivot_index: Arc<PivotIndex>, +} + +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 { + /// `pivot_values` are the pivot list in output order. The serde only sends a non-empty list + /// without duplicates, of a pivot column type that [`PivotIndex`] supports. + pub fn try_new(value_type: DataType, pivot_values: Vec<ScalarValue>) -> Result<Self> { + let pivot_index = Arc::new(PivotIndex::try_new(&pivot_values)?); + Ok(Self { + signature: Signature::user_defined(Immutable), + value_type, + pivot_values, + 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]) -> Result<DataType> { + Ok(DataType::List(Arc::new(Field::new_list_field( + self.value_type.clone(), + true, + )))) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<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) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(PivotFirstAccumulator::new( + self.value_type.clone(), + Arc::clone(&self.pivot_index), + self.pivot_values.len(), + ))) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + // Every value type in Spark's `PivotFirst.supportsDataType`. + matches!( + self.value_type, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Float32 + | DataType::Float64 + | DataType::Decimal128(_, _) + ) + } + + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + let index = Arc::clone(&self.pivot_index); + let num_slots = self.pivot_values.len(); + let value_type = self.value_type.clone(); + Ok(match &self.value_type { + // Booleans are stored as 0/1 bytes, so that one primitive implementation covers + // every value type. + DataType::Boolean | DataType::Int8 => Box::new( + PivotFirstGroupsAccumulator::<Int8Type>::new(value_type, index, num_slots), + ), + DataType::Int16 => Box::new(PivotFirstGroupsAccumulator::<Int16Type>::new( + value_type, index, num_slots, + )), + DataType::Int32 => Box::new(PivotFirstGroupsAccumulator::<Int32Type>::new( + value_type, index, num_slots, + )), + DataType::Int64 => Box::new(PivotFirstGroupsAccumulator::<Int64Type>::new( + value_type, index, num_slots, + )), + DataType::Float32 => Box::new(PivotFirstGroupsAccumulator::<Float32Type>::new( + value_type, index, num_slots, + )), + DataType::Float64 => Box::new(PivotFirstGroupsAccumulator::<Float64Type>::new( + value_type, index, num_slots, + )), + DataType::Decimal128(_, _) => Box::new( + PivotFirstGroupsAccumulator::<Decimal128Type>::new(value_type, index, num_slots), + ), + other => return internal_err!("PivotFirst does not support value type {other}"), + }) + } +} + +/// Marks a row whose pivot column value is in no slot. +const NO_SLOT: u32 = u32::MAX; + +/// Maps a pivot column value to its output slot the way Spark's `PivotFirst` matches it. +/// +/// For the pivot column types the serde admits, Spark looks the value up in a Scala +/// `HashMap[Any, Int]` built from the pivot list. Matching therefore goes through Scala's `==` +/// on the boxed Catalyst value, which is plain value equality except for floats, where it is +/// IEEE 754 equality: `-0.0` and `0.0` are one key, and NaN matches nothing, not even a NaN in +/// the pivot list. A null pivot column value matches a null in the pivot list, which is what +/// Spark's `pivotIndex.getOrElse(null, -1)` does. +#[derive(Debug)] +struct PivotIndex { + /// The Arrow type of the pivot values. A pivot column with another layout, such as a + /// dictionary-encoded string column, is cast to it before the lookup. + key_type: DataType, + keys: PivotKeys, + null_slot: Option<u32>, +} + +/// The slot of each non-null pivot value. +#[derive(Debug)] +enum PivotKeys { + /// Slots of `false` and `true`. + Boolean([Option<u32>; 2]), + /// Integers, dates, timestamps and decimals, keyed on their value as an `i128`. + Integer(HashMap<i128, u32>), + /// Floats, keyed on the bits of their [`ieee_equality_key`] as an `f64`. A NaN in the pivot + /// list has no key, so its slot is never written. + Float(HashMap<u64, u32>), + Utf8(HashMap<Box<str>, u32>), +} + +impl PivotIndex { + fn try_new(pivot_values: &[ScalarValue]) -> Result<Self> { + let Some(first) = pivot_values.first() else { + return internal_err!("PivotFirst requires at least one pivot value"); + }; + let key_type = first.data_type(); + let mut keys = match &key_type { + DataType::Boolean => PivotKeys::Boolean([None, None]), + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Date32 + | DataType::Timestamp(TimeUnit::Microsecond, _) + | DataType::Decimal128(_, _) => PivotKeys::Integer(HashMap::new()), + DataType::Float32 | DataType::Float64 => PivotKeys::Float(HashMap::new()), + DataType::Utf8 => PivotKeys::Utf8(HashMap::new()), + other => return internal_err!("PivotFirst does not support pivot column type {other}"), + }; + if pivot_values.len() >= NO_SLOT as usize { + return internal_err!("PivotFirst supports fewer than {NO_SLOT} pivot values"); + } + let mut null_slot = None; + for (slot, value) in pivot_values.iter().enumerate() { + let slot = slot as u32; + // `Some(true)` for a new key, `Some(false)` for a repeated one, `None` for a NaN. + let inserted = match (&mut keys, value) { + (_, value) if value.is_null() => Some(null_slot.replace(slot).is_none()), + (PivotKeys::Boolean(slots), ScalarValue::Boolean(Some(b))) => { + Some(slots[*b as usize].replace(slot).is_none()) + } + (PivotKeys::Integer(map), value) => match integer_key(value) { + Some(key) => Some(map.insert(key, slot).is_none()), + None => return mismatched_pivot_value(value, &key_type), + }, + (PivotKeys::Float(map), ScalarValue::Float32(Some(v))) => { + float_key(*v).map(|key| map.insert(key, slot).is_none()) + } + (PivotKeys::Float(map), ScalarValue::Float64(Some(v))) => { + float_key(*v).map(|key| map.insert(key, slot).is_none()) + } + (PivotKeys::Utf8(map), ScalarValue::Utf8(Some(s))) => { + Some(map.insert(s.as_str().into(), slot).is_none()) + } + (_, value) => return mismatched_pivot_value(value, &key_type), + }; + if inserted == Some(false) { + // Spark keeps a repeated value's last position and sizes its buffer from the + // distinct values, so the slot layout here would no longer match its schema. + return internal_err!( + "PivotFirst pivot value {value} is repeated; the serde should have fallen back" + ); + } + } + Ok(Self { + key_type, + keys, + null_slot, + }) + } + + /// The slot of every row of `pivots`, or [`NO_SLOT`] for a row whose value is not in the + /// pivot list. + fn slots(&self, pivots: &ArrayRef) -> Result<Vec<u32>> { + // A timestamp's value does not depend on its time zone label, so only the unit has to + // agree for the column to be read as it is. + let same_layout = match (pivots.data_type(), &self.key_type) { + (DataType::Timestamp(a, _), DataType::Timestamp(b, _)) => a == b, + (a, b) => a == b, + }; + let pivots = if same_layout { + Arc::clone(pivots) + } else { + cast(pivots, &self.key_type)? + }; + let null_slot = self.null_slot.unwrap_or(NO_SLOT); + Ok(match &self.keys { + PivotKeys::Boolean(slots) => { + lookup(pivots.as_boolean().iter(), null_slot, |b| slots[b as usize]) + } + PivotKeys::Integer(map) => match pivots.data_type() { + DataType::Int8 => integer_slots::<Int8Type>(&pivots, null_slot, map), + DataType::Int16 => integer_slots::<Int16Type>(&pivots, null_slot, map), + DataType::Int32 => integer_slots::<Int32Type>(&pivots, null_slot, map), + DataType::Int64 => integer_slots::<Int64Type>(&pivots, null_slot, map), + DataType::Date32 => integer_slots::<Date32Type>(&pivots, null_slot, map), + DataType::Timestamp(_, _) => { + integer_slots::<TimestampMicrosecondType>(&pivots, null_slot, map) + } + DataType::Decimal128(_, _) => { + integer_slots::<Decimal128Type>(&pivots, null_slot, map) + } + other => return internal_err!("PivotFirst cannot read pivot column type {other}"), + }, + PivotKeys::Float(map) => match pivots.data_type() { + DataType::Float32 => float_slots::<Float32Type>(&pivots, null_slot, map), + DataType::Float64 => float_slots::<Float64Type>(&pivots, null_slot, map), + other => return internal_err!("PivotFirst cannot read pivot column type {other}"), + }, + PivotKeys::Utf8(map) => lookup(pivots.as_string::<i32>().iter(), null_slot, |s| { + map.get(s).copied() + }), + }) + } +} + +fn mismatched_pivot_value<T>(value: &ScalarValue, key_type: &DataType) -> Result<T> { + internal_err!("PivotFirst pivot value {value:?} does not have the pivot column type {key_type}") +} + +fn integer_key(value: &ScalarValue) -> Option<i128> { + Some(match value { + ScalarValue::Int8(Some(v)) => i128::from(*v), + ScalarValue::Int16(Some(v)) => i128::from(*v), + ScalarValue::Int32(Some(v)) | ScalarValue::Date32(Some(v)) => i128::from(*v), + ScalarValue::Int64(Some(v)) | ScalarValue::TimestampMicrosecond(Some(v), _) => { + i128::from(*v) + } + ScalarValue::Decimal128(Some(v), _, _) => *v, + _ => return None, + }) +} + +fn float_key(v: impl Into<f64>) -> Option<u64> { + ieee_equality_key(v.into()).map(f64::to_bits) +} + +fn lookup<T>( + values: impl Iterator<Item = Option<T>>, + null_slot: u32, + slot_of: impl Fn(T) -> Option<u32>, +) -> Vec<u32> { + values + .map(|v| match v { + Some(v) => slot_of(v).unwrap_or(NO_SLOT), + None => null_slot, + }) + .collect() +} + +fn integer_slots<T>(pivots: &ArrayRef, null_slot: u32, map: &HashMap<i128, u32>) -> Vec<u32> +where + T: ArrowPrimitiveType, + T::Native: Into<i128>, +{ + lookup(pivots.as_primitive::<T>().iter(), null_slot, |v| { + map.get(&v.into()).copied() + }) +} + +fn float_slots<T>(pivots: &ArrayRef, null_slot: u32, map: &HashMap<u64, u32>) -> Vec<u32> +where + T: ArrowPrimitiveType, + T::Native: Into<f64>, +{ + lookup(pivots.as_primitive::<T>().iter(), null_slot, |v| { + float_key(v).and_then(|key| map.get(&key).copied()) + }) +} + +/// Grouped `PivotFirst` state. Each group owns `num_slots` cells, stored group-major in one flat +/// buffer, so that a group's cells are contiguous and `evaluate` hands the buffer to the output +/// list as it is. +#[derive(Debug)] +struct PivotFirstGroupsAccumulator<T: ArrowPrimitiveType> { + /// The value type Spark declared. `T` stores it, as `Int8` for `Boolean`. + value_type: DataType, + pivot_index: Arc<PivotIndex>, + num_slots: usize, + /// `values[group * num_slots + slot]`. Only the cells set in `is_set` hold a value. + values: Vec<T::Native>, + /// Which cells have been written. A cell nothing has written to is null, as in Spark. + is_set: BooleanBufferBuilder, +} + +impl<T: ArrowPrimitiveType> PivotFirstGroupsAccumulator<T> { + fn new(value_type: DataType, pivot_index: Arc<PivotIndex>, num_slots: usize) -> Self { + Self { + value_type, + pivot_index, + num_slots, + values: Vec::new(), + is_set: BooleanBufferBuilder::new(0), + } + } + + fn resize(&mut self, total_num_groups: usize) { + let num_cells = total_num_groups * self.num_slots; + if num_cells > self.values.len() { + self.values.resize(num_cells, T::Native::default()); + self.is_set.append_n(num_cells - self.is_set.len(), false); + } + } + + /// Writes `values[row]` into `cell` for each `(row, cell)` pair, skipping a null value, which + /// Spark never writes. A later row overwrites an earlier one, as Spark's `update` and `merge` + /// do. + fn write( + &mut self, + values: &ArrayRef, + writes: impl Iterator<Item = (usize, usize)>, + ) -> Result<()> { + let values = if self.value_type == DataType::Boolean { + cast(values, &DataType::Int8)? + } else { + Arc::clone(values) + }; + let Some(values) = values.as_primitive_opt::<T>() else { + return internal_err!( + "PivotFirst expected {} values, got {}", + self.value_type, + values.data_type() + ); + }; + for (row, cell) in writes { + if values.is_valid(row) { + self.values[cell] = values.value(row); + self.is_set.set_bit(cell, true); + } + } + Ok(()) + } + + /// Removes the cells of the groups that `emit_to` covers, returning their values and which of + /// them are set. The remaining groups move to the front. + fn take_cells(&mut self, emit_to: EmitTo) -> (Vec<T::Native>, BooleanBuffer) { + let num_cells = match emit_to { + EmitTo::All => self.values.len(), + EmitTo::First(n) => n * self.num_slots, + }; + let rest = self.values.split_off(num_cells); Review Comment: Could we use `EmitTo::take_needed` here? With 100K groups and 128 Int32 slots, `First(8192)` retains the full buffer and copies the remainder, nearly doubling live memory. I reproduced this through DataFusion’s partially sorted final aggregate. -- 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]
