From 0047452b375a8b0336d26427fdaf9bd5113979f3 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Mon, 26 Oct 2026 15:28:05 -0700 Subject: [PATCH 1/3] feat: run min and max over struct inputs natively Allow min/max over a struct whose fields, at any depth, are all min/max-supported types other than float and double. A struct with no nullable fields is Compatible. A struct with nullable fields is Incompatible (opt-in via allowIncompatible), because DataFusion skips a null field when comparing structs while Spark orders null first. The native planner checks that a struct input matches the declared type apart from field nullability, then relabels it with Comet's Cast, which now returns an error instead of panicking when a non-nullable field holds a null. Closes #7024 --- .../core/src/execution/planner/aggregate.rs | 107 +++++++++++++- .../spark-expr/src/conversion_funcs/cast.rs | 68 ++++++++- .../org/apache/comet/serde/aggregates.scala | 66 ++++++++- .../aggregate/min_max_collation.sql | 5 + .../expressions/aggregate/min_max_struct.sql | 133 ++++++++++++++++++ .../aggregate/min_max_struct_many_batches.sql | 34 +++++ .../aggregate/min_max_struct_mixed_stages.sql | 55 ++++++++ .../aggregate/min_max_struct_nullable.sql | 74 ++++++++++ .../aggregate/min_max_struct_timezone.sql | 56 ++++++++ 9 files changed, 587 insertions(+), 11 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql diff --git a/native/core/src/execution/planner/aggregate.rs b/native/core/src/execution/planner/aggregate.rs index 4437141b302..77afce648b5 100644 --- a/native/core/src/execution/planner/aggregate.rs +++ b/native/core/src/execution/planner/aggregate.rs @@ -51,10 +51,10 @@ use datafusion_comet_proto::{ spark_operator::{AggregateMode as ProtoAggregateMode, Operator}, }; use datafusion_comet_spark_expr::{ - ApproxPercentile, Avg, AvgDecimal, BloomFilterAgg, CometCollectList, CometCollectSet, - Correlation, Covariance, HllPlusPlus, HllSketchAgg, HllUnionAgg, MaxMinBy, Mode, Regr, - RegrType, SparkBloomFilterVersion, SparkListAgg, SparkMinMax, SparkPercentile, SparkPivotFirst, - Stddev, SumDecimal, SumInteger, Variance, ZeroOrder, + ApproxPercentile, Avg, AvgDecimal, BloomFilterAgg, Cast, CometCollectList, CometCollectSet, + Correlation, Covariance, EvalMode, HllPlusPlus, HllSketchAgg, HllUnionAgg, MaxMinBy, Mode, + Regr, RegrType, SparkBloomFilterVersion, SparkCastOptions, SparkListAgg, SparkMinMax, + SparkPercentile, SparkPivotFirst, Stddev, SumDecimal, SumInteger, Variance, ZeroOrder, }; use jni::objects::{Global, JObject}; @@ -339,7 +339,7 @@ impl PhysicalPlanner { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap()); let func = min_max_udaf(&datatype, false); - let child = Arc::new(CastExpr::new(child, datatype.clone(), None)); + let child = Self::cast_min_max_child(child, &datatype, &schema)?; AggregateExprBuilder::new(func, vec![child]) .schema(schema) @@ -353,7 +353,7 @@ impl PhysicalPlanner { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap()); let func = min_max_udaf(&datatype, true); - let child = Arc::new(CastExpr::new(child, datatype.clone(), None)); + let child = Self::cast_min_max_child(child, &datatype, &schema)?; AggregateExprBuilder::new(func, vec![child]) .schema(schema) @@ -807,6 +807,38 @@ impl PhysicalPlanner { Ok(Arc::new(CastExpr::new(child, nullable_type, None))) } + /// Cast a min/max input to the aggregate's declared type. Native expressions such as + /// `struct(...)` can report nullable fields where Spark declares them non-nullable, and + /// DataFusion's `CastExpr` rejects that narrowing. For a struct, only nullability may differ: + /// Comet's `Cast` then relabels the fields without touching the data (and errors if a field + /// declared non-nullable holds a null), so the input matches the state that shuffle and scans + /// stamp with Spark's type. Any other difference is a planning error. + fn cast_min_max_child( + child: Arc, + datatype: &DataType, + schema: &SchemaRef, + ) -> Result, ExecutionError> { + if !matches!(datatype, DataType::Struct(_)) { + return Ok(Arc::new(CastExpr::new(child, datatype.clone(), None))); + } + let child_type = child.data_type(schema.as_ref())?; + if &child_type == datatype { + return Ok(child); + } + if !equal_ignoring_nullability(&child_type, datatype) { + return Err(GeneralError(format!( + "min/max input type {child_type} does not match the declared type {datatype}" + ))); + } + Ok(Arc::new(Cast::new( + child, + datatype.clone(), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), + None, + None, + ))) + } + pub(super) fn create_aggr_func_expr( name: &str, schema: SchemaRef, @@ -823,6 +855,19 @@ impl PhysicalPlanner { } } +/// True if two types match apart from struct field nullability and metadata. +fn equal_ignoring_nullability(a: &DataType, b: &DataType) -> bool { + match (a, b) { + (DataType::Struct(fa), DataType::Struct(fb)) => { + fa.len() == fb.len() + && fa.iter().zip(fb.iter()).all(|(x, y)| { + x.name() == y.name() && equal_ignoring_nullability(x.data_type(), y.data_type()) + }) + } + _ => a == b, + } +} + /// Return a copy of `data_type` with every nested field marked nullable. Map key fields are left /// non-nullable to preserve Arrow's map invariant. Primitive types are returned unchanged. fn make_all_fields_nullable(data_type: &DataType) -> DataType { @@ -898,3 +943,53 @@ fn percentile_value(expr: &spark_expression::Expr) -> Result>>()?; - Ok(Arc::new(StructArray::new( + Ok(Arc::new(StructArray::try_new( to_fields.clone(), cast_fields, array.nulls().cloned(), - ))) + )?)) } _ => unreachable!(), } @@ -1248,6 +1248,70 @@ mod tests { } } + /// Struct min/max casts its input this way to relabel nullable fields as non-nullable, + /// so the values, including a UTC timestamp, must come through unchanged. + #[test] + fn test_cast_struct_relabels_nullable_fields_as_non_nullable() { + let ts_type = + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())); + let ts: ArrayRef = Arc::new( + PrimitiveArray::::from(vec![-1, 0, 1]).with_timezone("UTC"), + ); + let s: ArrayRef = Arc::new(StringArray::from(vec!["a", "中文", ""])); + let inner_from = Fields::from(vec![Field::new("s", DataType::Utf8, true)]); + let inner_to = Fields::from(vec![Field::new("s", DataType::Utf8, false)]); + let inner: ArrayRef = Arc::new(StructArray::new(inner_from.clone(), vec![s], None)); + let from_fields = Fields::from(vec![ + Field::new("ts", ts_type.clone(), true), + Field::new("inner", DataType::Struct(inner_from), true), + ]); + let to_fields = Fields::from(vec![ + Field::new("ts", ts_type, false), + Field::new("inner", DataType::Struct(inner_to), false), + ]); + // The outer struct itself may be null. + let outer: ArrayRef = Arc::new(StructArray::new( + from_fields, + vec![ts, inner], + Some(vec![true, false, true].into()), + )); + let to_type = DataType::Struct(to_fields); + + let result = cast_array( + Arc::clone(&outer), + &to_type, + &SparkCastOptions::new_without_timezone(EvalMode::Legacy), + ) + .unwrap(); + + assert_eq!(result.data_type(), &to_type); + assert_eq!(result.nulls(), outer.nulls()); + let result = result.as_struct(); + let original = outer.as_struct(); + assert_eq!(result.column(0).as_ref(), original.column(0).as_ref()); + assert_eq!( + result.column(1).as_struct().column(0).as_ref(), + original.column(1).as_struct().column(0).as_ref() + ); + } + + /// A null under a field relabeled as non-nullable must be an error, not a panic. + #[test] + fn test_cast_struct_to_non_nullable_field_with_null_returns_error() { + let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None])); + let from_fields = Fields::from(vec![Field::new("v", DataType::Int64, true)]); + let to_fields = Fields::from(vec![Field::new("v", DataType::Int64, false)]); + let array: ArrayRef = Arc::new(StructArray::new(from_fields, vec![values], None)); + + let result = cast_array( + array, + &DataType::Struct(to_fields), + &SparkCastOptions::new_without_timezone(EvalMode::Legacy), + ); + + assert!(result.is_err()); + } + #[test] fn test_cast_nested_struct_to_struct_ansi_overflow_returns_error() { let inner_values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), Some(128), None])); diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 17ab15d1810..c72b0c1a0a0 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, isSpark42Plus, withFallbackReason} import org.apache.comet.expressions.CometEvalMode -import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, isStringCollationType, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType} import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil, CometTypeShim} object CometMin extends CometAggregateExpressionSerde[Min] { @@ -45,6 +45,8 @@ object CometMin extends CometAggregateExpressionSerde[Min] { override def getUnsupportedReasons(): Seq[String] = AggSerde.minMaxUnsupportedReasons + override def getIncompatibleReasons(): Seq[String] = AggSerde.minMaxIncompatibleReasons + override def getSupportLevel(expr: Min): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -86,6 +88,8 @@ object CometMax extends CometAggregateExpressionSerde[Max] { override def getUnsupportedReasons(): Seq[String] = AggSerde.minMaxUnsupportedReasons + override def getIncompatibleReasons(): Seq[String] = AggSerde.minMaxIncompatibleReasons + override def getSupportLevel(expr: Max): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -1589,7 +1593,7 @@ object CometPivotFirst extends CometAggregateExpressionSerde[PivotFirst] with Co } } -object AggSerde { +object AggSerde extends CometTypeShim { import org.apache.spark.sql.types._ def minMaxDataTypeSupported(dt: DataType): Boolean = { @@ -1639,14 +1643,70 @@ object AggSerde { Unsupported(Some(ComparisonUtils.nonDefaultCollationReason("min/max"))) // DataFusion's grouped hash aggregate under-reports the memory these values hold (#6936). case _: StringType | BinaryType => Compatible() + case st: StructType => minMaxStructSupportLevel(st) case _ if minMaxDataTypeSupported(dt) => Compatible() case _ => Unsupported(Some(s"Unsupported data type: $dt")) } + /** + * DataFusion compares struct fields with Arrow's total order for floats, which differs from + * Spark for -0.0 and negative NaN. It also skips a null field when comparing two structs, while + * Spark orders null as the smallest value, so a struct is only Compatible when none of its + * fields (at any depth) can be null. + */ + private def minMaxStructSupportLevel(st: StructType): SupportLevel = { + if (st.fields.isEmpty) { + Unsupported(Some("Struct min/max over an empty struct")) + } else if (hasNonDefaultStringCollation(st)) { + Unsupported(Some(ComparisonUtils.nonDefaultCollationReason("min/max"))) + } else if (SupportLevel.containsType(st, classOf[FloatType], classOf[DoubleType])) { + Unsupported( + Some( + "Struct min/max over a float or double field. DataFusion orders -0.0 below 0.0 and" + + " a NaN with the sign bit set below every value, while Spark treats -0.0 and 0.0" + + " as equal and every NaN as the largest value")) + } else { + unsupportedStructLeaf(st) match { + case Some(dt) => Unsupported(Some(s"Struct min/max over a field of type $dt")) + case None if structHasNullableField(st) => Incompatible(Some(minMaxNullableStructReason)) + case None => Compatible() + } + } + } + + /** The first leaf type of a struct that native min/max cannot compare, if any. */ + private def unsupportedStructLeaf(dt: DataType): Option[DataType] = dt match { + case st: StructType if st.fields.isEmpty => Some(st) + case st: StructType => + st.fields.view.flatMap(f => unsupportedStructLeaf(f.dataType)).headOption + case _ => + minMaxSupportLevel(dt) match { + case _: Compatible => None + case _ => Some(dt) + } + } + + private def structHasNullableField(dt: DataType): Boolean = dt match { + case st: StructType => st.fields.exists(f => f.nullable || structHasNullableField(f.dataType)) + case _ => false + } + + private val minMaxNullableStructReason = + "Struct min/max with nullable fields (at any depth). DataFusion skips a null field when" + + " comparing structs, while Spark orders null as the smallest value. When a field that" + + " decides the order is null, the result can differ from Spark, and can also differ" + + " between runs because the skipping comparison does not give a consistent order." + + /** Reasons a struct min/max is not Compatible, for the generated compatibility guide. */ + val minMaxIncompatibleReasons: Seq[String] = Seq(minMaxNullableStructReason) + /** The inputs that `minMaxSupportLevel` rejects, for the generated compatibility guide. */ val minMaxUnsupportedReasons: Seq[String] = Seq( "Input of a type other than boolean, integral, floating-point, decimal, date, timestamp," + - " string or binary, such as timestamp without time zone, intervals, arrays and structs.", + " string or binary, such as timestamp without time zone, intervals, arrays and maps." + + " A struct is supported only if all of its fields are of supported types other than" + + " float and double, since DataFusion orders -0.0 and NaN with the sign bit set" + + " differently from Spark inside a struct.", "A string with a non-default collation (Spark 4.0+), which Spark orders by its collation" + " while Comet compares raw bytes.") diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql index 02d96805fb1..97a43f236ea 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql @@ -31,3 +31,8 @@ INSERT INTO test_min_max_collation VALUES ('x', 'a'), ('x', 'B'), ('x', NULL), ( query expect_fallback(min/max does not support non-UTF8_BINARY collated operands) SELECT grp, min(CAST(s AS STRING COLLATE UTF8_LCASE)), max(CAST(s AS STRING COLLATE UTF8_LCASE)) FROM test_min_max_collation GROUP BY grp ORDER BY grp + +-- A collated string inside a struct falls back for the same reason. +query expect_fallback(min/max does not support non-UTF8_BINARY collated operands) +SELECT grp, max(struct(coalesce(grp, '') AS g, CAST(s AS STRING COLLATE UTF8_LCASE) AS s)) +FROM test_min_max_collation GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql new file mode 100644 index 00000000000..ee4374fd9b7 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql @@ -0,0 +1,133 @@ +-- 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. + +-- min/max over a struct. Compatible only when no field can be null and no field is float/double. +-- Structs with nullable fields are opt-in because DataFusion skips null fields when comparing. + +statement +CREATE TABLE test_mm_struct(g int, ts bigint, s string, d decimal(38,0), dt date, t timestamp, n string) +USING parquet + +statement +INSERT INTO test_mm_struct VALUES + (1, 5, 'b', 99999999999999999999999999999999999999, DATE'1980-01-01', TIMESTAMP'1980-01-01 00:00:00', 'x'), + (1, 9, '日本語', -99999999999999999999999999999999999999, DATE'9999-12-31', TIMESTAMP'2100-01-01 00:00:00', 'y'), + (1, 9, '中文', 0, DATE'2000-02-29', TIMESTAMP'2000-02-29 12:00:00', NULL), + (2, 9223372036854775807, 'a', 1, DATE'1980-01-01', TIMESTAMP'1999-12-31 23:59:59', 'z'), + (2, -9223372036854775808, 'c', 2, DATE'1980-01-02', TIMESTAMP'1999-12-31 23:59:58', 'w'), + (3, 1, 'only', 3, DATE'2020-01-01', TIMESTAMP'2020-01-01 00:00:00', NULL) + +-- Non-nullable fields via coalesce: native. +query +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct GROUP BY g ORDER BY g + +query +SELECT g, min(struct(coalesce(d, 0BD) AS d, coalesce(dt, DATE'1970-01-01') AS dt, coalesce(t, TIMESTAMP'1970-01-01 00:00:00') AS t)), + max(struct(coalesce(d, 0BD) AS d, coalesce(dt, DATE'1970-01-01') AS dt, coalesce(t, TIMESTAMP'1970-01-01 00:00:00') AS t)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Nested struct, and a whole-struct null from a guarding CASE WHEN. +query +SELECT g, max(CASE WHEN ts > 0 THEN struct(coalesce(ts, 0L) AS ts, struct(coalesce(s, '') AS s, true AS b) AS inner) END) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- No grouping. +query +SELECT min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct + +-- Nullable fields: opt-in, falls back by default. +query expect_fallback(Struct min/max with nullable fields) +SELECT g, max(struct(ts, s, n)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Float and double fields: NaN ordering differs, so fall back. +query expect_fallback(float or double field) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, 1.5D AS d)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Array field: unsupported. +query expect_fallback(Struct min/max over a field of type ArrayType) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, array(1, 2) AS a)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Spark cannot order maps, so it rejects a struct holding one during analysis. +query expect_error(INVALID_ORDERING_TYPE) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, map('k', 1) AS m)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- timestamp_ntz is not supported by native min/max, so a struct holding one falls back too. +query expect_fallback(Struct min/max over a field of type TimestampNTZType) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, TIMESTAMP_NTZ'2020-01-01 00:00:00' AS t)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Boolean, integral and binary fields. Spark marks narrowing casts nullable, so each one is +-- wrapped in coalesce to keep the struct Compatible. +query +SELECT g, + min(struct(coalesce(ts, 0L) > 5 AS b, coalesce(CAST(ts % 100 AS TINYINT), 0Y) AS i8, + coalesce(CAST(ts % 30000 AS SMALLINT), 0S) AS i16, coalesce(CAST(ts % 100000 AS INT), 0) AS i32, + coalesce(CAST(s AS BINARY), X'') AS bin)), + max(struct(coalesce(ts, 0L) > 5 AS b, coalesce(CAST(ts % 100 AS TINYINT), 0Y) AS i8, + coalesce(CAST(ts % 30000 AS SMALLINT), 0S) AS i16, coalesce(CAST(ts % 100000 AS INT), 0) AS i32, + coalesce(CAST(s AS BINARY), X'') AS bin)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Ties on the first field are decided by the next one. Group 1 has two rows with ts = 9, and +-- '中文' sorts before '日本語' as raw UTF-8 bytes. +query +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct WHERE ts = 9 GROUP BY g ORDER BY g + +-- count(DISTINCT) makes Spark plan a PartialMerge stage, which merges the struct buffers again. +query +SELECT g, count(DISTINCT s), max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- A VALUES input keeps the partial aggregate on Spark, while the final one merges Spark's struct +-- buffers natively. The plan is not fully native, so only the answer is checked. +query spark_answer_only +SELECT k, min(struct(o, v)), max(struct(o, v)) +FROM VALUES (1, 5L, 'b'), (1, 9L, 'ｱ'), (2, -1L, '😀'), (2, 3L, 'a') AS t(k, o, v) +GROUP BY k ORDER BY k + +-- Empty input: one NULL row without grouping, no rows with grouping. +query +SELECT max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct WHERE g < 0 + +query +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct WHERE g < 0 GROUP BY g + +-- FILTER turns filtered-out rows into nulls before they reach the struct accumulator. +query +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FILTER (WHERE ts < 9) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- A nested struct that can be null makes the outer struct opt-in. +query expect_fallback(Struct min/max with nullable fields) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, CASE WHEN ts > 5 THEN struct(coalesce(s, '') AS s) END AS inner)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Spark reads file-source schemas as nullable, so min/max over a struct column read from a table +-- is opt-in even when every stored value is non-null. +statement +CREATE TABLE test_mm_struct_col(g int, st struct) USING parquet + +statement +INSERT INTO test_mm_struct_col VALUES + (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', '日本語')), + (2, named_struct('ts', 1L, 's', 'a')), (2, NULL) + +query expect_fallback(Struct min/max with nullable fields) +SELECT g, min(st), max(st) FROM test_mm_struct_col GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql new file mode 100644 index 00000000000..59cb56b18f7 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql @@ -0,0 +1,34 @@ +-- 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. + +-- Struct min/max over many small batches and many groups, so the accumulator carries each group's +-- winner from batch to batch instead of seeing all rows at once. range() produces non-nullable +-- values, so the struct stays Compatible. + +-- Config: spark.comet.batchSize=64 + +query +SELECT g, min(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)), + max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) +FROM (SELECT id, id % 1000 AS g FROM range(0, 200000, 1, 8)) GROUP BY g ORDER BY g + +query +SELECT max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) FROM range(0, 200000, 1, 8) + +query +SELECT g, max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) FILTER (WHERE id % 3 = 0) +FROM (SELECT id, id % 1000 AS g FROM range(0, 200000, 1, 8)) GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql new file mode 100644 index 00000000000..57503e070f8 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql @@ -0,0 +1,55 @@ +-- 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. + +-- Struct min/max with its partial and final stages on different engines. The test-only configs +-- below keep one stage on Spark, so the struct buffer crosses between Spark and Comet in both +-- directions: native partial -> Spark final, and Spark partial -> native final. Min and Max allow +-- both. The plan is mixed by design, so only the answer is checked. + +-- ConfigMatrix: spark.comet.testing.aggregate.partialMode.enabled=true,false +-- ConfigMatrix: spark.comet.testing.aggregate.finalMode.enabled=true,false + +statement +CREATE TABLE test_mm_struct_mixed(g int, ts bigint, s string, x int) USING parquet + +statement +INSERT INTO test_mm_struct_mixed VALUES + (1, 5, 'b', 1), (1, 9, '日本語', 2), (1, 9, '中文', 3), + (2, 9223372036854775807, 'a', 4), (2, -9223372036854775808, 'c', 5), + (3, NULL, NULL, NULL) + +query spark_answer_only +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +query spark_answer_only +SELECT g, max(CASE WHEN ts > 0 THEN struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s) END) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +query spark_answer_only +SELECT max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct_mixed + +-- count(DISTINCT) adds PartialMerge stages between the partial and final ones. +query spark_answer_only +SELECT g, count(DISTINCT s), min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +-- collect_list makes Spark plan an ObjectHashAggregate instead of a sort aggregate. +query spark_answer_only +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), sort_array(collect_list(x)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql new file mode 100644 index 00000000000..f060093988d --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql @@ -0,0 +1,74 @@ +-- 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. + +-- min/max over a struct with nullable fields, opted in with allowIncompatible. +-- DataFusion skips a field when either side is null and moves on to the next field, while Spark +-- orders null below every value. The answers match as long as the field that decides the order is +-- not null, which is what these queries check. + +-- Config: spark.comet.expression.Min.allowIncompatible=true +-- Config: spark.comet.expression.Max.allowIncompatible=true + +statement +CREATE TABLE test_mm_struct_null(g int, ts bigint, s string, n string, inner_v int) USING parquet + +statement +INSERT INTO test_mm_struct_null VALUES + (1, 5, 'b', 'x', 1), + (1, 9, NULL, NULL, NULL), + (1, 7, '日本語', NULL, 3), + (2, 1, 'a', 'z', NULL), + (2, 2, NULL, 'w', 2), + (3, 4, NULL, NULL, NULL), + (4, NULL, NULL, NULL, NULL) + +-- The deciding field ts is distinct and non-null within each group; later fields may be null. +query +SELECT g, min(struct(ts, s, n)), max(struct(ts, s, n)) +FROM test_mm_struct_null WHERE ts IS NOT NULL GROUP BY g ORDER BY g + +-- Nested struct with nullable inner fields. +query +SELECT g, min(struct(ts, struct(s, inner_v) AS inner)), max(struct(ts, struct(s, inner_v) AS inner)) +FROM test_mm_struct_null WHERE ts IS NOT NULL GROUP BY g ORDER BY g + +-- Whole-struct nulls are skipped by both engines. Group 4 has only a NULL struct, so the result +-- is NULL. +query +SELECT g, max(CASE WHEN ts IS NOT NULL THEN struct(ts, s) END) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- Known difference: in group 1 Spark returns (7, 'a') because a null ts sorts first. DataFusion +-- skips the null ts and compares the second field, where '~' is larger than 'a', so it returns +-- (NULL, '~'). +query ignore(DataFusion skips null struct fields when comparing, Spark orders null first) +SELECT max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) +FROM test_mm_struct_null WHERE g = 1 + +-- A struct column read from a table. Spark reads it with nullable fields, so it goes through the +-- opt-in path and the cast that relabels its fields. +statement +CREATE TABLE test_mm_struct_null_col(g int, st struct) USING parquet + +statement +INSERT INTO test_mm_struct_null_col VALUES + (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', NULL)), + (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL) + +query +SELECT g, min(st), max(st) FROM test_mm_struct_null_col GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql new file mode 100644 index 00000000000..1677acb61ae --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql @@ -0,0 +1,56 @@ +-- 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. + +-- min/max over a struct with timestamp and date fields under several session time zones, +-- including values near the ends of the range and DST transitions. 1901 stays clear of the +-- INT96 rebase cutoff in every zone listed. + +-- ConfigMatrix: spark.sql.session.timeZone=UTC,Etc/UTC,America/Los_Angeles,Asia/Kolkata,Australia/Lord_Howe + +statement +CREATE TABLE test_mm_struct_tz(g int, t timestamp, dt date, s string) USING parquet + +statement +INSERT INTO test_mm_struct_tz VALUES + (1, TIMESTAMP'1901-01-01 00:00:00', DATE'1901-01-01', 'a'), + (1, TIMESTAMP'9999-12-30 23:59:59.999999', DATE'9999-12-31', 'b'), + (1, TIMESTAMP'1970-01-01 00:00:00', DATE'1970-01-01', 'c'), + (2, TIMESTAMP'2024-03-10 01:59:59', DATE'2024-03-10', 'before'), + (2, TIMESTAMP'2024-03-10 03:00:00', DATE'2024-03-10', 'after'), + (2, TIMESTAMP'2024-11-03 01:30:00', DATE'2024-11-03', 'overlap'), + (3, TIMESTAMP'1969-12-31 23:59:59.999999', DATE'1969-12-31', '한국어'), + (3, NULL, NULL, NULL) + +query +SELECT g, + min(struct(coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t, coalesce(dt, DATE'2000-01-01') AS dt, coalesce(s, '') AS s)), + max(struct(coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t, coalesce(dt, DATE'2000-01-01') AS dt, coalesce(s, '') AS s)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g + +-- The date leads, so ties on the date are decided by the timestamp. +query +SELECT g, + min(struct(coalesce(dt, DATE'2000-01-01') AS dt, coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t)), + max(struct(coalesce(dt, DATE'2000-01-01') AS dt, coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g + +-- The timestamp field is computed natively from a date, which depends on the session time zone. +query +SELECT g, + min(struct(CAST(coalesce(dt, DATE'2000-01-01') AS TIMESTAMP) AS t, coalesce(s, '') AS s)), + max(struct(CAST(coalesce(dt, DATE'2000-01-01') AS TIMESTAMP) AS t, coalesce(s, '') AS s)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g From ea79a1f364f0254a292c661506b86a5ed8d62b9d Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Wed, 28 Oct 2026 15:26:04 -0700 Subject: [PATCH 2/3] Address review comments --- .../latest/compatibility/operators.md | 9 +- .../core/src/execution/planner/aggregate.rs | 9 +- .../src/agg_funcs/min_max_struct.rs | 589 ++++++++++++++++++ native/spark-expr/src/agg_funcs/mod.rs | 2 + .../org/apache/comet/serde/aggregates.scala | 31 +- .../expressions/aggregate/min_max_struct.sql | 22 +- .../aggregate/min_max_struct_many_batches.sql | 3 +- .../aggregate/min_max_struct_nullable.sql | 50 +- .../comet/exec/CometAggregateSuite.scala | 37 ++ 9 files changed, 683 insertions(+), 69 deletions(-) create mode 100644 native/spark-expr/src/agg_funcs/min_max_struct.rs diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 8785874c6f6..a1a74a38fa1 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -167,11 +167,14 @@ incorrect result. When any single window expression in a `WindowExec` falls back (`stddev`, `stddev_pop`, `stddev_samp`, `var_pop`, `var_samp`, `corr`, `covar_pop`, `covar_samp`). These run natively as plain aggregations but not as window functions ([#4766](https://github.com/apache/datafusion-comet/issues/4766)). -- `min` / `max` on timestamp-without-time-zone, interval, or nested (array / struct) input types, and `sum` / - `avg` on year-month or day-time interval input types. Windowed aggregates inherit the same input-type support as - the batch aggregates, so these fall back in both contexts. +- `min` / `max` on timestamp-without-time-zone, interval, or array input types, and `sum` / `avg` on year-month + or day-time interval input types. Windowed aggregates inherit the same input-type support as the batch + aggregates, so these fall back in both contexts. - `min` / `max` on string or binary input types. These run natively as plain aggregations but not as window functions. +- `min` / `max` on struct input types. A struct whose fields are all supported `min` / `max` types other than + `float` and `double` runs natively as a plain aggregation, but not as a window function. Other structs fall + back in both contexts. - `sum` or `avg` on `DECIMAL` with a sliding (non ever-expanding) frame, because the sliding path would wrap on overflow instead of returning Spark's `NULL`. - `avg` or `try_avg` on `DECIMAL(p, s)` with an expanding frame when `p >= 28`, because Spark can preserve diff --git a/native/core/src/execution/planner/aggregate.rs b/native/core/src/execution/planner/aggregate.rs index 77afce648b5..c1dd517da85 100644 --- a/native/core/src/execution/planner/aggregate.rs +++ b/native/core/src/execution/planner/aggregate.rs @@ -54,7 +54,8 @@ use datafusion_comet_spark_expr::{ ApproxPercentile, Avg, AvgDecimal, BloomFilterAgg, Cast, CometCollectList, CometCollectSet, Correlation, Covariance, EvalMode, HllPlusPlus, HllSketchAgg, HllUnionAgg, MaxMinBy, Mode, Regr, RegrType, SparkBloomFilterVersion, SparkCastOptions, SparkListAgg, SparkMinMax, - SparkPercentile, SparkPivotFirst, Stddev, SumDecimal, SumInteger, Variance, ZeroOrder, + SparkPercentile, SparkPivotFirst, SparkStructMinMax, Stddev, SumDecimal, SumInteger, Variance, + ZeroOrder, }; use jni::objects::{Global, JObject}; @@ -915,13 +916,15 @@ fn make_all_fields_nullable(data_type: &DataType) -> DataType { } } -/// `min` or `max` over `data_type`: Spark's version for floats, which DataFusion orders -/// differently, and DataFusion's for every other type. +/// `min` or `max` over `data_type`: Spark's version for floats and structs, which DataFusion +/// orders differently, and DataFusion's for every other type. pub(super) fn min_max_udaf(data_type: &DataType, is_max: bool) -> Arc { // Spark has only 32- and 64-bit floating-point types. Keep Float16 on DataFusion's path so // an internal plan using it does not reach SparkMinMax, whose accumulators reject that type. if matches!(data_type, DataType::Float32 | DataType::Float64) { Arc::new(AggregateUDF::new_from_impl(SparkMinMax::new(is_max))) + } else if matches!(data_type, DataType::Struct(_)) { + Arc::new(AggregateUDF::new_from_impl(SparkStructMinMax::new(is_max))) } else if is_max { max_udaf() } else { diff --git a/native/spark-expr/src/agg_funcs/min_max_struct.rs b/native/spark-expr/src/agg_funcs/min_max_struct.rs new file mode 100644 index 00000000000..ad88640750c --- /dev/null +++ b/native/spark-expr/src/agg_funcs/min_max_struct.rs @@ -0,0 +1,589 @@ +// 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. + +use arrow::array::{new_null_array, Array, ArrayRef, BooleanArray}; +use arrow::compute::{not, nullif, prep_null_mask_filter, SortOptions}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::row::{OwnedRow, RowConverter, Rows, SortField}; +use datafusion::common::{internal_err, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::{ + Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature, Volatility, +}; +use datafusion::physical_expr::expressions::format_state_name; +use std::fmt::Debug; +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; + +/// Spark's `max` or `min` over a struct. +/// +/// Spark orders structs field by field, with a null field below every value. Each input batch is +/// encoded once with Arrow's row format, ascending with nulls first, whose bytes sort in exactly +/// that order. Each group keeps the bytes of its current winner, and memory is accounted by the +/// size of those buffers. Float and double fields are not supported: the row format orders `-0.0` +/// below `0.0`, while Spark treats them as equal. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStructMinMax { + signature: Signature, + is_max: bool, +} + +impl SparkStructMinMax { + pub fn new(is_max: bool) -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + is_max, + } + } +} + +impl AggregateUDFImpl for SparkStructMinMax { + fn name(&self) -> &str { + if self.is_max { + "max" + } else { + "min" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, args: AccumulatorArgs) -> Result> { + Ok(Box::new(StructMinMaxAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + Ok(vec![Arc::new(Field::new( + format_state_name(args.name, self.name()), + args.return_field.data_type().clone(), + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result> { + Ok(Box::new(StructMinMaxGroupsAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } +} + +/// Encodes struct values to Spark-ordered row bytes and decodes them back. +#[derive(Debug)] +struct StructRows { + converter: RowConverter, + /// The encoding of a null struct, emitted for a group that saw no value. + null_row: OwnedRow, + is_max: bool, +} + +impl StructRows { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + if !matches!(data_type, DataType::Struct(_)) { + return internal_err!("Spark struct min/max expects a struct, got {data_type}"); + } + let options = SortOptions { + descending: false, + nulls_first: true, + }; + let converter = RowConverter::new(vec![SortField::new_with_options( + data_type.clone(), + options, + )])?; + let null_row = converter + .convert_columns(&[new_null_array(data_type, 1)])? + .row(0) + .owned(); + Ok(Self { + converter, + null_row, + is_max, + }) + } + + fn encode(&self, values: &ArrayRef) -> Result { + Ok(self.converter.convert_columns(&[Arc::clone(values)])?) + } + + /// Whether `candidate` replaces `current`. Only a strictly better value does. + #[inline] + fn replaces(&self, candidate: &[u8], current: &[u8]) -> bool { + if self.is_max { + candidate > current + } else { + candidate < current + } + } + + /// Keeps the better of `value` and `slot`, reusing the slot's buffer. Returns the change in + /// the slot's allocated bytes. + #[inline] + fn update(&self, slot: &mut Option>, value: &[u8]) -> isize { + match slot { + None => { + let bytes = value.to_vec(); + let added = bytes.capacity() as isize; + *slot = Some(bytes); + added + } + Some(current) if self.replaces(value, current) => { + let before = current.capacity() as isize; + current.clear(); + current.extend_from_slice(value); + current.capacity() as isize - before + } + Some(_) => 0, + } + } + + fn decode<'a>(&self, values: impl Iterator>) -> Result { + let parser = self.converter.parser(); + let rows: Vec<_> = values + .map(|value| match value { + Some(bytes) => parser.parse(bytes), + None => self.null_row.row(), + }) + .collect(); + let mut columns = self.converter.convert_rows(rows)?; + Ok(columns.remove(0)) + } + + fn size(&self) -> usize { + self.converter.size() + self.null_row.as_ref().len() + } +} + +/// Whether row `i` takes part: a null struct or a row the filter drops does not. +#[inline] +fn is_skipped(values: &ArrayRef, filter: Option<&BooleanArray>, i: usize) -> bool { + values.is_null(i) || filter.is_some_and(|f| !f.is_valid(i) || !f.value(i)) +} + +#[derive(Debug)] +struct StructMinMaxAccumulator { + rows: StructRows, + value: Option>, +} + +impl StructMinMaxAccumulator { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + Ok(Self { + rows: StructRows::try_new(data_type, is_max)?, + value: None, + }) + } +} + +impl Accumulator for StructMinMaxAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let values = &values[0]; + let encoded = self.rows.encode(values)?; + for i in 0..values.len() { + if !is_skipped(values, None, i) { + self.rows.update(&mut self.value, encoded.row(i).as_ref()); + } + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.update_batch(states) + } + + fn state(&mut self) -> Result> { + Ok(vec![self.evaluate()?]) + } + + fn evaluate(&mut self) -> Result { + let array = self.rows.decode(std::iter::once(self.value.as_deref()))?; + ScalarValue::try_from_array(&array, 0) + } + + fn size(&self) -> usize { + size_of_val(self) + + self.rows.size() + + self.value.as_ref().map_or(0, |value| value.capacity()) + } +} + +#[derive(Debug)] +struct StructMinMaxGroupsAccumulator { + rows: StructRows, + /// The row bytes of each group's current winner, or `None` before it sees a value. + values: Vec>>, + /// The bytes allocated by the buffers in `values`. + value_bytes: usize, +} + +impl StructMinMaxGroupsAccumulator { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + Ok(Self { + rows: StructRows::try_new(data_type, is_max)?, + values: vec![], + value_bytes: 0, + }) + } +} + +impl GroupsAccumulator for StructMinMaxGroupsAccumulator { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let values = &values[0]; + self.values.resize(total_num_groups, None); + let encoded = self.rows.encode(values)?; + let mut delta = 0isize; + for (i, &group) in group_indices.iter().enumerate() { + if !is_skipped(values, opt_filter, i) { + delta += self + .rows + .update(&mut self.values[group], encoded.row(i).as_ref()); + } + } + self.value_bytes = self.value_bytes.wrapping_add_signed(delta); + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let emitted = emit_to.take_needed(&mut self.values); + let freed: usize = emitted.iter().flatten().map(|value| value.capacity()).sum(); + self.value_bytes -= freed; + self.rows + .decode(emitted.iter().map(|value| value.as_deref())) + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + Ok(vec![self.evaluate(emit_to)?]) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + self.update_batch(values, group_indices, None, total_num_groups) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + let values = Arc::clone(&values[0]); + let Some(filter) = opt_filter else { + return Ok(vec![values]); + }; + // A row the filter drops, or whose filter value is null, contributes nothing. + let dropped = not(&prep_null_mask_filter(filter))?; + Ok(vec![nullif(&values, &dropped)?]) + } + + fn size(&self) -> usize { + size_of_val(self) + + self.rows.size() + + self.values.capacity() * size_of::>>() + + self.value_bytes + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, Int64Array, StringArray, StructArray}; + use arrow::datatypes::{Fields, Int64Type}; + + fn fields(nullable: bool) -> Fields { + Fields::from(vec![ + Field::new("k", DataType::Int64, nullable), + Field::new("s", DataType::Utf8, nullable), + ]) + } + + fn structs(rows: &[Option<(Option, Option<&str>)>]) -> ArrayRef { + let k: Int64Array = rows.iter().map(|r| r.and_then(|(k, _)| k)).collect(); + let s: StringArray = rows.iter().map(|r| r.and_then(|(_, s)| s)).collect(); + let validity = rows.iter().map(Option::is_some).collect::>(); + Arc::new(StructArray::new( + fields(true), + vec![Arc::new(k), Arc::new(s)], + Some(validity.into()), + )) + } + + fn decoded(array: &ArrayRef) -> Vec, Option)>> { + let array = array.as_struct(); + let k = array.column(0).as_primitive::(); + let s = array.column(1).as_string::(); + (0..array.len()) + .map(|i| { + array.is_valid(i).then(|| { + ( + k.is_valid(i).then(|| k.value(i)), + s.is_valid(i).then(|| s.value(i).to_string()), + ) + }) + }) + .collect() + } + + fn groups(is_max: bool) -> StructMinMaxGroupsAccumulator { + StructMinMaxGroupsAccumulator::try_new(&DataType::Struct(fields(true)), is_max).unwrap() + } + + /// Spark orders a null field below every value, at any position, and skips a null struct. + #[test] + fn orders_null_fields_first_like_spark() -> Result<()> { + let input = structs(&[ + Some((Some(1), Some("z"))), + Some((None, Some("zz"))), + Some((Some(2), Some("a"))), + None, + Some((Some(2), None)), + ]); + for (is_max, expected) in [ + (true, Some((Some(2), Some("a".to_string())))), + (false, Some((None, Some("zz".to_string())))), + ] { + let mut acc = groups(is_max); + acc.update_batch(&[Arc::clone(&input)], &[0, 0, 0, 0, 0], None, 1)?; + assert_eq!(decoded(&acc.evaluate(EmitTo::All)?), vec![expected]); + } + Ok(()) + } + + /// A group with no value, or whose rows the filter drops, emits a null struct. + #[test] + fn empty_and_filtered_groups_are_null() -> Result<()> { + let input = structs(&[Some((Some(1), Some("a"))), Some((Some(5), Some("b")))]); + let filter = BooleanArray::from(vec![Some(true), None]); + let mut acc = groups(true); + acc.update_batch(&[input], &[0, 1], Some(&filter), 3)?; + assert_eq!( + decoded(&acc.evaluate(EmitTo::All)?), + vec![Some((Some(1), Some("a".to_string()))), None, None] + ); + Ok(()) + } + + /// Emitting a prefix of the groups frees exactly the bytes those groups held, so ordered + /// aggregation does not keep reserving memory for groups it has already emitted. + #[test] + fn emitting_groups_releases_their_bytes() -> Result<()> { + let payload = "x".repeat(64 * 1024); + let mut acc = groups(true); + let empty = acc.size(); + for round in 0..3 { + let rows: Vec<_> = (0..64) + .map(|k| Some((Some(round * 64 + k), Some(payload.as_str())))) + .collect(); + let group_indices: Vec = (0..64).collect(); + acc.update_batch(&[structs(&rows)], &group_indices, None, 64)?; + assert!(acc.size() > 64 * payload.len()); + acc.evaluate(EmitTo::First(32))?; + acc.evaluate(EmitTo::First(32))?; + assert_eq!(acc.value_bytes, 0); + assert!(acc.size() < empty + 64 * size_of::>>()); + } + Ok(()) + } + + /// Partial states merge into the same answer as one pass over every row. + #[test] + fn merge_matches_single_pass() -> Result<()> { + let first = structs(&[Some((Some(3), Some("b"))), Some((Some(1), None))]); + let second = structs(&[Some((Some(3), Some("c"))), None]); + let mut left = groups(true); + left.update_batch(&[first], &[0, 1], None, 2)?; + let mut right = groups(true); + right.update_batch(&[second], &[0, 1], None, 2)?; + let mut merged = groups(true); + merged.merge_batch(&left.state(EmitTo::All)?, &[0, 1], 2)?; + merged.merge_batch(&right.state(EmitTo::All)?, &[0, 1], 2)?; + assert_eq!( + decoded(&merged.evaluate(EmitTo::All)?), + vec![ + Some((Some(3), Some("c".to_string()))), + Some((Some(1), None)) + ] + ); + Ok(()) + } + + #[test] + fn ungrouped_accumulator_matches_grouped() -> Result<()> { + let input = structs(&[ + Some((Some(2), Some("a"))), + None, + Some((Some(2), Some("b"))), + Some((None, Some("zzz"))), + ]); + let mut acc = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + acc.update_batch(&[Arc::clone(&input)])?; + let state = acc.state()?; + let mut merged = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + merged.merge_batch(&[state[0].to_array()?])?; + let result = merged.evaluate()?.to_array()?; + assert_eq!( + decoded(&result), + vec![Some((Some(2), Some("b".to_string())))] + ); + + let mut empty = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + assert_eq!(decoded(&empty.evaluate()?.to_array()?), vec![None]); + Ok(()) + } + + /// Non-nullable fields decode with the declared type. + #[test] + fn keeps_non_nullable_field_types() -> Result<()> { + let data_type = DataType::Struct(fields(false)); + let input: ArrayRef = Arc::new(StructArray::new( + fields(false), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["a", "b"])), + ], + None, + )); + let mut acc = StructMinMaxGroupsAccumulator::try_new(&data_type, false)?; + acc.update_batch(&[input], &[0, 0], None, 1)?; + let result = acc.evaluate(EmitTo::All)?; + assert_eq!(result.data_type(), &data_type); + Ok(()) + } + + /// Sorted aggregation emits finished groups after every batch. Struct max over 4,096 groups + /// of 4 KiB values (16 MiB in all), read in 64-row batches, must fit a 2 MiB pool, which it + /// only does if emitting a group releases the memory charged for it. DataFusion's own struct + /// max fails this with `ResourcesExhausted`, because it subtracts rows rather than bytes. + #[tokio::test] + async fn sorted_aggregation_fits_a_small_pool() -> Result<()> { + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::Schema; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::execution::TaskContext; + use datafusion::logical_expr::AggregateUDF; + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_expr::expressions::{col, Column}; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, + }; + use datafusion::physical_plan::{collect, ExecutionPlan, InputOrderMode}; + use datafusion::prelude::SessionConfig; + + const GROUPS: i64 = 4096; + const BATCH_ROWS: i64 = 64; + let payload = "x".repeat(4096); + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Struct(fields(true)), true), + ])); + let batches = (0..GROUPS / BATCH_ROWS) + .map(|b| { + let groups = (b * BATCH_ROWS..(b + 1) * BATCH_ROWS).collect::>(); + let rows: Vec<_> = groups + .iter() + .map(|&g| Some((Some(g), Some(payload.as_str())))) + .collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(groups)), structs(&rows)], + ) + }) + .collect::, _>>()?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("g", 0), + ))]) + .unwrap(); + let source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input: Arc = DataSourceExec::from_data_source(source); + + let max = Arc::new( + AggregateExprBuilder::new( + Arc::new(AggregateUDF::new_from_impl(SparkStructMinMax::new(true))), + vec![col("v", &schema)?], + ) + .schema(Arc::clone(&schema)) + .alias("max") + .build()?, + ); + // Comet splits a sort aggregate into a Partial and a Final stage. Both see sorted input, + // so DataFusion runs its ordered streams, which emit finished groups after every batch. + let group_by = PhysicalGroupBy::new_single(vec![(col("g", &schema)?, "g".to_string())]); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&max)], + vec![None], + input, + Arc::clone(&schema), + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![max], + vec![None], + Arc::clone(&partial) as Arc, + Arc::clone(&schema), + )?); + assert_eq!(partial.input_order_mode(), &InputOrderMode::Sorted); + // No disk, so an over-reserving accumulator fails instead of spilling. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(2 * 1024 * 1024, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(BATCH_ROWS as usize)) + .with_runtime(runtime), + ); + let output = collect(aggregate, context).await?; + assert_eq!( + output.iter().map(|b| b.num_rows()).sum::(), + GROUPS as usize + ); + Ok(()) + } +} diff --git a/native/spark-expr/src/agg_funcs/mod.rs b/native/spark-expr/src/agg_funcs/mod.rs index e7a72d60ec8..7133a9cff90 100644 --- a/native/spark-expr/src/agg_funcs/mod.rs +++ b/native/spark-expr/src/agg_funcs/mod.rs @@ -29,6 +29,7 @@ mod hll_union_agg; mod list_agg; mod max_min_by; mod min_max; +mod min_max_struct; mod mode; mod percentile; mod pivot_first; @@ -53,6 +54,7 @@ pub use hll_union_agg::HllUnionAgg; pub use list_agg::SparkListAgg; pub use max_min_by::MaxMinBy; pub use min_max::SparkMinMax; +pub use min_max_struct::SparkStructMinMax; pub use mode::Mode; pub use percentile::{SparkPercentile, ZeroOrder}; pub use pivot_first::SparkPivotFirst; diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index c72b0c1a0a0..a521d410ade 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -45,8 +45,6 @@ object CometMin extends CometAggregateExpressionSerde[Min] { override def getUnsupportedReasons(): Seq[String] = AggSerde.minMaxUnsupportedReasons - override def getIncompatibleReasons(): Seq[String] = AggSerde.minMaxIncompatibleReasons - override def getSupportLevel(expr: Min): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -88,8 +86,6 @@ object CometMax extends CometAggregateExpressionSerde[Max] { override def getUnsupportedReasons(): Seq[String] = AggSerde.minMaxUnsupportedReasons - override def getIncompatibleReasons(): Seq[String] = AggSerde.minMaxIncompatibleReasons - override def getSupportLevel(expr: Max): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -1649,10 +1645,10 @@ object AggSerde extends CometTypeShim { } /** - * DataFusion compares struct fields with Arrow's total order for floats, which differs from - * Spark for -0.0 and negative NaN. It also skips a null field when comparing two structs, while - * Spark orders null as the smallest value, so a struct is only Compatible when none of its - * fields (at any depth) can be null. + * The native accumulator compares structs by their Arrow row encoding, ascending with nulls + * first, which orders fields the way Spark does, a null field below every value included. The + * encoding orders floats by IEEE total order, which differs from Spark for -0.0 and NaN with + * the sign bit set, so float and double fields fall back. */ private def minMaxStructSupportLevel(st: StructType): SupportLevel = { if (st.fields.isEmpty) { @@ -1662,13 +1658,12 @@ object AggSerde extends CometTypeShim { } else if (SupportLevel.containsType(st, classOf[FloatType], classOf[DoubleType])) { Unsupported( Some( - "Struct min/max over a float or double field. DataFusion orders -0.0 below 0.0 and" + + "Struct min/max over a float or double field. Comet orders -0.0 below 0.0 and" + " a NaN with the sign bit set below every value, while Spark treats -0.0 and 0.0" + " as equal and every NaN as the largest value")) } else { unsupportedStructLeaf(st) match { case Some(dt) => Unsupported(Some(s"Struct min/max over a field of type $dt")) - case None if structHasNullableField(st) => Incompatible(Some(minMaxNullableStructReason)) case None => Compatible() } } @@ -1686,26 +1681,12 @@ object AggSerde extends CometTypeShim { } } - private def structHasNullableField(dt: DataType): Boolean = dt match { - case st: StructType => st.fields.exists(f => f.nullable || structHasNullableField(f.dataType)) - case _ => false - } - - private val minMaxNullableStructReason = - "Struct min/max with nullable fields (at any depth). DataFusion skips a null field when" + - " comparing structs, while Spark orders null as the smallest value. When a field that" + - " decides the order is null, the result can differ from Spark, and can also differ" + - " between runs because the skipping comparison does not give a consistent order." - - /** Reasons a struct min/max is not Compatible, for the generated compatibility guide. */ - val minMaxIncompatibleReasons: Seq[String] = Seq(minMaxNullableStructReason) - /** The inputs that `minMaxSupportLevel` rejects, for the generated compatibility guide. */ val minMaxUnsupportedReasons: Seq[String] = Seq( "Input of a type other than boolean, integral, floating-point, decimal, date, timestamp," + " string or binary, such as timestamp without time zone, intervals, arrays and maps." + " A struct is supported only if all of its fields are of supported types other than" + - " float and double, since DataFusion orders -0.0 and NaN with the sign bit set" + + " float and double, since Comet orders -0.0 and NaN with the sign bit set" + " differently from Spark inside a struct.", "A string with a non-default collation (Spark 4.0+), which Spark orders by its collation" + " while Comet compares raw bytes.") diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql index ee4374fd9b7..9818b0b2938 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql @@ -15,8 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- min/max over a struct. Compatible only when no field can be null and no field is float/double. --- Structs with nullable fields are opt-in because DataFusion skips null fields when comparing. +-- min/max over a struct. Spark orders structs field by field, with a null field below every +-- value. Float and double fields fall back. Null fields are covered in min_max_struct_nullable.sql. statement CREATE TABLE test_mm_struct(g int, ts bigint, s string, d decimal(38,0), dt date, t timestamp, n string) @@ -31,7 +31,7 @@ INSERT INTO test_mm_struct VALUES (2, -9223372036854775808, 'c', 2, DATE'1980-01-02', TIMESTAMP'1999-12-31 23:59:58', 'w'), (3, 1, 'only', 3, DATE'2020-01-01', TIMESTAMP'2020-01-01 00:00:00', NULL) --- Non-nullable fields via coalesce: native. +-- Non-nullable fields. query SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) @@ -51,8 +51,8 @@ FROM test_mm_struct GROUP BY g ORDER BY g query SELECT min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct --- Nullable fields: opt-in, falls back by default. -query expect_fallback(Struct min/max with nullable fields) +-- Nullable fields, including a null in the deciding field. +query SELECT g, max(struct(ts, s, n)) FROM test_mm_struct GROUP BY g ORDER BY g -- Float and double fields: NaN ordering differs, so fall back. @@ -71,8 +71,7 @@ SELECT g, max(struct(coalesce(ts, 0L) AS ts, map('k', 1) AS m)) FROM test_mm_str query expect_fallback(Struct min/max over a field of type TimestampNTZType) SELECT g, max(struct(coalesce(ts, 0L) AS ts, TIMESTAMP_NTZ'2020-01-01 00:00:00' AS t)) FROM test_mm_struct GROUP BY g ORDER BY g --- Boolean, integral and binary fields. Spark marks narrowing casts nullable, so each one is --- wrapped in coalesce to keep the struct Compatible. +-- Boolean, integral and binary fields. query SELECT g, min(struct(coalesce(ts, 0L) > 5 AS b, coalesce(CAST(ts % 100 AS TINYINT), 0Y) AS i8, @@ -114,13 +113,12 @@ query SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FILTER (WHERE ts < 9) FROM test_mm_struct GROUP BY g ORDER BY g --- A nested struct that can be null makes the outer struct opt-in. -query expect_fallback(Struct min/max with nullable fields) +-- A nested struct that can be null orders below every non-null nested struct. +query SELECT g, max(struct(coalesce(ts, 0L) AS ts, CASE WHEN ts > 5 THEN struct(coalesce(s, '') AS s) END AS inner)) FROM test_mm_struct GROUP BY g ORDER BY g --- Spark reads file-source schemas as nullable, so min/max over a struct column read from a table --- is opt-in even when every stored value is non-null. +-- A struct column read from a table, whose fields Spark reads as nullable. statement CREATE TABLE test_mm_struct_col(g int, st struct) USING parquet @@ -129,5 +127,5 @@ INSERT INTO test_mm_struct_col VALUES (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL) -query expect_fallback(Struct min/max with nullable fields) +query SELECT g, min(st), max(st) FROM test_mm_struct_col GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql index 59cb56b18f7..87df5c13d09 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql @@ -16,8 +16,7 @@ -- under the License. -- Struct min/max over many small batches and many groups, so the accumulator carries each group's --- winner from batch to batch instead of seeing all rows at once. range() produces non-nullable --- values, so the struct stays Compatible. +-- winner from batch to batch instead of seeing all rows at once. -- Config: spark.comet.batchSize=64 diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql index f060093988d..2d7247d8cfd 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql @@ -15,13 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- min/max over a struct with nullable fields, opted in with allowIncompatible. --- DataFusion skips a field when either side is null and moves on to the next field, while Spark --- orders null below every value. The answers match as long as the field that decides the order is --- not null, which is what these queries check. - --- Config: spark.comet.expression.Min.allowIncompatible=true --- Config: spark.comet.expression.Max.allowIncompatible=true +-- min/max over structs with nullable fields. Spark orders structs field by field, with a null +-- field (or a null nested struct) below every value, and skips a null struct entirely. statement CREATE TABLE test_mm_struct_null(g int, ts bigint, s string, n string, inner_v int) USING parquet @@ -34,41 +29,48 @@ INSERT INTO test_mm_struct_null VALUES (2, 1, 'a', 'z', NULL), (2, 2, NULL, 'w', 2), (3, 4, NULL, NULL, NULL), - (4, NULL, NULL, NULL, NULL) + (4, NULL, NULL, NULL, NULL), + (5, NULL, 'z', 'p', 1), + (5, 1, 'a', 'q', 2), + (5, 1, NULL, 'r', 3), + (6, 3, 'a', NULL, NULL), + (6, 3, 'a', 'b', NULL) --- The deciding field ts is distinct and non-null within each group; later fields may be null. +-- Nulls in every position, including the deciding field: group 5 orders (NULL, ...) first and +-- breaks the tie on ts = 1 with the null s first. query SELECT g, min(struct(ts, s, n)), max(struct(ts, s, n)) -FROM test_mm_struct_null WHERE ts IS NOT NULL GROUP BY g ORDER BY g +FROM test_mm_struct_null GROUP BY g ORDER BY g --- Nested struct with nullable inner fields. +-- Nested struct with nullable inner fields, and a nested struct that is itself null. query -SELECT g, min(struct(ts, struct(s, inner_v) AS inner)), max(struct(ts, struct(s, inner_v) AS inner)) -FROM test_mm_struct_null WHERE ts IS NOT NULL GROUP BY g ORDER BY g +SELECT g, + min(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)), + max(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)) +FROM test_mm_struct_null GROUP BY g ORDER BY g --- Whole-struct nulls are skipped by both engines. Group 4 has only a NULL struct, so the result --- is NULL. +-- Whole-struct nulls are skipped. Group 4 has only a NULL struct, so the result is NULL. query SELECT g, max(CASE WHEN ts IS NOT NULL THEN struct(ts, s) END) FROM test_mm_struct_null GROUP BY g ORDER BY g --- Known difference: in group 1 Spark returns (7, 'a') because a null ts sorts first. DataFusion --- skips the null ts and compares the second field, where '~' is larger than 'a', so it returns --- (NULL, '~'). -query ignore(DataFusion skips null struct fields when comparing, Spark orders null first) -SELECT max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, - CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) +-- A null deciding field sorts first: Spark returns (7, 'a') for max and (NULL, '~') for min. +query +SELECT min(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)), + max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) FROM test_mm_struct_null WHERE g = 1 --- A struct column read from a table. Spark reads it with nullable fields, so it goes through the --- opt-in path and the cast that relabels its fields. +-- A struct column read from a table, whose fields Spark reads as nullable. statement CREATE TABLE test_mm_struct_null_col(g int, st struct) USING parquet statement INSERT INTO test_mm_struct_null_col VALUES (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', NULL)), - (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL) + (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL), + (3, named_struct('ts', NULL, 's', 'x')), (3, named_struct('ts', 2L, 's', NULL)) query SELECT g, min(st), max(st) FROM test_mm_struct_null_col GROUP BY g ORDER BY g diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 5c73433a552..deaeb7c2bd3 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -1416,6 +1416,43 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("struct min/max buffers cross engines in both directions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + // JVM columnar shuffle can carry Spark's partial rows to a native final. + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withTable("struct_min_max") { + sql("CREATE TABLE struct_min_max(g int, ts bigint, s string) USING parquet") + // Group 1 has a null deciding field, which sorts first on both engines. + sql( + "INSERT INTO struct_min_max VALUES (1, 5, 'b'), (1, 9, NULL), (1, NULL, 'z'), " + + "(2, 1, '日本語'), (2, 1, 'a'), (3, NULL, NULL)") + for { + (disabled, nativeMode) <- Seq( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE -> Partial, + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE -> Final) + } { + withSQLConf(disabled.key -> "false") { + val df = sql( + "SELECT g, min(struct(ts, s)), max(struct(ts, s)) FROM struct_min_max GROUP BY g") + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan + assert( + collect(plan) { case a: CometBaseAggregateExec => a.modes }.flatten == + Seq(nativeMode), + plan) + assert( + collect(plan) { case a: SortAggregateExec => + a.aggregateExpressions.map(_.mode).distinct + }.flatten == Seq(if (nativeMode == Partial) Final else Partial), + plan) + } + } + } + } + } + test("count, avg with null") { Seq(false, true).foreach { dictionary => withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { From f62550fc518a26a2a4d0dab8947ce7276fd27f09 Mon Sep 17 00:00:00 2001 From: Parth Chandra Date: Thu, 29 Oct 2026 09:55:17 -0700 Subject: [PATCH 3/3] address more review comments --- .../src/agg_funcs/min_max_struct.rs | 214 +++++++++++++++--- .../expressions/aggregate/min_max_struct.sql | 68 +++++- .../aggregate/min_max_struct_nullable.sql | 76 ------- .../CometAggregateExpressionBenchmark.scala | 41 ++++ 4 files changed, 278 insertions(+), 121 deletions(-) delete mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql diff --git a/native/spark-expr/src/agg_funcs/min_max_struct.rs b/native/spark-expr/src/agg_funcs/min_max_struct.rs index ad88640750c..d50caee0554 100644 --- a/native/spark-expr/src/agg_funcs/min_max_struct.rs +++ b/native/spark-expr/src/agg_funcs/min_max_struct.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{new_null_array, Array, ArrayRef, BooleanArray}; -use arrow::compute::{not, nullif, prep_null_mask_filter, SortOptions}; +use arrow::array::{ + make_comparator, new_null_array, Array, ArrayRef, BooleanArray, DynComparator, UInt32Array, +}; +use arrow::compute::{not, nullif, prep_null_mask_filter, take, SortOptions}; use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::row::{OwnedRow, RowConverter, Rows, SortField}; use datafusion::common::{internal_err, Result, ScalarValue}; @@ -25,17 +27,20 @@ use datafusion::logical_expr::{ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature, Volatility, }; use datafusion::physical_expr::expressions::format_state_name; +use std::cmp::Ordering; use std::fmt::Debug; use std::mem::{size_of, size_of_val}; use std::sync::Arc; /// Spark's `max` or `min` over a struct. /// -/// Spark orders structs field by field, with a null field below every value. Each input batch is -/// encoded once with Arrow's row format, ascending with nulls first, whose bytes sort in exactly -/// that order. Each group keeps the bytes of its current winner, and memory is accounted by the -/// size of those buffers. Float and double fields are not supported: the row format orders `-0.0` -/// below `0.0`, while Spark treats them as equal. +/// Spark orders structs field by field, with a null field below every value. Within a batch, each +/// group's best row is found with Arrow's struct comparator, which stops at the first field that +/// differs, so a wide trailing field that never decides costs nothing. Only those candidates are +/// encoded with Arrow's row format (ascending, nulls first, the same order), and each group keeps +/// the bytes of its current winner. Memory is accounted by the size of those buffers. Float and +/// double fields are not supported: both orders put `-0.0` below `0.0`, while Spark treats them +/// as equal. #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkStructMinMax { signature: Signature, @@ -98,6 +103,12 @@ impl AggregateUDFImpl for SparkStructMinMax { } } +/// Spark's struct order: ascending, with nulls (of the struct or of any field) first. +const SPARK_ORDER: SortOptions = SortOptions { + descending: false, + nulls_first: true, +}; + /// Encodes struct values to Spark-ordered row bytes and decodes them back. #[derive(Debug)] struct StructRows { @@ -112,13 +123,9 @@ impl StructRows { if !matches!(data_type, DataType::Struct(_)) { return internal_err!("Spark struct min/max expects a struct, got {data_type}"); } - let options = SortOptions { - descending: false, - nulls_first: true, - }; let converter = RowConverter::new(vec![SortField::new_with_options( data_type.clone(), - options, + SPARK_ORDER, )])?; let null_row = converter .convert_columns(&[new_null_array(data_type, 1)])? @@ -135,32 +142,42 @@ impl StructRows { Ok(self.converter.convert_columns(&[Arc::clone(values)])?) } - /// Whether `candidate` replaces `current`. Only a strictly better value does. + /// Compares rows of `values` with each other in Spark's struct order. + fn comparator(&self, values: &ArrayRef) -> Result { + Ok(make_comparator( + values.as_ref(), + values.as_ref(), + SPARK_ORDER, + )?) + } + + /// Whether a value that compares to the current one as `ordering` replaces it. Only a strictly + /// better value does, so of equal values the first one seen is kept. #[inline] - fn replaces(&self, candidate: &[u8], current: &[u8]) -> bool { + fn beats(&self, ordering: Ordering) -> bool { if self.is_max { - candidate > current + ordering == Ordering::Greater } else { - candidate < current + ordering == Ordering::Less } } - /// Keeps the better of `value` and `slot`, reusing the slot's buffer. Returns the change in - /// the slot's allocated bytes. + /// Keeps the better of `value` and `slot`, reusing the slot's buffer. Returns the bytes the + /// slot's buffer grew by, which is never negative since `clear` keeps the capacity. #[inline] - fn update(&self, slot: &mut Option>, value: &[u8]) -> isize { + fn update(&self, slot: &mut Option>, value: &[u8]) -> usize { match slot { None => { let bytes = value.to_vec(); - let added = bytes.capacity() as isize; + let added = bytes.capacity(); *slot = Some(bytes); added } - Some(current) if self.replaces(value, current) => { - let before = current.capacity() as isize; + Some(current) if self.beats(value.cmp(current.as_slice())) => { + let before = current.capacity(); current.clear(); current.extend_from_slice(value); - current.capacity() as isize - before + current.capacity() - before } Some(_) => 0, } @@ -207,12 +224,20 @@ impl StructMinMaxAccumulator { impl Accumulator for StructMinMaxAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let values = &values[0]; - let encoded = self.rows.encode(values)?; + let cmp = self.rows.comparator(values)?; + let mut best: Option = None; for i in 0..values.len() { - if !is_skipped(values, None, i) { - self.rows.update(&mut self.value, encoded.row(i).as_ref()); + if is_skipped(values, None, i) { + continue; + } + if best.is_none_or(|b| self.rows.beats(cmp(i, b))) { + best = Some(i); } } + if let Some(best) = best { + let encoded = self.rows.encode(&values.slice(best, 1))?; + self.rows.update(&mut self.value, encoded.row(0).as_ref()); + } Ok(()) } @@ -243,14 +268,22 @@ struct StructMinMaxGroupsAccumulator { values: Vec>>, /// The bytes allocated by the buffers in `values`. value_bytes: usize, + /// Scratch space for `update_batch`: each group's best row in the current batch, or + /// `NO_ROW`, and the groups that have one. Reset after every batch. + batch_best: Vec, + touched: Vec, } +const NO_ROW: u32 = u32::MAX; + impl StructMinMaxGroupsAccumulator { fn try_new(data_type: &DataType, is_max: bool) -> Result { Ok(Self { rows: StructRows::try_new(data_type, is_max)?, values: vec![], value_bytes: 0, + batch_best: vec![], + touched: vec![], }) } } @@ -265,16 +298,40 @@ impl GroupsAccumulator for StructMinMaxGroupsAccumulator { ) -> Result<()> { let values = &values[0]; self.values.resize(total_num_groups, None); - let encoded = self.rows.encode(values)?; - let mut delta = 0isize; + self.batch_best.resize(total_num_groups, NO_ROW); + + // Find each group's best row in this batch without encoding anything. + let cmp = self.rows.comparator(values)?; for (i, &group) in group_indices.iter().enumerate() { - if !is_skipped(values, opt_filter, i) { - delta += self - .rows - .update(&mut self.values[group], encoded.row(i).as_ref()); + if is_skipped(values, opt_filter, i) { + continue; } + let best = &mut self.batch_best[group]; + if *best == NO_ROW { + *best = i as u32; + self.touched.push(group); + } else if self.rows.beats(cmp(i, *best as usize)) { + *best = i as u32; + } + } + if self.touched.is_empty() { + return Ok(()); + } + + // Encode only those rows, and keep each one that beats its group's stored winner. + let indices: UInt32Array = self + .touched + .iter() + .map(|&group| self.batch_best[group]) + .collect(); + let encoded = self.rows.encode(&take(values.as_ref(), &indices, None)?)?; + for (k, &group) in self.touched.iter().enumerate() { + self.value_bytes += self + .rows + .update(&mut self.values[group], encoded.row(k).as_ref()); + self.batch_best[group] = NO_ROW; } - self.value_bytes = self.value_bytes.wrapping_add_signed(delta); + self.touched.clear(); Ok(()) } @@ -318,6 +375,8 @@ impl GroupsAccumulator for StructMinMaxGroupsAccumulator { + self.rows.size() + self.values.capacity() * size_of::>>() + self.value_bytes + + self.batch_best.capacity() * size_of::() + + self.touched.capacity() * size_of::() } } @@ -487,6 +546,95 @@ mod tests { Ok(()) } + /// Nested structs with nulls at every level, used to check the two orders agree. + fn nested_structs(n: usize) -> ArrayRef { + // Small value ranges so that ties on the leading fields are common. + let k: Int64Array = (0..n) + .map(|i| (i % 7 != 0).then_some((i * 7919 % 5) as i64 - 2)) + .collect(); + let s: StringArray = (0..n) + .map(|i| (i % 5 != 0).then(|| ["", "a", "ab", "日本", "\u{7f}"][i * 31 % 5])) + .collect(); + let inner_v: Int64Array = (0..n) + .map(|i| (i % 3 != 0).then_some((i * 104729 % 3) as i64)) + .collect(); + let inner_fields = Fields::from(vec![Field::new("v", DataType::Int64, true)]); + let inner_validity: Vec = (0..n).map(|i| i % 11 != 0).collect(); + let inner = StructArray::new( + inner_fields.clone(), + vec![Arc::new(inner_v)], + Some(inner_validity.into()), + ); + let outer_fields = Fields::from(vec![ + Field::new("k", DataType::Int64, true), + Field::new("s", DataType::Utf8, true), + Field::new("inner", DataType::Struct(inner_fields), true), + ]); + let outer_validity: Vec = (0..n).map(|i| i % 13 != 0).collect(); + Arc::new(StructArray::new( + outer_fields, + vec![Arc::new(k), Arc::new(s), Arc::new(inner)], + Some(outer_validity.into()), + )) + } + + /// Each batch picks candidates with Arrow's comparator, and winners are kept as row bytes. + /// The two must order every pair of values the same way. + #[test] + fn comparator_and_row_bytes_agree() -> Result<()> { + let values = nested_structs(200); + let rows = StructRows::try_new(values.data_type(), true)?; + let cmp = rows.comparator(&values)?; + let encoded = rows.encode(&values)?; + for i in 0..values.len() { + for j in 0..values.len() { + assert_eq!( + cmp(i, j), + encoded.row(i).cmp(&encoded.row(j)), + "rows {i} and {j}" + ); + } + } + Ok(()) + } + + /// Many rows per group in each batch, across several batches, against a brute-force answer. + #[test] + fn matches_brute_force_over_batches() -> Result<()> { + let values = nested_structs(600); + let data_type = values.data_type().clone(); + let groups: Vec = (0..values.len()).map(|i| i * 17 % 9).collect(); + let reference = StructRows::try_new(&data_type, true)?; + let encoded = reference.encode(&values)?; + for is_max in [true, false] { + let mut acc = StructMinMaxGroupsAccumulator::try_new(&data_type, is_max)?; + for start in (0..values.len()).step_by(100) { + acc.update_batch( + &[values.slice(start, 100)], + &groups[start..start + 100], + None, + 9, + )?; + } + let result = acc.evaluate(EmitTo::All)?; + let result_rows = reference.encode(&result)?; + for group in 0..9 { + let members = + (0..values.len()).filter(|&i| groups[i] == group && values.is_valid(i)); + let expected = if is_max { + members.max_by(|&a, &b| encoded.row(a).cmp(&encoded.row(b))) + } else { + members.min_by(|&a, &b| encoded.row(a).cmp(&encoded.row(b))) + }; + match expected { + Some(i) => assert_eq!(result_rows.row(group), encoded.row(i), "group {group}"), + None => assert!(result.is_null(group), "group {group}"), + } + } + } + Ok(()) + } + /// Sorted aggregation emits finished groups after every batch. Struct max over 4,096 groups /// of 4 KiB values (16 MiB in all), read in 64-row batches, must fit a 2 MiB pool, which it /// only does if emitting a group releases the memory charged for it. DataFusion's own struct diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql index 9818b0b2938..d5ffc6334cf 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql @@ -15,8 +15,9 @@ -- specific language governing permissions and limitations -- under the License. --- min/max over a struct. Spark orders structs field by field, with a null field below every --- value. Float and double fields fall back. Null fields are covered in min_max_struct_nullable.sql. +-- min/max over a struct. Spark orders structs field by field, with a null field (or a null +-- nested struct) below every value, and skips a null struct entirely. Float and double fields +-- fall back. statement CREATE TABLE test_mm_struct(g int, ts bigint, s string, d decimal(38,0), dt date, t timestamp, n string) @@ -51,11 +52,8 @@ FROM test_mm_struct GROUP BY g ORDER BY g query SELECT min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct --- Nullable fields, including a null in the deciding field. -query -SELECT g, max(struct(ts, s, n)) FROM test_mm_struct GROUP BY g ORDER BY g - --- Float and double fields: NaN ordering differs, so fall back. +-- Float and double fields fall back: Comet orders -0.0 below 0.0, and a NaN with the sign bit +-- set below every value. query expect_fallback(float or double field) SELECT g, max(struct(coalesce(ts, 0L) AS ts, 1.5D AS d)) FROM test_mm_struct GROUP BY g ORDER BY g @@ -118,14 +116,60 @@ query SELECT g, max(struct(coalesce(ts, 0L) AS ts, CASE WHEN ts > 5 THEN struct(coalesce(s, '') AS s) END AS inner)) FROM test_mm_struct GROUP BY g ORDER BY g +-- Null fields at every level. +statement +CREATE TABLE test_mm_struct_null(g int, ts bigint, s string, n string, inner_v int) USING parquet + +statement +INSERT INTO test_mm_struct_null VALUES + (1, 5, 'b', 'x', 1), + (1, 9, NULL, NULL, NULL), + (1, 7, '日本語', NULL, 3), + (2, 1, 'a', 'z', NULL), + (2, 2, NULL, 'w', 2), + (3, 4, NULL, NULL, NULL), + (4, NULL, NULL, NULL, NULL), + (5, NULL, 'z', 'p', 1), + (5, 1, 'a', 'q', 2), + (5, 1, NULL, 'r', 3), + (6, 3, 'a', NULL, NULL), + (6, 3, 'a', 'b', NULL) + +-- Nulls in every position, including the deciding field: group 5 orders (NULL, ...) first and +-- breaks the tie on ts = 1 with the null s first. +query +SELECT g, min(struct(ts, s, n)), max(struct(ts, s, n)) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- Nested struct with nullable inner fields, and a nested struct that is itself null. +query +SELECT g, + min(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)), + max(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- Whole-struct nulls are skipped. Group 4 has only a NULL struct, so the result is NULL. +query +SELECT g, max(CASE WHEN ts IS NOT NULL THEN struct(ts, s) END) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- A null deciding field sorts first: Spark returns (7, 'a') for max and (NULL, '~') for min. +query +SELECT min(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)), + max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) +FROM test_mm_struct_null WHERE g = 1 + -- A struct column read from a table, whose fields Spark reads as nullable. statement -CREATE TABLE test_mm_struct_col(g int, st struct) USING parquet +CREATE TABLE test_mm_struct_null_col(g int, st struct) USING parquet statement -INSERT INTO test_mm_struct_col VALUES - (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', '日本語')), - (2, named_struct('ts', 1L, 's', 'a')), (2, NULL) +INSERT INTO test_mm_struct_null_col VALUES + (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', NULL)), + (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL), + (3, named_struct('ts', NULL, 's', 'x')), (3, named_struct('ts', 2L, 's', NULL)) query -SELECT g, min(st), max(st) FROM test_mm_struct_col GROUP BY g ORDER BY g +SELECT g, min(st), max(st) FROM test_mm_struct_null_col GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql deleted file mode 100644 index 2d7247d8cfd..00000000000 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_nullable.sql +++ /dev/null @@ -1,76 +0,0 @@ --- 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. - --- min/max over structs with nullable fields. Spark orders structs field by field, with a null --- field (or a null nested struct) below every value, and skips a null struct entirely. - -statement -CREATE TABLE test_mm_struct_null(g int, ts bigint, s string, n string, inner_v int) USING parquet - -statement -INSERT INTO test_mm_struct_null VALUES - (1, 5, 'b', 'x', 1), - (1, 9, NULL, NULL, NULL), - (1, 7, '日本語', NULL, 3), - (2, 1, 'a', 'z', NULL), - (2, 2, NULL, 'w', 2), - (3, 4, NULL, NULL, NULL), - (4, NULL, NULL, NULL, NULL), - (5, NULL, 'z', 'p', 1), - (5, 1, 'a', 'q', 2), - (5, 1, NULL, 'r', 3), - (6, 3, 'a', NULL, NULL), - (6, 3, 'a', 'b', NULL) - --- Nulls in every position, including the deciding field: group 5 orders (NULL, ...) first and --- breaks the tie on ts = 1 with the null s first. -query -SELECT g, min(struct(ts, s, n)), max(struct(ts, s, n)) -FROM test_mm_struct_null GROUP BY g ORDER BY g - --- Nested struct with nullable inner fields, and a nested struct that is itself null. -query -SELECT g, - min(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)), - max(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)) -FROM test_mm_struct_null GROUP BY g ORDER BY g - --- Whole-struct nulls are skipped. Group 4 has only a NULL struct, so the result is NULL. -query -SELECT g, max(CASE WHEN ts IS NOT NULL THEN struct(ts, s) END) -FROM test_mm_struct_null GROUP BY g ORDER BY g - --- A null deciding field sorts first: Spark returns (7, 'a') for max and (NULL, '~') for min. -query -SELECT min(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, - CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)), - max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, - CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) -FROM test_mm_struct_null WHERE g = 1 - --- A struct column read from a table, whose fields Spark reads as nullable. -statement -CREATE TABLE test_mm_struct_null_col(g int, st struct) USING parquet - -statement -INSERT INTO test_mm_struct_null_col VALUES - (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', NULL)), - (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL), - (3, named_struct('ts', NULL, 's', 'x')), (3, named_struct('ts', 2L, 's', NULL)) - -query -SELECT g, min(st), max(st) FROM test_mm_struct_null_col GROUP BY g ORDER BY g diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala index 950ebaab7aa..be98e2dca9b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala @@ -199,6 +199,24 @@ object CometAggregateExpressionBenchmark extends CometBenchmarkBase { "min_by_high_card", "SELECT min_by(c_int, c_long) FROM parquetV1Table GROUP BY high_card_grp")) + // min / max over a struct of an ordering key and a payload. The key decides every comparison, + // so the narrow (64 byte) and wide (4 KiB) cases differ only in what has to be carried along. + private val structMinMaxAggregates = List( + AggExprConfig( + "max_struct_narrow", + "SELECT MAX(struct(k, narrow)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "max_struct_wide", + "SELECT MAX(struct(k, wide)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "min_struct_wide", + "SELECT MIN(struct(k, wide)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "max_struct_narrow_global", + "SELECT MAX(struct(k, narrow)) FROM parquetV1Table"), + AggExprConfig("max_struct_wide_global", "SELECT MAX(struct(k, wide)) FROM parquetV1Table"), + AggExprConfig("min_struct_wide_global", "SELECT MIN(struct(k, wide)) FROM parquetV1Table")) + override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024 @@ -231,5 +249,28 @@ object CometAggregateExpressionBenchmark extends CometBenchmarkBase { } } } + + // Fewer rows than above, since each wide payload is 4 KiB. + runBenchmarkWithTable("Struct min/max benchmarks", 200 * 1000) { v => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s""" + SELECT + CAST(value % 1000 AS INT) AS grp, + CAST(value AS LONG) AS k, + CASE WHEN value % 10 = 0 THEN NULL ELSE sha2(CAST(value AS STRING), 256) END AS narrow, + CASE WHEN value % 10 = 0 THEN NULL + ELSE repeat(sha2(CAST(value AS STRING), 256), 64) END AS wide + FROM $tbl + """)) + + structMinMaxAggregates.foreach { config => + runExpressionBenchmark(config.name, v.toLong, config.query, config.extraCometConfigs) + } + } + } + } } }