viirya commented on code in PR #4798:
URL: https://github.com/apache/datafusion-comet/pull/4798#discussion_r4252564318
##########
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:
This comment sits above the reason helper rather than the
`PivotFirst.supportsDataType` call it describes. Could you move it to
`getSupportLevel`? Also, `getUnsupportedReasons` below hard-codes the
value-type list, and so does `groups_accumulator_supported` on the native side.
The claim that "the two lists cannot drift" only holds for the check itself.
Could you reword it so it says which places would need updating if Spark adds a
type?
##########
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:
Spark has no `PivotTransformer`. The two-phase plan comes from
`Analyzer.ResolvePivot`, which the audit notes already name correctly. The same
name appears in the `PivotFirst.sql` header and the `CometAggregateSuite`
comment, so could you fix all three?
##########
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:
This gate has no test. `Config` and `MinSparkVersion` are file-level
directives, so a small separate SQL file would do. Give it `-- MinSparkVersion:
4.0` and a `STRING COLLATE UTF8_LCASE` pivot column queried with
`expect_fallback`. That would cover it on the versions where it matters.
Spark's `DataFramePivotSuite` "pivot with explicit values under UTF8_LCASE
collation" from SPARK-56306 is a good shape to borrow.
##########
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:
There is no threshold here. `ResolvePivot` picks the `PivotFirst` plan
whenever every aggregate type is in `supportsDataType`, and the query below has
no 1..10 values. Could you drop this sentence?
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -3861,4 +3861,87 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ // Optimized-pivot fast path: Spark's PivotTransformer emits a two-phase
aggregate whose
+ // second phase uses PivotFirst whenever every aggregate output type is in
+ // PivotFirst.supportsDataType. We assert:
+ // 1. The physical plan actually contains a PivotFirst (Spark still picked
the fast path).
+ // 2. Comet converted the second-phase aggregate to CometHashAggregateExec
(i.e. our
+ // CometPivotFirst serde accepted the aggregate).
+ // 3. Results match Spark exactly.
+ // If any of the three fails, the test surfaces the specific breakage.
+ test("PivotFirst runs natively on the optimized pivot fast path") {
Review Comment:
These two tests look like duplicates of `PivotFirst.sql`. The default
`query` mode there already asserts that every operator runs natively and
compares with Spark. The first test is the same query as the first SQL case,
and the double-pivot section covers the same `-0.0`/NaN data in both
directions. Could you drop them, or say what they catch that the SQL file
doesn't?
##########
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 path isn't exercised natively by any Comet test. Everything in
`PivotFirst.sql` has a grouping column, so only `grouped_and_ungrouped_agree`
covers this accumulator. Could you add a `PivotFirst.sql` case with no grouping
column, e.g. `SELECT * FROM (SELECT course, earnings FROM pf_sales) PIVOT
(sum(earnings) FOR course IN ('dotNET', 'Java'))`? Please also add one over
empty input (`... WHERE false`). Spark returns a single row of nulls there, and
it would be good to pin down that the native final aggregate does too.
##########
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.
Review Comment:
The query below doesn't cast. The table column is already `float`, and `min`
keeps the type. Could you reword this to say that?
--
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]
