comphead commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r4260679224
########## 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); + let values = std::mem::replace(&mut self.values, rest); + let is_set = self.is_set.finish(); + self.is_set + .append_buffer(&is_set.slice(num_cells, is_set.len() - num_cells)); + (values, is_set.slice(0, num_cells)) + } + + /// Gives stored cells the declared value type: the precision and scale of a decimal, or + /// booleans in place of their `Int8` storage. + fn to_value_type(&self, cells: PrimitiveArray<T>) -> Result<ArrayRef> { + if self.value_type == DataType::Boolean { + Ok(cast(&cells, &DataType::Boolean)?) + } else { + Ok(Arc::new(cells.with_data_type(self.value_type.clone()))) + } + } +} + +impl<T: ArrowPrimitiveType> GroupsAccumulator for PivotFirstGroupsAccumulator<T> { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let [pivots, values] = values else { + return internal_err!( + "PivotFirst expects 2 inputs (pivot, value); got {}", + values.len() + ); + }; + self.resize(total_num_groups); + let slots = self.pivot_index.slots(pivots)?; + let num_slots = self.num_slots; + let writes = group_indices + .iter() + .zip(slots) + .enumerate() + .filter(|&(row, (_, slot))| { + slot != NO_SLOT && opt_filter.is_none_or(|f| f.is_valid(row) && f.value(row)) + }) + .map(|(row, (&group, slot))| (row, group * num_slots + slot as usize)); + self.write(values, writes) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + if values.len() != self.num_slots { + return internal_err!( + "PivotFirst merge expects {} state columns; got {}", + self.num_slots, + values.len() + ); + } + self.resize(total_num_groups); + // Each column is one slot, and any non-null cell overwrites the group's slot (matches + // Spark's `PivotFirst.merge`). A group's state holds only the pivot values its rows had, + // so the columns are mostly null, and only their valid rows are visited. + let num_slots = self.num_slots; + for (slot, column) in values.iter().enumerate() { + let cell = |row: usize| (row, group_indices[row] * num_slots + slot); + match column.logical_nulls() { + Some(nulls) => self.write(column, nulls.valid_indices().map(cell))?, + None => self.write(column, (0..column.len()).map(cell))?, + } + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + let (values, is_set) = self.take_cells(emit_to); + let num_groups = values.len() / self.num_slots; + let cells = PrimitiveArray::new(values.into(), Some(NullBuffer::new(is_set))); + let offsets = + OffsetBuffer::<i32>::from_lengths(std::iter::repeat_n(self.num_slots, num_groups)); + let field = Arc::new(Field::new_list_field(self.value_type.clone(), true)); + Ok(Arc::new(ListArray::try_new( + field, + offsets, + self.to_value_type(cells)?, + None, + )?)) + } + + fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let (values, is_set) = self.take_cells(emit_to); + let num_slots = self.num_slots; + let num_groups = values.len() / num_slots; + (0..num_slots) + .map(|slot| { + let column: ScalarBuffer<T::Native> = (0..num_groups) + .map(|group| values[group * num_slots + slot]) + .collect(); + let nulls = BooleanBuffer::collect_bool(num_groups, |group| { + is_set.value(group * num_slots + slot) + }); + self.to_value_type(PrimitiveArray::new(column, Some(NullBuffer::new(nulls)))) + }) + .collect() + } + + fn convert_to_state( + &self, + _values: &[ArrayRef], + _opt_filter: Option<&BooleanArray>, + ) -> Result<Vec<ArrayRef>> { + not_impl_err!("Input batch conversion to state not implemented") + } + + fn size(&self) -> usize { + size_of_val(self) + + self.values.capacity() * size_of::<T::Native>() + + self.is_set.capacity() / 8 + } +} + +/// Ungrouped 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 { Review Comment: The ungrouped `PivotFirstAccumulator` repeats the update, merge and state logic that `PivotFirstGroupsAccumulator` already has. It could be a thin `Accumulator` over a one-group `PivotFirstGroupsAccumulator`. `update_batch` and `merge_batch` would pass an all-zero group index vector with one group, and `state` and `evaluate` would call the groups accumulator with `EmitTo::All` and read row 0. That would remove the `ScalarValue::try_from_array` in `update_batch`, which runs for every matching input row, and an ungrouped pivot has no group column to bound the number of rows. It would also remove the second `slot_or_null` and merge implementation, and most of the ungrouped Rust tests (`update_writes_by_pivot_index`, `update_ignores_unmatched_pivot`, `update_ignores_null_value`, `evaluate_returns_list_with_nulls_for_empty_slots`, `state_returns_one_scalar_per_slot`, `merge_pastes_non_null_slots`, `grouped_and_ungrouped_agree`). Those repeat what the grouped tests assert. The ungrouped `PivotFirst.sql` case viirya asked for would then run the same code as the grouped path for decimal and boolean values, which only the grouped Rust tests cover today. I have not tried the wrapper. `take_cells` consumes the state, so it relies on `state` and `evaluate` being called once each on the ungrouped accumulator. ########## spark/src/test/resources/sql-tests/expressions/aggregate/PivotFirst.sql: ########## @@ -0,0 +1,525 @@ +-- 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. + +-- PivotFirst is Spark's internal aggregate emitted by the two-phase optimized pivot plan +-- (PivotTransformer). It runs whenever every aggregate output type is in +-- PivotFirst.supportsDataType (Boolean, Byte, Short, Int, Long, Float, Double, Decimal). +-- These SQL tests exercise that path via SQL PIVOT clauses, one per supported value type and +-- per supported pivot column type, plus a handful of edge cases (unmatched pivot values, +-- all-null groups, null pivot column). The pivot column types and pivot lists whose matching +-- Comet does not reproduce come last and must fall back. + +-- ============================================================ +-- Setup: sales-like table used for most cases +-- ============================================================ + +statement +CREATE TABLE pf_sales (year int, course string, earnings int) USING parquet + +statement +INSERT INTO pf_sales VALUES + (2012, 'dotNET', 15000), + (2012, 'Java', 20000), + (2013, 'dotNET', 48000), + (2013, 'Java', 30000) + +-- ============================================================ +-- Int values: canonical optimized-pivot query. Extra pivot values +-- (1..10) push the aggregate over the threshold that turns on the +-- PivotFirst plan. +-- ============================================================ + +query +SELECT * FROM pf_sales + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Long values: sum of a bigint aggregates to bigint. +-- ============================================================ + +statement +CREATE TABLE pf_sales_bi (year int, course string, earnings bigint) USING parquet + +statement +INSERT INTO pf_sales_bi VALUES + (2012, 'dotNET', 15000000000), + (2012, 'Java', 20000000000), + (2013, 'dotNET', 48000000000), + (2013, 'Java', 30000000000) + +query +SELECT * FROM pf_sales_bi + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Double values: avg(earnings) is double. +-- ============================================================ + +query +SELECT * FROM pf_sales + PIVOT (avg(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Float values: cast the sum result to float via a wrapping min. +-- ============================================================ + +statement +CREATE TABLE pf_sales_f (year int, course string, earnings float) USING parquet + +statement +INSERT INTO pf_sales_f VALUES + (2012, 'dotNET', CAST(15000 AS FLOAT)), + (2012, 'Java', CAST(20000 AS FLOAT)), + (2013, 'dotNET', CAST(48000 AS FLOAT)), + (2013, 'Java', CAST(30000 AS FLOAT)) + +query +SELECT * FROM pf_sales_f + PIVOT (min(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Short and Byte value types (sum promotes them to bigint, so use min +-- to keep the aggregate output type equal to the input type). +-- ============================================================ + +statement +CREATE TABLE pf_sales_small ( + year int, course string, s smallint, b tinyint) USING parquet + +statement +INSERT INTO pf_sales_small VALUES + (2012, 'dotNET', 150, 10), + (2012, 'Java', 200, 20), + (2013, 'dotNET', 480, 30), + (2013, 'Java', 300, 40) + +query +SELECT * FROM pf_sales_small + PIVOT (min(s) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +query +SELECT * FROM pf_sales_small + PIVOT (min(b) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Boolean values: bool_or's output type is boolean. +-- ============================================================ + +statement +CREATE TABLE pf_flags (year int, course string, flag boolean) USING parquet + +statement +INSERT INTO pf_flags VALUES + (2012, 'dotNET', true), + (2012, 'Java', false), + (2013, 'dotNET', false), + (2013, 'Java', true) + +query +SELECT * FROM pf_flags + PIVOT (bool_or(flag) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Decimal values. +-- ============================================================ + +statement +CREATE TABLE pf_sales_dec (year int, course string, earnings decimal(10,2)) USING parquet + +statement +INSERT INTO pf_sales_dec VALUES + (2012, 'dotNET', 15000.00), + (2012, 'Java', 20000.00), + (2013, 'dotNET', 48000.00), + (2013, 'Java', 30000.00) + +query +SELECT * FROM pf_sales_dec + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Multiple aggregates in a single pivot. +-- ============================================================ + +query +SELECT * FROM pf_sales + PIVOT (sum(earnings) AS s, avg(earnings) AS a + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Pivot values that do not appear in the data: the corresponding +-- output slots must be NULL. +-- ============================================================ + +query +SELECT * FROM pf_sales + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java', 'MissingA', 'MissingB')) + ORDER BY year + +-- ============================================================ +-- Rows whose pivot column is not in the pivot list must be ignored: +-- 'Python' below should not contribute to any output slot. +-- ============================================================ + +statement +CREATE TABLE pf_sales_extra (year int, course string, earnings int) USING parquet + +statement +INSERT INTO pf_sales_extra VALUES + (2012, 'dotNET', 15000), + (2012, 'Java', 20000), + (2012, 'Python', 12345), + (2013, 'dotNET', 48000), + (2013, 'Java', 30000), + (2013, 'Python', 67890) + +query +SELECT * FROM pf_sales_extra + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Group with no matching rows: every output slot is NULL for that +-- pivot combination. Ensures NULL initialization matches Spark. +-- ============================================================ + +statement +CREATE TABLE pf_sales_partial (year int, course string, earnings int) USING parquet + +statement +INSERT INTO pf_sales_partial VALUES + (2012, 'dotNET', 15000), + (2012, 'Java', 20000), + (2013, 'Python', 42000) + +query +SELECT * FROM pf_sales_partial + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- NULL value in an aggregated column: the null must not clobber a +-- previously written slot (Spark PivotFirst.update ignores nulls). +-- ============================================================ + +statement +CREATE TABLE pf_sales_nullable (year int, course string, earnings int) USING parquet + +statement +INSERT INTO pf_sales_nullable VALUES + (2012, 'dotNET', 15000), + (2012, 'dotNET', NULL), + (2012, 'Java', NULL), + (2013, 'dotNET', 48000), + (2013, 'Java', 30000) + +query +SELECT * FROM pf_sales_nullable + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- NULL in the pivot COLUMN. Spark tracks this with a dedicated +-- test ("pivot with null should not throw NPE") because early +-- versions of PivotFirst NPE'd on null keys with a TreeMap. A null +-- row matches no slot unless the pivot list holds NULL too. +-- ============================================================ + +statement +CREATE TABLE pf_sales_nullpivot (year int, course string, earnings int) USING parquet + +statement +INSERT INTO pf_sales_nullpivot VALUES + (2012, 'dotNET', 15000), + (2012, NULL, 9999), + (2013, 'Java', 30000), + (2013, NULL, 42) + +query +SELECT * FROM pf_sales_nullpivot + PIVOT (sum(earnings) + FOR course IN ('dotNET', 'Java')) + ORDER BY year + +-- ============================================================ +-- Non-string pivot column: integer. Exercises the pivot value +-- serialization path with a numeric literal. +-- ============================================================ + +statement +CREATE TABLE pf_sales_intpivot (region string, year int, earnings int) USING parquet + +statement +INSERT INTO pf_sales_intpivot VALUES + ('NA', 2012, 15000), + ('NA', 2013, 48000), + ('EU', 2012, 20000), + ('EU', 2013, 30000) + +query +SELECT * FROM pf_sales_intpivot + PIVOT (sum(earnings) + FOR year IN (2012, 2013)) + ORDER BY region + +-- ============================================================ +-- Non-string pivot column: date. Exercises the pivot value +-- serialization path with a DateType literal. +-- ============================================================ + +statement +CREATE TABLE pf_sales_datepivot (region string, d date, earnings int) USING parquet + +statement +INSERT INTO pf_sales_datepivot VALUES + ('NA', DATE '2024-01-01', 15000), + ('NA', DATE '2024-06-15', 48000), + ('EU', DATE '2024-01-01', 20000), + ('EU', DATE '2024-06-15', 30000) + +query +SELECT * FROM pf_sales_datepivot + PIVOT (sum(earnings) + FOR d IN (DATE '2024-01-01', DATE '2024-06-15')) + ORDER BY region + +-- ============================================================ +-- Non-string pivot column: double, including the two values whose +-- key matching differs between Spark and a plain hash of the native +-- scalar. Spark keys the pivot list on the boxed Catalyst value, so +-- -0.0 and 0.0 are the same key while NaN matches nothing at all. +-- +-- The zero rows are written as -0.0D / 0.0D rather than +-- CAST(-0.0 AS DOUBLE), because the cast constant-folds to +0.0 and +-- would make the case vacuous. +-- ============================================================ + +statement +CREATE TABLE pf_sales_dblpivot (region string, k double, earnings int) USING parquet + +statement +INSERT INTO pf_sales_dblpivot VALUES + ('NA', 0.0D, 15000), + ('NA', -0.0D, 48000), + ('NA', double('NaN'), 11000), + ('NA', 1.5D, 20000), + ('EU', 0.0D, 30000), + ('EU', double('NaN'), 12000), + ('EU', 1.5D, 40000) + +-- -0.0 in the pivot list has to collect both signs of zero (63000 / 30000), +-- and the NaN column has to stay NULL for every group. +query +SELECT * FROM pf_sales_dblpivot + PIVOT (sum(earnings) + FOR k IN (-0.0D, double('NaN'), 1.5D)) + ORDER BY region + +-- The mirror: a 0.0 entry in the pivot list must also catch the -0.0 rows. +query +SELECT * FROM pf_sales_dblpivot + PIVOT (sum(earnings) + FOR k IN (0.0D, 1.5D)) + ORDER BY region + +-- ============================================================ +-- Non-string pivot column: float. Same two rules. +-- ============================================================ + +statement +CREATE TABLE pf_sales_fltpivot (region string, k float, earnings int) USING parquet + +statement +INSERT INTO pf_sales_fltpivot VALUES + ('NA', 0.0F, 15000), + ('NA', -0.0F, 48000), + ('NA', float('NaN'), 11000), + ('EU', -0.0F, 30000) + +query +SELECT * FROM pf_sales_fltpivot + PIVOT (sum(earnings) + FOR k IN (-0.0F, float('NaN'))) + ORDER BY region + +-- ============================================================ +-- NULL in the pivot LIST collects the rows whose pivot column is +-- NULL, as Spark's HashMap lookup of a null key does. +-- ============================================================ + +query +SELECT * FROM pf_sales_nullpivot + PIVOT (sum(earnings) + FOR course IN ('dotNET', NULL)) + ORDER BY year + +-- ============================================================ +-- Two NaN entries are two distinct keys for Spark (NaN equals +-- nothing), so the list has no duplicates and stays native, and +-- both columns stay NULL. +-- ============================================================ + +query +SELECT * FROM pf_sales_dblpivot + PIVOT (sum(earnings) + FOR k IN (double('NaN'), 1.5D, double('NaN'))) + ORDER BY region + +-- ============================================================ +-- The other pivot column types that run natively: boolean, +-- tinyint, smallint, bigint, decimal, timestamp and timestamp_ntz. +-- ============================================================ + +statement +CREATE TABLE pf_sales_types ( + region string, flag boolean, b tinyint, s smallint, l bigint, d decimal(10,2), + ts timestamp, ntz timestamp_ntz, earnings int) USING parquet + +statement +INSERT INTO pf_sales_types VALUES + ('NA', true, 1, 100, 10000000000, 1.50, + TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', 15000), + ('NA', false, 2, 200, 20000000000, 2.25, + TIMESTAMP '2024-06-15 12:30:00', TIMESTAMP_NTZ '2024-06-15 12:30:00', 48000), + ('EU', true, 2, 200, 20000000000, 2.25, + TIMESTAMP '2024-06-15 12:30:00', TIMESTAMP_NTZ '2024-06-15 12:30:00', 20000), + ('EU', NULL, 3, 300, 30000000000, 3.00, + TIMESTAMP '2025-01-01 00:00:00', TIMESTAMP_NTZ '2025-01-01 00:00:00', 30000) + +query +SELECT region, `true`, `false` FROM ( + SELECT region, flag, earnings FROM pf_sales_types) + PIVOT (sum(earnings) FOR flag IN (true, false)) + ORDER BY region + +query +SELECT * FROM (SELECT region, b, earnings FROM pf_sales_types) + PIVOT (sum(earnings) FOR b IN (CAST(1 AS TINYINT), CAST(2 AS TINYINT))) + ORDER BY region + +query +SELECT * FROM (SELECT region, s, earnings FROM pf_sales_types) + PIVOT (sum(earnings) FOR s IN (CAST(100 AS SMALLINT), CAST(300 AS SMALLINT))) + ORDER BY region + +query +SELECT * FROM (SELECT region, l, earnings FROM pf_sales_types) + PIVOT (sum(earnings) FOR l IN (10000000000, 20000000000)) + ORDER BY region + +query +SELECT * FROM (SELECT region, d, earnings FROM pf_sales_types) + PIVOT (sum(earnings) FOR d IN (2.25, 3.00)) + ORDER BY region + +query +SELECT * FROM (SELECT region, ts, earnings FROM pf_sales_types) + PIVOT (sum(earnings) + FOR ts IN (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP '2024-06-15 12:30:00')) + ORDER BY region + +query +SELECT * FROM (SELECT region, ntz, earnings FROM pf_sales_types) + PIVOT (sum(earnings) + FOR ntz IN (TIMESTAMP_NTZ '2024-06-15 12:30:00', TIMESTAMP_NTZ '2025-01-01 00:00:00')) + ORDER BY region + +-- ============================================================ +-- Pivot column types that fall back. Spark compares binary pivot +-- values by reference before 4.2, so no row matches there, and by +-- content from 4.2. It matches arrays and structs through their +-- interpreted ordering. A pivot on two columns pivots on a struct. +-- ============================================================ + +statement +CREATE TABLE pf_sales_complex ( + region string, k binary, a array<int>, course string, year int, earnings int) + USING parquet + +statement +INSERT INTO pf_sales_complex VALUES + ('NA', X'61', array(1, 1), 'dotNET', 2012, 15000), + ('NA', X'62', array(2, 2), 'Java', 2013, 48000), + ('EU', X'61', array(1, 1), 'dotNET', 2013, 20000) + +query expect_fallback(Unsupported pivot column data type: BinaryType) +SELECT * FROM (SELECT region, k, earnings FROM pf_sales_complex) + PIVOT (sum(earnings) FOR k IN (X'61', X'62')) + ORDER BY region + +query expect_fallback(Unsupported pivot column data type: ArrayType) +SELECT * FROM (SELECT region, a, earnings FROM pf_sales_complex) + PIVOT (sum(earnings) FOR a IN (array(1, 1), array(2, 2))) + ORDER BY region + +query expect_fallback(Unsupported pivot column data type: StructType) +SELECT * FROM (SELECT region, course, year, earnings FROM pf_sales_complex) + PIVOT (sum(earnings) FOR (course, year) IN (('dotNET', 2012), ('Java', 2013))) + ORDER BY region + +-- ============================================================ +-- Duplicate pivot values fall back. Spark keeps one buffer slot per +-- distinct value, so its partial aggregate emits fewer columns than +-- the list has entries, and the projection reads the missing ones +-- past the end of the array: NULL with ANSI off, an error with it +-- on. 0.0 and -0.0 are duplicates under Spark's key equality. Only +-- rows matching the first entry are used, because Spark writes a +-- repeated value's row past the end of its buffer. +-- ============================================================ + +statement +SET spark.sql.ansi.enabled=false + +query expect_fallback(Pivot values list contains duplicates) +SELECT * FROM pf_sales_intpivot + PIVOT (sum(earnings) FOR year IN (2012, 2014, 2014)) + ORDER BY region + +statement +CREATE TABLE pf_sales_dblnozero (region string, k double, earnings int) USING parquet + +statement +INSERT INTO pf_sales_dblnozero VALUES + ('NA', 1.5D, 15000), + ('NA', 2.5D, 48000), + ('EU', 1.5D, 20000) + +query expect_fallback(Pivot values list contains duplicates) +SELECT * FROM pf_sales_dblnozero + PIVOT (sum(earnings) FOR k IN (1.5D, 0.0D, -0.0D)) + ORDER BY region Review Comment: Every case in this file has a handful of groups, one partition and one output batch, and I found no `range(` or `REPARTITION` in it. So the grouped accumulator's paths for emitting more groups than one batch (`EmitTo::First`), merging one group's state from several partial aggregates, and growing the buffers across many `update_batch` calls are only reached by the Rust unit tests and by `CometPivotBenchmark`, which CI does not run. Those Rust tests call the accumulator directly, so they do not go through DataFusion's emission or the shuffle. Could you add one case over more groups than one batch (for example 20000 groups from `range(60000)`, which is above the 8192 in rich7420's comment) with a `/*+ REPARTITION(4) */` hint so each group has several partial states? Giving each group and pivot value pair a single non-null value keeps the result deterministic. Other fixtures in this directory already use `range(N)` for size, for example `approx_count_distinct.sql`. -- 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]
