andygrove commented on code in PR #4798: URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r4261918095
########## 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: Added both in 24ae599fc6. The empty one filters with `WHERE 1 = 0` rather than `WHERE false`. A literal `false` lets the optimizer replace the scan with an empty `LocalTableScan`, which Comet doesn't convert by default, so that query ran entirely in Spark. The fixtures run without constant folding, so `1 = 0` keeps the scan, the filter and both `PivotFirst` aggregates native, and the row of nulls comes from the native final aggregate. Since cafc8d1f65 the ungrouped accumulator is the grouped one holding a single group, so these run the same code as the grouped cases. ########## spark/src/main/scala/org/apache/comet/serde/aggregates.scala: ########## @@ -1434,6 +1434,96 @@ object CometMode extends CometAggregateExpressionSerde[Mode] with CometTypeShim } } +object CometPivotFirst extends CometAggregateExpressionSerde[PivotFirst] with CometTypeShim { + + // 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 def unsupportedPivotTypeReason(dt: DataType): String = + s"Unsupported pivot column data type: $dt" + + private val emptyPivotValuesReason = "Pivot values list is empty" + + private val duplicatePivotValuesReason = "Pivot values list contains duplicates" + + override def getUnsupportedReasons(): Seq[String] = Seq( + "Value data type outside PivotFirst.supportsDataType " + + "(Boolean, Byte, Short, Int, Long, Float, Double, Decimal)", + "Pivot column of a type other than Boolean, Byte, Short, Int, Long, Float, Double, " + + "Decimal, String (UTF8_BINARY collation), Date, Timestamp or TimestampNTZ, which " + + "includes the struct that a pivot on more than one column builds", + emptyPivotValuesReason, + s"$duplicatePivotValuesReason, including `0.0` next to `-0.0`") + + // The pivot column types for which Spark finds a value's slot in a `HashMap` keyed by plain + // value equality, which the native lookup reproduces. Spark keys arrays and structs in a + // `TreeMap` by their interpreted ordering. It keys binary values by reference identity before + // 4.2 and by ordering from 4.2, which orders strings with a non-UTF8_BINARY collation too. + private def isSupportedPivotType(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType => true + case FloatType | DoubleType => true + case _: DecimalType => true + case DateType | TimestampType | TimestampNTZType => true + case st: StringType => !isStringCollationType(st) Review Comment: Added `PivotFirst_collation.sql` in 24ae599fc6, with `MinSparkVersion: 4.0` and the UTF8_LCASE shape from SPARK-56306. On 4.1 the query falls back before it reaches `PivotFirst` anyway, since the native scan declines the collated schema and Comet shuffle declines a collated hash key. What the fixture pins is the gate's own reason, which nothing else records. With the collation check removed, `expect_fallback` fails on the missing reason. It passes locally on 4.0.4, 4.1.3 and 4.2.0. ########## spark/src/main/scala/org/apache/comet/serde/aggregates.scala: ########## @@ -1434,6 +1434,96 @@ object CometMode extends CometAggregateExpressionSerde[Mode] with CometTypeShim } } +object CometPivotFirst extends CometAggregateExpressionSerde[PivotFirst] with CometTypeShim { + + // Delegate to Spark's own PivotFirst.supportsDataType so the two lists cannot drift if a Review Comment: Reworded in cafc8d1f65, and the check changed with it. Now that the ungrouped accumulator is built on the grouped one, native code has nothing that could store a value type it doesn't know. With `PivotFirst.supportsDataType` as the check, a type that a later Spark version adds would get through to an internal error instead of a fallback. So `getSupportLevel` calls `isSupportedValueType`, which lists the same types Spark accepts on 3.4.3 through 4.2.0. The comment now sits on that list and names the three places a new type needs: the list itself, `getUnsupportedReasons`, and `SparkPivotFirst::groups_accumulator`. `groups_accumulator_supported` no longer lists types. ########## 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 Review Comment: Fixed in all three places. The module doc (cafc8d1f65) and the `PivotFirst.sql` header (24ae599fc6) now name `ResolvePivot`, and the `CometAggregateSuite` comment went away with the tests. ########## 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 Review Comment: Dropped in 24ae599fc6. -- 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]
