From 88a5d406b65fb20f2ee01c1d7bafbf64bed39ace Mon Sep 17 00:00:00 2001 From: rich7420 Date: Thu, 24 Sep 2026 20:37:56 +0800 Subject: [PATCH 01/89] feat: support native wide-decimal hashing (WIP) --- .../expression-audits/hash_funcs.md | 4 +- docs/source/user-guide/latest/expressions.md | 4 +- native/spark-expr/src/hash_funcs/utils.rs | 118 ++++++++++++++---- native/spark-expr/src/hash_funcs/xxhash64.rs | 4 +- .../src/hash_funcs/xxhash64_diff.rs | 43 ++++++- .../scala/org/apache/comet/serde/hash.scala | 7 +- .../shuffle/CometShuffleExchangeExec.scala | 4 - .../comet/CometHashExpressionSuite.scala | 71 +++++++---- .../comet/exec/CometNativeShuffleSuite.scala | 21 ++++ .../CometHashExpressionBenchmark.scala | 63 ++++++++++ 10 files changed, 274 insertions(+), 65 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index 989d1d20668..14336cc290a 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -34,7 +34,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `Murmur3Hash(children, seed) extends HashExpression[Int]`; produces a Murmur3 hash with a configurable Int seed and `IntegerType` result. Comet routes via `CometMurmur3Hash` to the native `murmur3_hash` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Known limitation: `DecimalType` children with precision > 18 fall back because Spark hashes them through Java `BigDecimal`; `TimeType` (Spark 4.0+) is also unsupported. The same limitations apply to `xxhash64`, `sha1`, `sha2` through the shared `HashUtils`. +- Decimals with precision 19–38, including nested decimals, hash natively using the minimal signed big-endian bytes of the unscaled value, matching Spark's `BigInteger.toByteArray()`. Precision ≤18 retains unscaled-long hashing. `TimeType` remains unsupported through the shared `HashUtils`. ## md5 @@ -70,6 +70,6 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Upstream (2026-09-15): Comet's native `xxhash64` UDF delegates compatible arguments at Spark's default seed (`42`) to `datafusion-spark::SparkXxhash64`. Differential tests in `native/spark-expr/src/hash_funcs/xxhash64_diff.rs` compare Comet's kernel against `SparkXxhash64` for primitives, both Decimal128 widths, dictionaries, lists, maps, and nested combinations. The Comet kernel is retained for a non-default seed (`SparkXxhash64` hardcodes 42), `Struct` (upstream does not push a parent null mask into children; see #5753), a `Dictionary` nested in a list/map (upstream restarts those hashes from 42), and `Time64`. `murmur3` is unchanged; `create_xxhash64_hashes` remains for `approx_count_distinct` and the fallback path. +- Native routing: Comet's native `xxhash64` UDF delegates compatible arguments at Spark's default seed (`42`) to `datafusion-spark::SparkXxhash64`. Differential tests in `native/spark-expr/src/hash_funcs/xxhash64_diff.rs` compare Comet's kernel against `SparkXxhash64` for primitives, narrow decimals, dictionaries, lists, maps, and nested combinations. Wide-decimal tests verify that the expression stays on Comet's kernel: the pinned DataFusion implementation uses fixed-width little-endian bytes rather than Spark's minimal signed big-endian bytes. The Comet kernel is retained for wide decimals (also within nested types), a non-default seed (`SparkXxhash64` hardcodes 42), `Struct` (upstream does not push a parent null mask into children; see #5753), a `Dictionary` nested in a list/map (upstream restarts those hashes from 42), and `Time64`. `create_xxhash64_hashes` remains for `approx_count_distinct` and the non-delegated native path. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index cb28013a77c..b15abb9ab77 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -357,12 +357,12 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `crc32` | ✅ | Native | | -| `hash` | ✅ | Native | | +| `hash` | ✅ | Native | Includes decimals with precision 19–38 and nested decimals | | `md5` | ✅ | Native | | | `sha` | ✅ | Native | | | `sha1` | ✅ | Native | | | `sha2` | ✅ | Native | | -| `xxhash64` | ✅ | Native | | +| `xxhash64` | ✅ | Native | Includes decimals with precision 19–38 and nested decimals | --- diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index 55beb6cd378..b11e065025f 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -237,30 +237,15 @@ macro_rules! hash_array_small_decimal { #[macro_export] macro_rules! hash_array_decimal { ($array_type:ident, $column: ident, $hashes: ident, $hash_method: ident) => { - let array = $column - .as_any() - .downcast_ref::<$array_type>() - .unwrap_or_else(|| { - panic!( - "Failed to downcast column to {}. Actual data type: {:?}.", - stringify!($array_type), - $column.data_type() - ) - }); - - if array.null_count() == 0 { - // Fast path: no nulls, use direct indexing - for i in 0..$hashes.len() { - $hashes[i] = $hash_method(array.value(i).to_le_bytes(), $hashes[i]); - } - } else { - // Slow path: check nulls - for i in 0..$hashes.len() { - if !array.is_null(i) { - $hashes[i] = $hash_method(array.value(i).to_le_bytes(), $hashes[i]); - } - } - } + let hash_decimal = |value: &i128, seed| { + let bytes = value.to_be_bytes(); + // Match BigInteger.toByteArray(): drop redundant sign bytes, retaining + // one sign bit. Zero and -1 still occupy one byte. + let sign_bits = value.leading_zeros().max(value.leading_ones()); + let start = (sign_bits.saturating_sub(1) / 8) as usize; + $hash_method(&bytes[start..], seed) + }; + $crate::hash_array!($array_type, $column, $hashes, hash_decimal); }; } @@ -1323,3 +1308,88 @@ pub(crate) mod test_utils { }; } } + +#[cfg(test)] +mod tests { + use super::super::{spark_murmur3_hash, spark_xxhash64}; + use arrow::array::{Array, ArrayRef, BinaryArray, Decimal128Array, Int32Array}; + use datafusion::common::{Result, ScalarValue}; + use datafusion::logical_expr::ColumnarValue; + use num::BigInt; + use std::sync::Arc; + + type HashFunction = fn(&[ColumnarValue]) -> Result; + + #[test] + fn wide_decimal_hashes_match_signed_big_endian_bytes() { + for precision in 19..=38 { + let max = 10_i128.pow(precision as u32) - 1; + let mut values = vec![Some(0), Some(max), Some(-max), None]; + // Both sides of every signed byte-length boundary, including 127/128 + // and -128/-129, as well as values that fit in i64 despite the type. + for bit in (7..127).step_by(8) { + for value in [(1_i128 << bit) - 1, 1_i128 << bit, (1_i128 << bit) + 1] { + if value <= max { + values.extend([Some(value), Some(-value)]); + } + } + } + let bytes: Vec<_> = values + .iter() + .map(|v| v.map(|v| BigInt::from(v).to_signed_bytes_be())) + .collect(); + let binary: ArrayRef = + Arc::new(BinaryArray::from_iter(bytes.iter().map(|v| v.as_deref()))); + for scale in [0, 2, precision as i8] { + let decimal: ArrayRef = Arc::new( + Decimal128Array::from(values.clone()) + .with_precision_and_scale(precision, scale) + .unwrap(), + ); + // Exercise both null-aware and null-free loops, including a slice. + for (offset, len) in [(0, values.len()), (4, values.len() - 4)] { + let decimal = decimal.slice(offset, len); + let binary = binary.slice(offset, len); + let prefix: ArrayRef = Arc::new(Int32Array::from_iter_values(0..len as i32)); + for seed in [0_i64, 42, -1, i32::MIN as i64] { + for chained in [false, true] { + let args = |value: ArrayRef, seed: ScalarValue| { + let mut args = vec![ColumnarValue::Array(value)]; + if chained { + args.insert(0, ColumnarValue::Array(Arc::clone(&prefix))); + args.push(ColumnarValue::Array(Arc::clone(&prefix))); + } + args.push(ColumnarValue::Scalar(seed)); + args + }; + for (hash, seed) in [ + ( + spark_murmur3_hash as HashFunction, + ScalarValue::Int32(Some(seed as i32)), + ), + ( + spark_xxhash64 as HashFunction, + ScalarValue::Int64(Some(seed)), + ), + ] { + let actual = hash(&args(Arc::clone(&decimal), seed.clone())) + .unwrap() + .into_array(len) + .unwrap(); + let expected = hash(&args(Arc::clone(&binary), seed)) + .unwrap() + .into_array(len) + .unwrap(); + assert_eq!( + actual.to_data(), + expected.to_data(), + "precision={precision}, scale={scale}, chained={chained}" + ); + } + } + } + } + } + } + } +} diff --git a/native/spark-expr/src/hash_funcs/xxhash64.rs b/native/spark-expr/src/hash_funcs/xxhash64.rs index bdda7c85f1d..5ed08933a74 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64.rs @@ -47,6 +47,8 @@ const SPARK_DEFAULT_SEED: i64 = 42; /// - `Struct` (and anything containing one): `SparkXxhash64` does not push a parent null /// mask into children, so hidden values of a NULL struct would affect the hash /// - a `Dictionary` nested in a list/map: `SparkXxhash64` restarts those hashes from 42 +/// - wide decimals: `SparkXxhash64` hashes fixed-width little-endian bytes instead +/// of Spark's minimal signed big-endian representation /// - `Time64`, which `SparkXxhash64` does not dispatch pub fn spark_xxhash64(args: &[ColumnarValue]) -> Result { let length = args.len(); @@ -129,7 +131,7 @@ fn type_compatible_with_spark_xxhash64(dt: &DataType, in_list_or_map: bool) -> b Boolean | Int8 | Int16 | Int32 | Int64 | Float32 | Float64 => true, Utf8 | LargeUtf8 | Binary | LargeBinary | FixedSizeBinary(_) => true, Date32 | Date64 | Timestamp(_, _) => true, - Decimal128(_, _) => true, + Decimal128(precision, _) => *precision <= 18, Dictionary(_, value) if !in_list_or_map => { type_compatible_with_spark_xxhash64(value.as_ref(), true) } diff --git a/native/spark-expr/src/hash_funcs/xxhash64_diff.rs b/native/spark-expr/src/hash_funcs/xxhash64_diff.rs index 3c51a0e4862..73f0239babe 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64_diff.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64_diff.rs @@ -29,11 +29,11 @@ //! | Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary | yes | //! | Date32, Date64, Timestamp | yes | //! | Decimal128 precision ≤ 18 | yes | -//! | Decimal128 precision > 18 | yes | +//! | Decimal128 precision > 18 | **no** — upstream uses fixed-width little-endian bytes | //! | Dictionary (top-level) | yes | //! | List / LargeList / FixedSizeList of primitives | yes | //! | Map<Utf8, Int32> / Map<Int32, Utf8> / Map<Utf8, Utf8> / Map<Int32, Int32> | yes | -//! | Map<Utf8, Decimal128> | yes | +//! | Map<Utf8, Decimal128> | only for precision ≤ 18 | //! | Struct (non-null parent) | yes, but not routed (see below) | //! | Struct NULL with hidden children | **no** — `SparkXxhash64` hashes hidden children | //! | List<Dictionary> | **no** — upstream restarts from seed 42 | @@ -124,6 +124,18 @@ fn assert_compatible(label: &str, arrays: &[ArrayRef]) { assert_eq!(expr, upstream, "{label}: expression mismatch"); } +// The pinned DataFusion implementation still uses the old decimal encoding. +// Spark parity is checked separately; routing must not reintroduce that encoding. +fn assert_wide_decimal_uses_comet(label: &str, arrays: &[ArrayRef]) { + let kernel = comet_kernel(arrays, SPARK_DEFAULT_SEED).unwrap(); + let upstream = spark_xxhash64_upstream(arrays).unwrap(); + assert_ne!( + kernel, upstream, + "{label}: expected different decimal encodings" + ); + assert_eq!(comet_expr(arrays, 42).unwrap(), kernel, "{label}: routing"); +} + fn col(array: impl Array + 'static) -> Vec { vec![Arc::new(array) as ArrayRef] } @@ -584,7 +596,7 @@ fn decimal128_precision_18_fits_i64() { #[test] fn decimal128_precision_20_does_not_fit_i64() { let too_big = 10_000_000_000_000_000_000i128; // 1e19, beyond i64::MAX - assert_compatible( + assert_wide_decimal_uses_comet( "Decimal128(20,2)", &col(decimal128( 20, @@ -604,7 +616,7 @@ fn decimal128_precision_20_does_not_fit_i64() { #[test] fn decimal128_precision_38() { let wide = 10i128.pow(28); - assert_compatible( + assert_wide_decimal_uses_comet( "Decimal128(38,10)", &col(decimal128( 38, @@ -621,6 +633,27 @@ fn decimal128_precision_38() { ); } +#[test] +fn wide_decimal_list_and_dictionary_use_comet() { + let values: ArrayRef = Arc::new(decimal128(38, 0, vec![Some(128), Some(-129), None])); + let dict: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(1), Some(0), Some(2), None]), + Arc::clone(&values), + ) + .unwrap(), + ); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 3].into()), + values, + None, + )); + for array in [dict, list] { + assert_wide_decimal_uses_comet("nested wide decimal", &[array]); + } +} + // --------------------------------------------------------------------------- // Dictionary // --------------------------------------------------------------------------- @@ -923,7 +956,7 @@ fn map_utf8_decimal128_small() { #[test] fn map_utf8_decimal128_large() { let wide = 10_000_000_000_000_000_000i128; - assert_compatible( + assert_wide_decimal_uses_comet( "Map", &[map_utf8_decimal( 20, diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index ee3e80059d5..04cd184fe41 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -20,7 +20,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} -import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructType} import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} @@ -116,12 +116,10 @@ object CometSha1 extends CometExpressionSerde[Sha1] { private object HashUtils { - private val unsupportedDecimalReason = - "`DecimalType` with precision > 18 is not supported (Spark hashes via Java `BigDecimal`)" private val unsupportedTimeTypeReason = "`TimeType` is not supported" val unsupportedReasons: Seq[String] = - Seq(unsupportedDecimalReason, unsupportedTimeTypeReason, "Unsupported child data type") + Seq(unsupportedTimeTypeReason, "Unsupported child data type") def supportLevelForChildren(expr: Expression): SupportLevel = { expr.children.iterator @@ -134,7 +132,6 @@ private object HashUtils { } private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { - case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason) case s: StructType => s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption case a: ArrayType => unsupportedReasonFor(a.elementType) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 7a0d231a141..5c82ffd6688 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -556,10 +556,6 @@ object CometShuffleExchangeExec _: TimestampNTZType | _: DateType => true case _: DecimalType => - // TODO enforce this check - // https://github.com/apache/datafusion-comet/issues/3079 - // Decimals with precision > 18 require Java BigDecimal conversion before hashing - // d.precision <= 18 true case dt if isTimeType(dt) => true diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index bbf6c4289ad..a7b9250477d 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** * Test suite for Spark `hash` (murmur3) and `xxhash64` compatibility between Spark and Comet. @@ -35,6 +36,17 @@ import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGener */ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { + private def checkNativeHashes(query: String): Unit = { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + CometScalaUDFCodegen.resetStats() + checkSparkAnswerAndImpl(sql(query), native = Seq("hash", "xxhash64")) + val stats = CometScalaUDFCodegen.stats() + assert( + stats.compileCount + stats.cacheHitCount == 0, + s"unexpected dispatcher activity: $stats") + } + } + test("hash - boolean") { withTable("t") { sql("CREATE TABLE t(c BOOLEAN) USING parquet") @@ -134,49 +146,64 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") + checkNativeHashes("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") } } } test("hash - decimal (precision > 18)") { - Seq((20, 2), (38, 10)).foreach { case (precision, scale) => + Seq((19, 0), (20, 2), (38, 0), (38, 10), (38, 38)).foreach { case (precision, scale) => withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") - sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - // Large decimals may fall back to Spark, so just check the answer - checkSparkAnswer("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") + val max = BigInt(10).pow(precision) - 1 + val unscaled = Seq(BigInt(0), max, -max) ++ + Seq(1, 127, 128, 129, 255, 256, 32767, 32768).flatMap(v => + Seq(BigInt(v), -BigInt(v))) ++ + Seq( + BigInt(Long.MaxValue), + BigInt(Long.MaxValue) + 1, + BigInt(Long.MinValue), + BigInt(Long.MinValue) - 1) + val values = unscaled.map { v => + val decimal = new java.math.BigDecimal(v.bigInteger, scale).toPlainString + s"(CAST('$decimal' AS DECIMAL($precision, $scale)))" + } + sql(s"INSERT INTO t VALUES ${(values :+ "(null)").mkString(",")}") + checkNativeHashes("SELECT c, hash(c), xxhash64(c) FROM t") + checkNativeHashes("SELECT hash(1, c, c), xxhash64(1, c, c) FROM t") } } } - test("hash - array of decimal (precision > 18) falls back to Spark") { + test("hash - array of decimal (precision > 18)") { withTable("t") { - sql("CREATE TABLE t(c ARRAY) USING parquet") - sql("INSERT INTO t VALUES (array(1.23, 2.34)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c), xxhash64(c) FROM t", "precision > 18") + sql("CREATE TABLE t(c ARRAY) USING parquet") + sql("""INSERT INTO t VALUES (array(127, 128, -128, -129)), + (array(99999999999999999999999999999999999999BD, null, -1)), + (array()), (array(null)), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") } } - test("hash - struct with decimal (precision > 18) falls back to Spark") { + test("hash - struct with decimal (precision > 18)") { withTable("t") { - sql("CREATE TABLE t(c STRUCT) USING parquet") - sql("INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c), xxhash64(c) FROM t", "precision > 18") + sql("CREATE TABLE t(c STRUCT) USING parquet") + sql("""INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.28)), + (named_struct('a', 2, 'b', -1.29)), (named_struct('a', null, 'b', null)), + (named_struct('a', 3, 'b', 999999999999999999999999999999999999.99BD)), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") + checkNativeHashes("SELECT hash(array(c, c)), xxhash64(array(c, c)) FROM t") } } - test("hash - map with decimal (precision > 18) value falls back to Spark") { + test("hash - map with decimal (precision > 18) keys and values") { withSQLConf("spark.sql.legacy.allowHashOnMapType" -> "true") { withTable("t") { - sql("CREATE TABLE t(c MAP) USING parquet") - sql("INSERT INTO t VALUES (map('a', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason( - "SELECT c, hash(c), xxhash64(c) FROM t", - "precision > 18") + sql("CREATE TABLE t(c MAP) USING parquet") + sql("""INSERT INTO t VALUES (map(128, -129, -128, 127)), + (map(99999999999999999999999999999999999999BD, null)), + (map()), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index cbacf4368a4..052a6cb23c8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -771,6 +771,27 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("native shuffle on wide decimal hash partitioning keys") { + withNestedHashPartitioning { + withTable("wide_decimals") { + sql("CREATE TABLE wide_decimals(id INT, c DECIMAL(38, 0)) USING parquet") + sql("""INSERT INTO wide_decimals VALUES (0, null), (1, 0), (2, 128), (3, -129), + (4, 99999999999999999999999999999999999999BD), + (5, -99999999999999999999999999999999999999BD), (6, 128)""") + Seq(Seq("c"), Seq("a"), Seq("s"), Seq("id", "c")).foreach { keys => + def shuffled = sql("""SELECT id, c, array(c, c) AS a, + named_struct('d', c) AS s FROM wide_decimals""") + .repartition(10, keys.map(col): _*) + checkCometExchange(shuffled, 1, native = true) + // This pins the common hash helper's shuffle caller, including chained and nested + // inputs. The old fixed-width encoding produced different partition assignments. + checkSparkAnswer(shuffled.selectExpr("id", "spark_partition_id()")) + checkSparkAnswer(shuffled.groupBy("c").count()) + } + } + } + } + test("native shuffle on struct hash partitioning key") { withNestedHashPartitioning { Seq(10, 201).foreach { numPartitions => diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala index 238d8b662b1..d1d3e206a4e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala @@ -19,6 +19,11 @@ package org.apache.spark.sql.benchmark +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.Row + +import org.apache.comet.{CometConf, ExtendedExplainInfo} + case class HashExprConfig( name: String, query: String, @@ -46,6 +51,10 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024 + if (mainArgs.contains("--wide-decimal")) { + runWideDecimalBenchmarks(values) + return + } runBenchmarkWithTable("Hash expression benchmarks", values) { v => withTempPath { dir => @@ -74,6 +83,60 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { runMurmur3HashBenchmarks(values) } + // Run only these cases with --wide-decimal. The fallback arm is the path used by + // wide-decimal hash on main before native support; codegen dispatch is tracked in #5835. + private def runWideDecimalBenchmarks(values: Int): Unit = { + runBenchmarkWithTable("Wide decimal hash", values) { rows => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s""" + SELECT d, array(d, d) AS a FROM ( + SELECT CASE WHEN value % 100 = 0 THEN NULL + WHEN value % 2 = 0 THEN CAST(value AS DECIMAL(38, 0)) + ELSE CAST(value AS DECIMAL(38, 0)) * 1000000000000000000BD + END AS d FROM $tbl) + """)) + for (function <- Seq("hash", "xxhash64"); column <- Seq("d", "a")) { + val exprName = if (function == "hash") "Murmur3Hash" else "XxHash64" + // A checksum consumes every hash result and avoids collecting a million rows. + val query = + s"SELECT sum(CAST($function($column) AS DECIMAL(38, 0))) FROM parquetV1Table" + val cases = Seq( + ("Spark", false, false), + ("Comet with Spark hash fallback", true, false), + ("Comet native", true, true)) + var expected: Seq[Row] = Seq.empty + val benchmark = + new Benchmark(s"$function($column): decimal(38,0)", rows, output = output) + cases.foreach { case (name, comet, native) => + val conf = Seq( + CometConf.COMET_ENABLED.key -> comet.toString, + CometConf.COMET_EXEC_ENABLED.key -> comet.toString, + CometConf.getExprEnabledConfigKey(exprName) -> native.toString) + withSQLConf(conf: _*) { + val df = spark.sql(query) + val actual = df.collect().toSeq + if (!comet) expected = actual else assert(actual == expected, name) + if (comet) { + val plan = df.queryExecution.executedPlan + val info = new ExtendedExplainInfo() + assert(info.getNativeExpressions(plan).contains(function) == native, name) + assert(!info.getCodegenDispatchExpressions(plan).contains(function), name) + } + } + benchmark.addCase(name) { _ => + withSQLConf(conf: _*) { spark.sql(query).collect() } + } + } + benchmark.run() + } + } + } + } + } + /** * Comprehensive benchmarks for murmur3 hash function across different data types. These * benchmarks cover primitive types, complex types (arrays, structs), and nested structures to From 70c6190183c38fab7c1f13227ce420f5ccf58966 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Mon, 28 Sep 2026 23:53:45 +0800 Subject: [PATCH 02/89] perf: avoid per-element slices when hashing wide decimal arrays --- .../expression-audits/hash_funcs.md | 3 + native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/decimal_hash.rs | 117 +++++++++++++ native/spark-expr/src/hash_funcs/utils.rs | 156 +++++++++++++++++- 4 files changed, 272 insertions(+), 8 deletions(-) create mode 100644 native/spark-expr/benches/decimal_hash.rs diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index 14336cc290a..0177b3488aa 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -35,6 +35,7 @@ - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. - Decimals with precision 19–38, including nested decimals, hash natively using the minimal signed big-endian bytes of the unscaled value, matching Spark's `BigInteger.toByteArray()`. Precision ≤18 retains unscaled-long hashing. `TimeType` remains unsupported through the shared `HashUtils`. +- Performance (tuned 2026-09-28, issue #5994): typed wide-decimal list scans avoid per-element Arrow slices and recursive dispatch. Local Criterion measurements were 5.0–14.2x faster across 8,192-row `decimal(38,0)` List, LargeList and FixedSizeList inputs with 2/32 elements and no/sparse/dense nulls; paired scalar remeasurement found no significant slowdown. Benchmark: `benches/decimal_hash.rs`. ## md5 @@ -66,6 +67,8 @@ ## xxhash64 +- Performance (tuned 2026-09-28, issue #5994): the shared typed wide-decimal list path was 2.1–15.8x faster in local Criterion measurements across 8,192-row `decimal(38,0)` List, LargeList and FixedSizeList inputs with 2/32 elements and no/sparse/dense nulls; paired scalar remeasurement found no significant slowdown. Benchmark: `benches/decimal_hash.rs`. + - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. - Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index a97334b56f1..bb9af7e291f 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -74,6 +74,10 @@ path = "src/lib.rs" name = "hash" harness = false +[[bench]] +name = "decimal_hash" +harness = false + [[bench]] name = "hash_alloc" harness = false diff --git a/native/spark-expr/benches/decimal_hash.rs b/native/spark-expr/benches/decimal_hash.rs new file mode 100644 index 00000000000..01e7580cc87 --- /dev/null +++ b/native/spark-expr/benches/decimal_hash.rs @@ -0,0 +1,117 @@ +// 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. + +//! Wide-decimal scalar controls and typed list hashing, including sliced inputs, +//! short/long arrays and no/sparse/dense nulls. Both public SQL hash paths are measured. + +use arrow::array::{ + Array, ArrayRef, Decimal128Array, FixedSizeListArray, LargeListArray, ListArray, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::Field; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::ColumnarValue; +use datafusion_comet_spark_expr::{spark_murmur3_hash, spark_xxhash64}; +use std::{hint::black_box, sync::Arc}; + +fn bench(c: &mut Criterion) { + let rows = 8192; + let mut group = c.benchmark_group("decimal_hash"); + for (null_every, tag) in [(0, "no_null"), (100, "sparse_null"), (2, "dense_null")] { + for width in [2, 32] { + let values: ArrayRef = Arc::new( + (0..(rows + 2) * width) + .map(|i| { + if null_every != 0 && i % null_every == 0 { + None + } else { + // Mix short signed encodings and values occupying all 16 bytes. + let magnitude = if i % 3 == 0 { 10_i128.pow(37) } else { 1 }; + Some((i as i128 % 7 - 3) * magnitude) + } + }) + .collect::() + .with_precision_and_scale(38, 0) + .unwrap(), + ); + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let nulls = (null_every != 0) + .then(|| NullBuffer::from_iter((0..rows + 2).map(|i| (i + 1) % null_every != 0))); + let cases: Vec<(&str, ArrayRef)> = vec![ + ("scalar", values.slice(width, rows)), + ( + "list", + (Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows + 2)), + Arc::clone(&values), + nulls.clone(), + )) as ArrayRef) + .slice(1, rows), + ), + ( + "large_list", + (Arc::new(LargeListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows + 2)), + Arc::clone(&values), + nulls.clone(), + )) as ArrayRef) + .slice(1, rows), + ), + ( + "fixed_list", + (Arc::new(FixedSizeListArray::new(field, width as i32, values, nulls)) + as ArrayRef) + .slice(1, rows), + ), + ]; + for (shape, array) in cases { + // A scalar control needs only one width. + if shape == "scalar" && width != 2 { + continue; + } + for name in ["murmur3", "xxhash64"] { + let args = vec![ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(if name == "murmur3" { + ScalarValue::Int32(Some(42)) + } else { + ScalarValue::Int64(Some(42)) + }), + ]; + group.bench_with_input( + BenchmarkId::new(name, format!("{shape}/{width}/{tag}")), + &args, + |b, args| { + if name == "murmur3" { + b.iter(|| black_box(spark_murmur3_hash(black_box(args)).unwrap())); + } else { + b.iter(|| black_box(spark_xxhash64(black_box(args)).unwrap())); + } + }, + ); + } + } + } + } + group.finish(); +} + +criterion_group!(benches, bench); +criterion_main!(benches); diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index b11e065025f..a44102c5bc4 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -234,17 +234,24 @@ macro_rules! hash_array_small_decimal { }; } +/// Hash one wide decimal without allocating BigInteger's signed byte representation. +#[macro_export] +macro_rules! hash_decimal { + ($value:expr, $seed:expr, $hash_method:ident) => {{ + let value: i128 = $value; + let bytes = value.to_be_bytes(); + // Match BigInteger.toByteArray(): drop redundant sign bytes, retaining + // one sign bit. Zero and -1 still occupy one byte. + let sign_bits = value.leading_zeros().max(value.leading_ones()); + let start = (sign_bits.saturating_sub(1) / 8) as usize; + $hash_method(&bytes[start..], $seed) + }}; +} + #[macro_export] macro_rules! hash_array_decimal { ($array_type:ident, $column: ident, $hashes: ident, $hash_method: ident) => { - let hash_decimal = |value: &i128, seed| { - let bytes = value.to_be_bytes(); - // Match BigInteger.toByteArray(): drop redundant sign bytes, retaining - // one sign bit. Zero and -1 still occupy one byte. - let sign_bits = value.leading_zeros().max(value.leading_ones()); - let start = (sign_bits.saturating_sub(1) / 8) as usize; - $hash_method(&bytes[start..], seed) - }; + let hash_decimal = |value: &i128, seed| $crate::hash_decimal!(*value, seed, $hash_method); $crate::hash_array!($array_type, $column, $hashes, hash_decimal); }; } @@ -408,6 +415,12 @@ macro_rules! hash_list_with_primitive_elements { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, $hash_method, |v: i64| v.to_le_bytes()); } + DataType::Decimal128(precision, _) if *precision > 18 => { + // Downcast once per list column instead of slicing and dispatching each element. + let elem_array = $values.as_any().downcast_ref::().unwrap(); + let hash_decimal = |value: i128, seed| $crate::hash_decimal!(value, seed, $hash_method); + $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, hash_decimal, |v| v); + } _ => { // Fall back to recursive approach for complex element types $crate::hash_list_array!($list_array_type, $fallback_offset_type, $col, $hashes_buffer, $recursive_hash_method); @@ -503,6 +516,11 @@ macro_rules! hash_list_with_primitive_elements { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, $hash_method, |v: i64| v.to_le_bytes()); } + DataType::Decimal128(precision, _) if *precision > 18 => { + let elem_array = $values.as_any().downcast_ref::().unwrap(); + let hash_decimal = |value: i128, seed| $crate::hash_decimal!(value, seed, $hash_method); + $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, hash_decimal, |v| v); + } _ => { // Fall back to recursive approach for complex element types if $list_array.null_count() == 0 { @@ -1392,4 +1410,126 @@ mod tests { } } } + + #[test] + fn decimal_lists_match_independent_binary_encoding() { + use arrow::array::{FixedSizeListArray, LargeListArray, ListArray}; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::Field; + + for precision in [18, 19, 38] { + let max = 10_i128.pow(precision as u32) - 1; + let values = vec![ + 0, -1, 127, 128, -128, -129, max, -max, 255, 256, -256, -257, 32767, 32768, -32768, + -32769, + ]; + for scale in [0, 2] { + for parent_nulls in [false, true] { + for child_nulls in [false, true] { + let nulls = child_nulls + .then(|| NullBuffer::from_iter((0..values.len()).map(|i| i % 3 != 1))); + // Keep nonzero values under null slots, and values under null parents. + let decimal: ArrayRef = Arc::new( + Decimal128Array::new(values.clone().into(), nulls.clone()) + .with_precision_and_scale(precision, scale) + .unwrap(), + ); + let bytes: Vec<_> = values + .iter() + .map(|&v| { + if precision <= 18 { + (v as i64).to_le_bytes().to_vec() + } else { + BigInt::from(v).to_signed_bytes_be() + } + }) + .collect(); + let binary: ArrayRef = Arc::new(BinaryArray::from_iter( + bytes.iter().enumerate().map(|(i, v)| { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + None + } else { + Some(v.as_slice()) + } + }), + )); + for kind in ["list", "large_list", "fixed_list"] { + let wrap = |values: ArrayRef| -> ArrayRef { + let field = + Arc::new(Field::new("item", values.data_type().clone(), true)); + let nulls = parent_nulls + .then(|| NullBuffer::from_iter((0..8).map(|i| i % 3 != 1))); + match kind { + "list" => Arc::new(ListArray::new( + field, + OffsetBuffer::new( + vec![0_i32, 2, 2, 5, 8, 11, 12, 14, 16].into(), + ), + values, + nulls, + )), + "large_list" => Arc::new(LargeListArray::new( + field, + OffsetBuffer::new( + vec![0_i64, 2, 2, 5, 8, 11, 12, 14, 16].into(), + ), + values, + nulls, + )), + _ => Arc::new(FixedSizeListArray::new(field, 2, values, nulls)), + } + }; + let decimal = wrap(Arc::clone(&decimal)); + let binary = wrap(Arc::clone(&binary)); + for (offset, len) in [(0, 8), (1, 6), (3, 1), (8, 0)] { + let prefix: ArrayRef = + Arc::new(Int32Array::from_iter_values(0..len as i32)); + for seed in [0_i64, 42, -1, i32::MIN as i64] { + for chained in [false, true] { + let args = |array: &ArrayRef, seed: ScalarValue| { + let mut args = vec![ColumnarValue::Array( + array.slice(offset, len), + )]; + if chained { + args.insert( + 0, + ColumnarValue::Array(Arc::clone(&prefix)), + ); + args.push(ColumnarValue::Array(Arc::clone( + &prefix, + ))); + } + args.push(ColumnarValue::Scalar(seed)); + args + }; + for (hash, seed) in [ + ( + spark_murmur3_hash as HashFunction, + ScalarValue::Int32(Some(seed as i32)), + ), + ( + spark_xxhash64 as HashFunction, + ScalarValue::Int64(Some(seed)), + ), + ] { + let actual = hash(&args(&decimal, seed.clone())) + .unwrap() + .into_array(len) + .unwrap(); + let expected = hash(&args(&binary, seed)) + .unwrap() + .into_array(len) + .unwrap(); + assert_eq!(actual.to_data(), expected.to_data(), + "{kind}, precision={precision}, scale={scale}, parent_nulls={parent_nulls}, child_nulls={child_nulls}, offset={offset}, len={len}, chained={chained}"); + } + } + } + } + } + } + } + } + } + } } From 36902dae1d9c1edfd969f3c042cbaf97a6653cf8 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 3 Oct 2026 00:40:09 +0800 Subject: [PATCH 03/89] test: cover mixed wide-decimal shuffle joins --- .../comet/exec/CometNativeShuffleSuite.scala | 46 +++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 052a6cb23c8..6ca5721cbf2 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -39,13 +39,14 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.plans.logical.LocalRelation import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning -import org.apache.spark.sql.comet.{CometExec, CometLocalTableScanExec, CometMetricNode, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.comet.{CometExec, CometLocalTableScanExec, CometMetricNode, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.comet.execution.arrow.CometArrowStream -import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} -import org.apache.spark.sql.execution.LocalTableScanExec +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, LocalTableScanExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.functions.{col, count, sum} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DataType, LongType, MapType, StructField, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -792,6 +793,45 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("wide decimal joins keep native and Spark shuffle inputs copartitioned") { + withSQLConf( + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "false", + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet,json", + SQLConf.SHUFFLE_PARTITIONS.key -> "7", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + withTable("decimal_native", "decimal_spark") { + Seq("decimal_native" -> "parquet", "decimal_spark" -> "json").foreach { + case (table, format) => + sql(s"CREATE TABLE $table(id INT, k DECIMAL(38, 0)) USING $format") + sql(s"""INSERT INTO $table VALUES (1, 1), (2, -1), (3, 128), (4, -129), + (5, 99999999999999999999999999999999999999BD), + (6, -99999999999999999999999999999999999999BD)""") + } + for (mode <- Seq("native", "auto"); adaptive <- Seq(false, true)) { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> mode, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val df = sql("""SELECT n.id, s.id FROM decimal_native n + JOIN decimal_spark s ON n.k = s.k""") + val (_, plan) = checkSparkAnswer(df) + val exchanges = collect(plan) { case e: CometShuffleExchangeExec => e } + assert(exchanges.count(_.shuffleType == CometNativeShuffle) == 1, plan.treeString) + if (mode == "auto") { + assert(exchanges.count(_.shuffleType == CometColumnarShuffle) == 1, plan.treeString) + } else { + assert(collect(plan) { case _: ShuffleExchangeExec => 1 }.sum == 1, plan.treeString) + } + assert(collect(plan) { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) + assert(collect(plan) { case _: FileSourceScanExec => 1 }.sum == 1, plan.treeString) + } + } + } + } + } + test("native shuffle on struct hash partitioning key") { withNestedHashPartitioning { Seq(10, 201).foreach { numPartitions => From 0d058fd73f1b71c915b790074c4bc6f6bbb9a0c2 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 3 Oct 2026 16:12:23 +0800 Subject: [PATCH 04/89] fix: explicitly widen wide-decimal benchmark row count --- .../spark/sql/benchmark/CometHashExpressionBenchmark.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala index d1d3e206a4e..6450b822fe2 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala @@ -109,7 +109,7 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { ("Comet native", true, true)) var expected: Seq[Row] = Seq.empty val benchmark = - new Benchmark(s"$function($column): decimal(38,0)", rows, output = output) + new Benchmark(s"$function($column): decimal(38,0)", rows.toLong, output = output) cases.foreach { case (name, comet, native) => val conf = Seq( CometConf.COMET_ENABLED.key -> comet.toString, From 070951d2a77fbaa597056d8de48d29812b10f2e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 13:25:52 +0000 Subject: [PATCH 05/89] feat: convert OneRowRelation to Arrow by default (#6666) * feat: convert OneRowRelation to Arrow by default Turn on spark.comet.convert.oneRowRelation.enabled by default. The row a query without a FROM clause reads has no columns, so converting it costs almost nothing, and it lets the operators above it run in Comet. Where the parent cannot run in Comet, EliminateRedundantTransitions drops the conversion again, so no extra transition is added. On Spark 4.2 the TPC-DS q77a golden now records the remaining blocker, the struct-typed scalar subquery result the projection above each OneRowRelation reads (#5834), instead of the unsupported leaf. * fix: handle expressions exposed by converting OneRowRelation by default Folded array literals are UnsafeArrayData and have no backing array, so serialize them through toArray. might_contain accepts an untyped NULL for either argument. parse_url raises Spark's INVALID_URL error under ANSI instead of a generic native error. An Expand with no output columns falls back to Spark. * fix: decline codegen dispatch for expressions that contain a subquery * test: adapt Spark explain tests to the Comet plan for OneRowRelation queries --- dev/diffs/3.4.3.diff | 11 ++++- dev/diffs/3.5.9.diff | 11 ++++- dev/diffs/4.0.4.diff | 23 +++++++++- dev/diffs/4.1.3.diff | 23 +++++++++- dev/diffs/4.2.0.diff | 23 +++++++++- .../latest/compatibility/spark-versions.md | 12 ++--- docs/source/user-guide/latest/datasources.md | 3 +- .../user-guide/latest/migration-guide.md | 5 ++- native/common/src/error.rs | 13 ++++++ .../bloom_filter_might_contain.rs | 27 +++++++++++- native/spark-expr/src/url_funcs/parse_url.rs | 13 +++--- .../scala/org/apache/comet/CometConf.scala | 5 ++- .../codegen/CometBatchKernelCodegen.scala | 12 ++--- .../org/apache/comet/serde/literals.scala | 4 +- .../apache/spark/sql/comet/operators.scala | 7 +++ .../comet/shims/ShimSparkErrorConverter.scala | 6 +++ .../comet/shims/ShimSparkErrorConverter.scala | 6 +++ .../comet/shims/ShimSparkErrorConverter.scala | 6 +++ .../sql-tests/operators/one_row_relation.sql | 44 +++++++++++++++++++ .../q77a/extended.txt | 12 ++--- .../org/apache/comet/CometCodegenSuite.scala | 14 +++++- .../apache/comet/CometExpressionSuite.scala | 24 ++++++++++ .../comet/exec/CometExec3_4PlusSuite.scala | 10 +++++ .../apache/comet/exec/CometExecSuite.scala | 9 ++++ .../comet/rules/CometExecRuleSuite.scala | 14 ++++++ .../org/apache/spark/sql/CometTestBase.scala | 6 +-- 26 files changed, 300 insertions(+), 43 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/operators/one_row_relation.sql diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 5c23919a2ee..af31251b32e 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -593,7 +593,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index a6b295578d6..1167bbe6554 100644 +index a6b295578d6..1fa898bf840 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -260,7 +260,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -606,6 +606,15 @@ index a6b295578d6..1167bbe6554 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -331,7 +332,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -463,7 +464,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index d83fe66edce..801199ddb8e 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -574,7 +574,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 739557bef30..718fbc5b786 100644 +index 739557bef30..51682036f38 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -265,7 +265,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -587,6 +587,15 @@ index 739557bef30..718fbc5b786 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -336,7 +337,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -468,7 +469,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index f12fdc9e352..21ef06bdfec 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -711,7 +711,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index ef9c301a66a..49fad8cbbf0 100644 +index ef9c301a66a..d97e9f76dc3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -724,6 +724,15 @@ index ef9c301a66a..49fad8cbbf0 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1192,7 +1201,7 @@ index ad424b3a7cc..4ece0117a34 100644 protected val baseResourcePath = { // use the same way as `SQLQueryTestSuite` to get the resource path diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index f294ff81021..02d72be8d29 100644 +index f294ff81021..62ca9a9bcd5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -1524,7 +1524,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark @@ -1222,6 +1231,16 @@ index f294ff81021..02d72be8d29 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -4744,7 +4752,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 575a4ae69d1..129d9f27232 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index c99c1c242d2..d034d20b990 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -813,7 +813,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 4c62c47971a..ecc7ed21546 100644 +index 4c62c47971a..fe788debf70 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -826,6 +826,15 @@ index 4c62c47971a..ecc7ed21546 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1313,7 +1322,7 @@ index cb9d0909554..084d6515e8b 100644 } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 74cdee49e55..f7452c9abb7 100644 +index 74cdee49e55..1a7bc4567b4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark @@ -1343,6 +1352,16 @@ index 74cdee49e55..f7452c9abb7 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -4741,7 +4749,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 23f0144dcec..40d536bb23a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 01a29c48ae8..3335dfb3713 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -842,7 +842,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 618f0ab675e..be186f99b66 100644 +index 618f0ab675e..efcd4853251 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -855,6 +855,15 @@ index 618f0ab675e..be186f99b66 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1373,7 +1382,7 @@ index 291aa7cab72..7783c37683e 100644 super.test(testName, testTags: _*) { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index da6f6aca2ad..c02b7c99490 100644 +index da6f6aca2ad..12f48c4bb5d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -1529,7 +1529,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper @@ -1403,6 +1412,16 @@ index da6f6aca2ad..c02b7c99490 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -4755,7 +4763,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 395cb67f441..33ac6ed19af 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala diff --git a/docs/source/user-guide/latest/compatibility/spark-versions.md b/docs/source/user-guide/latest/compatibility/spark-versions.md index 337252cb794..84baf6194fa 100644 --- a/docs/source/user-guide/latest/compatibility/spark-versions.md +++ b/docs/source/user-guide/latest/compatibility/spark-versions.md @@ -89,11 +89,13 @@ Spark 4.2.0 is supported with Java 17 and Scala 2.13. - **`FILTER` on window aggregates**: Spark 4.2 accepts a `FILTER (WHERE ...)` clause on an aggregate function used as a window function. Comet does not support the clause there, so the window operator falls back to Spark. -- **`OneRowRelation` in `UNION` branches** - ([#4949](https://github.com/apache/datafusion-comet/issues/4949)): Spark 4.2 plans a - `OneRowRelation` into the `UNION` branches of some queries, for example TPC-DS q77a. Comet does - not convert `OneRowRelation` by default, so the union and the aggregates above it fall back to - Spark. +- **Merged subqueries in `UNION` branches** + ([#4949](https://github.com/apache/datafusion-comet/issues/4949)): Spark 4.2 turns a single-row + aggregate in some queries, for example TPC-DS q77a, into a scalar subquery that returns a struct, + and reads its fields in a projection over a `OneRowRelation`. Comet does not support a + struct-typed scalar subquery result yet + ([#5834](https://github.com/apache/datafusion-comet/issues/5834)), so that projection, and the + union and the aggregates above it, fall back to Spark. - **ANSI arithmetic overflow** ([#4967](https://github.com/apache/datafusion-comet/issues/4967)): Spark 4.2 changed some ANSI-mode arithmetic overflow behavior, and Comet does not match these changes yet. See diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index 30da3df12e1..eb56d1e8970 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -52,7 +52,8 @@ converted into Arrow format, allowing the Comet pipeline to take over after that ### Other Spark inputs Comet can also convert the output of these Spark inputs to Arrow format, so that the operators -above them can run in Comet. Each conversion is off by default. +above them can run in Comet. Only `spark.comet.convert.oneRowRelation.enabled` is on by default, +because the row it converts has no columns. - `spark.comet.convert.range.enabled`: `spark.range` and SQL `range()`, for ranges that Comet does not generate natively with `spark.comet.exec.range.enabled`. diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index b55cd94d97c..3ce3e88d1da 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -79,8 +79,9 @@ set, `spark.comet.sparkToColumnar.enabled=true` still converts the first four, w by default, and a list that names any of the five still converts it, as before. In both cases the driver logs a warning that names the config to use instead. Before Spark 4.1, Spark plans a query without a `FROM` clause as an `RDDScan`, so on those versions `RDDScan` in the list also converted -it. The list remains the way to convert other leaf operators, such as the scan of a Data Source V2 -connector. +it. `spark.comet.convert.oneRowRelation.enabled` is on by default, so such a query is now converted +without either setting, and logs no warning. The list remains the way to convert other leaf +operators, such as the scan of a Data Source V2 connector. ## Upgrading to Comet 1.1.0 diff --git a/native/common/src/error.rs b/native/common/src/error.rs index 1346b398061..23f92b46393 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -215,6 +215,9 @@ pub enum SparkError { group_index: i32, }, + #[error("[INVALID_URL] The url is invalid: {url}. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] + InvalidUrl { url: String }, + #[error("[DATATYPE_CANNOT_ORDER] Cannot order by type: {data_type}.")] DatatypeCannotOrder { data_type: String }, @@ -379,6 +382,7 @@ impl SparkError { SparkError::UnexpectedPositiveValue { .. } => "UnexpectedPositiveValue", SparkError::UnexpectedNegativeValue { .. } => "UnexpectedNegativeValue", SparkError::InvalidRegexGroupIndex { .. } => "InvalidRegexGroupIndex", + SparkError::InvalidUrl { .. } => "InvalidUrl", SparkError::DatatypeCannotOrder { .. } => "DatatypeCannotOrder", SparkError::ScalarSubqueryTooManyRows => "ScalarSubqueryTooManyRows", SparkError::MergeCardinalityViolation => "MergeCardinalityViolation", @@ -613,6 +617,11 @@ impl SparkError { "groupIndex": group_index, }) } + SparkError::InvalidUrl { url } => { + serde_json::json!({ + "url": url, + }) + } SparkError::DatatypeCannotOrder { data_type } => { serde_json::json!({ "dataType": data_type, @@ -742,6 +751,7 @@ impl SparkError { // IllegalArgumentException SparkError::DatatypeCannotOrder { .. } + | SparkError::InvalidUrl { .. } | SparkError::InvalidUtf8String { .. } | SparkError::IllegalDayOfWeek { .. } | SparkError::SequenceIllegalBoundaries { .. } => { @@ -853,6 +863,9 @@ impl SparkError { // Regex errors SparkError::InvalidRegexGroupIndex { .. } => Some("INVALID_PARAMETER_VALUE"), + // URL errors + SparkError::InvalidUrl { .. } => Some("INVALID_URL"), + // Unsupported operation errors SparkError::DatatypeCannotOrder { .. } => Some("DATATYPE_CANNOT_ORDER"), diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs b/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs index 66168444d9a..e60e6f462c1 100644 --- a/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs +++ b/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs @@ -57,6 +57,8 @@ fn evaluate_bloom_filter( ColumnarValue::Scalar(ScalarValue::Binary(v)) => { Ok(v.map(|v| SparkBloomFilter::from(v.as_slice()))) } + // An untyped NULL literal, as in `might_contain(null, 1L)` + ColumnarValue::Scalar(ScalarValue::Null) => Ok(None), _ => internal_err!("Bloom filter expression should be evaluated as a scalar binary value"), } } @@ -90,6 +92,10 @@ fn execute_might_contain( .and_then(|filter| optional_value.map(|value| filter.might_contain_long(value))); Ok(ColumnarValue::Scalar(ScalarValue::Boolean(result))) } + // An untyped NULL literal, as in `might_contain(bf, null)` + ColumnarValue::Scalar(ScalarValue::Null) => { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None))) + } ColumnarValue::Array(values_array) => { let values = values_array .as_any() @@ -109,7 +115,7 @@ fn execute_might_contain( }) .unwrap_or_else(|| Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)))) } - _ => internal_err!("Expected Int64Array or Int64 Scalar as arguments"), + other => internal_err!("Expected Int64Array or Int64 Scalar as arguments, got {other:?}"), } } @@ -117,6 +123,7 @@ fn execute_might_contain( mod tests { use super::*; use arrow::array::BooleanArray; + use datafusion::physical_expr::expressions::Literal; fn assert_result_eq>>(result: ColumnarValue, expected: Vec) { let array = result.to_array(1).unwrap(); @@ -171,6 +178,24 @@ mod tests { assert_all_null(result); } + #[test] + fn test_execute_untyped_null_value() { + let mut filter = SparkBloomFilter::from((4, 64)); + filter.put_long(123); + + let args = [ColumnarValue::Scalar(ScalarValue::Null)]; + + let result = execute_might_contain(&Some(filter), &args).unwrap(); + assert_all_null(result); + } + + #[test] + fn test_untyped_null_bloom_filter() { + let bloom_filter_expr: Arc = Arc::new(Literal::new(ScalarValue::Null)); + let expr = BloomFilterMightContain::try_new(bloom_filter_expr).unwrap(); + assert!(expr.bloom_filter.is_none()); + } + #[test] fn test_execute_array() { let mut filter = SparkBloomFilter::from((4, 64)); diff --git a/native/spark-expr/src/url_funcs/parse_url.rs b/native/spark-expr/src/url_funcs/parse_url.rs index 7e072b68957..ec1b33ab353 100644 --- a/native/spark-expr/src/url_funcs/parse_url.rs +++ b/native/spark-expr/src/url_funcs/parse_url.rs @@ -29,12 +29,14 @@ use arrow::array::{ }; use arrow::datatypes::DataType; use datafusion::common::cast::{as_large_string_array, as_string_array, as_string_view_array}; -use datafusion::common::{exec_datafusion_err, exec_err, Result, ScalarValue}; +use datafusion::common::{exec_err, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; use regex::Regex; +use crate::SparkError; + // RFC 3986 Appendix B decomposition regex. // Groups: 2=scheme, 4=authority, 5=path, 7=query, 9=fragment static URI_REGEX: LazyLock = LazyLock::new(|| { @@ -109,11 +111,10 @@ fn has_invalid_uri_chars(s: &str) -> bool { } fn invalid_url_err(value: &str) -> datafusion::common::DataFusionError { - exec_datafusion_err!( - "[INVALID_URL] The provided URL '{}' is not valid. Use `try_parse_url` to tolerate \ - malformed URLs and return NULL instead. SQLSTATE: 22P02", - value - ) + SparkError::InvalidUrl { + url: value.to_string(), + } + .into() } fn parse_url_component(value: &str, part: &str, key: Option<&str>) -> Result> { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 09167333deb..6dee9e3c23e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -244,9 +244,10 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "When enabled, the single row that a query without a FROM clause, such as `SELECT 1`, " + - "reads will be converted to Arrow format.") + "reads will be converted to Arrow format. The row has no columns, so converting it " + + "costs almost nothing, and it lets the operators above it run in Comet.") .booleanConf - .createWithDefault(false) + .createWithDefault(true) val COMET_CONVERT_FROM_ROW_DATA_SOURCE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.convert.rowDataSource.enabled") diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 2052539efec..e478057ceef 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -26,6 +26,7 @@ import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, Unevaluable} import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.execution.ExecSubqueryExpression import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -173,11 +174,9 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // instance with a single `init(partitionIndex)` call, so `Rand` / `MonotonicallyIncreasingID` // state advances correctly across batches. // - // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`) is accepted: the surrounding - // Comet operator's inherited `SparkPlan.waitForSubqueries` populates the subquery's - // `result` field before evaluation. The closure serializer captures that value into the - // arg-0 bytes, and the dispatcher keys its compile cache on those bytes, so distinct subquery - // results produce distinct cache entries. + // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`): rejected. The tree is + // closure-serialized at plan time, before Spark has run the subquery, so the deserialized + // copy never holds a result and `ScalarSubquery.doGenCode` fails with "has not finished". // // `Unevaluable`: rejected by default. `isCodegenInertUnevaluable` exempts version-specific // leaves that are `Unevaluable` but never invoked by codegen (e.g. Spark 4.0's @@ -188,6 +187,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come boundExpr.find { case _: org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction => true case _: org.apache.spark.sql.catalyst.expressions.Generator => true + case _: ExecSubqueryExpression => true case u: Unevaluable if isCodegenInertUnevaluable(u) => false case _: Unevaluable => true case _: org.apache.comet.udf.NativeUdfCall => true @@ -196,7 +196,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case Some(bad) => return Some( s"codegen dispatch: expression ${bad.getClass.getSimpleName} not supported " + - "(aggregate, generator, unevaluable, or native UDF)") + "(aggregate, generator, subquery, unevaluable, or native UDF)") case None => } val badRef = boundExpr.collectFirst { diff --git a/spark/src/main/scala/org/apache/comet/serde/literals.scala b/spark/src/main/scala/org/apache/comet/serde/literals.scala index 621dd99ebe8..ef6fa1fb466 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -113,7 +113,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit case arr: ArrayType => val listLiteralBuilder: ListLiteral.Builder = - makeListLiteral(value.asInstanceOf[ArrayData].array, arr) + makeListLiteral(value.asInstanceOf[ArrayData].toArray[Any](arr.elementType), arr) exprBuilder.setListVal(listLiteralBuilder.build()) exprBuilder.setDatatype(serializeDataType(dataType).get) case dt => @@ -225,7 +225,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit array.foreach(v => { val casted = v.asInstanceOf[ArrayData] listLiteralBuilder.addListValues(if (casted != null) { - makeListLiteral(casted.array, a) + makeListLiteral(casted.toArray[Any](a.elementType), a) } else ListLiteral.newBuilder()) listLiteralBuilder.addNullMask(casted != null) }) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 909539bf2b6..fdb07c9acff 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1714,6 +1714,13 @@ object CometExpandExec extends CometOperatorSerde[ExpandExec] { op: ExpandExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { + // The native Expand regroups the flattened expressions by their count per projection, so a + // projection with no columns, which column pruning leaves when nothing above the Expand reads + // its output, would reach it as no projections at all. + if (op.projections.isEmpty || op.projections.head.isEmpty) { + withFallbackReason(op, "Expand without output columns is not supported") + return None + } val projExprs = op.projections.flatMap(_.map(e => exprToProto(e, op.child.output))) if (projExprs.forall(_.isDefined) && childOp.nonEmpty) { diff --git a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 4fd3606203b..f47c7682389 100644 --- a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -294,6 +294,12 @@ trait ShimSparkErrorConverter { messageParameters = params.map { case (k, v) => (k, v.toString) }, cause = null)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => // orderedOperationUnsupportedByDataTypeError takes DataType in Spark 3.4, not String Some( diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 39a76f386c5..ee4dbea9725 100644 --- a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -291,6 +291,12 @@ trait ShimSparkErrorConverter { params("groupCount").toString.toInt, params("groupIndex").toString.toInt)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => Some( QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 67a4363870d..5e5641fbb2e 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -318,6 +318,12 @@ trait ShimSparkErrorConverter { params("groupCount").toString.toInt, params("groupIndex").toString.toInt)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => Some( QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( diff --git a/spark/src/test/resources/sql-tests/operators/one_row_relation.sql b/spark/src/test/resources/sql-tests/operators/one_row_relation.sql new file mode 100644 index 00000000000..e28d80cdf89 --- /dev/null +++ b/spark/src/test/resources/sql-tests/operators/one_row_relation.sql @@ -0,0 +1,44 @@ +-- 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. + +-- Queries without a FROM clause run over a OneRowRelation, which Comet converts by default. + +-- A folded array literal is an UnsafeArrayData, not a GenericArrayData. +query +SELECT sequence(1, 5) + +query +SELECT sequence(1, 5), array(array(1, 2), array(3)), array('a', NULL, 'c') + +query +SELECT explode(sequence(1, 3)) + +-- GROUPING SETS over a single row. +statement +CREATE TEMPORARY VIEW one_row_user_view AS SELECT current_user + +query +SELECT count(*) FROM one_row_user_view GROUP BY current_user GROUPING SETS ((current_user)) + +query +SELECT current_user, grouping(current_user) FROM one_row_user_view GROUP BY ROLLUP(current_user) + +query +SELECT count(*) FROM (SELECT 1 AS a) GROUP BY GROUPING SETS ((a), ()) + +statement +DROP VIEW one_row_user_view diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt index c97afc95a3f..30de4b735f5 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt @@ -65,7 +65,7 @@ CometColumnarToRow : : : +- CometProject : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- Project + : : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : : :- Subquery : : : : +- CometColumnarToRow : : : : +- CometProject @@ -81,7 +81,7 @@ CometColumnarToRow : : : : +- CometFilter : : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : : +- ReusedSubquery - : : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : : +- Scan OneRowRelation : +- CometColumnarToRow : +- CometProject : +- CometBroadcastHashJoin @@ -184,7 +184,7 @@ CometColumnarToRow : : : +- CometProject : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- Project + : : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : : :- Subquery : : : : +- CometColumnarToRow : : : : +- CometProject @@ -200,7 +200,7 @@ CometColumnarToRow : : : : +- CometFilter : : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : : +- ReusedSubquery - : : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : : +- Scan OneRowRelation : +- CometColumnarToRow : +- CometProject : +- CometBroadcastHashJoin @@ -303,7 +303,7 @@ CometColumnarToRow : : +- CometProject : : +- CometFilter : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- Project + : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : :- Subquery : : : +- CometColumnarToRow : : : +- CometProject @@ -319,7 +319,7 @@ CometColumnarToRow : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : +- ReusedSubquery - : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : +- Scan OneRowRelation +- CometColumnarToRow +- CometProject +- CometBroadcastHashJoin diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index c9371fbf6da..e9105230d7f 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.{ByteArray, UTF8String} -import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.serde.{CometInvokeTargets, CometScalaUDF, QueryPlanSerde} @@ -1502,6 +1502,18 @@ class CometCodegenSuite checkSparkAnswerAndOperator(df) } + test("scalar subquery inside a dispatched expression falls back to Spark") { + assume(isSpark40Plus, "collations are Spark 4.0+") + withTable("t") { + sql("CREATE TABLE t (c1 STRING) USING parquet") + sql("INSERT INTO t VALUES ('a')") + checkSparkAnswer(sql("SELECT COUNT(CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE))")) + checkSparkAnswer(sql("SELECT CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE) = 'A'")) + checkSparkAnswer( + sql("SELECT COUNT(CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE)) FROM t")) + } + } + test("ScalaUDF composed with reused scalar subquery across projection and filter") { // The same scalar subquery appears in two sites: the projection (which the dispatcher // compiles into a fused kernel) and the filter (a separate operator). Each site holds its diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 098b8dd3877..5d2ca1d91de 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -565,6 +565,30 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("ANSI parse_url of an invalid URL raises a Spark error") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("parse_url_invalid") { + sql("CREATE TABLE parse_url_invalid (u STRING) USING PARQUET") + sql("INSERT INTO parse_url_invalid VALUES ('inva lid://user:pass@host/file')") + + val df = sql("SELECT parse_url(u, 'HOST') FROM parse_url_invalid") + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + checkSparkAnswerMaybeThrows(df) match { + case (Some(sparkException), Some(cometException)) => + def sparkError(error: Throwable): SparkThrowable = + causeChain(error) + .collectFirst { case error: SparkThrowable => error } + .getOrElse(fail(s"Expected a SparkThrowable, got $error")) + val expected = sparkError(sparkException) + val actual = sparkError(cometException) + assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getSqlState == expected.getSqlState) + case errors => fail(s"Expected Spark and Comet invalid URL errors, got $errors") + } + } + } + } + test("ANSI decimal divide by zero raises a Spark error") { withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { withTable("decimal_div_zero") { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala index c6924eb91c4..516ece40be0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala @@ -158,6 +158,16 @@ class CometExec3_4PlusSuite extends CometTestBase { } } + test("test NULL inputs for BloomFilterMightContain without a FROM clause") { + assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4968") + checkSparkAnswer(""" + |SELECT might_contain(null, null) both_null, + | might_contain(null, 1L) null_bf, + | might_contain((SELECT bloom_filter_agg(cast(id as long)) FROM range(1, 10000)), + | null) null_value + |""".stripMargin) + } + test("test NULL inputs for BloomFilterMightContain") { assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4968") val table = "test" diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index d042992c696..22e76a9abd1 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, CatalogTable} import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, Hex, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate, Final} +import org.apache.spark.sql.catalyst.plans.logical.Expand import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog @@ -2854,6 +2855,14 @@ class CometExecSuite extends CometTestBase { } } + test("expand operator with no output columns falls back to Spark") { + // Column pruning leaves this shape when nothing above the Expand reads its output. + val child = spark.range(3).queryExecution.analyzed + val expand = Expand(Seq(Seq.empty[Expression], Seq.empty[Expression]), Seq.empty, child) + val df = datasetOfRows(spark, expand).groupBy().count() + checkSparkAnswerAndFallbackReason(df, "Expand without output columns is not supported") + } + test("multiple distinct multiple columns sets") { withTable("agg2") { val data2 = Seq[(Integer, Integer, Integer)]( diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 006cc1ae126..d8f40f2b635 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -2356,6 +2356,20 @@ class CometExecRuleSuite extends CometTestBase { } } + test("only the OneRowRelation conversion is on by default") { + withConversionQueries { queries => + // CometTestBase turns every conversion on, so set each back to its default. + val defaults = queries.map { case (entry, _) => entry.key -> entry.defaultValueString } + for ((entry, query) <- queries) { + withClue(s"$query: ") { + assert( + convertedLeaves(query, defaults: _*).nonEmpty == + (entry eq CometConf.COMET_CONVERT_FROM_ONE_ROW_RELATION_ENABLED)) + } + } + } + } + test("the deprecated sparkToColumnar settings still convert what they converted before") { val switch = CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true" val list = CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 07864c6fec8..60b646fd74e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -79,9 +79,9 @@ abstract class CometTestBase } /** - * Sets each Spark-to-Arrow conversion that [[sparkConf]] turns on, which are off by default, to - * `enabled`. These are the conversions `spark.comet.sparkToColumnar.enabled` stood for before - * each had a config of its own. + * Sets each Spark-to-Arrow conversion that [[sparkConf]] turns on to `enabled`. All but the + * `OneRowRelation` conversion are off by default. These are the conversions + * `spark.comet.sparkToColumnar.enabled` stood for before each had a config of its own. */ protected def sparkToArrowConversionConfs(enabled: Boolean): Seq[(String, String)] = Seq( From 011cce3916cbc99556c66b6368b1ddfaae30e2fe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 13:51:33 +0000 Subject: [PATCH 06/89] fix: rethrow the exception a JVM input throws instead of a CometNativeException (#6243) * fix: rethrow the exception a JVM input throws instead of a CometNativeException A native plan pulls each batch of a JVM input through the Arrow C stream exported by CometArrowStream.stream. When the JVM producer throws while building a batch, Arrow Java's exported stream catches the throwable and hands native only its text, so the task failed with a CometNativeException ("C Data interface error: ...") instead of the exception Spark throws. For inputs of Arrow-backed batches this happened past the first batch only, because schema reconciliation reads the first one on the JVM. Row inputs through CometSparkToColumnarExec hit it on every batch. CometArrowStream.stream now wraps each exported reader to keep the first throwable its loadNextBatch throws, keyed by the exported stream until the task completes. When executePlan fails, CometExecIterator rethrows that throwable if one of its input streams recorded one. Closes #6234. * fix: drop unneeded string interpolator flagged by scalafix * refactor: match the recorded input failure once, share the failed-write test helper CometExecIterator looked up the recorded input failure twice, in the guard and in the body of the catch case. An extractor does it once, and the comment at the catch site shrinks to a line now that CometArrowStream.inputFailure carries the explanation. The native-versus-JVM failed-write scaffolding in CometIcebergWriteActionSuite was copied three times. failedWrites runs a write against a native-writer table and a JVM-writer table, checks that the native run went through CometIcebergWriteExec, and returns both failures. --- docs/source/contributor-guide/ffi.md | 10 ++ .../org/apache/comet/CometExecIterator.scala | 17 ++++ .../arrow/CometNativeArrowSource.scala | 66 ++++++++++++- .../comet/CometIcebergWriteActionSuite.scala | 97 ++++++++++++++----- .../apache/comet/exec/CometExecSuite.scala | 51 +++++++++- .../CometExecIteratorLifecycleSuite.scala | 55 ++++++++++- 6 files changed, 267 insertions(+), 29 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 7027080bbf3..36328468f34 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -119,6 +119,16 @@ realigns under-aligned buffers on import ([apache/arrow-rs#10030](https://github the `realigns_under_aligned_decimal128` test in `scan.rs` guards against an arrow downgrade that would bring back the panic ([apache/arrow-rs#10028](https://github.com/apache/arrow-rs/issues/10028)). +### Errors from the JVM Producer + +Arrow Java's exported stream catches whatever the reader throws in `get_next` and hands native only its text, so +native fails the plan with a `CometNativeException` built from that text. `CometArrowStream.stream` wraps every +reader to keep the throwable itself until the task completes, and `CometExecIterator` rethrows it in place of the +native error. The task then fails with the exception Spark would have thrown, such as a `SparkArithmeticException` +from an upstream plan. For an input of Arrow-backed `ColumnarBatch`es, the first batch never takes this path: schema +reconciliation reads it on the JVM before the stream is exported, so what it throws propagates directly. A test of +this path therefore has to fail a later batch. + ### Schema Reconciliation `CometArrowStream.reconcileStreamSchema` advertises the stream's schema from the actual `CometVector` types in the diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 22143b4f188..545c01fafd4 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -26,6 +26,7 @@ import java.util.concurrent.atomic.AtomicBoolean import scala.util.control.NonFatal +import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.memory.BufferAllocator import org.apache.hadoop.conf.Configuration import org.apache.spark._ @@ -33,6 +34,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.comet.CometMetricNode +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized._ import org.apache.spark.util.SerializableConfiguration @@ -224,6 +226,17 @@ class CometExecIterator( CometExecIterator.startMemoryUsageLog() + /** + * Matches any throwable once a JVM input of this plan has failed, and extracts what that input + * threw. Native saw only its text. See `CometArrowStream.inputFailure`. + */ + private object InputFailure { + def unapply(nativeFailure: Throwable): Option[Throwable] = + inputObjects.iterator + .collect { case stream: ArrowArrayStream => CometArrowStream.inputFailure(stream) } + .collectFirst { case Some(failure) => failure } + } + private def getNextBatch: Option[ColumnarBatch] = { assert(partitionIndex >= 0 && partitionIndex < numParts) @@ -246,6 +259,10 @@ class CometExecIterator( result } catch { + // Native saw only the text of what a JVM input threw, so rethrow the exception itself. + case InputFailure(failure) => + throw failure + // Handle CometQueryExecutionException with JSON payload first case e: CometQueryExecutionException => logError(s"Native execution for task $taskAttemptId failed", e) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index 19532cefa14..316afe68a22 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -19,11 +19,15 @@ package org.apache.spark.sql.comet.execution.arrow +import java.util.concurrent.ConcurrentHashMap + import scala.jdk.CollectionConverters._ import scala.reflect.ClassTag import org.apache.arrow.c.{ArrowArrayStream, Data} import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema} import org.apache.spark.TaskContext @@ -233,16 +237,19 @@ object CometArrowStream extends Logging { * the JNI global ref that arrow-java holds on the stream's private data. Taking the stream * leaves a null release callback in the JVM's struct, and arrow-java's `release` skips a null * callback, so a stream native took is still released only once. + * + * The reader is wrapped so that what it throws while native pulls a batch stays available to + * [[inputFailure]] until the task completes. */ def stream( name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = { val context = TaskContext.get() val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) - var reader: ArrowReader = null + var reader: FailureRecordingReader = null var arrowStream: ArrowArrayStream = null try { - reader = readerFactory(allocator) + reader = new FailureRecordingReader(allocator, readerFactory(allocator)) arrowStream = ArrowArrayStream.allocateNew(allocator) Data.exportArrayStream(allocator, reader, arrowStream) } catch { @@ -256,7 +263,9 @@ object CometArrowStream extends Logging { } if (context != null) { val streamRef = arrowStream + exportedReaders.put(streamRef, reader) context.addTaskCompletionListener[Unit] { _ => + exportedReaders.remove(streamRef) // Release before close: `close` frees the struct that `release` reads. streamRef.release() streamRef.close() @@ -266,6 +275,59 @@ object CometArrowStream extends Logging { Iterator.single(arrowStream) } + /** + * The reader behind each stream exported by [[stream]] whose task has not completed yet, keyed + * by the stream object that its consumer holds as an input. + */ + private val exportedReaders = new ConcurrentHashMap[ArrowArrayStream, FailureRecordingReader]() + + /** + * What the reader behind `stream` threw while native pulled a batch from it, if anything. Arrow + * Java's exported stream catches that throwable and hands native only its text, so the native + * plan fails with a `CometNativeException` built from the text. The throwable itself is what + * the task would have thrown without Comet, so the plan's consumer rethrows it instead. + */ + def inputFailure(stream: ArrowArrayStream): Option[Throwable] = + Option(exportedReaders.get(stream)).flatMap(reader => Option(reader.failure)) + + /** + * Forwards everything to `reader`, keeping the first throwable its `loadNextBatch` throws for + * [[inputFailure]]. Its own `ArrowReader` state is never initialized: every method that the + * stream exporter calls is forwarded, and so is `close`. + */ + private final class FailureRecordingReader(allocator: BufferAllocator, reader: ArrowReader) + extends ArrowReader(allocator) { + + @volatile var failure: Throwable = _ + + override def loadNextBatch(): Boolean = + try reader.loadNextBatch() + catch { + case t: Throwable => + if (failure == null) failure = t + throw t + } + + override def getVectorSchemaRoot(): VectorSchemaRoot = reader.getVectorSchemaRoot() + + override def getDictionaryVectors(): java.util.Map[java.lang.Long, Dictionary] = + reader.getDictionaryVectors() + + override def lookup(id: Long): Dictionary = reader.lookup(id) + + override def getDictionaryIds(): java.util.Set[java.lang.Long] = reader.getDictionaryIds() + + override def bytesRead(): Long = reader.bytesRead() + + override def close(): Unit = reader.close() + + override def close(closeReadSource: Boolean): Unit = reader.close(closeReadSource) + + override protected def closeReadSource(): Unit = () + + override protected def readSchema(): Schema = reader.getVectorSchemaRoot().getSchema + } + /** Close a resource, swallowing any error, for use in rollback paths. Null-safe. */ private def closeQuietly(c: AutoCloseable): Unit = { if (c != null) { diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index e13a2ec21ad..53388f1b491 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -38,7 +38,7 @@ import org.json4s.jackson.JsonMethods.parse import org.apache.hadoop.fs.{Path => HadoopPath} import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile -import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext} +import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, SparkThrowable, Success, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase @@ -2235,18 +2235,8 @@ class CometIcebergWriteActionSuite "region", "amount") - var nativeFailure: Throwable = null - val nativePlans = withNativeEnabled { - capturePlans(spark, includeFailures = true) { - nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows)) - } - } - val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows)) - assert( - nativePlans.exists(p => - collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), - "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" + - nativePlans.mkString("\n--\n")) + val (nativeFailure, jvmFailure) = + failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows)) val nativeError = clusteredWriterError(nativeFailure) val jvmError = clusteredWriterError(jvmFailure) @@ -2293,18 +2283,8 @@ class CometIcebergWriteActionSuite (2, new Timestamp(0L), Array[Byte](0)), (3, new Timestamp(-1500L), Array[Byte](0, 1, -1))).toDF("id", "ts", "bin") - var nativeFailure: Throwable = null - val nativePlans = withNativeEnabled { - capturePlans(spark, includeFailures = true) { - nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows)) - } - } - val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows)) - assert( - nativePlans.exists(p => - collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), - "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" + - nativePlans.mkString("\n--\n")) + val (nativeFailure, jvmFailure) = + failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows)) val (_, nativeMessage) = clusteredWriterError(nativeFailure) assert( @@ -2319,6 +2299,51 @@ class CometIcebergWriteActionSuite } } + // https://github.com/apache/datafusion-comet/issues/6234. The native writer pulls its input + // through an Arrow C stream, and Arrow Java hands native only the text of an exception thrown + // while producing a batch, so the user got a CometNativeException instead of Spark's exception. + // The first batch is read on the JVM to derive the stream's schema, which is why the overflow + // has to land past it: row 90000 of a single 100000-row data file, so one task reads it all. + test("native acceleration: an input error past the first batch fails like the JVM writer") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + val (nativeTable, jvmTable) = ("overflow_native", "overflow_jvm") + Seq("overflow_src", nativeTable, jvmTable).foreach { table => + spark.sql(s"CREATE TABLE $catalog.$ns.$table (v INT) USING iceberg") + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 100000, 1, 1) + .selectExpr("CAST(id AS INT) AS v") + .writeTo(s"$catalog.$ns.overflow_src") + .append() + } + def insert(table: String): Unit = + spark.sql( + s"INSERT INTO $catalog.$ns.$table " + + s"SELECT v + ${Int.MaxValue - 90000} FROM $catalog.$ns.overflow_src") + + val (nativeFailure, jvmFailure) = failedWrites(nativeTable, jvmTable)(insert) + + def overflow(failure: Throwable): (Int, Class[_], String, String) = { + val chain = causeChain(failure) + val depth = chain.indexWhere(_.isInstanceOf[ArithmeticException]) + assert(depth >= 0, s"no ArithmeticException in the cause chain of $failure") + val error = chain(depth).asInstanceOf[Throwable with SparkThrowable] + (depth, error.getClass, error.getErrorClass, error.getSqlState) + } + assert(overflow(jvmFailure)._3 == "ARITHMETIC_OVERFLOW", s"JVM writer: $jvmFailure") + assert( + overflow(nativeFailure) == overflow(jvmFailure), + s"native writer: $nativeFailure\nJVM writer: $jvmFailure") + Seq(nativeTable, jvmTable).foreach { table => + assert(spark.table(s"$catalog.$ns.$table").count() == 0, s"$table committed rows") + } + } + } + } + test("native acceleration: target-file-size rolls on iceberg-java's 1000-row cadence") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => @@ -3646,6 +3671,28 @@ class CometIcebergWriteActionSuite (depth, chain(depth).getMessage) } + /** + * Runs `write` against `nativeTable` with the native writer and against `jvmTable` with the JVM + * writer, expects both to fail, and returns the two failures. The native write has to go + * through [[CometIcebergWriteExec]], so a fallback to the JVM writer cannot pass as parity. + */ + private def failedWrites(nativeTable: String, jvmTable: String)( + write: String => Unit): (Throwable, Throwable) = { + var nativeFailure: Throwable = null + val nativePlans = withNativeEnabled { + capturePlans(spark, includeFailures = true) { + nativeFailure = intercept[Exception](write(nativeTable)) + } + } + val jvmFailure = intercept[Exception](write(jvmTable)) + assert( + nativePlans.exists(p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), + "expected the failed write to have gone through CometIcebergWriteExec. Plans:\n" + + nativePlans.mkString("\n--\n")) + (nativeFailure, jvmFailure) + } + private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) val plans = capturePlans(spark)(action) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 22e76a9abd1..5071c40cb61 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -52,7 +52,7 @@ import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo} +import org.apache.comet.{CometConf, CometExecIterator, CometNativeException, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} import org.apache.comet.rules.CometCoalesceShufflePartitions import org.apache.comet.serde.Config.ConfigMap @@ -3979,6 +3979,55 @@ class CometExecSuite extends CometTestBase { }) } + // https://github.com/apache/datafusion-comet/issues/6234. The final TopK reads its child's + // batches through an Arrow C stream, and Arrow Java hands native only the text of an exception + // thrown while producing one. The first batch is read on the JVM to derive the stream's schema, + // so the overflow has to land past it: row 90000 of a single file, read by one task. + test("TakeOrderedAndProjectExec: an input error past the first batch keeps Spark's exception") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTempPath { dir => + spark + .range(0, 100000, 1, 1) + .selectExpr("CAST(id AS INT) AS id") + .write + .parquet(dir.getCanonicalPath) + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("overflow_src") + val df = + sql(s"SELECT id + ${Int.MaxValue - 90000} AS v FROM overflow_src ORDER BY v LIMIT 5") + checkSparkError(df, "ARITHMETIC_OVERFLOW") + } + } + } + + // https://github.com/apache/datafusion-comet/issues/6234. CometSparkToColumnarExec exports its + // rows to native without reading a batch on the JVM first, so even an exception thrown for the + // first batch reached the user only as the text inside a CometNativeException. + test("SparkToColumnar keeps the exception its row input throws") { + val rows = spark.sparkContext.parallelize(1 to 10, 1).map { i => + if (i == 5) throw new IllegalStateException("injected row input failure") + Row(i) + } + val df = spark + .createDataFrame(rows, StructType(Seq(StructField("a", IntegerType)))) + .groupBy() + .sum("a") + assert( + collect(df.queryExecution.executedPlan) { case c: CometSparkToColumnarExec => c }.nonEmpty, + "expected the rows to reach native through CometSparkToColumnarExec:\n" + + df.queryExecution.executedPlan) + + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + def injected(error: Option[Throwable], engine: String): (Class[_], Int) = { + val chain = causeChain(error.getOrElse(fail(s"$engine did not fail"))) + assert(!chain.exists(_.isInstanceOf[CometNativeException]), s"$engine: ${chain.head}") + val depth = chain.indexWhere(t => + t.isInstanceOf[IllegalStateException] && t.getMessage == "injected row input failure") + assert(depth >= 0, s"$engine did not surface the injected failure: ${chain.head}") + (chain.head.getClass, depth) + } + assert(injected(cometError, "Comet") == injected(sparkError, "Spark")) + } + // Arrow Java ignores ArrowArray.offset on import, and a sliced boolean is the one array arrow-rs // exports with a non-zero offset, so native has to zero it at every level, not only the top. // https://github.com/apache/datafusion-comet/issues/6288 diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index d6237f76b8e..0485fb61818 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -29,6 +29,7 @@ import scala.reflect.ClassTag import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.BigIntVector import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} import org.apache.spark.executor.TaskMetrics @@ -40,11 +41,13 @@ import org.apache.spark.sql.comet.{CometExec, CometExecUtils, CometMetricNode, C import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.types.{LongType, StructField, StructType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.JsonProtocol -import org.apache.comet.{CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native} +import org.apache.comet.{CometArrowAllocator, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native} import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.vector.CometPlainVector /** * Regression tests for the native plan lifecycle: every `createPlan` must be balanced by exactly @@ -237,6 +240,56 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { } } + test("an input that throws after its first batch fails the plan with its own exception") { + val allocator = CometArrowAllocator.newChildAllocator("input-failure", 0, Long.MaxValue) + try { + withTaskContext(4600000L) { + val vector = new BigIntVector("test", allocator) + vector.allocateNew(2) + vector.set(0, 1L) + vector.set(1, 2L) + vector.setValueCount(2) + val first = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), 2) + val boom = new IllegalStateException("injected input failure") + // The stream's schema is derived from the first batch on the JVM, so a failure there + // propagates unchanged. Native pulls every later batch through the exported stream, and + // Arrow Java hands native only the text of what that pull threw. + val input = Iterator.single(first) ++ new Iterator[ColumnarBatch] { + override def hasNext: Boolean = throw boom + override def next(): ColumnarBatch = throw boom + } + val schema = StructType(Seq(StructField("test", LongType, nullable = false))) + val stream = CometArrowStream.fromColumnarBatchIter( + input, + schema, + CometArrowStream.NATIVE_TIMEZONE, + "input-failure") + val limitOp = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val iter = CometExec.getCometIterator( + Array(stream.asInstanceOf[Object]), + 1, + limitOp, + CometMetricNode(Map.empty), + 1, + 0, + None, + Seq.empty) + + assert(iter.hasNext && iter.next().numRows() == 2) + val thrown = intercept[Throwable](iter.hasNext) + assert(thrown eq boom, s"expected the input's own exception, got: $thrown") + + TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None) + assert( + CometArrowStream.inputFailure(stream).isEmpty, + "a completed task must not keep its input streams' failures") + } + } finally { + allocator.close() + } + } + test("a plan fed only by native scans returns its memory when its consumer stops early") { // See issue #2453. The sort keeps its sorted runs in memory, and the merge that reads them // spawns a Tokio task for each run. From 45e91beb705caae6afbc5116fe48a913b6551a6b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 13:52:03 +0000 Subject: [PATCH 07/89] docs: check timezone handling against the contributor guide in the PR review skills (#6340) * docs: check timezone handling against the contributor guide in the PR review skills Point the review skills at docs/source/contributor-guide/timezones.md (added in #6337) when a PR deals with timestamps or the session timezone. - review-comet-pr: read the page up front, and a cross-cutting "Timestamps and timezones" check covering the UTC label invariant, where the timezone comes from, Etc/UTC, NTZ values, tests, and which parts of the page a PR can make stale. - review-comet-expression-pr: the page in the guide table, serde and kernel checklist items, the page's testing guidance, and a common finding for mislabelled timestamp results. - review-comet-ffi-pr: timestamps cross the boundary unconverted and new producers label with NATIVE_TIMEZONE. - review-comet-iceberg-write-pr: timestamp partition values are UTC, and iceberg-rust's years/months depend on the +00:00 relabel. * docs: update the timezone review notes for CometTimeZone and the Iceberg partition kernels --- .../review-comet-expression-pr/SKILL.md | 30 ++++++++++---- .ai/skills/review-comet-ffi-pr/SKILL.md | 5 +++ .../review-comet-iceberg-write-pr/SKILL.md | 10 +++++ .ai/skills/review-comet-pr/SKILL.md | 40 +++++++++++++++++++ 4 files changed, 78 insertions(+), 7 deletions(-) diff --git a/.ai/skills/review-comet-expression-pr/SKILL.md b/.ai/skills/review-comet-expression-pr/SKILL.md index 2c83d9d15ba..17f8e2f04ce 100644 --- a/.ai/skills/review-comet-expression-pr/SKILL.md +++ b/.ai/skills/review-comet-expression-pr/SKILL.md @@ -30,11 +30,12 @@ bar, and the output format. This skill only covers what is specific to expressio ## Read the Contributor Guide First -| Doc | What you need from it | -| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------- | -| `docs/source/contributor-guide/adding_a_new_expression.md` | The serde contract, support levels, when to set the return type explicitly, shimming | -| `docs/source/contributor-guide/sql-file-tests.md` | The test framework expression PRs are expected to use, and every directive it supports | -| `docs/source/contributor-guide/optimizing_expressions.md` | The benchmark workflow and the no-regression rule, for PRs that change an existing kernel | +| Doc | What you need from it | +| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `docs/source/contributor-guide/adding_a_new_expression.md` | The serde contract, support levels, when to set the return type explicitly, shimming | +| `docs/source/contributor-guide/sql-file-tests.md` | The test framework expression PRs are expected to use, and every directive it supports | +| `docs/source/contributor-guide/optimizing_expressions.md` | The benchmark workflow and the no-regression rule, for PRs that change an existing kernel | +| `docs/source/contributor-guide/timezones.md` | The timestamp label rule, where the session timezone comes from, and how to test it, for datetime expressions, casts, and anything that returns a timestamp | Hold the diff against these. If the PR does something a guide says to do differently, either the PR is wrong or the guide is out of date, and you need to say which. @@ -77,6 +78,10 @@ Location: `spark/src/main/scala/org/apache/comet/serde/` - [ ] A change that routes a whole class of expressions to the JVM codegen dispatcher lists the new shapes it admits and tests one of each, including decimal results whose scale differs from the declared type and boolean inputs from sliced batches (#6424, #6425) +- [ ] A timezone-aware expression serializes `CometTimeZone.nativeId(expr.timeZoneId)`, the + timezone Spark stamped on it in a form native code can parse, rather than + `SQLConf.get.sessionLocalTimeZone` or the JVM default, and returns + `CometTimeZone.supportLevel` from `getSupportLevel` when `nativeId` gives `None` ### Registration in `QueryPlanSerde.scala` @@ -112,6 +117,13 @@ Location: `native/spark-expr/src/`, registered in `comet_scalar_funcs.rs`. - [ ] Comparisons, grouping and set operations treat `-0.0` and `0.0`, and NaN payloads, as Spark does, including inside arrays and structs. Where a DataFusion kernel treats them differently, the serde falls back (#5507, #5701). +- [ ] A `TimestampType` result is labelled `"UTC"`, in `return_type()` or `data_type()` and in the + arrays the kernel builds. Not the session timezone, and not `None`. +- [ ] Local-time work uses the session timezone from the proto. DataFusion's own datetime functions + take the timezone from the input's label, so wiring one in for a timezone-aware Spark + expression evaluates it in UTC unless the session timezone is passed explicitly. +- [ ] A UTC fast path matches a fixed list of UTC aliases, as `is_utc_timezone` in + `extract_date_part.rs` does, and sends every other timezone ID down the general path Before accepting a hand-written kernel, ask whether the function already exists upstream in DataFusion or the `datafusion-spark` crate. Comet prefers wiring an upstream function over carrying @@ -169,8 +181,10 @@ single file by appending a substring of its name to the suite argument. - [ ] Edge cases tested: empty input, overflow, boundary values, negative values - [ ] Both literal and column arguments tested, in every combination for multi-argument expressions. They take different code paths. -- [ ] Timezone handling tested for timestamp and datetime expressions, including a non-UTC session - timezone and timestamps with and without timezone +- [ ] Timestamp and datetime expressions are tested as "Testing timezone-sensitive code" in + `timezones.md` describes. That means several session timezones through `ConfigMatrix`, + including `Etc/UTC` and a zone with DST, both `TIMESTAMP` and `TIMESTAMP_NTZ` inputs, and a + result that is compared or fed to another expression rather than only projected. - [ ] SQL syntax gated with `MinSparkVersion` when it only parses on newer Spark - [ ] For a function from another project, such as an Iceberg transform, the expected values come from that project's Java implementation, which is what Spark runs, not from iceberg-rust @@ -254,3 +268,5 @@ reference in that doc is the only place they are documented. 5. **Missing `getSupportLevel`**, divergences left undeclared rather than marked `Incompatible` 6. **Version-specific Spark behavior implemented once**, with no shim 7. **Name collides with a DataFusion built-in** and no explicit return type +8. **Timestamp result mislabelled**, with the session timezone or no timezone instead of `"UTC"`. + It only shows once the result is compared or fed to another expression. diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md index b42ad33d1e3..88d47c1683c 100644 --- a/.ai/skills/review-comet-ffi-pr/SKILL.md +++ b/.ai/skills/review-comet-ffi-pr/SKILL.md @@ -138,6 +138,11 @@ a new subclass needs no new case as long as `getValueVector` returns an Arrow ve `ColumnarBatchArrowReader` decodes dictionary-encoded columns before export, and `reconcileStreamSchema` advertises the value type. A new reader that exports a dictionary would reach `ScanExec` as is, and only the cast in `build_record_batch` would unpack it. +- [ ] **Timestamps cross unconverted.** Both directions pass the raw microseconds. JVM producers + label `TimestampType` with `CometArrowStream.NATIVE_TIMEZONE`, which is `"UTC"`, and a new + producer must use it too rather than the session timezone. A change that shifts values by a + timezone at the boundary gives wrong answers. See "How Comet represents timestamps" in + `docs/source/contributor-guide/timezones.md`. ## 5. Memory Accounting diff --git a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md index 9752f865a5d..68e787ff1a6 100644 --- a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md +++ b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md @@ -111,6 +111,16 @@ only the latest. `Float.compare` does not ([#6138](https://github.com/apache/datafusion-comet/issues/6138)). Check equality, hashing, ordering and rendering for float, double, timestamp, timestamptz, binary and decimal. +- [ ] **Timestamp partition values are UTC.** Iceberg's `years`, `months`, `days` and `hours` + ignore the session timezone. For date and timestamp sources, `PartitionValueCalculator` + (`iceberg_partition_value.rs`) computes them with the same kernels in + `iceberg_funcs/temporal.rs` that key the sort in front of a clustered write, and those follow + iceberg-java's `DateTimeUtil`, including its pre-1970 rounding. A change that moves one of + these transforms back to iceberg-rust, or that makes the kernels read the column's timezone + label, breaks that agreement. iceberg-rust's `years` and `months` follow the label + ([apache/iceberg-rust#3142](https://github.com/apache/iceberg-rust/issues/3142)). Ask for a + test in a non-UTC session for a change here. "Partition values" in `iceberg-writes.md` lists + where iceberg-rust differs, and `timezones.md` covers how Comet labels timestamps. - [ ] **Partition paths use the Java renderers.** Directory names come from `CometLocationGenerator` / `partition_to_path` in `iceberg_partition_path.rs`, which follow iceberg-java's `partitionToPath`, with `java_float_string` for floats. A PR that calls diff --git a/.ai/skills/review-comet-pr/SKILL.md b/.ai/skills/review-comet-pr/SKILL.md index 050a60802ce..e93361b69f9 100644 --- a/.ai/skills/review-comet-pr/SKILL.md +++ b/.ai/skills/review-comet-pr/SKILL.md @@ -47,6 +47,11 @@ change. For a new operator, read `docs/source/contributor-guide/adding_a_new_operator.md` alongside this skill. There is no dedicated operator review skill yet. +For a PR that deals with timestamps or the session timezone, read +`docs/source/contributor-guide/timezones.md` as well, whichever areas it touches. Timezone handling +cuts across the serde, the native kernels, the scans and the JVM/native boundary, so no area skill +owns it. "Timestamps and timezones" in step 5 says what to check. + If the PR falls outside all of these, for example build, CI, docs only, or release tooling, this skill on its own is the review. @@ -176,6 +181,41 @@ the minor version is wrong for every earlier patch. CI builds only the newest pa it can't catch that (#6042, #5701). The pull request CI also runs only the default Spark profile, so logic that depends on the Spark version needs the matching `run-spark-*` labels (#6156). +### Timestamps and timezones + +Timezone bugs are easy to miss in review and in tests. A mislabelled timestamp column passes any +test that only projects it, and a result computed in the wrong timezone looks plausible. A PR deals +with timezones if it touches a datetime expression, a cast to or from a timestamp, how a scan reads +timestamps, the timestamp type at the JVM/native boundary, or anything that reads the session +timezone. Searching the diff flags most of these PRs: + +```shell +gh pr diff --repo apache/datafusion-comet | grep -inE 'time_?zone|zoneid|chrono_tz|timestamp(ntz)?type|timestampmicro|timestamp\(' +``` + +For such a PR, hold the diff against "The invariant" and "Guidelines" in +`docs/source/contributor-guide/timezones.md`. Look for these first: + +- A `TimestampType` value labelled with the session timezone, or with no timezone. Inside a native + plan every `TimestampType` value is labelled exactly `"UTC"`, and every `TimestampNTZType` value + has no timezone. Check the declared type and the arrays the code builds, not only the values. +- A timezone taken from the JVM default or the host. An expression uses the `timeZoneId` Spark + stamped on it, passed through `CometTimeZone.nativeId`, not `SQLConf.get.sessionLocalTimeZone`. +- A path gated on a UTC session. `Etc/UTC` is the session default on Ubuntu and Debian images, so + check what the gate does with it, and that the output there is still labelled `"UTC"` rather than + `"Etc/UTC"`. +- A timezone applied to a `TimestampNTZType` value, other than to convert it to `TimestampType`. +- Tests that use a single session timezone, or only project the result. "Testing timezone-sensitive + code" in the same page says what to ask for. + +`timezones.md` also describes specific code: where the `"UTC"` label is set, which serdes serialize +a timezone, how `CometTimeZone` rewrites timezone IDs, what `array_with_timezone` does, how the +scans adapt timestamps, and which expressions go through the codegen dispatcher. A PR that changes +any of these updates the page in the same PR. The page also documents some known limitations as +current behavior, such as chrono-tz's DST horizon and the timezone database versions. A fix for one +of the bugs tracked in [#6335](https://github.com/apache/datafusion-comet/issues/6335) usually +changes that text too. + ### Configuration New configs go in `CometConf.scala` and must follow From 8c783aa88104616dcf0f7876b4f7a9ba71e2bf31 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 13:56:14 +0000 Subject: [PATCH 08/89] feat: charge native write buffers to the task's memory pool (#6247) * feat: charge native write buffers to the task's memory pool Neither native writer reserved what it buffers, so a native Parquet or Iceberg write held its row groups in memory that no budget accounted for. A fanout Iceberg write multiplies that by the number of partitions a task writes, and when it outgrew the executor's memory overhead the container was killed. Both writers now register one memory consumer per task and resize it after every batch. ParquetWriterExec reserves parquet-rs's estimate of its in-progress row group. The Iceberg writer wraps iceberg-rust's ParquetWriterBuilder so that every open file reports what it holds, capped at the row-group size because iceberg-rust exposes only the bytes written plus the in-progress row group, and adds the rows its pacers hold back. Neither writer can spill, so a resize the pool refuses fails the task with a native out-of-memory error that Spark can retry. Closes #5648. Closes #6115. * fix: release an Iceberg file's reservation once its row groups leave memory A data file's share was its written size capped at one row group, so every fanout file that had flushed stayed charged a whole row group until it closed. A local file keeps nothing it has flushed, so a pool of eight 128 KiB row groups failed a twelve-partition write whose files held almost nothing. Each file's OutputFile is now wrapped so that its writer counts the bytes that leave memory on their way to storage, and the file reports what it has written less those. A local file releases a row group as it is handed on. S3 and GCS do not: OpenDAL collects what a file flushes into parts of at least 5 MiB and holds each part until the next one is complete or the file closes, so a file there keeps what it has flushed since its last upload charged. --- .../review-comet-iceberg-write-pr/SKILL.md | 28 +- .../contributor-guide/iceberg-writes.md | 34 +- .../contributor-guide/memory_management.md | 37 +- .../user-guide/latest/iceberg-writes.md | 12 + native/core/src/execution/memory_pools/mod.rs | 2 + .../src/execution/memory_pools/testing.rs | 81 ++ .../src/execution/operators/iceberg_common.rs | 2 +- .../src/execution/operators/iceberg_write.rs | 971 +++++++++++++++++- .../src/execution/operators/parquet_writer.rs | 98 +- .../comet/CometIcebergWriteActionSuite.scala | 68 ++ .../parquet/CometParquetWriterSuite.scala | 35 + 11 files changed, 1330 insertions(+), 38 deletions(-) create mode 100644 native/core/src/execution/memory_pools/testing.rs diff --git a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md index 68e787ff1a6..a835f45441a 100644 --- a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md +++ b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md @@ -152,6 +152,14 @@ Read the ownership table in the contributor guide before reviewing any change ne `SparkWrite.abort` keeps working. - [ ] The user-visible exception type is still what Spark's own write path throws on each Spark version ([#6143](https://github.com/apache/datafusion-comet/issues/6143)). +- [ ] **Every open file and every held-back row is reserved.** Files must be opened through + `MeteredParquetWriterBuilder`, and rows the writer holds outside iceberg-rust (the + `PartitionFeed`s, or anything a new feed holds back) must be added to what `run_write_task` reserves. A + builder that constructs `ParquetWriterBuilder` directly leaves its files out of the task's + memory reservation, so a wide fanout write grows past the pool instead of failing its task. + A storage scheme newly supported for writes needs its entry in + `StorageWrites::for_location`: a file reports its flushed row groups until its storage has + written them out, which a local file does at once and an object store only part by part. ## 5. Plan Shape @@ -189,16 +197,16 @@ Read the ownership table in the contributor guide before reviewing any change ne ## 7. Tests -| Suite | Covers | -| ----------------------------------- | ------------------------------------------------------------------------------------ | -| `CometIcebergWriteActionSuite` | Both layers end to end, parity with iceberg-java, DML, failure cleanup, AQE re-plans | -| `CometIcebergWriteDetectionSuite` | One case per eligibility rule | -| `CometIcebergSystemFunctionSuite` | Native partition transforms keeping partitioned writes native | -| `CometIcebergRewriteActionSuite` | `rewrite_data_files` through the split plan and the native writer | -| `IcebergWriteProtoTranslationSuite` | Property to `IcebergParquetWriteSettings` translation | -| Rust tests in `iceberg_write.rs` | Rolling, fanout order, clustered checks, cleanup guard, manifest round trip | -| `iceberg_partition_path.rs` tests | Partition path rendering against iceberg-java | -| `CometIcebergWriteBenchmark` | Native versus iceberg-java throughput | +| Suite | Covers | +| ----------------------------------- | ----------------------------------------------------------------------------------------------- | +| `CometIcebergWriteActionSuite` | Both layers end to end, parity with iceberg-java, DML, failure cleanup, AQE re-plans | +| `CometIcebergWriteDetectionSuite` | One case per eligibility rule | +| `CometIcebergSystemFunctionSuite` | Native partition transforms keeping partitioned writes native | +| `CometIcebergRewriteActionSuite` | `rewrite_data_files` through the split plan and the native writer | +| `IcebergWriteProtoTranslationSuite` | Property to `IcebergParquetWriteSettings` translation | +| Rust tests in `iceberg_write.rs` | Rolling, fanout order, clustered checks, cleanup guard, manifest round trip, memory reservation | +| `iceberg_partition_path.rs` tests | Partition path rendering against iceberg-java | +| `CometIcebergWriteBenchmark` | Native versus iceberg-java throughput | Ask specifically: diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 073b7c836d9..a90cff08075 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -272,10 +272,22 @@ Points where Comet adapts iceberg-rust to match iceberg-java: ([#5776](https://github.com/apache/datafusion-comet/issues/5776)). Manifest order becomes the row order of an unordered read, so any map iteration that reaches the output needs the same care. +iceberg-rust keeps each open file writer private inside its rolling and partitioning writers, so +Comet wraps `ParquetWriterBuilder` in `MeteredParquetWriterBuilder`, whose files report what they +hold in memory. It hands `ParquetWriter` each file's `OutputFile` behind a `CountedOutput`, whose +writer counts the bytes that leave memory on their way to storage, and a file reports what it has +written less those. When they leave depends on the storage (`StorageWrites`). After every batch +`run_write_task` resizes the task's reservation to what the open files report plus the rows each +`PartitionFeed` holds, for the dictionary choice or for pacing, and a resize the pool refuses fails +the task. What that figure covers, and what it misses, is described under +[Native writers](memory_management.md#native-writers). + `FileIO` comes from `load_file_io` in `iceberg_common.rs`, shared with the native scan. It picks the storage backend from the data location's scheme and wires in Comet's S3 credential bridge when one is configured. For writes the bridge fails closed: if a configured provider cannot initialize, -the task fails rather than writing with the default credential chain. +the task fails rather than writing with the default credential chain. A storage scheme newly +supported for writes also needs its entry in `StorageWrites::for_location`, which treats an unknown +scheme as one that uploads in parts. The JVM eligibility gate fails closed before building that `FileIO` for an `s3` or `s3a` data location. It checks effective `fs.s3a.*` Hadoop settings against the six keys translated by @@ -396,16 +408,16 @@ the writer: run the write suites and the Iceberg Spark tests. The pin policy is ## Testing -| Suite | What it covers | -| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, AQE re-planning. | -| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | -| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | -| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | -| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | -| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, partition path rendering, partition values past `chrono`'s calendar. | -| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | -| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | +| Suite | What it covers | +| ---------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, a fanout write outgrowing the memory pool, AQE re-planning. | +| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | +| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | +| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | +| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | +| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, memory reservation, partition path rendering, partition values past `chrono`'s calendar. | +| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | +| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | The Comet suites run against the Iceberg version each Spark profile pins in `spark/pom.xml`: 1.5.2 for Spark 3.4, 1.8.1 for 3.5, 1.10.0 for 4.0 and 4.2, and 1.11.0 for 4.1. Only the default profile diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 2a48baa3f18..7bd16d880f6 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -375,6 +375,40 @@ Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservat An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. +### Native writers + +Both native writers reserve what they hold between batches through a single consumer per task, +`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch. Neither can spill, so +when the pool refuses a resize the task fails with a `CometNativeException` whose message starts +`Additional allocation failed for` and names the consumer. That is a task failure Spark can retry. +Unreserved, the same memory would count only toward the container limit, where exceeding it kills +the executor. + +- `ParquetWriterExec` writes one file per task and reserves parquet-rs's estimate of the file's + in-progress row group (`ArrowWriter::memory_size`). The estimate counts encoded pages, encoder + buffers, dictionaries and Bloom filters. The writer keeps parquet-rs's default row-group limit of + 1Mi rows and sets no byte limit, so a wide schema can hold a large row group + ([#5304](https://github.com/apache/datafusion-comet/issues/5304)). +- The Iceberg writer keeps one file open per partition in a fanout write, so its reservation grows + with the number of partitions a task writes. iceberg-rust's `ParquetWriter` does not expose + parquet-rs's `memory_size`, only every byte the file has written, its in-progress row group's + encoded size included. So each file's writer counts the bytes that leave memory on their way to + storage, and the file reports what it has written less those. A local file writes a flushed row + group out at once, so it reports its in-progress row group and the few KiB parquet-rs buffers in + front of storage. S3 and GCS take a file in parts of at least 5 MiB, and OpenDAL holds each part + in memory until the next one is complete or the file closes, so there a file also reports what + it has flushed since its last part was uploaded. At the default row-group size + (`write.parquet.row-group-size-bytes`, 128 MiB) that is the last row group it flushed. The + reservation also covers the rows each partition holds back, first for its dictionary choice and + then until they fill the 1000-row unit the rolling writer is fed in. It does not cover what + parquet-rs holds beyond the encoded size: dictionary hash tables, unencoded dictionary indices + and buffer capacity. Native writes decline Bloom filters today, and neither figure would include + them. + +The writers register one consumer per task rather than one per open file because every consumer +registered with `fair_unified` lowers the share of every other consumer in the task. A consumer per +partition would shrink the task's other operators' shares as a fanout write widened. + ## Crossing the FFI boundary Batches move between the JVM and native over the Arrow C Data and C Stream interfaces, which are @@ -425,7 +459,8 @@ The pool tracks _declared reservations_. Container RSS counts _pages the process diverge for several structural reasons: - **Undeclared allocations.** Arrow array builders, expression kernels producing intermediate - arrays, decompression buffers, Parquet metadata structures, `object_store` request buffers, and + arrays, decompression buffers, Parquet metadata structures, the part of a native writer's buffers + its estimate misses (see [Native writers](#native-writers)), `object_store` request buffers, and tokio's own machinery all allocate without reserving. Only operators that were explicitly written to reserve show up in the pool. - **Rounding and padding.** Arrow buffers are padded to 64-byte boundaries and builders grow by diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index a89c397d493..831dcc20386 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -299,6 +299,18 @@ failure), the error propagates as an ordinary Spark task failure and Spark's tas re-executes it — through the native writer again. Retries cannot collide: each attempt's task attempt id is embedded in its data file names. +The native writer's buffers are charged to Comet's memory pool, the off-heap budget Comet's other +native operators draw on, where iceberg-java's buffers sit on the JVM heap. A fanout write keeps a +data file open for every partition a task writes to. Each open file holds the row group it is +writing in memory, up to `write.parquet.row-group-size-bytes`, and on S3 or GCS also the last row +group it flushed, which is uploaded once the next one is complete or the file closes. So a task +writing to many partitions needs memory in proportion to them. When the pool cannot grant it, the +task fails with a `CometNativeException` reading `Additional allocation failed for IcebergWriteExec` +instead of exceeding the executor's memory, and Spark retries it like any other task failure. Such a +write fits in less memory with the fanout writer disabled (`write.spark.fanout.enabled=false`): +Spark then sorts each task's rows by partition, and the task keeps one file open at a time. A +smaller row-group size also helps. Otherwise the write needs a larger `spark.memory.offHeap.size`. + Partial results are never committed. The commit set is exactly the commit messages returned by successful tasks — a failed task contributes none — and if the job fails, the driver-side commit operator aborts without committing anything. A failed task attempt also deletes the diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 1fc4db2021c..e2fd2e58a73 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -22,6 +22,8 @@ mod plan_pool; mod spark_memory; mod spill_replay; mod task_shared; +#[cfg(test)] +pub(crate) mod testing; mod unified_pool; use datafusion::execution::memory_pool::{MemoryPool, TrackConsumersPool, UnboundedMemoryPool}; diff --git a/native/core/src/execution/memory_pools/testing.rs b/native/core/src/execution/memory_pools/testing.rs new file mode 100644 index 00000000000..ce1603a26f5 --- /dev/null +++ b/native/core/src/execution/memory_pools/testing.rs @@ -0,0 +1,81 @@ +// 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. + +//! Memory pools for tests of operators that reserve memory. + +use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::Result; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, MemoryReservation}; + +/// A [`GreedyMemoryPool`] that also records the most it ever had reserved, so a test can check +/// what an operator reserved while it ran, not only that it gave everything back. +#[derive(Debug)] +pub(crate) struct PeakMemoryPool { + inner: GreedyMemoryPool, + peak: AtomicUsize, +} + +impl PeakMemoryPool { + pub(crate) fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + peak: AtomicUsize::new(0), + } + } + + pub(crate) fn peak(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + fn record(&self) { + self.peak + .fetch_max(self.inner.reserved(), Ordering::Relaxed); + } +} + +impl fmt::Display for PeakMemoryPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "PeakMemoryPool({}, peak {})", self.inner, self.peak()) + } +} + +impl MemoryPool for PeakMemoryPool { + fn name(&self) -> &str { + "PeakMemoryPool" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record(); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } +} diff --git a/native/core/src/execution/operators/iceberg_common.rs b/native/core/src/execution/operators/iceberg_common.rs index f4ea3f7983b..262325c30cb 100644 --- a/native/core/src/execution/operators/iceberg_common.rs +++ b/native/core/src/execution/operators/iceberg_common.rs @@ -574,7 +574,7 @@ fn env_region_present() -> bool { /// The JVM write gate (`CometIcebergNativeWrite.storageScheme`) mirrors this rule exactly, case /// included. Change both together, and keep the cases in /// `scheme_of_extracts_scheme_from_all_uri_forms` in step with its `storageScheme` test. -fn scheme_of(path: &str) -> &str { +pub(crate) fn scheme_of(path: &str) -> &str { match path.split_once(':') { Some((scheme, _)) if !scheme.is_empty() && !scheme.contains('/') => scheme, _ => "file", diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index f7d5312fd92..ba9b08bd9da 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -27,6 +27,7 @@ use std::collections::HashMap; use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use arrow::array::{ @@ -35,8 +36,10 @@ use arrow::array::{ }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef}; +use bytes::Bytes; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -48,12 +51,13 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, }; +use futures::stream::BoxStream; use futures::TryStreamExt; use iceberg::arrow::arrow_struct_to_literal; -use iceberg::io::FileIO; +use iceberg::io::{FileIO, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage}; use iceberg::spec::{ - DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey, PartitionSpec, - PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, + DataFile, DataFileBuilder, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey, + PartitionSpec, PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, Struct as IcebergStruct, StructType, }; use iceberg::writer::base_writer::data_file_writer::{DataFileWriter, DataFileWriterBuilder}; @@ -61,16 +65,19 @@ use iceberg::writer::file_writer::location_generator::{ DefaultFileNameGenerator, LocationGenerator, }; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; -use iceberg::writer::file_writer::ParquetWriterBuilder; +use iceberg::writer::file_writer::{ + FileWriter, FileWriterBuilder, ParquetWriter, ParquetWriterBuilder, +}; use iceberg::writer::partitioning::clustered_writer::ClusteredWriter; use iceberg::writer::partitioning::fanout_writer::FanoutWriter; use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter; -use iceberg::writer::IcebergWriterBuilder; +use iceberg::writer::{CurrentFileStatus, IcebergWriterBuilder}; use iceberg::ErrorKind; #[cfg(test)] use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use serde::{Deserialize, Serialize}; use datafusion_comet_proto::spark_operator::{ CompressionCodec as ProtoCompressionCodec, IcebergParquetWriteSettings, IcebergWrite, @@ -79,16 +86,20 @@ use datafusion_comet_proto::spark_operator::{ use crate::cloud::s3::credential_bridge::AccessMode; use crate::errors::CometError; -use crate::execution::operators::iceberg_common::load_file_io; +use crate::execution::operators::iceberg_common::{load_file_io, scheme_of}; use crate::execution::operators::iceberg_dictionary::DictionaryChooser; use crate::execution::operators::iceberg_partition_path::{ partition_to_path, CometLocationGenerator, }; use crate::execution::operators::iceberg_partition_value::PartitionValueCalculator; -/// The writer every partition gets: iceberg-rust's own data file writer, rolling Parquet files. -type PartitionDataFileWriter = - DataFileWriter; +/// The writer every partition gets: iceberg-rust's own data file writer, rolling Parquet files +/// that report what they hold in memory. +type PartitionDataFileWriter = DataFileWriter< + MeteredParquetWriterBuilder, + TrackingLocationGenerator, + DefaultFileNameGenerator, +>; /// Opens each partition's data file writer with the Parquet writer properties chosen for that /// partition. @@ -99,7 +110,8 @@ type PartitionDataFileWriter = /// `ParquetWriterBuilder` -> `RollingFileWriterBuilder` -> `DataFileWriterBuilder` chain once per /// partition, with that partition's properties. Every file the partition rolls into keeps them. /// The location and file name generators are shared, so file names still count up across the -/// task and every location is still tracked for cleanup. +/// task and every location is still tracked for cleanup, and the files all report to one +/// [`OpenFileMemory`]. struct PartitionWriterBuilder { schema: IcebergSchemaRef, target_file_size: usize, @@ -107,6 +119,8 @@ struct PartitionWriterBuilder { location_generator: TrackingLocationGenerator, file_name_generator: DefaultFileNameGenerator, properties: Arc, + open_files: OpenFileMemory, + storage: StorageWrites, } #[async_trait::async_trait] @@ -117,8 +131,13 @@ impl IcebergWriterBuilder for PartitionWriterBuilder { let properties = self .properties .take(partition_key.as_ref().map(PartitionKey::data)); + let parquet_builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new(properties, Arc::clone(&self.schema)), + open_files: self.open_files.clone(), + storage: self.storage, + }; let rolling_builder = RollingFileWriterBuilder::new( - ParquetWriterBuilder::new(properties, Arc::clone(&self.schema)), + parquet_builder, self.target_file_size, self.file_io.clone(), self.location_generator.clone(), @@ -175,6 +194,335 @@ impl PartitionProperties { } } +/// What a task's open data files hold in memory between them. +/// +/// iceberg-rust keeps each open file writer private inside its rolling and partitioning writers, +/// so the files report their own shares here through [`MeteredParquetWriter`], and `run_write_task` +/// reserves the total. A fanout write keeps one file open per partition, so this is what grows +/// with the partition count. +#[derive(Clone, Debug, Default)] +struct OpenFileMemory(Arc); + +impl OpenFileMemory { + fn bytes(&self) -> usize { + self.0.load(Ordering::Relaxed) + } + + fn share(&self) -> OpenFileShare { + OpenFileShare { + total: self.clone(), + bytes: 0, + } + } +} + +/// One open file's part of [`OpenFileMemory`], given back when the file closes, or when it is +/// dropped because the task failed mid-write. +#[derive(Debug)] +struct OpenFileShare { + total: OpenFileMemory, + bytes: usize, +} + +impl OpenFileShare { + fn set(&mut self, bytes: usize) { + if bytes > self.bytes { + self.total + .0 + .fetch_add(bytes - self.bytes, Ordering::Relaxed); + } else { + self.total + .0 + .fetch_sub(self.bytes - bytes, Ordering::Relaxed); + } + self.bytes = bytes; + } +} + +impl Drop for OpenFileShare { + fn drop(&mut self) { + self.set(0); + } +} + +/// [`ParquetWriterBuilder`] whose files report what they hold in memory to the task's +/// [`OpenFileMemory`]. +#[derive(Clone, Debug)] +struct MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder, + open_files: OpenFileMemory, + /// How the files' storage takes the row groups they flush. + storage: StorageWrites, +} + +impl FileWriterBuilder for MeteredParquetWriterBuilder { + type R = MeteredParquetWriter; + + async fn build(&self, output_file: OutputFile) -> iceberg::Result { + let released = Arc::new(AtomicUsize::new(0)); + let output_file = CountedOutput::wrap(output_file, self.storage, Arc::clone(&released)); + Ok(MeteredParquetWriter { + inner: self.inner.build(output_file).await?, + share: self.open_files.share(), + released, + }) + } +} + +/// iceberg-rust's [`ParquetWriter`], reporting after every write how much of its file it still +/// holds in memory. +/// +/// parquet-rs keeps a file's in-progress row group in memory and hands it to storage once it +/// reaches the row-group size. `ParquetWriter` does not expose parquet-rs's estimate of that +/// memory, only `current_written_size`: every byte the file has written, the in-progress row +/// group's encoded size included. So the file's writer counts the bytes that leave memory on +/// their way to storage (see [`CountedOutput`]), and the share is what the file has written less +/// those: its in-progress row group, the few KiB parquet-rs buffers in front of storage, and what +/// the storage still holds of the row groups the file has flushed (see [`StorageWrites`]). +/// +/// The share does not see what parquet-rs keeps beyond the encoded estimate: dictionary +/// encoders' hash tables and unencoded indices, and buffer capacity past what is used. Nor would +/// it see Bloom filters, which parquet-rs sizes when a file opens and the eligibility gate +/// declines today. +struct MeteredParquetWriter { + inner: ParquetWriter, + share: OpenFileShare, + /// The bytes of the file that have left memory, counted by its [`CountingFileWrite`]. + released: Arc, +} + +impl FileWriter for MeteredParquetWriter { + async fn write(&mut self, batch: &RecordBatch) -> iceberg::Result<()> { + self.inner.write(batch).await?; + let released = self.released.load(Ordering::Relaxed); + self.share + .set(self.inner.current_written_size().saturating_sub(released)); + Ok(()) + } + + async fn close(self) -> iceberg::Result> { + // Closing flushes the last row group, and the file's buffers go with the writer. The share + // is given back when it drops at the end of this call, whether or not the close succeeded. + let Self { + inner, + share: _share, + .. + } = self; + inner.close().await + } +} + +impl CurrentFileStatus for MeteredParquetWriter { + fn current_file_path(&self) -> String { + self.inner.current_file_path() + } + + fn current_row_num(&self) -> usize { + self.inner.current_row_num() + } + + fn current_written_size(&self) -> usize { + self.inner.current_written_size() + } +} + +/// How the storage behind a task's data files takes the row groups a file flushes, which decides +/// when they leave memory. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum StorageWrites { + /// A local file is written as it goes, so a flushed row group leaves memory at once. + Through, + /// S3 and GCS take a file in parts. OpenDAL collects what a file flushes until it makes a part + /// of at least [`MIN_PART_BYTES`], and holds each part until the next one is complete or the + /// file closes, because only then does it know whether the part is the last. So a file whose + /// row groups are at least a part each keeps the last one it flushed in memory. + InParts, +} + +/// The smallest part OpenDAL uploads, its `write_multi_min_size` for both S3 and GCS. +const MIN_PART_BYTES: usize = 5 * 1024 * 1024; + +impl StorageWrites { + /// How the storage `load_file_io` builds for `location` takes a file. Native writes support + /// local storage (`file`, `memory`), S3 and GCS, and any other scheme is taken to upload in + /// parts, the way that reserves more. + fn for_location(location: &str) -> Self { + match scheme_of(location) { + "file" | "memory" => StorageWrites::Through, + _ => StorageWrites::InParts, + } + } +} + +/// What a file's storage holds of the bytes the file has handed it, following [`StorageWrites`]. +#[derive(Debug)] +struct HeldByStorage { + writes: StorageWrites, + /// Handed over since the last part was complete. + collecting: usize, + /// The last complete part, held until the next one. + last_part: usize, +} + +impl HeldByStorage { + fn new(writes: StorageWrites) -> Self { + Self { + writes, + collecting: 0, + last_part: 0, + } + } + + /// Takes `len` more bytes and returns how many bytes leave memory as it does: for a local + /// file these bytes, and for a file taken in parts the part held until now, if these complete + /// the next one. + fn hand_on(&mut self, len: usize) -> usize { + match self.writes { + StorageWrites::Through => len, + StorageWrites::InParts => { + self.collecting += len; + if self.collecting < MIN_PART_BYTES { + return 0; + } + std::mem::replace(&mut self.last_part, std::mem::take(&mut self.collecting)) + } + } + } +} + +/// The [`Storage`] behind one data file's [`OutputFile`], which opens the file's writer as a +/// [`CountingFileWrite`]. +/// +/// `ParquetWriter` opens its file's writer through the `OutputFile` it is built with, and an +/// `OutputFile` can only be made over a `Storage`. This one serves the one file it was made for, +/// through the `OutputFile` iceberg-rust made for it, and refuses what `ParquetWriter` never does +/// with its file. Serde exists only to satisfy the `Storage` typetag supertraits; Comet never +/// serializes it. +#[derive(Debug, Default, Serialize, Deserialize)] +struct CountedOutput { + #[serde(skip)] + file: Option, +} + +#[derive(Debug)] +struct CountedFile { + output: OutputFile, + writes: StorageWrites, + released: Arc, +} + +impl CountedOutput { + /// `output`, its writer counting into `released` the bytes that leave memory. + fn wrap(output: OutputFile, writes: StorageWrites, released: Arc) -> OutputFile { + let path = output.location().to_string(); + let file = CountedFile { + output, + writes, + released, + }; + OutputFile::new(Arc::new(CountedOutput { file: Some(file) }), path) + } + + fn file(&self, path: &str) -> iceberg::Result<&CountedFile> { + self.file + .as_ref() + .filter(|file| file.output.location() == path) + .ok_or_else(|| { + iceberg::Error::new( + ErrorKind::Unexpected, + format!("{path} is not the data file this output was made for"), + ) + }) + } + + fn unsupported(operation: &str, path: &str) -> iceberg::Error { + iceberg::Error::new( + ErrorKind::FeatureUnsupported, + format!("{operation} is not supported on the output of data file {path}"), + ) + } +} + +#[async_trait::async_trait] +#[typetag::serde(name = "CometCountedOutputStorage")] +impl Storage for CountedOutput { + async fn exists(&self, path: &str) -> iceberg::Result { + self.file(path)?.output.exists().await + } + + async fn metadata(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("metadata", path)) + } + + async fn read(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("read", path)) + } + + async fn reader(&self, path: &str) -> iceberg::Result> { + Err(Self::unsupported("reader", path)) + } + + async fn write(&self, path: &str, _bs: Bytes) -> iceberg::Result<()> { + Err(Self::unsupported("write", path)) + } + + async fn writer(&self, path: &str) -> iceberg::Result> { + let file = self.file(path)?; + Ok(Box::new(CountingFileWrite { + inner: file.output.writer().await?, + storage: HeldByStorage::new(file.writes), + released: Arc::clone(&file.released), + })) + } + + async fn delete(&self, path: &str) -> iceberg::Result<()> { + // `ParquetWriter` deletes a file it closes without having written a row. + self.file(path)?.output.delete().await + } + + async fn delete_prefix(&self, path: &str) -> iceberg::Result<()> { + Err(Self::unsupported("delete_prefix", path)) + } + + async fn delete_stream(&self, _paths: BoxStream<'static, String>) -> iceberg::Result<()> { + Err(Self::unsupported("delete_stream", "")) + } + + fn new_input(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("new_input", path)) + } + + fn new_output(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("new_output", path)) + } +} + +/// A data file's writer, counting into `released` the bytes that leave memory as the file hands +/// them to the storage's own writer, `inner`. +/// +/// parquet-rs hands a file's bytes on when it flushes a row group: everything it has written +/// since the last flush, except what is still in its write buffer. +struct CountingFileWrite { + inner: Box, + storage: HeldByStorage, + released: Arc, +} + +#[async_trait::async_trait] +impl FileWrite for CountingFileWrite { + async fn write(&mut self, bs: Bytes) -> iceberg::Result<()> { + let len = bs.len(); + self.inner.write(bs).await?; + self.released + .fetch_add(self.storage.hand_on(len), Ordering::Relaxed); + Ok(()) + } + + async fn close(&mut self) -> iceberg::Result { + self.inner.close().await + } +} + /// [`CometLocationGenerator`] that records every location it hands to a file writer. /// /// iceberg-rust's writers keep the `DataFile`s they have finalized private until `close`, and @@ -471,6 +819,10 @@ impl ExecutionPlan for IcebergWriteExec { // Time spent inside the iceberg-rust writer stack (write + close), excluding time spent // waiting on the upstream input stream. Surfaced on the JVM exec's SQL metrics by name. let write_time = MetricBuilder::new(&self.metrics).subset_time("write_time", partition); + // One consumer for the whole task, however many files it opens: a consumer per file + // would shrink every other consumer's share of the fair pool as a fanout write widened. + let reservation = MemoryConsumer::new(format!("IcebergWriteExec[{partition}]")) + .register(context.memory_pool()); let input_stream = self.input.execute(partition, context)?; let common = Arc::clone(&self.common); let iceberg_schema = Arc::clone(&self.iceberg_schema); @@ -492,6 +844,7 @@ impl ExecutionPlan for IcebergWriteExec { partition_id, task_attempt_id, write_time, + reservation, ) .await?; // The guard is still armed: until the output batch reaches the JVM nothing else knows @@ -561,6 +914,10 @@ impl DisplayAs for IcebergWriteExec { /// Iceberg field IDs, and routes through `UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter` /// depending on `writer_mode`. /// +/// After every batch, `reservation` is resized to what the task's writers hold between batches: +/// the open files' shares (see [`MeteredParquetWriter`]) and the rows waiting in their +/// [`PartitionFeed`]s. The writer cannot spill, so a reservation the pool refuses fails the task. +/// /// On success the still-armed [`AbortOnDrop`] is returned along with the data files: the caller /// owns cleanup until the JVM acknowledges the output after recording its locations. #[allow(clippy::too_many_arguments)] @@ -574,6 +931,7 @@ async fn run_write_task( partition_id: Option, task_attempt_id: Option, write_time: Time, + reservation: MemoryReservation, ) -> DFResult<(Vec, AbortOnDrop)> { // The JVM exec wrapper stamps both ids per task; a missing id means the plan template was // executed directly, and defaulting would make every task collide on the same file names. @@ -623,6 +981,7 @@ async fn run_write_task( writer_properties, &target_schema, )?)); + let open_files = OpenFileMemory::default(); let data_file_builder = PartitionWriterBuilder { schema: Arc::clone(&iceberg_schema), target_file_size: common.target_file_size_bytes as usize, @@ -630,6 +989,8 @@ async fn run_write_task( location_generator: location_generator.clone(), file_name_generator, properties: Arc::clone(&properties), + open_files: open_files.clone(), + storage: StorageWrites::for_location(&common.data_location), }; let mut abort_guard = AbortOnDrop { file_io, @@ -672,8 +1033,10 @@ async fn run_write_task( if nests_floats { decorated = drop_unwritten_values(decorated)?; } - let _timer = write_time.timer(); + let timer = write_time.timer(); writer.write(decorated, &properties).await?; + timer.done(); + reservation.try_resize(open_files.bytes() + writer.pending_bytes())?; } let _timer = write_time.timer(); writer.close(&properties).await @@ -718,6 +1081,23 @@ enum InnerWriter { } impl InnerWriter { + /// Memory held by the rows waiting in the feeds: those held back for a dictionary choice, and + /// up to `ROWS_DIVISOR - 1` rows for the rest of a unit, for every partition the task has + /// seen, in a fanout write. + fn pending_bytes(&self) -> usize { + match self { + InnerWriter::Unpartitioned(_, feed) => feed.pending_bytes(), + InnerWriter::Fanout(_, _, fanout) => fanout + .feeds + .values() + .map(|(_, feed)| feed.pending_bytes()) + .sum(), + InnerWriter::Clustered(_, _, live) => { + live.as_ref().map_or(0, |(_, feed)| feed.pending_bytes()) + } + } + } + /// Writes `batch` through the writer this task built, in the [`ROWS_DIVISOR`]-row units the /// rolling writer needs to re-check the target file size on iceberg-java's cadence. Rows are /// fed after partition splitting, so each partition's dictionary choice is made from its own @@ -1187,6 +1567,8 @@ struct RowPacer { /// Rows handed in but not yet handed over; fewer than [`ROWS_DIVISOR`] in total. pending: Vec, pending_rows: usize, + /// Memory held by `pending`. + pending_bytes: usize, } impl RowPacer { @@ -1195,6 +1577,7 @@ impl RowPacer { slicer, pending: Vec::new(), pending_rows: 0, + pending_bytes: 0, } } @@ -1217,6 +1600,7 @@ impl RowPacer { if offset < rows { let rest = self.slicer.detach(&batch, offset, rows - offset)?; self.pending_rows += rest.num_rows(); + self.pending_bytes += rest.get_array_memory_size(); self.pending.push(rest); } Ok(units) @@ -1243,6 +1627,7 @@ impl RowPacer { fn concat_pending(&mut self) -> DFResult { self.pending_rows = 0; + self.pending_bytes = 0; if self.pending.len() == 1 { return Ok(self.pending.pop().expect("pending has one batch")); } @@ -1295,6 +1680,12 @@ impl PartitionFeed { self.held.as_ref().map_or(0, |held| held.bytes) } + /// Memory taken by every row the feed has not handed to the writer: those held back for the + /// dictionary choice and those waiting in the pacer for the rest of their unit. + fn pending_bytes(&self) -> usize { + self.held_bytes() + self.pacer.pending_bytes + } + /// The complete units `batch` makes available, in row order. fn push( &mut self, @@ -2230,6 +2621,75 @@ mod tests { } } + /// The task reserves what the pacers hold back, so the count has to follow the rows: up while + /// they wait, and back to nothing once they are handed over. + #[test] + fn pacer_counts_the_memory_of_the_rows_it_holds_back() { + for gather in [false, true] { + let mut pacer = RowPacer::new(RowSlicer { gather }); + assert!(pacer.push(int_batch(800)).unwrap().is_empty()); + assert!( + pacer.pending_bytes >= 800 * std::mem::size_of::(), + "gather={gather}: {} bytes for 800 held rows", + pacer.pending_bytes + ); + // Completing the unit hands every held row over. + assert_eq!(pacer.push(int_batch_from(800, 200)).unwrap().len(), 1); + assert_eq!(pacer.pending_bytes, 0, "gather={gather}"); + assert!(pacer.push(int_batch(300)).unwrap().is_empty()); + assert!(pacer.pending_bytes > 0, "gather={gather}"); + assert!(pacer.flush().unwrap().is_some()); + assert_eq!(pacer.pending_bytes, 0, "gather={gather}"); + } + } + + /// A local file keeps nothing it hands on. S3 and GCS take a file in parts: OpenDAL collects + /// what it is handed until it makes a part of at least `MIN_PART_BYTES`, and holds each part + /// until the next one is complete, so a part leaves memory only when the next one is made. + #[test] + fn storage_taken_in_parts_holds_the_last_part_until_the_next() { + let mib = 1024 * 1024; + let mut local = HeldByStorage::new(StorageWrites::Through); + assert_eq!(local.hand_on(mib), mib); + + let mut parts = HeldByStorage::new(StorageWrites::InParts); + // Small row groups are collected until they make a part, which is then held. + for _ in 0..5 { + assert_eq!(parts.hand_on(mib), 0); + } + // A row group of a part or more is a part of its own, and sends the one held before it. + assert_eq!(parts.hand_on(6 * mib), 5 * mib); + assert_eq!(parts.hand_on(mib), 0); + assert_eq!(parts.hand_on(4 * mib), 6 * mib); + assert_eq!(parts.hand_on(8 * mib), 5 * mib); + } + + #[test] + fn only_local_storage_writes_through() { + for location in [ + "file:///tmp/warehouse", + "/tmp/warehouse", + "memory:///warehouse", + ] { + assert_eq!( + StorageWrites::for_location(location), + StorageWrites::Through, + "{location}" + ); + } + for location in [ + "s3://bucket/warehouse", + "s3a://bucket/warehouse", + "gs://bucket/w", + ] { + assert_eq!( + StorageWrites::for_location(location), + StorageWrites::InParts, + "{location}" + ); + } + } + // -- Integration tests against the real iceberg-rust writer stack --------- mod integration { @@ -2237,15 +2697,18 @@ mod tests { use arrow::array::{BinaryArray, Int32Array, StringArray, TimestampMicrosecondArray}; use arrow::datatypes::TimeUnit; use datafusion::common::Result as DFResult; + use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_comet_proto::spark_operator::{ CompressionCodec as ProtoCodec, IcebergParquetWriteSettings, IcebergWriteCommon, IcebergWriterMode as ProtoIcebergWriterMode, }; use futures::StreamExt; + use iceberg::io::FileIOBuilder; use iceberg::spec::{ Manifest, NestedField, PartitionSpec, PrimitiveType, Schema, Transform, Type, }; + use iceberg_storage_opendal::OpenDalStorageFactory; use parquet::file::properties::WriterProperties; use std::collections::HashMap; use std::path::PathBuf; @@ -2253,6 +2716,8 @@ mod tests { use std::time::Duration; use tempfile::TempDir; + use crate::execution::memory_pools::testing::PeakMemoryPool; + fn user_schema() -> SchemaRef { Arc::new(ArrowSchema::new(vec![ Field::new("id", DataType::Int32, false), @@ -2296,6 +2761,12 @@ mod tests { )) } + /// A reservation no write can outgrow, for the tests that are not about memory. + fn unbounded_reservation() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("IcebergWriteExec[0]").register(&pool) + } + fn common( data_location: String, spec_json: String, @@ -2362,6 +2833,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await?; // These tests assert on the written files, so they stand in for the JVM taking @@ -2461,6 +2933,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -2508,6 +2981,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3784,6 +4258,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3912,6 +4387,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3979,6 +4455,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -4071,6 +4548,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -4087,6 +4565,475 @@ mod tests { ); } + // -- Memory accounting --------------------------------------------------------------- + + const TARGET_FILE_SIZE: u64 = 512 * 1024 * 1024; + + /// parquet-rs writes a file through a `BufWriter` of the default capacity, so up to this + /// much of a row group it has flushed can still be in memory after the flush. + const PARQUET_WRITE_BUFFER_BYTES: usize = 8 * 1024; + + /// Rows `first..first + rows`, spread round robin over the regions `r0`..`r{partitions}`. + fn round_robin_batch_from(first: usize, rows: usize, partitions: usize) -> RecordBatch { + let ids: Vec = (first..first + rows).map(|id| id as i32).collect(); + let regions: Vec = (0..rows) + .map(|row| format!("r{}", row % partitions)) + .collect(); + let regions: Vec<&str> = regions.iter().map(String::as_str).collect(); + batch(&ids, ®ions) + } + + fn round_robin_batch(rows: usize, partitions: usize) -> RecordBatch { + round_robin_batch_from(0, rows, partitions) + } + + /// As many rows for each of the regions `r0`..`r{partitions}`, one region after another, + /// the way a clustered write receives them. + fn clustered_batch(rows: usize, partitions: usize) -> RecordBatch { + let ids: Vec = (0..rows as i32).collect(); + let regions: Vec = (0..rows) + .map(|row| format!("r{}", row * partitions / rows)) + .collect(); + let regions: Vec<&str> = regions.iter().map(String::as_str).collect(); + batch(&ids, ®ions) + } + + /// [`iceberg_user_schema`] with a `payload` column, for rows large enough to fill row + /// groups well past what parquet-rs buffers in front of storage. + fn iceberg_payload_schema() -> Schema { + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "region", Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::required(3, "payload", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build() + .unwrap() + } + + /// Rows `first..first + rows` of `region`, each with its own `payload_bytes`-byte payload. + fn payload_batch_from( + region: &str, + first: usize, + rows: usize, + payload_bytes: usize, + ) -> RecordBatch { + let ids: Vec = (first..first + rows).map(|id| id as i32).collect(); + let payloads: Vec = (first..first + rows) + .map(|id| format!("{id:0>payload_bytes$}")) + .collect(); + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(vec![region; rows])), + Arc::new(StringArray::from(payloads)), + ], + ) + .unwrap() + } + + /// Writes `batches` into a fresh table of `schema`, partitioned by `region` unless the + /// mode is unpartitioned, reserving from `pool` the way `IcebergWriteExec::execute` + /// reserves from the task's pool. A successful write's files are left for the caller to + /// inspect. + async fn write_reserving_from( + pool: &Arc, + schema: Schema, + writer_mode: ProtoIcebergWriterMode, + batches: Vec, + writer_properties: WriterProperties, + target_file_size_bytes: u64, + ) -> (TempDir, DFResult>) { + let temp_dir = TempDir::new().unwrap(); + let spec = match writer_mode { + ProtoIcebergWriterMode::IcebergWriterUnpartitioned => { + PartitionSpec::builder(Arc::new(schema.clone())) + .build() + .unwrap() + } + _ => identity_region_spec(&schema), + }; + let common = with_target_file_size( + common( + format!("file://{}", temp_dir.path().display()), + serde_json::to_string(&spec).unwrap(), + serde_json::to_string(&schema).unwrap(), + writer_mode, + ), + target_file_size_bytes, + ); + let pool: Arc = Arc::::clone(pool); + let reservation = MemoryConsumer::new("IcebergWriteExec[0]").register(&pool); + let written = run_write_task( + input_stream(batches), + common, + Arc::new(schema), + Arc::new(spec), + writer_mode, + writer_properties, + Some(0), + Some(0), + Time::default(), + reservation, + ) + .await + .map(|(data_files, mut abort_guard)| { + abort_guard.disarm(); + data_files + }); + (temp_dir, written) + } + + /// How many files one batch was written to, and the most the write had reserved. Every + /// write has to give back what it reserved. + async fn peak_reservation( + writer_mode: ProtoIcebergWriterMode, + batch: RecordBatch, + target_file_size_bytes: u64, + ) -> (usize, usize) { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + // A partition's first page of rows is held back for its dictionary choice. One unit + // is that page here, so every unit goes straight to its file. + let properties = WriterProperties::builder() + .set_data_page_row_count_limit(ROWS_DIVISOR) + .build(); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + writer_mode, + vec![batch], + properties, + target_file_size_bytes, + ) + .await; + let files = written.unwrap().len(); + assert_eq!(pool.reserved(), 0, "{writer_mode:?} kept a reservation"); + (files, pool.peak()) + } + + fn parquet_files_under(dir: &std::path::Path) -> Vec { + let mut found = Vec::new(); + let mut dirs = vec![dir.to_path_buf()]; + while let Some(dir) = dirs.pop() { + for entry in std::fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + dirs.push(path); + } else if path.extension().is_some_and(|ext| ext == "parquet") { + found.push(path); + } + } + } + found + } + + /// A fanout partition holds rows back until they fill a unit, and every partition value + /// has its own. Here no partition fills one, so nothing reaches a file before the close: + /// the reservation is the held-back rows alone, and it grows with the partitions holding + /// them. + #[tokio::test] + async fn a_fanout_write_reserves_the_rows_its_partitions_hold_back() { + let rows_per_partition = 900; + let mut peaks = Vec::new(); + for partitions in [1, 8] { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + vec![round_robin_batch( + rows_per_partition * partitions, + partitions, + )], + WriterProperties::builder().build(), + TARGET_FILE_SIZE, + ) + .await; + assert_eq!( + record_counts(&written.unwrap()), + vec![rows_per_partition as u64; partitions] + ); + assert_eq!(pool.reserved(), 0, "the write kept a reservation"); + peaks.push(pool.peak()); + } + assert!(peaks[0] > 0, "held-back rows were not reserved: {peaks:?}"); + assert!( + peaks[1] > 4 * peaks[0], + "the reservation did not grow with the partitions holding rows back: {peaks:?}" + ); + } + + /// An open file reserves what it holds and gives it back when it closes. A fanout write + /// keeps every partition's file open, a clustered write one partition's at a time, and a + /// rolling writer one file however many it rolls through. + #[tokio::test] + async fn the_reservation_follows_the_files_that_are_open() { + // Eight partitions of exactly one unit each: no rows are waiting for a full unit, so + // after the batch every row is in a file. + let rows = 8 * ROWS_DIVISOR; + let (fanout_files, fanout) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterFanout, + round_robin_batch(rows, 8), + TARGET_FILE_SIZE, + ) + .await; + let (clustered_files, clustered) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterClustered, + clustered_batch(rows, 8), + TARGET_FILE_SIZE, + ) + .await; + // A one-byte target rolls to a new file for every unit. + let (rolled_files, rolled) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, + round_robin_batch(rows, 8), + 1, + ) + .await; + assert_eq!((fanout_files, clustered_files, rolled_files), (8, 8, 8)); + assert!( + clustered > 0 && rolled > 0, + "open files were not reserved: clustered {clustered}, rolled {rolled}" + ); + // One file open against eight. A closed file that kept its share would leave the + // clustered and rolled writes holding all eight too. + assert!( + 4 * clustered < fanout, + "clustered {clustered}, fanout {fanout}" + ); + assert!(4 * rolled < fanout, "rolled {rolled}, fanout {fanout}"); + } + + /// parquet-rs holds only a file's in-progress row group and what it buffers in front of + /// storage; a local file keeps none of the row groups it has flushed. A file written + /// through many row groups reserves no more than one and that buffer, although its written + /// size keeps growing. + #[tokio::test] + async fn a_file_reserves_no_more_than_one_row_group() { + let row_group_bytes = 16 * 1024; + let properties = WriterProperties::builder() + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(); + // Whole units, so nothing is held back between batches. + let batches = (0..50) + .map(|unit| round_robin_batch_from(unit * ROWS_DIVISOR, ROWS_DIVISOR, 8)) + .collect(); + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, + batches, + properties, + TARGET_FILE_SIZE, + ) + .await; + let data_files = written.unwrap(); + assert_eq!(data_files.len(), 1); + let file_size = data_files[0].file_size_in_bytes() as usize; + assert!( + file_size > 8 * row_group_bytes, + "the file never outgrew a row group: {file_size} bytes" + ); + assert!(pool.peak() > 0, "the open file was not reserved"); + assert!( + pool.peak() <= row_group_bytes + PARQUET_WRITE_BUFFER_BYTES, + "reserved {} bytes for a {row_group_bytes}-byte row group", + pool.peak() + ); + } + + /// A local file keeps nothing of the row groups it has flushed, so a file that flushes + /// gives back what it reserved for them. Here every partition's unit is larger than a row + /// group and flushes as soon as it is written, so twelve open files together hold less + /// than one row group, and a pool of eight row groups holds the write. Charging a flushed + /// file a whole row group until it closed failed this write on its ninth partition. + #[tokio::test] + async fn files_that_flush_give_back_what_they_reserved() { + let row_group_bytes = 128 * 1024; + let properties = WriterProperties::builder() + // One unit is one page, so no rows wait for the dictionary choice. + .set_data_page_row_count_limit(ROWS_DIVISOR) + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(); + // One batch per partition, each a unit of distinct 256-byte payloads: about 256 KB. + let batches = (0..12) + .map(|partition| payload_batch_from(&format!("r{partition}"), 0, ROWS_DIVISOR, 256)) + .collect(); + let pool = Arc::new(PeakMemoryPool::new(8 * row_group_bytes)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_payload_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches, + properties, + TARGET_FILE_SIZE, + ) + .await; + let data_files = written.expect("flushed row groups stayed reserved"); + assert_eq!(data_files.len(), 12); + assert!( + data_files + .iter() + .all(|file| file.file_size_in_bytes() > row_group_bytes as u64), + "a partition's file did not outgrow a row group" + ); + assert!( + pool.peak() < row_group_bytes, + "twelve flushed files reserved {} bytes", + pool.peak() + ); + assert_eq!(pool.reserved(), 0, "the write kept a reservation"); + } + + /// A write whose open files outgrow what the pool grants fails with the pool's error, + /// rather than holding memory nothing accounts for, and deletes the files it had opened. + #[tokio::test] + async fn a_write_the_pool_cannot_hold_fails_and_deletes_its_files() { + let batches = || vec![round_robin_batch(16 * ROWS_DIVISOR, 16)]; + let properties = || WriterProperties::builder().build(); + + // With room to spare, every partition's unit reaches its own file. + let roomy = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (dir, written) = write_reserving_from( + &roomy, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + assert_eq!(written.unwrap().len(), 16); + assert_eq!(parquet_files_under(dir.path()).len(), 16); + + let tight = Arc::new(PeakMemoryPool::new(roomy.peak() / 2)); + let (dir, written) = write_reserving_from( + &tight, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + let error = written.expect_err("a write that outgrows the pool must fail"); + assert!( + matches!(error, DataFusionError::ResourcesExhausted(_)), + "{error:?}" + ); + assert!(error.to_string().contains("IcebergWriteExec[0]"), "{error}"); + assert_eq!( + parquet_files_under(dir.path()), + Vec::::new(), + "the failed write left files behind" + ); + assert_eq!(tight.reserved(), 0, "the failed write kept a reservation"); + } + + /// A task that fails mid-write drops its open files without closing them, so a file's + /// share has to come back either way. + #[tokio::test] + async fn an_open_file_gives_back_its_share_when_it_closes_or_is_dropped() { + let schema = Arc::new(iceberg_user_schema()); + let target = Arc::new(iceberg::arrow::schema_to_arrow_schema(&schema).unwrap()); + let rows = + decorate_batch_with_field_ids(round_robin_batch(ROWS_DIVISOR, 1), &target).unwrap(); + let file_io = FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build(); + let open_files = OpenFileMemory::default(); + let builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new(WriterProperties::builder().build(), schema), + open_files: open_files.clone(), + storage: StorageWrites::Through, + }; + let mut closed = builder + .build(file_io.new_output("memory:/t/closed.parquet").unwrap()) + .await + .unwrap(); + let mut dropped = builder + .build(file_io.new_output("memory:/t/dropped.parquet").unwrap()) + .await + .unwrap(); + + closed.write(&rows).await.unwrap(); + let one = open_files.bytes(); + assert!(one > 0, "a written file reported nothing"); + dropped.write(&rows).await.unwrap(); + assert_eq!(open_files.bytes(), 2 * one); + + closed.close().await.unwrap(); + assert_eq!(open_files.bytes(), one); + drop(dropped); + assert_eq!(open_files.bytes(), 0); + } + + /// A file's share is what it has written and still holds. A local file keeps nothing it + /// has flushed, so its share falls back whenever a row group flushes and never passes a + /// row group and parquet-rs's write buffer. A file taken in parts keeps what it flushes + /// until that makes a part, so here, short of a part, it holds everything it has written. + #[tokio::test] + async fn a_files_share_is_what_it_still_holds() { + let row_group_bytes = 128 * 1024; + let schema = Arc::new(iceberg_payload_schema()); + let target = Arc::new(iceberg::arrow::schema_to_arrow_schema(&schema).unwrap()); + let file_io = FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build(); + for storage in [StorageWrites::Through, StorageWrites::InParts] { + let open_files = OpenFileMemory::default(); + let builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new( + WriterProperties::builder() + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(), + Arc::clone(&schema), + ), + open_files: open_files.clone(), + storage, + }; + let location = format!("memory:/t/{storage:?}.parquet"); + let mut file = builder + .build(file_io.new_output(&location).unwrap()) + .await + .unwrap(); + let mut fell_back = false; + let mut previous = 0; + // Units of about 40 KB, so a row group flushes every three or four. + for unit in 0..20 { + let rows = payload_batch_from("r0", unit * ROWS_DIVISOR, ROWS_DIVISOR, 32); + let rows = decorate_batch_with_field_ids(rows, &target).unwrap(); + file.write(&rows).await.unwrap(); + let share = open_files.bytes(); + match storage { + StorageWrites::Through => assert!( + share <= row_group_bytes + PARQUET_WRITE_BUFFER_BYTES, + "unit {unit}: a local file reserved {share} bytes" + ), + StorageWrites::InParts => { + assert_eq!(share, file.current_written_size(), "unit {unit}") + } + } + fell_back |= share < previous; + previous = share; + } + assert!( + file.current_written_size() > 4 * row_group_bytes, + "{storage:?}: the file flushed too few row groups" + ); + if storage == StorageWrites::Through { + assert!(fell_back, "the share never fell back as row groups flushed"); + } + file.close().await.unwrap(); + assert_eq!(open_files.bytes(), 0, "{storage:?}"); + } + } + /// Regression for apache/datafusion-comet#6145. Past `chrono`'s calendar (year 262142) the /// `day` and `identity(date)` directory names used to panic the task while iceberg-rust /// rendered them, and `year` and `month` came back as NULL partition values. The expected diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 9897ed626c6..c96cb17020e 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -39,7 +39,7 @@ use async_trait::async_trait; use datafusion::{ common::tree_node::TreeNodeRecursion, error::{DataFusionError, Result}, - execution::context::TaskContext, + execution::{context::TaskContext, memory_pool::MemoryConsumer}, physical_expr::{EquivalenceProperties, PhysicalExpr}, physical_plan::{ execution_plan::{Boundedness, EmissionType}, @@ -157,6 +157,20 @@ impl ParquetWriter { } } + /// Memory the writer holds between batches: parquet-rs's estimate for the in-progress row + /// group, which counts encoded pages, encoder buffers, dictionaries and any Bloom filters. + /// The remote writer flushes a row group after every batch and stages it in a buffer that is + /// cleared after each upload but keeps its capacity, so that buffer counts too. + fn memory_size(&self) -> usize { + match self { + ParquetWriter::LocalFile(writer) => writer.memory_size(), + #[cfg(feature = "hdfs-opendal")] + ParquetWriter::Remote(writer, ..) => { + writer.memory_size() + writer.inner().get_ref().capacity() + } + } + } + /// Close the writer and finalize the file async fn close(self) -> std::result::Result<(), parquet::errors::ParquetError> { match self { @@ -492,6 +506,8 @@ impl ExecutionPlan for ParquetWriterExec { let rows_written = MetricBuilder::new(&self.metrics).counter("rows_written", partition); let runtime_env = context.runtime_env(); + let reservation = MemoryConsumer::new(format!("ParquetWriterExec[{partition}]")) + .register(context.memory_pool()); let input = self.input.execute(partition, context)?; let input_schema = self.input.schema(); let work_dir = self.work_dir.clone(); @@ -577,6 +593,8 @@ impl ExecutionPlan for ParquetWriterExec { writer.write(&renamed_batch).await.map_err(|e| { DataFusionError::Execution(format!("Failed to write batch: {}", e)) })?; + // The writer cannot spill, so a row group the pool will not grant fails the task. + reservation.try_resize(writer.memory_size())?; } writer.close().await.map_err(|e| { @@ -619,11 +637,14 @@ impl ExecutionPlan for ParquetWriterExec { #[cfg(test)] mod tests { use super::*; + use crate::execution::memory_pools::testing::PeakMemoryPool; use arrow::array::{Array, Int32Array, ListArray, StringArray}; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::prelude::SessionContext; + use datafusion::execution::memory_pool::MemoryPool; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::prelude::{SessionConfig, SessionContext}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::basic::Repetition; use parquet::file::reader::{FileReader, SerializedFileReader}; @@ -845,7 +866,6 @@ mod tests { /// Helper function to create a test RecordBatch with 1000 rows of (int, string) data /// Example batch_id 1 -> 0..1000, 2 -> 1001..2000 - #[allow(dead_code)] fn create_test_record_batch(batch_id: i32) -> Result { assert!(batch_id > 0, "batch_id must be greater than 0"); let num_rows = batch_id * 1000; @@ -868,6 +888,78 @@ mod tests { .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } + /// Writes `batches` to a local file, reserving from `pool` as a task reserves from its own, + /// and returns the size of the file written. + async fn write_reserving_from( + pool: &Arc, + batches: Vec, + ) -> Result { + let schema = batches[0].schema(); + let column_names = schema.fields().iter().map(|f| f.name().clone()).collect(); + let memory_source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let temp_dir = tempfile::tempdir()?; + let file = temp_dir.path().join("part-00000.parquet"); + let writer = ParquetWriterExec::try_new( + input, + format!("file://{}", file.display()), + None, + None, + None, + ParquetCompression::None, + 0, + column_names, + None, + HashMap::new(), + )?; + let pool: Arc = Arc::::clone(pool); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + let context = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let mut stream = writer.execute(0, context.task_ctx())?; + while stream.try_next().await?.is_some() {} + Ok(std::fs::metadata(file)?.len()) + } + + /// The writer reserves its in-progress row group, here the whole file, until the file is + /// closed. + #[tokio::test] + async fn test_parquet_writer_reserves_its_row_group_until_the_file_closes() -> Result<()> { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let batches = (1..=5) + .map(create_test_record_batch) + .collect::>()?; + let file_size = write_reserving_from(&pool, batches).await?; + assert!( + pool.peak() as u64 >= file_size / 2, + "reserved {} bytes for a row group written as a {file_size}-byte file", + pool.peak() + ); + assert_eq!(pool.reserved(), 0, "the writer kept its reservation"); + Ok(()) + } + + /// A row group the pool will not grant fails the write with the pool's error, rather than + /// holding memory nothing accounts for. + #[tokio::test] + async fn test_parquet_writer_fails_when_the_pool_refuses_its_row_group() -> Result<()> { + let pool = Arc::new(PeakMemoryPool::new(1024)); + let error = write_reserving_from(&pool, vec![create_test_record_batch(1)?]) + .await + .expect_err("a row group larger than the pool must fail the write"); + assert!( + matches!(error, DataFusionError::ResourcesExhausted(_)), + "{error:?}" + ); + assert!( + error.to_string().contains("ParquetWriterExec[0]"), + "{error}" + ); + assert_eq!(pool.reserved(), 0, "the failed write kept its reservation"); + Ok(()) + } + #[tokio::test] #[cfg(feature = "hdfs-opendal")] #[ignore = "This test requires a running HDFS cluster"] diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 53388f1b491..e9ab624c562 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -2850,6 +2850,74 @@ class CometIcebergWriteActionSuite } } + // The native writer reserves what its open files hold from the task's memory pool. A fanout + // write keeps a file open for every partition, so over enough partitions it outgrows the pool + // and fails its task, which Spark can retry, instead of growing in memory that no budget + // accounts for until the executor is killed. + test("native acceleration: a fanout write that outgrows the memory pool fails its task") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + val session = spark + import session.implicits._ + // One task writes 64 partitions of 1000 random 500-byte strings, a run of rows per + // partition. Each open file holds about 500 KB, so the pool of about 4 MiB (0.002 of the + // suite's 2 GiB off-heap size) is outgrown long before all 64 files are open. + val rows = 64000 + (0 until rows) + .map(i => (i, s"r${i / 1000}", new scala.util.Random(i).alphanumeric.take(500).mkString)) + .toDF("id", "region", "payload") + .coalesce(1) + .createOrReplaceTempView("fanout_oom_src") + Seq("fanout_oom", "fanout_oom_control").foreach { table => + // No distribution, so one task writes every partition. + spark.sql(s""" + CREATE TABLE $catalog.$ns.$table (id INT, region STRING, payload STRING) + USING iceberg + PARTITIONED BY (region) + TBLPROPERTIES ( + 'write.spark.fanout.enabled'='true', + 'write.distribution-mode'='none') + """) + } + + withNativeEnabled { + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { + val (failedPlans, error) = captureFailedPlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom SELECT * FROM fanout_oom_src") + } + assert( + error.toSeq + .flatMap(exceptionChain) + .exists(t => + Option(t.getMessage) + .exists(_.contains("Additional allocation failed for IcebergWriteExec"))), + s"expected the native writer's out-of-memory error, got $error") + assert( + failedPlans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + s"the failed write did not run natively:\n${failedPlans.mkString("\n--\n")}") + } + assert(countSnapshots("fanout_oom") == 0L, "the failed write must not commit") + assert( + parquetFiles(dataDir("fanout_oom")).isEmpty, + s"the failed task left data files behind: ${parquetFiles(dataDir("fanout_oom"))}") + + // The executor outlived the failure, and with the whole pool the same write succeeds. + val controlPlans = capturePlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom_control SELECT * FROM fanout_oom_src") + } + assert( + controlPlans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + "the control write did not run natively") + assert( + spark.sql(s"SELECT count(*) FROM $catalog.$ns.fanout_oom_control").head().getLong(0) + == rows) + assert(parquetFiles(dataDir("fanout_oom_control")).size == 64) + } + } + } + test("native acceleration: a pre-native handoff failure cleans up task files") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index babe86b3c72..ce981240bd8 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -1442,6 +1442,41 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a row group the pool cannot hold fails its task with a native out-of-memory error") { + // The native writer reserves its in-progress row group from the task's memory pool, so a row + // group that outgrows the pool fails the task, which Spark can retry, instead of growing in + // memory that no budget accounts for until the executor is killed. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + // One task writes about 10 MB of random strings, which snappy cannot shrink, into a single + // row group, against a pool of about 4 MiB: 0.002 of the suite's 2 GiB off-heap size. + val rows = 20000 + val df = materializeAsCometSource( + (0 until rows) + .map(i => (i, new Random(i).alphanumeric.take(500).mkString)) + .toDF("id", "payload") + .repartition(1), + sourcePath) + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { + val e = intercept[Exception] { + df.write.parquet(outputPath) + } + val outOfMemory = causeChain(e) + .flatMap(t => Option(t.getMessage)) + .exists(_.contains("Additional allocation failed for ParquetWriterExec")) + assert(outOfMemory, s"Expected the native writer's out-of-memory error, got: $e") + } + + // The executor outlived the failure, and with the whole pool the same write succeeds. + val plan = captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + assert(spark.read.parquet(outputPath).count() == rows) + } + } + } + test("the Spark 3.x opt-in key still enables native writes on Spark 4.0+") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // The opt-in moved from DataWritingCommandExec to WriteFilesExec with the operator. Jobs that From 7821e4579eae2dc9d7379f7e78bdaa6ae690b161 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:41:41 +0000 Subject: [PATCH 09/89] build(deps): bump rustls from 0.23.44 to 0.23.45 in /native (#6701) Bumps [rustls](https://github.com/rustls/rustls) from 0.23.44 to 0.23.45. - [Release notes](https://github.com/rustls/rustls/releases) - [Changelog](https://github.com/rustls/rustls/blob/main/CHANGELOG.md) - [Commits](https://github.com/rustls/rustls/compare/v/0.23.44...v/0.23.45) --- updated-dependencies: - dependency-name: rustls dependency-version: 0.23.45 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- native/Cargo.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index e777e80d00c..30de7e06387 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -5960,9 +5960,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.44" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "aws-lc-rs", "once_cell", @@ -6586,7 +6586,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", From a32529c4d84f5157be7a23aa0742ecdc15090eb7 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:48:12 +0000 Subject: [PATCH 10/89] test: move mixed-type shift coverage to SQL fixtures (#6652) --- .../expressions/bitwise/shift_mixed_types.sql | 46 +++++++++++++++++++ .../comet/CometBitwiseExpressionSuite.scala | 20 -------- 2 files changed, 46 insertions(+), 20 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql diff --git a/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql b/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql new file mode 100644 index 00000000000..a8100c48392 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql @@ -0,0 +1,46 @@ +-- 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. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- BIGINT values with INT amounts exercise Comet's cast of the shift amount. +-- Migrated from CometBitwiseExpressionSuite's mixed-type shift test (#6634). +statement +CREATE TABLE test_shift_mixed_types(v bigint, amount int) USING parquet + +statement +INSERT INTO test_shift_mixed_types +SELECT /*+ COALESCE(1) */ v, amount +FROM VALUES + (1111, 2), + (1111, 2), + (3333, 4), + (5555, 6), + (4294967296, 1), + (-4294967296, 1), + (9223372036854775807, 1), + (-9223372036854775808, 1), + (0, 0), + (NULL, 2), + (1111, NULL) +AS input(v, amount) + +query +SELECT shiftright(v, 2), shiftright(v, amount) FROM test_shift_mixed_types + +query +SELECT shiftleft(v, 2), shiftleft(v, amount) FROM test_shift_mixed_types diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index 99a57b15757..db72532ee50 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -54,26 +54,6 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe } } - test("bitwise shift with different left/right types") { - Seq(false, true).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "test" - withTable(table) { - sql(s"create table $table(col1 long, col2 int) using parquet") - sql(s"insert into $table values(1111, 2)") - sql(s"insert into $table values(1111, 2)") - sql(s"insert into $table values(3333, 4)") - sql(s"insert into $table values(5555, 6)") - - checkSparkAnswerAndOperator( - s"SELECT shiftright(col1, 2), shiftright(col1, col2) FROM $table") - checkSparkAnswerAndOperator( - s"SELECT shiftleft(col1, 2), shiftleft(col1, col2) FROM $table") - } - } - } - } - test("bitwise_get - throws exceptions") { def checkSparkAndCometEqualThrows(query: String): Unit = { checkSparkAnswerMaybeThrows(sql(query)) match { From 6940c2e584b96ae9e6666228ab14d5d8bcc693f6 Mon Sep 17 00:00:00 2001 From: Yu Hsun Hung <75130467+sam-1112@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:51:47 +0000 Subject: [PATCH 11/89] feat: Use DataFusion date_trunc for scalar-format timestamp truncation (#5956) * feat: Use DataFusion date_trunc for compatible trunc and date_trunc paths Reuse the upstream physical date_trunc for scalar formats while keeping Comet per-row format, dictionary, and session-timezone handling, with a fallback outside DataFusion's TimestampNanosecond range. * perf: retain direct Date32 truncation kernel * fix: preserve timezone semantics for minute truncation * fix: take the earlier midnight for non-UTC coarse truncation * fix: align timestamp truncation with Spark DST semantics and optimize UTC * fix: preserve timestamp truncation nulls and stabilize DST tests * fix: truncate extreme UTC timestamps with integer calendar arithmetic * fix: handle dictionary and fine-grained timestamp truncation boundaries * perf: Avoid unnecessary dictionary key scans in timestamp truncation --- .../expression-audits/datetime_funcs.md | 4 + native/Cargo.lock | 1 + native/Cargo.toml | 1 + native/spark-expr/Cargo.toml | 1 + native/spark-expr/benches/date_trunc.rs | 36 +- native/spark-expr/benches/timestamp_trunc.rs | 132 +- native/spark-expr/src/kernels/temporal.rs | 1503 ++++++++++++++++- .../expressions/datetime/trunc_date.sql | 68 +- .../expressions/datetime/trunc_timestamp.sql | 145 +- .../trunc_timestamp_dst_ambiguous.sql | 74 +- .../comet/CometTemporalExpressionSuite.scala | 46 +- 11 files changed, 1867 insertions(+), 144 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/datetime_funcs.md b/docs/source/contributor-guide/expression-audits/datetime_funcs.md index 242fd7e6c35..dc88f63915f 100644 --- a/docs/source/contributor-guide/expression-audits/datetime_funcs.md +++ b/docs/source/contributor-guide/expression-audits/datetime_funcs.md @@ -33,6 +33,10 @@ - Constant-folded to a literal by Spark's `ComputeCurrentTime` rule before Comet sees the plan. +## date_trunc + +- Performance (tuned 2026-10-06, PR [#5956](https://github.com/apache/datafusion-comet/pull/5956)): scalar timestamp truncation skips dictionary key masking for infallible fine units and safe UTC/NTZ values with substantial repetition. A values-only range check retains masking near the lower timestamp bound; non-UTC calendar paths and high-cardinality coarse batches retain masking to avoid errors and unnecessary calendar work on unused entries. About 19x faster for MICROSECOND with 65,536 keys and 32 distinct values. Matched Criterion coverage includes sparse/dense/all-NULL keys, high-cardinality YEAR batches, and used/unused overflow values. Benchmark: `benches/timestamp_trunc.rs`. + ## dayname - Spark 4.0+. Implemented natively: maps a `DateType` value to a fixed US-English abbreviated day name (`DayOfWeek.getDisplayName(TextStyle.SHORT, Locale.US)`), with no session-locale or timezone dependence. diff --git a/native/Cargo.lock b/native/Cargo.lock index 30de7e06387..c73688fe1d8 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2143,6 +2143,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-functions", "datafusion-spark", "datasketches", "futures", diff --git a/native/Cargo.toml b/native/Cargo.toml index a2e01fd75fa..dac17e5a289 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -46,6 +46,7 @@ bytes = { version = "1.11.1" } parquet = { version = "59.2.0", default-features = false, features = ["experimental"] } datafusion = { version = "55.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } datafusion-datasource = { version = "55.1.0" } +datafusion-functions = { version = "55.1.0", default-features = false, features = ["datetime_expressions"] } datafusion-physical-expr-adapter = { version = "55.1.0" } datafusion-spark = { version = "55.1.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index d177a42fb8c..7020f8aab75 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -30,6 +30,7 @@ edition = { workspace = true } arrow = { workspace = true } chrono = { workspace = true } datafusion = { workspace = true } +datafusion-functions = { workspace = true } datafusion-spark = { workspace = true } num = { workspace = true } regex = { workspace = true } diff --git a/native/spark-expr/benches/date_trunc.rs b/native/spark-expr/benches/date_trunc.rs index a3dd0900c61..b8342c247df 100644 --- a/native/spark-expr/benches/date_trunc.rs +++ b/native/spark-expr/benches/date_trunc.rs @@ -18,31 +18,57 @@ use arrow::array::{ArrayRef, Date32Array}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion_comet_spark_expr::date_trunc_dyn; -use std::sync::Arc; +use std::{hint::black_box, sync::Arc}; fn criterion_benchmark(c: &mut Criterion) { - let date_array = create_date_array(); + let modern_dates: ArrayRef = Arc::new(create_modern_date_array()); + let mixed_range_null_dates: ArrayRef = Arc::new(create_mixed_range_null_date_array()); let mut group = c.benchmark_group("date_trunc"); // Benchmark each truncation format for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { - let array_ref: ArrayRef = Arc::new(date_array.clone()); group.bench_function(format!("date_trunc_{}", format.to_lowercase()), |b| { - b.iter(|| date_trunc_dyn(&array_ref, format.to_string()).unwrap()); + b.iter(|| { + black_box(date_trunc_dyn(black_box(&modern_dates), format.to_string()).unwrap()) + }); }); + group.bench_function( + format!("date_trunc_{}_mixed_range_null", format.to_lowercase()), + |b| { + b.iter(|| { + black_box( + date_trunc_dyn(black_box(&mixed_range_null_dates), format.to_string()) + .unwrap(), + ) + }); + }, + ); } group.finish(); } -fn create_date_array() -> Date32Array { +fn create_modern_date_array() -> Date32Array { // Create 10000 dates spanning several years (more realistic workload) // Days since Unix epoch: range from 0 (1970-01-01) to ~19000 (2022) let dates: Vec = (0..10000).map(|i| (i * 2) % 19000).collect(); Date32Array::from(dates) } +fn create_mixed_range_null_date_array() -> Date32Array { + // Mix ordinary modern dates with dates around year 3333, which are outside + // DataFusion's TimestampNanosecond range, and nulls. + let dates: Vec> = (0..10000) + .map(|i| match i % 10 { + 0 | 1 => None, + 2 | 3 => Some(497_826 + (i * 17) % 365), + _ => Some((i * 2) % 19000), + }) + .collect(); + Date32Array::from(dates) +} + fn config() -> Criterion { Criterion::default() } diff --git a/native/spark-expr/benches/timestamp_trunc.rs b/native/spark-expr/benches/timestamp_trunc.rs index 6dd5d9fad11..1ad974a01e1 100644 --- a/native/spark-expr/benches/timestamp_trunc.rs +++ b/native/spark-expr/benches/timestamp_trunc.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; -use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::array::{Array, DictionaryArray, Int32Array, RecordBatch, TimestampMicrosecondArray}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_expr::expressions::{lit, Column}; use datafusion::physical_expr::PhysicalExpr; @@ -26,37 +26,125 @@ use std::sync::Arc; #[path = "common/mod.rs"] mod common; -use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS}; +use common::{is_null, timestamp_micros_array, NULL_RATIOS, ROW_COUNTS}; -// Non-UTC session timezone so the benchmark exercises the `array_with_timezone` resolution path -// (the hot part of `TimestampTruncExpr`), not just the truncation kernel. -const TZ: &str = "America/Los_Angeles"; +// Cover the default UTC path and non-UTC timezone resolution with identical input shapes. const MICROS_PER_DAY: i64 = 86_400_000_000; +const BASE_MICROS: i64 = 1_704_067_200_123_456; // 2024-01-01T00:00:00.123456Z fn criterion_benchmark(c: &mut Criterion) { - let schema = Arc::new(Schema::new(vec![Field::new( - "a", - DataType::Timestamp(TimeUnit::Microsecond, Some(TZ.into())), - true, - )])); - let mut group = c.benchmark_group("timestamp_trunc"); - for format in ["YEAR", "MONTH", "DAY", "HOUR"] { - let expr = - TimestampTruncExpr::new(Arc::new(Column::new("a", 0)), lit(format), TZ.to_string()); - for rows in ROW_COUNTS { - for (null_ratio, tag) in NULL_RATIOS { - let ts = timestamp_micros_array(rows, null_ratio, Some(TZ), |i| { - (i as i64) * MICROS_PER_DAY - }); - let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ts]).unwrap(); + for timezone in ["UTC", "America/Los_Angeles"] { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.into())), + true, + )])); + for format in [ + "YEAR", + "QUARTER", + "MONTH", + "WEEK", + "DAY", + "HOUR", + "MINUTE", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ] { + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit(format), + timezone.to_string(), + ); + for rows in ROW_COUNTS { + for (null_ratio, tag) in NULL_RATIOS.into_iter().chain([(0.875, "dense")]) { + // Keep every row in the modern range, including the largest batch. Vary + // sub-day values so fine-unit truncation does actual work as well. + let ts = timestamp_micros_array(rows, null_ratio, Some(timezone), |i| { + BASE_MICROS + + (i % 366) as i64 * MICROS_PER_DAY + + (i as i64 * 1_234_567).rem_euclid(MICROS_PER_DAY) + }); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ts]).unwrap(); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{timezone}/{format}/{rows}/{tag}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + } + } + } + group.finish(); + + // Many repeated keys should cost only as much as truncating the distinct values. + // Keep NULL keys and unused/used overflowing entries in the same matched benchmark. + let mut group = c.benchmark_group("timestamp_trunc_dictionary"); + for rows in [8_192, 65_536] { + for (format, cardinality) in [ + ("MICROSECOND", 32), + ("SECOND", 32), + ("YEAR", 32), + ("YEAR", rows), + ] { + for (null_ratio, tag) in NULL_RATIOS.into_iter().chain([(0.875, "dense")]) { + let values = TimestampMicrosecondArray::from( + (0..cardinality) + .map(|i| BASE_MICROS + i as i64 * 1_234_567) + .collect::>(), + ) + .with_timezone("UTC"); + let keys = Int32Array::from_iter( + (0..rows) + .map(|i| (!is_null(i, null_ratio)).then_some((i % cardinality) as i32)), + ); + let input = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(input)]).unwrap(); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit(format), + "UTC".into(), + ); + let cardinality_tag = if cardinality == rows { "high" } else { "low" }; group.bench_with_input( - BenchmarkId::from_parameter(format!("{format}/{rows}/{tag}")), + BenchmarkId::from_parameter(format!("{format}/{rows}/{cardinality_tag}/{tag}")), &batch, |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), ); } } + for (tag, referenced) in [("unused_overflow", false), ("used_overflow", true)] { + let values = + TimestampMicrosecondArray::from(vec![BASE_MICROS, i64::MIN]).with_timezone("UTC"); + let keys = Int32Array::from_iter((0..rows).map(|i| { + if referenced && i == rows - 1 { + 1 + } else { + 0 + } + })); + let input = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(input)]).unwrap(); + let expr = + TimestampTruncExpr::new(Arc::new(Column::new("a", 0)), lit("YEAR"), "UTC".into()); + assert_eq!(expr.evaluate(&batch).is_err(), referenced); + group.bench_with_input( + BenchmarkId::from_parameter(format!("YEAR/{rows}/{tag}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)))), + ); + } } group.finish(); } diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index a7332d2099d..c3735d741e6 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -18,7 +18,8 @@ //! temporal kernels use chrono::{ - DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc, + DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, Offset, TimeZone, + Timelike, Utc, }; use std::sync::Arc; @@ -32,11 +33,54 @@ use arrow::array::{ }; use arrow::{ array::*, - datatypes::{DataType, TimeUnit}, + datatypes::{DataType, Field, TimeUnit}, }; +use datafusion::{ + common::{config::ConfigOptions, ScalarValue}, + logical_expr::{ColumnarValue, ScalarFunctionArgs}, +}; +use datafusion_functions::datetime; use crate::SparkError; +/// Invoke DataFusion's physical `date_trunc` implementation with a scalar granularity. +/// +/// Spark syntax normalization and compatibility fallback decisions deliberately live outside +/// this helper so the upstream execution boundary stays obvious. +fn datafusion_date_trunc( + array: ArrayRef, + granularity: &'static str, +) -> Result { + let data_type = array.data_type().clone(); + let number_rows = array.len(); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string()))), + ColumnarValue::Array(array), + ]; + let arg_fields = args + .iter() + .enumerate() + .map(|(index, value)| { + Arc::new(Field::new( + format!("date_trunc_arg_{index}"), + value.data_type(), + true, + )) + }) + .collect(); + + datetime::date_trunc() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Arc::new(Field::new("date_trunc", data_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .and_then(|value| value.to_array(number_rows)) + .map_err(|error| SparkError::Internal(error.to_string())) +} + // Copied from arrow_arith/temporal.rs macro_rules! return_compute_error_with { ($msg:expr, $param:expr) => { @@ -47,6 +91,7 @@ macro_rules! return_compute_error_with { // The number of days between the beginning of the proleptic gregorian calendar (0001-01-01) // and the beginning of the Unix Epoch (1970-01-01) const DAYS_TO_UNIX_EPOCH: i32 = 719_163; +const MICROS_PER_MINUTE: i64 = 60_000_000; // Optimized date truncation functions that work directly with days since epoch // These avoid the overhead of converting to/from NaiveDateTime @@ -353,19 +398,100 @@ fn date_trunc_fn_for_format(format: &str) -> Result { }) } -/// Optimized date truncation for Date32 arrays -/// Works directly with days since epoch instead of converting to/from NaiveDateTime -fn date_trunc_date32(array: &Date32Array, format: String) -> Result { - // Select the truncation function based on format - let trunc_fn = date_trunc_fn_for_format(&format)?; +const MICROS_PER_DAY: i64 = 86_400_000_000; - // Apply truncation to each element - let result: Date32Array = array - .iter() - .map(|opt_days| opt_days.and_then(trunc_fn)) - .collect(); +#[inline] +fn fits_timestamp_nanosecond(micros: i64) -> bool { + micros.checked_mul(1_000).is_some() +} - Ok(result) +/// DataFusion's coarse truncation first converts the input to nanoseconds. Although an input near +/// the lower TimestampNanosecond bound can itself be represented, truncating it may move the +/// result before that bound: for example, `1677-09-22` truncated to YEAR becomes `1677-01-01`. +/// The result can move backward by 365 days (366 when starting from December 31 in a leap year), +/// and timezone gap handling can shift it by a few more hours. Round that worst case up to 370 +/// days so both DataFusion's input and coarse-truncation result remain representable. The +/// effective microsecond interval is approximately `1678-09-26T00:12:43.145225Z` through +/// `2262-04-11T23:47:16.854775Z`; because Date32 values are UTC midnight, its first upstream date +/// is 1678-09-27 and its last is 2262-04-11. +#[inline] +fn fits_datafusion_coarse_trunc_range(micros: i64) -> bool { + const LOWER_NANOSECOND_MICROS: i64 = i64::MIN / 1_000; + const COARSE_TRUNC_MARGIN_MICROS: i64 = 370 * MICROS_PER_DAY; + + fits_timestamp_nanosecond(micros) + && micros >= LOWER_NANOSECOND_MICROS + COARSE_TRUNC_MARGIN_MICROS +} + +// Integer proleptic Gregorian calendar conversion, adapted from DataFusion 55.1.0's +// datetime/date_trunc.rs (Howard Hinnant's algorithms): +// https://howardhinnant.github.io/date_algorithms.html +// These helpers only receive days derived from i64 microseconds (about +/-107 million), +// so the calendar intermediates fit in i64 even outside chrono's year range. +fn civil_from_days(days: i64) -> (i64, i64, i64) { + let z = days + 719_468; + let era = z.div_euclid(146_097); + let day_of_era = z.rem_euclid(146_097); + let year_of_era = + (day_of_era - day_of_era / 1460 + day_of_era / 36524 - day_of_era / 146_096) / 365; + let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + let month_index = (5 * day_of_year + 2) / 153; + let day = day_of_year - (153 * month_index + 2) / 5 + 1; + let month = if month_index < 10 { + month_index + 3 + } else { + month_index - 9 + }; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { + let year = year - i64::from(month <= 2); + let era = year.div_euclid(400); + let year_of_era = year.rem_euclid(400); + let month_index = if month > 2 { month - 3 } else { month + 9 }; + let day_of_year = (153 * month_index + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + era * 146_097 + day_of_era - 719_468 +} + +/// Truncate UTC/NTZ coarse units directly in microseconds, avoiding both DataFusion's +/// intermediate nanosecond conversion and chrono's year limit. A representable input can +/// truncate below i64::MIN; report that overflow rather than returning NULL or wrapping. +fn timestamp_trunc_coarse_micros(micros: i64, granularity: &str) -> Result { + let days = micros.div_euclid(MICROS_PER_DAY); + let truncated_days = match granularity { + // The epoch is Thursday, three days after Monday. + "week" => days - (days + 3).rem_euclid(7), + "month" => { + let (_, _, day) = civil_from_days(days); + days - (day - 1) + } + "quarter" => { + let (year, month, _) = civil_from_days(days); + days_from_civil(year, 1 + 3 * ((month - 1) / 3), 1) + } + "year" => { + let (year, _, _) = civil_from_days(days); + days_from_civil(year, 1, 1) + } + _ => unreachable!("integer fallback only handles normalized coarse units"), + }; + truncated_days.checked_mul(MICROS_PER_DAY).ok_or_else(|| { + SparkError::Internal(format!( + "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" + )) + }) +} + +/// Truncate Date32 values directly in days since the epoch. +/// +/// Routing Date32 through DataFusion's timestamp kernel requires two casts and temporary arrays, +/// which is materially slower than this single-pass implementation. +fn date_trunc_date32(array: &Date32Array, format: String) -> Result { + let trunc_fn = date_trunc_fn_for_format(&format)?; + Ok(array.iter().map(|value| value.and_then(trunc_fn)).collect()) } /// @@ -548,7 +674,26 @@ pub(crate) fn timestamp_trunc_dyn( DataType::Dictionary(_, _) => { downcast_dictionary_array!( array => { - let truncated_values = timestamp_trunc_dyn(array.values(), format)?; + // Safe, low-cardinality truncation can operate only on distinct values, + // rather than scanning every key just to discover unused values. + if !timestamp_trunc_dictionary_needs_mask(array.values(), array.len(), &format)? { + let values = timestamp_trunc_dyn(array.values(), format)?; + return Ok(Arc::new(array.with_values(values))); + } + // Dictionary values can outlive the rows that reference them (e.g. after + // filtering). Unused entries, including entries hidden by NULL keys, + // must not raise errors in the fallible timestamp kernel. + let mut unused = vec![true; array.values().len()]; + for key in array.keys_iter().flatten() { + unused[key] = false; + } + let values = if unused.iter().any(|unused| *unused) { + arrow::compute::nullif(array.values(), &BooleanArray::from(unused)) + .map_err(|error| SparkError::Internal(error.to_string()))? + } else { + Arc::clone(array.values()) + }; + let truncated_values = timestamp_trunc_dyn(values.as_ref(), format)?; Ok(Arc::new(array.with_values(truncated_values))) } dt => return_compute_error_with!("timestamp_trunc does not support", dt), @@ -566,6 +711,40 @@ pub(crate) fn timestamp_trunc_dyn( } } +/// Keep key masking for fallible values and for coarse units with many distinct values, where +/// masking can avoid expensive calendar work on unused entries. Fine arithmetic processes the +/// physical values regardless of validity, so infallible fine units never need a key scan. +fn timestamp_trunc_dictionary_needs_mask( + values: &dyn Array, + keys_len: usize, + format: &str, +) -> Result { + let Some(values) = values.as_any().downcast_ref::() else { + return Ok(true); + }; + let granularity = normalize_timestamp_trunc_format(format)?; + if matches!(granularity, "microsecond" | "millisecond" | "second") { + return Ok(false); + } + if values.timezone().is_some_and(|tz| !is_utc_timezone(tz)) { + // Unused non-UTC values can also hit chrono boundary panics, not just Result errors. + return Ok(true); + } + if matches!(granularity, "week" | "month" | "quarter" | "year") && values.len() > keys_len / 64 + { + // Require substantial repetition before speculatively computing unused calendar values. + // Dense NULLs/high cardinality otherwise benefit from masking before truncation. + return Ok(true); + } + // Truncation moves backwards by at most a year. Above this conservative lower-bound margin, + // UTC/NTZ truncation cannot underflow i64 and all values can be evaluated without a key scan. + const LOWER_SAFE_MICROS: i64 = i64::MIN + 370 * MICROS_PER_DAY; + Ok(values + .iter() + .flatten() + .any(|micros| micros < LOWER_SAFE_MICROS)) +} + /// Convert microseconds since epoch to NaiveDateTime #[inline] fn micros_to_naive(micros: i64) -> Option { @@ -584,6 +763,25 @@ type NtzTruncFn = fn(NaiveDateTime) -> Option; /// Truncates a `DateTime`, returning `None` if the result is out of range. type TzTruncFn = fn(DateTime) -> Option>; +/// The Spark `date_trunc` spellings and their canonical DataFusion granularities. +const TIMESTAMP_TRUNC_ALIASES: [(&str, &str); 15] = [ + ("YEAR", "year"), + ("YYYY", "year"), + ("YY", "year"), + ("QUARTER", "quarter"), + ("MONTH", "month"), + ("MON", "month"), + ("MM", "month"), + ("WEEK", "week"), + ("DAY", "day"), + ("DD", "day"), + ("HOUR", "hour"), + ("MINUTE", "minute"), + ("SECOND", "second"), + ("MILLISECOND", "millisecond"), + ("MICROSECOND", "microsecond"), +]; + /// The `timestamp_trunc` formats Spark accepts for the NTZ path, and the truncation each one /// selects. All entries are ASCII, so `eq_ignore_ascii_case` on the raw input matches Spark /// without allocating. @@ -653,6 +851,19 @@ fn tz_trunc_fn_for_format(format: &str) -> Result { }) } +/// Normalize Spark `date_trunc` aliases without accepting additional DataFusion spellings. +fn normalize_timestamp_trunc_format(format: &str) -> Result<&'static str, SparkError> { + TIMESTAMP_TRUNC_ALIASES + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(format)) + .map(|(_, granularity)| *granularity) + .ok_or_else(|| { + SparkError::Internal(format!( + "Unsupported format: {format:?} for function 'timestamp_trunc'" + )) + }) +} + /// Truncate a TimestampNTZ array without any timezone conversion. /// NTZ values are timezone-independent; we treat the raw microseconds as a naive datetime. fn timestamp_trunc_ntz( @@ -680,6 +891,316 @@ where Ok(result) } +/// The zone-aware scalar-format fallback for HOUR/DAY outside DataFusion 55.1's internal +/// TimestampNanosecond range. UTC/NTZ coarse units use integer calendar arithmetic instead. +/// Row-format paths continue to call the same underlying helpers. +fn timestamp_trunc_legacy( + array: &TimestampMicrosecondArray, + format: &str, +) -> Result { + let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + let iter = ArrayIter::new(array); + match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, None) => { + timestamp_trunc_ntz(array, format.to_string()) + } + DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { + let trunc_fn = tz_trunc_fn_for_format(format)?; + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + iter, + builder, + tz, + |dt| as_micros_from_unix_epoch_utc(trunc_fn(dt)), + ) + } + dt => return_compute_error_with!( + "Unsupported input type '{:?}' for function 'timestamp_trunc'", + dt + ), + } +} + +fn datafusion_timestamp_trunc_requires_nanos(granularity: &str, has_timezone: bool) -> bool { + match granularity { + "microsecond" | "millisecond" | "second" | "minute" => false, + "hour" | "day" => has_timezone, + "week" | "month" | "quarter" | "year" => true, + _ => unreachable!("granularity was normalized before compatibility dispatch"), + } +} + +/// Returns whether a timezone has a zero UTC offset at every instant. +/// +/// Keep this list conservative: an unlisted timezone takes the zone-aware path, which is always +/// correct even when its current offset happens to be zero. +fn is_utc_timezone(timezone: &str) -> bool { + matches!( + timezone, + "UTC" | "Etc/UTC" | "Etc/GMT" | "GMT" | "Z" | "+00:00" | "-00:00" | "00:00" + ) +} + +/// Truncate timezone-aware timestamps to the local minute boundary. +/// +/// DataFusion floors the stored UTC microseconds directly. That is only equivalent to Spark's +/// local-time truncation when the timezone offset is a whole number of minutes. Historical offsets +/// can include seconds, so resolve the offset for each instant before finding the local remainder. +/// If truncation crosses an offset transition, re-resolve the truncated local datetime just like +/// `ZonedDateTime.truncatedTo` rather than retaining the input instant's offset. +fn timestamp_trunc_minute_tz( + array: &TimestampMicrosecondArray, + timezone: &str, +) -> Result { + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + ArrayIter::new(array), + TimestampMicrosecondBuilder::with_capacity(array.len()), + timezone, + |dt| { + let micros = dt.timestamp_micros(); + let timezone = dt.timezone(); + let original_offset_secs = dt.offset().fix().local_minus_utc(); + let offset_micros = i64::from(original_offset_secs) * 1_000_000; + let candidate = micros - (micros + offset_micros).rem_euclid(MICROS_PER_MINUTE); + let candidate_dt = + as_datetime_with_timezone::(candidate, timezone) + .expect("truncated minute candidate must be a valid datetime"); + let candidate_offset_secs = candidate_dt.offset().fix().local_minus_utc(); + + if candidate_offset_secs == original_offset_secs { + return candidate; + } + + let truncated_local = dt + .naive_local() + .with_second(0) + .and_then(|local| local.with_nanosecond(0)) + .expect("truncated local minute must be a valid datetime"); + match timezone.from_local_datetime(&truncated_local) { + LocalResult::Single(resolved) => resolved.timestamp_micros(), + LocalResult::Ambiguous(earlier, later) => { + // ZonedDateTime retains the original offset when it is valid in an overlap. + if earlier.offset().fix().local_minus_utc() == original_offset_secs { + earlier.timestamp_micros() + } else if later.offset().fix().local_minus_utc() == original_offset_secs { + later.timestamp_micros() + } else { + earlier.timestamp_micros() + } + } + LocalResult::None => { + // The candidate lies immediately before a forward transition. Java advances + // a nonexistent local time by the gap, which is equivalent to resolving it + // with the candidate's pre-transition offset. + naive_to_micros(truncated_local) - i64::from(candidate_offset_secs) * 1_000_000 + } + } + }, + ) +} + +/// Truncate timezone-aware timestamps to a coarse local-date boundary. +/// +/// Spark resolves these boundaries with `LocalDate.atStartOfDay`. Unlike +/// `ZonedDateTime.truncatedTo`, an overlap always selects the earlier instant, independent of +/// the input timestamp's offset. A midnight in a gap resolves to the transition instant, +/// the first valid local time after the gap, even when the gap starts before midnight. +fn timestamp_trunc_coarse_tz( + array: &TimestampMicrosecondArray, + format: &str, + timezone: &str, +) -> Result { + let trunc_fn = ntz_trunc_fn_for_format(format)?; + let tz: Tz = timezone.parse()?; + + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + ArrayIter::new(array), + TimestampMicrosecondBuilder::with_capacity(array.len()), + timezone, + |dt| { + let truncated_local = trunc_fn(dt.naive_local()) + .expect("truncated coarse local datetime must be a valid datetime"); + + match tz.from_local_datetime(&truncated_local) { + LocalResult::Single(resolved) => resolved.timestamp_micros(), + // `LocalDate.atStartOfDay` takes the earlier occurrence in an overlap. + LocalResult::Ambiguous(earlier, _) => earlier.timestamp_micros(), + LocalResult::None => { + // atStartOfDay returns the end of the gap, not midnight shifted by its + // length. Find a valid pre-gap instant, extending the usual three-hour + // probe for zones that skipped an entire date (e.g. Pacific/Apia). + let mut probe = truncated_local - Duration::hours(3); + let before_gap = loop { + match tz.from_local_datetime(&probe) { + LocalResult::Single(resolved) | LocalResult::Ambiguous(resolved, _) => { + break resolved; + } + LocalResult::None => probe -= Duration::hours(3), + } + }; + let pre_gap_offset_secs = before_gap.offset().fix().local_minus_utc(); + let mut low = before_gap.timestamp(); + let mut high = + truncated_local.and_utc().timestamp() - i64::from(pre_gap_offset_secs); + // TZ offsets and transitions have whole-second precision. The old offset + // holds at low, and has changed by high; find the first changed second. + while high - low > 1 { + let mid = low + (high - low) / 2; + let utc = DateTime::from_timestamp(mid, 0) + .expect("midnight gap probe must be a valid datetime") + .naive_utc(); + if tz.offset_from_utc_datetime(&utc).fix().local_minus_utc() + == pre_gap_offset_secs + { + low = mid; + } else { + high = mid; + } + } + high * 1_000_000 + } + } + }, + ) +} + +/// Handle the rare fine-unit batch whose physical values would underflow DataFusion's +/// unchecked subtraction. Iterate logical values so NULL slots never cause errors. +fn timestamp_trunc_fine_boundary( + array: &TimestampMicrosecondArray, + granularity: &str, + unit: i64, +) -> Result { + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for value in array.iter() { + match value { + None => builder.append_null(), + Some(micros) => { + let remainder = micros.rem_euclid(unit); + let truncated = if matches!(granularity, "second" | "millisecond") { + // Spark uses unchecked Long subtraction for these timezone-independent units. + micros.wrapping_sub(remainder) + } else { + // MINUTE/HOUR/DAY use Spark's exact instant-to-microseconds conversion. + micros.checked_sub(remainder).ok_or_else(|| { + SparkError::Internal(format!( + "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" + )) + })? + }; + builder.append_value(truncated); + } + } + } + Ok(builder.finish().with_timezone_opt(array.timezone())) +} + +fn timestamp_trunc_upstream( + array: &TimestampMicrosecondArray, + format: &str, +) -> Result { + let granularity = normalize_timestamp_trunc_format(format)?; + + if matches!( + granularity, + "hour" | "day" | "week" | "month" | "quarter" | "year" + ) && array.timezone().is_some_and(is_utc_timezone) + { + // UTC local time equals the instant. Removing its label lets DataFusion use arithmetic + // for HOUR/DAY (without a nanosecond limit) and integer calendar code for coarse units. + // The unlabelled call still applies the coarse-unit range guard and calendar fallback. + let input = array.clone().with_timezone_opt(None::>); + let result = timestamp_trunc_upstream(&input, format)?; + return Ok(result.with_timezone_opt(array.timezone())); + } + + if granularity == "minute" { + if let Some(timezone) = array.timezone().filter(|tz| !is_utc_timezone(tz)) { + return timestamp_trunc_minute_tz(array, timezone); + } + } + + if matches!(granularity, "week" | "month" | "quarter" | "year") { + if let Some(timezone) = array.timezone().filter(|tz| !is_utc_timezone(tz)) { + return timestamp_trunc_coarse_tz(array, format, timezone); + } + } + + let fine_unit = match granularity { + "millisecond" => Some(1_000), + "second" => Some(1_000_000), + "minute" => Some(MICROS_PER_MINUTE), + "hour" if array.timezone().is_none() => Some(3_600_000_000), + "day" if array.timezone().is_none() => Some(MICROS_PER_DAY), + _ => None, + }; + if let Some(unit) = fine_unit { + // First aligned microsecond value that can be represented after truncation. + let lower = i64::MIN + (unit - i64::MIN.rem_euclid(unit)).rem_euclid(unit); + // Check physical values too: DataFusion processes NULL slots before restoring validity. + if array.values().iter().any(|micros| *micros < lower) { + return timestamp_trunc_fine_boundary(array, granularity, unit); + } + } + + let requires_nanos = + datafusion_timestamp_trunc_requires_nanos(granularity, array.timezone().is_some()); + + if !requires_nanos + || array + .iter() + .flatten() + .all(fits_datafusion_coarse_trunc_range) + { + let result = datafusion_date_trunc(Arc::new(array.clone()), granularity)?; + return Ok(result + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch") + .clone()); + } + + let upstream_input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + let upstream = datafusion_date_trunc(Arc::new(upstream_input), granularity)?; + let upstream = upstream + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch"); + // Non-UTC HOUR/DAY still need the zone-aware fallback. UTC coarse units have already + // removed their timezone label and must not fall back through chrono. + let legacy = if array.timezone().is_some() { + let input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| !fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + Some(timestamp_trunc_legacy(&input, format)?) + } else { + None + }; + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for (index, value) in array.iter().enumerate() { + match value { + None => builder.append_null(), + Some(micros) if fits_datafusion_coarse_trunc_range(micros) => { + builder.append_option((!upstream.is_null(index)).then(|| upstream.value(index))); + } + Some(micros) => match &legacy { + Some(result) => { + builder.append_option((!result.is_null(index)).then(|| result.value(index))); + } + None => builder.append_value(timestamp_trunc_coarse_micros(micros, granularity)?), + }, + } + } + Ok(builder.finish().with_timezone_opt(array.timezone())) +} + /// Truncate a single NTZ value and append to builder fn timestamp_trunc_ntz_single( value: Option, @@ -711,19 +1232,14 @@ where T: ArrowTemporalType + ArrowNumericType, i64: From, { - let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); - let iter = ArrayIter::new(array); match array.data_type() { - DataType::Timestamp(TimeUnit::Microsecond, None) => { - // TimestampNTZ: operate directly on naive microsecond values without timezone - timestamp_trunc_ntz(array, format) - } - DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { - let trunc_fn = tz_trunc_fn_for_format(&format)?; - as_timestamp_tz_with_op::<&PrimitiveArray, T, _>(iter, builder, tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - }) - } + DataType::Timestamp(TimeUnit::Microsecond, _) => timestamp_trunc_upstream( + array + .as_any() + .downcast_ref::() + .expect("TimestampMicrosecond type mismatch"), + &format, + ), dt => return_compute_error_with!( "Unsupported input type '{:?}' for function 'timestamp_trunc'", dt @@ -896,40 +1412,250 @@ where #[cfg(test)] mod tests { + use super::{ + naive_to_micros, normalize_timestamp_trunc_format, ntz_trunc_fn_for_format, + timestamp_trunc_coarse_micros, timestamp_trunc_ntz, MICROS_PER_DAY, + TIMESTAMP_TRUNC_ALIASES, + }; use crate::kernels::temporal::{ - date_trunc, date_trunc_array_fmt_dyn, timestamp_trunc, timestamp_trunc_array_fmt_dyn, + date_trunc, date_trunc_array_fmt_dyn, date_trunc_dyn, timestamp_trunc, + timestamp_trunc_array_fmt_dyn, timestamp_trunc_dyn, }; + use crate::SparkError; use arrow::array::{ builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder}, iterator::ArrayIter, types::{Date32Type, Int32Type, TimestampMicrosecondType}, - Array, Date32Array, PrimitiveArray, StringArray, TimestampMicrosecondArray, + Array, Date32Array, DictionaryArray, Int32Array, PrimitiveArray, StringArray, + TimestampMicrosecondArray, }; + use arrow::datatypes::{DataType, TimeUnit}; + use chrono::{DateTime, Datelike, NaiveDate}; use std::sync::Arc; + fn epoch_days(date: &str) -> i32 { + NaiveDate::parse_from_str(date, "%Y-%m-%d") + .unwrap() + .num_days_from_ce() + - 719_163 + } + + fn assert_date_trunc(format: &str, input: &[Option<&str>], expected: &[Option<&str>]) { + let input = Date32Array::from( + input + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + let expected = Date32Array::from( + expected + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + assert_eq!(date_trunc(&input, format.to_string()).unwrap(), expected); + } + + fn instant_micros(instant: &str) -> i64 { + DateTime::parse_from_rfc3339(instant) + .unwrap() + .timestamp_micros() + } + + fn assert_timestamp_trunc( + format: &str, + timezone: Option<&str>, + input: &[Option<&str>], + expected: &[Option<&str>], + ) { + let input = TimestampMicrosecondArray::from( + input + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from( + expected + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.to_string()).unwrap(), + expected + ); + } + #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_date_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i32); + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-01-01")]); } - let array = Date32Array::from(vec); - for fmt in [ - "YEAR", "YYYY", "YY", "QUARTER", "MONTH", "MON", "MM", "WEEK", - ] { - match date_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) - } - } - _ => unreachable!(), - } + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-05-01")]); + } + + assert_date_trunc( + "QUARTER", + &[ + Some("2024-01-01"), + Some("2024-03-31"), + Some("2024-04-01"), + Some("2024-06-30"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + &[ + Some("2024-01-01"), + Some("2024-01-01"), + Some("2024-04-01"), + Some("2024-04-01"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + ); + + assert_date_trunc( + "week", + &[ + Some("2024-05-13"), + Some("2024-05-14"), + Some("2024-05-19"), + Some("2024-05-01"), + Some("2024-01-01"), + Some("2023-12-31"), + ], + &[ + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-04-29"), + Some("2024-01-01"), + Some("2023-12-25"), + ], + ); + + let dates = [ + Some("2024-02-29"), + Some("2000-02-29"), + Some("1900-02-28"), + Some("1969-12-31"), + Some("1960-02-29"), + Some("1900-01-01"), + // The input fits TimestampNanosecond, but truncating it to YEAR does not. + Some("1677-09-22"), + // Valid Spark Date32 outside TimestampNanosecond's range. DataFusion 55.1 + // date_trunc internally converts coarse granularities to nanoseconds. + Some("3333-05-17"), + None, + ]; + assert_date_trunc( + "YEAR", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-01-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-01-01"), + Some("3333-01-01"), + None, + ], + ); + assert_date_trunc( + "QUARTER", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-10-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-07-01"), + Some("3333-04-01"), + None, + ], + ); + assert_date_trunc( + "MONTH", + &dates, + &[ + Some("2024-02-01"), + Some("2000-02-01"), + Some("1900-02-01"), + Some("1969-12-01"), + Some("1960-02-01"), + Some("1900-01-01"), + Some("1677-09-01"), + Some("3333-05-01"), + None, + ], + ); + assert_date_trunc( + "WEEK", + &dates, + &[ + Some("2024-02-26"), + Some("2000-02-28"), + Some("1900-02-26"), + Some("1969-12-29"), + Some("1960-02-29"), + Some("1900-01-01"), + Some("1677-09-20"), + Some("3333-05-11"), + None, + ], + ); + + let input = Date32Array::from(vec![epoch_days("2024-05-17")]); + for format in ["DAY", "HOUR", "SECOND", "invalid", " YEAR ", ""] { + let SparkError::Internal(message) = date_trunc(&input, format.to_string()).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'date_trunc'") + ); } } + #[test] + fn test_date_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append(epoch_days("2024-06-30")).unwrap(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append_null(); + builder.append(epoch_days("1969-12-31")).unwrap(); + let input = builder.finish(); + let input_keys = input.keys().clone(); + + let result = date_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(epoch_days("2024-05-01")), + Some(epoch_days("2024-06-01")), + Some(epoch_days("2024-05-01")), + None, + Some(epoch_days("1969-12-01")), + ] + ); + } + #[test] // This test only verifies that the various input array types work. Actually correctness to // ensure this produces the same results as spark is verified in the JVM tests @@ -1047,42 +1773,679 @@ mod tests { } #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_timestamp_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i64); + let input = [Some("2024-05-17T12:34:56.123456Z"), None]; + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-01-01T00:00:00Z"), None], + ); } - let array = TimestampMicrosecondArray::from(vec).with_timezone_utc(); - for fmt in [ + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-05-01T00:00:00Z"), None], + ); + } + for (format, expected) in [ + ("QUARTER", "2024-04-01T00:00:00Z"), + ("WEEK", "2024-05-13T00:00:00Z"), + ("DAY", "2024-05-17T00:00:00Z"), + ("DD", "2024-05-17T00:00:00Z"), + ("HOUR", "2024-05-17T12:00:00Z"), + ("MINUTE", "2024-05-17T12:34:00Z"), + ("SECOND", "2024-05-17T12:34:56Z"), + ("MILLISECOND", "2024-05-17T12:34:56.123Z"), + ("MICROSECOND", "2024-05-17T12:34:56.123456Z"), + ] { + assert_timestamp_trunc(format, Some("UTC"), &input, &[Some(expected), None]); + } + + let invalid_input = + TimestampMicrosecondArray::from(vec![instant_micros("2024-05-17T12:34:56Z")]) + .with_timezone_utc(); + for format in ["MILLISECONDS", "invalid", " DAY ", ""] { + let SparkError::Internal(message) = + timestamp_trunc(&invalid_input, format.to_string()).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'timestamp_trunc'") + ); + } + } + + #[test] + fn test_timestamp_trunc_utc_aliases_and_wide_range() { + let input = [ + Some("1500-06-15T12:34:56.123456Z"), + Some("1678-06-01T12:34:56.123456Z"), + Some("1969-12-31T23:59:59.123456Z"), + Some("3333-05-17T12:34:56.123456Z"), + None, + ]; + for timezone in [ + "UTC", "Etc/UTC", "Etc/GMT", "GMT", "Z", "+00:00", "-00:00", "00:00", + ] { + for (format, expected) in [ + ( + "HOUR", + [ + Some("1500-06-15T12:00:00Z"), + Some("1678-06-01T12:00:00Z"), + Some("1969-12-31T23:00:00Z"), + Some("3333-05-17T12:00:00Z"), + None, + ], + ), + ( + "DAY", + [ + Some("1500-06-15T00:00:00Z"), + Some("1678-06-01T00:00:00Z"), + Some("1969-12-31T00:00:00Z"), + Some("3333-05-17T00:00:00Z"), + None, + ], + ), + ( + "YEAR", + [ + Some("1500-01-01T00:00:00Z"), + Some("1678-01-01T00:00:00Z"), + Some("1969-01-01T00:00:00Z"), + Some("3333-01-01T00:00:00Z"), + None, + ], + ), + ] { + assert_timestamp_trunc(format, Some(timezone), &input, &expected); + } + } + } + + #[test] + fn test_timestamp_trunc_wide_range_fallback() { + let input = [ + Some("2024-05-17T12:34:56.123456Z"), + Some("3333-05-17T12:34:56.123456Z"), + Some("1969-12-31T23:59:59.123456Z"), + Some("1677-09-22T00:00:00Z"), + None, + ]; + assert_timestamp_trunc( "YEAR", - "YYYY", - "YY", + Some("UTC"), + &input, + &[ + Some("2024-01-01T00:00:00Z"), + Some("3333-01-01T00:00:00Z"), + Some("1969-01-01T00:00:00Z"), + Some("1677-01-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "QUARTER", - "MONTH", - "MON", - "MM", + Some("UTC"), + &input, + &[ + Some("2024-04-01T00:00:00Z"), + Some("3333-04-01T00:00:00Z"), + Some("1969-10-01T00:00:00Z"), + Some("1677-07-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "WEEK", - "DAY", - "DD", - "HOUR", - "MINUTE", - "SECOND", - "MILLISECOND", - "MICROSECOND", - ] { - match timestamp_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) + Some("UTC"), + &input, + &[ + Some("2024-05-13T00:00:00Z"), + Some("3333-05-11T00:00:00Z"), + Some("1969-12-29T00:00:00Z"), + Some("1677-09-20T00:00:00Z"), + None, + ], + ); + } + + #[test] + fn test_timestamp_trunc_out_of_chrono_range() { + for timezone in [None, Some("UTC"), Some("Etc/UTC"), Some("+00:00")] { + let input = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MAX), + Some(i64::MAX - 1), + Some(i64::MAX - 2), + Some(instant_micros("3333-05-17T12:34:56Z")), + Some(-9_000_000_000_000_000_000), + None, + ]) + .with_timezone_opt(timezone); + // Expected extremes were computed independently with java.time.LocalDate: + // MAX, MAX - 1, and MAX - 2 are all +294247-01-10; + // the negative input is -283229-05-10. + for (format, recent, future, maximum, negative) in [ + ( + "YEAR", + "2024-01-01T00:00:00Z", + "3333-01-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_011_174_400_000_000, + ), + ( + "QUARTER", + "2024-04-01T00:00:00Z", + "3333-04-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_003_398_400_000_000, + ), + ( + "MONTH", + "2024-05-01T00:00:00Z", + "3333-05-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_000_806_400_000_000, + ), + ( + "WEEK", + "2024-05-13T00:00:00Z", + "3333-05-11T00:00:00Z", + 9_223_371_504_000_000_000, + -9_000_000_028_800_000_000, + ), + ] { + let expected = TimestampMicrosecondArray::from(vec![ + Some(instant_micros(recent)), + Some(maximum), + Some(maximum), + Some(maximum), + Some(instant_micros(future)), + Some(negative), + None, + ]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.to_string()).unwrap(), + expected + ); + for (alias, canonical) in TIMESTAMP_TRUNC_ALIASES { + if canonical == normalize_timestamp_trunc_format(format).unwrap() { + assert_eq!( + timestamp_trunc(&input, alias.to_lowercase()).unwrap(), + expected + ); + } + } + } + } + } + + #[test] + fn test_timestamp_trunc_coarse_overflow() { + for timezone in [None, Some("UTC"), Some("Etc/UTC"), Some("+00:00")] { + let input = + TimestampMicrosecondArray::from(vec![Some(i64::MIN)]).with_timezone_opt(timezone); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let error = timestamp_trunc(&input, format.to_string()).unwrap_err(); + assert!(error.to_string().contains("out of range"), "{error}"); + } + } + } + + #[test] + fn test_timestamp_trunc_integer_calendar_matches_chrono() { + // Exercise negative eras, year zero, Gregorian leap centuries, and month boundaries. + for year in [-2000, -400, -1, 0, 1, 1500, 1600, 1900, 2000, 2024, 3333] { + for month in 1..=12 { + for day in [1, 15, 28, 29, 30, 31] { + let Some(date) = NaiveDate::from_ymd_opt(year, month, day) else { + continue; + }; + let datetime = date.and_hms_micro_opt(23, 59, 59, 999_999).unwrap(); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let expected = naive_to_micros( + ntz_trunc_fn_for_format(format).unwrap()(datetime).unwrap(), + ); + assert_eq!( + timestamp_trunc_coarse_micros( + naive_to_micros(datetime), + normalize_timestamp_trunc_format(format).unwrap() + ) + .unwrap(), + expected + ); } } - _ => unreachable!(), } } } + #[test] + fn test_timestamp_trunc_coarse_guard_boundaries() { + let lower = i64::MIN / 1_000 + 370 * MICROS_PER_DAY; + let upper = i64::MAX / 1_000; + let input = TimestampMicrosecondArray::from(vec![ + Some(lower - 1), + Some(lower), + Some(lower + 1), + Some(upper - 1), + Some(upper), + Some(upper + 1), + None, + ]); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let expected = timestamp_trunc_ntz(&input, format.to_string()).unwrap(); + assert_eq!( + timestamp_trunc(&input, format.to_string()).unwrap(), + expected + ); + } + } + + #[test] + fn test_timestamp_trunc_extreme_dictionary() { + let values = TimestampMicrosecondArray::from(vec![Some(i64::MAX), None]); + let keys = Int32Array::from(vec![Some(0), None, Some(1), Some(0)]); + let input = + DictionaryArray::::try_new(keys.clone(), Arc::new(values)) + .unwrap(); + let expected_values = + TimestampMicrosecondArray::from(vec![Some(9_223_371_244_800_000_000), None]); + let expected = DictionaryArray::::try_new( + keys, + Arc::new(expected_values), + ) + .unwrap(); + let result = timestamp_trunc_dyn(&input, "YEAR".to_string()).unwrap(); + assert_eq!(result.as_ref(), &expected as &dyn Array); + } + + #[test] + fn test_timestamp_trunc_empty_and_null_batches() { + for timezone in [None, Some("UTC")] { + for values in [vec![], vec![None, None]] { + let input = TimestampMicrosecondArray::from(values).with_timezone_opt(timezone); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + assert_eq!(timestamp_trunc(&input, format.to_string()).unwrap(), input); + } + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_unused_overflow() { + for timezone in [None, Some("UTC")] { + let values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MIN), + None, + ]) + .with_timezone_opt(timezone); + for keys in [ + vec![Some(0), None, Some(2), Some(0)], + vec![None, None], + vec![], + vec![Some(0); 8192], + ] { + let keys = Int32Array::from(keys); + let input = + DictionaryArray::::try_new(keys.clone(), Arc::new(values.clone())) + .unwrap(); + let expected_values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-01-01T00:00:00Z")), + None, + None, + ]) + .with_timezone_opt(timezone); + // Slice past the first key to check that referenced-entry masking respects offsets. + let input = input.slice( + usize::from(!input.is_empty()), + input.len().saturating_sub(1), + ); + let expected = + DictionaryArray::::try_new(keys, Arc::new(expected_values)).unwrap(); + let expected = expected.slice( + usize::from(!expected.is_empty()), + expected.len().saturating_sub(1), + ); + let result = timestamp_trunc_dyn(&input, "YEAR".into()).unwrap(); + // Unused entries can be masked, so compare the decoded logical rows. + let result = arrow::compute::cast( + result.as_ref(), + &DataType::Timestamp(TimeUnit::Microsecond, timezone.map(Into::into)), + ) + .unwrap(); + let expected = arrow::compute::cast(&expected, result.data_type()).unwrap(); + assert_eq!(result.as_ref(), expected.as_ref()); + let used = DictionaryArray::::try_new( + Int32Array::from(vec![1]), + Arc::new(values.clone()), + ) + .unwrap(); + assert!(timestamp_trunc_dyn(&used, "YEAR".into()) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_unused_chrono_boundary() { + // These unused values can panic in timezone-aware calendar helpers, so they must still + // be masked before computing the distinct values, even in a low-cardinality dictionary. + for timezone in ["+01:00", "Asia/Tokyo"] { + let values = TimestampMicrosecondArray::from(vec![ + instant_micros("2024-05-17T12:34:56Z"), + chrono::NaiveDateTime::MIN.and_utc().timestamp_micros() + MICROS_PER_DAY, + ]) + .with_timezone(timezone); + let input = DictionaryArray::::try_new( + Int32Array::from(vec![Some(0); 8192]), + Arc::new(values.clone()), + ) + .unwrap(); + for format in ["YEAR", "WEEK"] { + let result = timestamp_trunc_dyn(&input, format.into()).unwrap(); + let decoded = arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); + let expected = timestamp_trunc(&values.slice(0, 1), format.into()).unwrap(); + let decoded = decoded + .as_any() + .downcast_ref::() + .unwrap(); + assert!(decoded.iter().all(|value| value == Some(expected.value(0)))); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_infallible_unused_extremes() { + for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { + let values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MIN), + None, + ]) + .with_timezone_opt(timezone); + for keys in [ + vec![Some(0), None, Some(2), Some(0)], + vec![None, None], + vec![], + ] { + let input = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(values.clone()), + ) + .unwrap(); + let decoded_input = arrow::compute::cast(&input, values.data_type()).unwrap(); + for format in ["MICROSECOND", "MILLISECOND", "SECOND"] { + let result = timestamp_trunc_dyn(&input, format.into()).unwrap(); + let decoded = + arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); + let expected = + timestamp_trunc_dyn(decoded_input.as_ref(), format.into()).unwrap(); + assert_eq!(decoded.as_ref(), expected.as_ref()); + } + } + } + } + + #[test] + fn test_timestamp_trunc_fine_poisoned_nulls() { + for timezone in [None, Some("UTC")] { + let input = TimestampMicrosecondArray::new( + vec![ + i64::MIN, + instant_micros("2024-05-17T12:34:56Z"), + i64::MIN + 1, + i64::MAX, + ] + .into(), + Some(arrow::buffer::NullBuffer::from(vec![ + false, true, false, true, + ])), + ) + .with_timezone_opt(timezone); + let clean = TimestampMicrosecondArray::from(vec![ + None, + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + Some(i64::MAX), + ]) + .with_timezone_opt(timezone); + for format in [ + "MINUTE", + "HOUR", + "DAY", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ] { + assert_eq!( + timestamp_trunc(&input, format.into()).unwrap(), + timestamp_trunc(&clean, format.into()).unwrap() + ); + let slice = input.slice(2, 2); + assert_eq!( + timestamp_trunc(&slice, format.into()).unwrap(), + timestamp_trunc(&clean.slice(2, 2), format.into()).unwrap() + ); + } + } + } + + #[test] + fn test_timestamp_trunc_fine_lower_bound() { + // First representable UTC boundary for each unit. + for (format, boundary) in [ + ("MINUTE", -9_223_372_036_800_000_000), + ("HOUR", -9_223_372_036_800_000_000), + ("DAY", -9_223_372_022_400_000_000), + ] { + for timezone in [None, Some("UTC")] { + for micros in [i64::MIN, i64::MIN + 1, boundary - 1] { + let input = + TimestampMicrosecondArray::from(vec![micros]).with_timezone_opt(timezone); + assert!(timestamp_trunc(&input, format.into()) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + let input = TimestampMicrosecondArray::from(vec![boundary, boundary + 1]) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from(vec![boundary, boundary]) + .with_timezone_opt(timezone); + assert_eq!(timestamp_trunc(&input, format.into()).unwrap(), expected); + } + } + } + + #[test] + fn test_timestamp_trunc_fine_spark_wrapping() { + // Spark uses unchecked Long subtraction for SECOND/MILLISECOND, even at MIN. + for (format, expected) in [ + ("SECOND", 9_223_372_036_854_551_616), + ("MILLISECOND", 9_223_372_036_854_775_616), + ] { + for timezone in [None, Some("UTC"), Some("Asia/Kolkata")] { + let input = TimestampMicrosecondArray::from(vec![Some(i64::MIN), None]) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from(vec![Some(expected), None]) + .with_timezone_opt(timezone); + assert_eq!(timestamp_trunc(&input, format.into()).unwrap(), expected); + } + } + } + + #[test] + fn test_timestamp_trunc_dst_gap_and_overlap() { + assert_timestamp_trunc( + "HOUR", + Some("America/Los_Angeles"), + &[ + Some("2024-11-03T08:30:15.123456Z"), + Some("2024-11-03T09:30:15.123456Z"), + ], + &[Some("2024-11-03T08:00:00Z"), Some("2024-11-03T09:00:00Z")], + ); + assert_timestamp_trunc( + "DAY", + Some("America/New_York"), + &[ + Some("2024-03-10T06:30:15.123456Z"), + Some("2024-03-10T07:30:15.123456Z"), + Some("2024-11-03T06:30:15.123456Z"), + ], + &[ + Some("2024-03-10T05:00:00Z"), + Some("2024-03-10T05:00:00Z"), + Some("2024-11-03T04:00:00Z"), + ], + ); + // Sao Paulo advanced from 23:59:59 on November 3 directly to 01:00 on November 4. + // Spark resolves DAY for a post-gap instant to that day's first valid local time. + assert_timestamp_trunc( + "DAY", + Some("America/Sao_Paulo"), + &[ + Some("2018-11-04T02:30:15.123456Z"), + Some("2018-11-04T03:30:15.123456Z"), + ], + &[Some("2018-11-03T03:00:00Z"), Some("2018-11-04T03:00:00Z")], + ); + } + + #[test] + fn test_timestamp_trunc_coarse_midnight_gap() { + // Toronto's gap starts on the previous date at 23:30 and ends at 00:30. + // atStartOfDay selects 00:30, rather than shifting midnight to 01:00. + for timezone in [ + "America/Toronto", + "Canada/Eastern", + "America/Montreal", + "America/Nassau", + ] { + assert_timestamp_trunc( + "WEEK", + Some(timezone), + &[ + Some("1919-04-02T16:00:00Z"), + Some("1919-03-31T04:30:00Z"), + None, + ], + &[ + Some("1919-03-31T04:30:00Z"), + Some("1919-03-31T04:30:00Z"), + None, + ], + ); + } + // Asuncion's gap starts exactly at midnight, at both a month and quarter boundary. + for format in ["MONTH", "QUARTER"] { + assert_timestamp_trunc( + format, + Some("America/Asuncion"), + &[Some("2023-10-15T12:00:00Z"), None], + &[Some("2023-10-01T04:00:00Z"), None], + ); + } + } + + #[test] + fn test_timestamp_trunc_coarse_midnight_overlap() { + // Havana repeated local midnight on 2026-11-01. MONTH must use the first occurrence + // (UTC-4), even though the input is after the fall-back transition and has UTC-5. + assert_timestamp_trunc( + "MONTH", + Some("America/Havana"), + &[Some("2026-11-15T12:00:00Z")], + &[Some("2026-11-01T04:00:00Z")], + ); + } + + #[test] + fn test_timestamp_trunc_minute_with_historical_offset_seconds() { + // Africa/Monrovia used UTC-00:44:30 in 1960. The input instant is local 10:30:45, + // which Spark truncates to local 10:30:00 (11:14:30 UTC), not a UTC minute boundary. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1960-06-15T11:15:15Z"), None], + &[Some("1960-06-15T11:14:30Z"), None], + ); + + // Asia/Aden changed from UTC+03:06:52 to UTC+03:00 within the local 23:53 minute. + // Truncating the post-transition 23:53:30 must re-resolve local 23:53:00 with the + // pre-transition offset, matching ZonedDateTime.truncatedTo. + assert_timestamp_trunc( + "MINUTE", + Some("Asia/Aden"), + &[Some("1947-03-13T20:53:30.123Z")], + &[Some("1947-03-13T20:46:08Z")], + ); + + // Monrovia's 1972 transition skipped local 00:00:00 through 00:44:29. Truncating + // 00:44:45 targets the gap, which ZonedDateTime shifts forward by 44 minutes 30 seconds. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1972-01-07T00:44:45Z")], + &[Some("1972-01-07T01:28:30Z")], + ); + } + + #[test] + fn test_timestamp_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder + .append(instant_micros("2024-06-30T23:59:59Z")) + .unwrap(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder.append_null(); + let input = builder.finish(); + let input = input.with_values(Arc::new( + input + .values() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + .with_timezone_utc(), + )); + let input_keys = input.keys().clone(); + + let result = timestamp_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(instant_micros("2024-05-01T00:00:00Z")), + Some(instant_micros("2024-06-01T00:00:00Z")), + Some(instant_micros("2024-05-01T00:00:00Z")), + None, + ] + ); + } + #[test] // test takes too long with miri #[cfg_attr(miri, ignore)] diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql index 672281627c9..0bc3b431563 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql @@ -15,21 +15,77 @@ -- specific language governing permissions and limitations -- under the License. +-- Dictionary-encoded dates reuse the scalar Date32 path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED + statement CREATE TABLE test_trunc_date(d date) USING parquet statement -INSERT INTO test_trunc_date VALUES (date('2024-06-15')), (date('2024-01-01')), (date('2024-12-31')), (NULL) +INSERT INTO test_trunc_date VALUES + (date('2024-05-17')), + -- Monday / Tuesday / Sunday for Monday-based WEEK truncation. + (date('2024-05-13')), + (date('2024-05-14')), + (date('2024-05-19')), + (date('2024-01-01')), + (date('2023-12-31')), + (date('2024-03-31')), + (date('2024-04-01')), + (date('2024-06-30')), + (date('2024-07-01')), + (date('2024-10-01')), + (date('2024-02-29')), + (date('2000-02-29')), + (date('1900-02-28')), + (date('1969-12-31')), + (date('1960-02-29')), + (date('1900-01-01')), + (date('1500-06-15')), + (date('1678-06-01')), + -- Valid Spark Date32 outside TimestampNanosecond's range. + (date('3333-05-17')), + (NULL) + +query +SELECT d, trunc(d, 'YEAR'), trunc(d, 'YYYY'), trunc(d, 'YY') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'QUARTER') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'MONTH'), trunc(d, 'MON'), trunc(d, 'MM') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'WEEK') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'year'), trunc(d, 'Year'), trunc(d, 'yEaR'), trunc(d, 'month'), trunc(d, 'Mon'), trunc(d, 'week') FROM test_trunc_date ORDER BY d + +-- Unsupported scalar formats fall back to Spark's codegen dispatcher and remain NULL-compatible. +query +SELECT d, trunc(d, 'DAY'), trunc(d, 'HOUR'), trunc(d, 'SECOND'), trunc(d, 'invalid'), trunc(d, ' YEAR '), trunc(d, '') FROM test_trunc_date ORDER BY d +-- NULL format is Incompatible on the native path, so this uses the codegen dispatcher. query -SELECT trunc(d, 'year') FROM test_trunc_date +SELECT d, trunc(d, NULL) FROM test_trunc_date ORDER BY d query -SELECT trunc(d, 'month') FROM test_trunc_date +SELECT trunc(NULL, 'YEAR'), trunc(NULL, NULL) +-- Literal arguments exercise the scalar date input branch. query -SELECT trunc(d, 'quarter') FROM test_trunc_date +SELECT + trunc(DATE '2024-05-17', 'YEAR'), + trunc(DATE '2024-05-17', 'QUARTER'), + trunc(DATE '2024-05-17', 'MONTH'), + trunc(DATE '2024-05-17', 'WEEK') --- literal arguments +-- Literal arguments exercise the scalar date input branch. query -SELECT trunc(date('2024-06-15'), 'year'), trunc(date('2024-06-15'), 'month'), trunc(date('2024-06-15'), 'quarter') +SELECT + trunc(DATE '3333-05-17', 'YEAR'), + trunc(DATE '3333-05-17', 'QUARTER'), + trunc(DATE '3333-05-17', 'MONTH'), + trunc(DATE '3333-05-17', 'WEEK') diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql index 1a8bbef9880..642777cd6ad 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql @@ -15,26 +15,155 @@ -- specific language governing permissions and limitations -- under the License. --- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Keep the wide-range fallback fixture independent of far-future JVM/native timezone rules. +-- Config: spark.sql.session.timeZone=UTC +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- Dictionary-encoded timestamps reuse the scalar timestamp path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true statement CREATE TABLE test_trunc_ts(ts timestamp) USING parquet statement -INSERT INTO test_trunc_ts VALUES (timestamp('2024-06-15 10:30:45')), (timestamp('2024-01-01 00:00:00')), (NULL) +INSERT INTO test_trunc_ts VALUES + (timestamp('2024-05-17 12:34:56.123456')), + (timestamp('2024-02-29 23:59:59.999999')), + (timestamp('2000-02-29 00:00:00')), + (timestamp('1900-02-28 00:00:00')), + (timestamp('1969-12-31 23:59:59.123456')), + -- Below the nanosecond lower bound and inside its 370-day truncation margin. + (timestamp('1500-06-15 12:34:56.123456')), + (timestamp('1678-06-01 12:34:56.123456')), + -- Valid Spark timestamp outside TimestampNanosecond's range. + (timestamp('3333-05-17 12:34:56.123456')), + (NULL) query -SELECT date_trunc('year', ts) FROM test_trunc_ts +SELECT ts, date_trunc('YEAR', ts), date_trunc('YYYY', ts), date_trunc('YY', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('month', ts) FROM test_trunc_ts +SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('day', ts) FROM test_trunc_ts +SELECT ts, date_trunc('MONTH', ts), date_trunc('MON', ts), date_trunc('MM', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('hour', ts) FROM test_trunc_ts +SELECT ts, date_trunc('WEEK', ts), date_trunc('DAY', ts), date_trunc('DD', ts) FROM test_trunc_ts ORDER BY ts --- literal arguments query -SELECT date_trunc('year', timestamp('2024-06-15 10:30:45')), date_trunc('month', timestamp('2024-06-15 10:30:45')), date_trunc('day', timestamp('2024-06-15 10:30:45')) +SELECT + ts, + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ts +ORDER BY ts + +query +SELECT + ts, + date_trunc('year', ts), + date_trunc('Year', ts), + date_trunc('yEaR', ts), + date_trunc('month', ts), + date_trunc('Mon', ts), + date_trunc('week', ts) +FROM test_trunc_ts +ORDER BY ts + +-- NULL format is Incompatible on the native path. Without allowIncompatible the +-- codegen dispatcher runs Spark's TruncTimestamp and returns NULL. +query +SELECT ts, date_trunc(NULL, ts) FROM test_trunc_ts ORDER BY ts + +query +SELECT date_trunc('YEAR', NULL), date_trunc(NULL, NULL) + +-- Literal arguments exercise the scalar timestamp input branch. +query +SELECT + date_trunc('YEAR', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('QUARTER', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MONTH', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('WEEK', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('DAY', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('HOUR', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MINUTE', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('SECOND', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MILLISECOND', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MICROSECOND', TIMESTAMP '2024-05-17 12:34:56.123456') + +-- Literal arguments exercise the scalar timestamp input branch. +query +SELECT + date_trunc('YEAR', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('QUARTER', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MONTH', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('WEEK', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('DAY', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('HOUR', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MINUTE', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('SECOND', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MILLISECOND', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MICROSECOND', TIMESTAMP '3333-05-17 12:34:56.123456') + +-- Long.MaxValue is used as an end-of-time sentinel and exceeds chrono's range. +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Direct date_trunc projections must execute natively and match Spark exactly. +-- Wrapping them in unix_micros would dispatch the whole subtree to the JVM. +statement +CREATE TABLE test_trunc_ts_extreme(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_ts_extreme VALUES + (timestamp_micros(9223372036854775807)), + (timestamp_micros(-9000000000000000000)), + (timestamp('2024-05-17 12:34:56.123456')), + (timestamp('3333-05-17 12:34:56.123456')), + (NULL) + +query expect_native(date_trunc) +SELECT + date_trunc('YEAR', ts), + date_trunc('QUARTER', ts), + date_trunc('MONTH', ts), + date_trunc('WEEK', ts) +FROM test_trunc_ts_extreme +ORDER BY ts + +-- A valid microsecond input can truncate below Long.MinValue. Spark raises an error. +statement +CREATE TABLE test_trunc_ts_overflow(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_ts_overflow VALUES (timestamp_micros(-9223372036854775808)) + +query expect_error(long overflow) +SELECT date_trunc('YEAR', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('QUARTER', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('MONTH', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('WEEK', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('MINUTE', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('HOUR', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('DAY', ts) FROM test_trunc_ts_overflow + +-- Spark intentionally wraps Long subtraction for SECOND/MILLISECOND at the lower bound. +query expect_native(date_trunc) +SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) +FROM test_trunc_ts_overflow diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql index a1059736b9c..74d2c09ab4c 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql @@ -15,36 +15,62 @@ -- specific language governing permissions and limitations -- under the License. --- Regression test for ambiguous local times during DST fall-back. --- On 2024-11-03 at 2:00 AM America/Los_Angeles, clocks fall back to 1:00 AM, --- so 1:30 AM occurs twice (once in PDT, once in PST). Truncating 01:30 to HOUR --- gives 01:00, which is ambiguous. chrono's DateTime::with_minute(0) returns --- None for ambiguous results, causing a panic in as_micros_from_unix_epoch_utc. +-- Differential coverage for scalar date_trunc around DST overlaps, gaps, and historical offsets +-- containing seconds. Explicit UTC offsets include both occurrences of the repeated US fall-back +-- hour. The 2018 Sao Paulo values exercise its historic midnight spring-forward gap: truncating a +-- valid 01:30 local timestamp to DAY targets the nonexistent local midnight. Africa/Monrovia used +-- UTC-00:44:30 until 1972, so its local minute boundaries do not align with UTC minute boundaries. +-- Asia/Aden covers minute truncation across a historical offset transition. Havana repeated +-- midnight on 2020-11-01, which makes MONTH select the earlier UTC occurrence. Toronto +-- skipped 1919-03-30 23:30 through 1919-03-31 00:30, so WEEK must select the gap end. +-- Asuncion skipped midnight on 2023-10-01, a MONTH and QUARTER boundary. -- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true --- Config: spark.sql.session.timeZone=America/Los_Angeles +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- ConfigMatrix: spark.sql.session.timeZone=America/Los_Angeles,America/New_York,America/Sao_Paulo,Africa/Monrovia,Asia/Aden,America/Havana,America/Toronto,America/Asuncion statement CREATE TABLE test_trunc_ambiguous(ts timestamp) USING parquet statement INSERT INTO test_trunc_ambiguous VALUES - (timestamp('2024-11-03 01:30:00')) + (timestamp('1919-04-02T16:00:00Z')), + (timestamp('1919-03-31T04:30:00Z')), + (timestamp('2023-10-15T12:00:00Z')), + (TIMESTAMP '1960-06-15 10:30:45'), + (timestamp('1947-03-13T20:53:30.123Z')), + (timestamp('1972-01-07T00:44:45Z')), + (timestamp('2018-11-04T01:30:15.123456Z')), + (timestamp('2018-11-04T02:30:15.123456Z')), + (timestamp('2018-11-04T03:30:15.123456Z')), + (timestamp('2018-11-04T04:30:15.123456Z')), + (timestamp('2024-03-10T06:30:15.123456Z')), + (timestamp('2024-03-10T07:30:15.123456Z')), + (timestamp('2024-03-10T08:30:15.123456Z')), + (timestamp('2024-03-10T09:30:15.123456Z')), + (timestamp('2024-03-10T10:30:15.123456Z')), + (timestamp('2024-03-10T11:30:15.123456Z')), + (timestamp('2024-11-03T05:30:15.123456Z')), + (timestamp('2024-11-03T06:30:15.123456Z')), + (timestamp('2024-11-03T07:30:15.123456Z')), + (timestamp('2024-11-03T08:30:15.123456Z')), + (timestamp('2024-11-03T09:30:15.123456Z')), + (timestamp('2024-11-03T10:30:15.123456Z')), + (timestamp('2020-11-15T12:00:00Z')), + (NULL) -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('DAY', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('HOUR', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('WEEK', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('MONTH', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('YEAR', ts) FROM test_trunc_ambiguous ORDER BY ts +query +SELECT + ts, + date_trunc('YEAR', ts), + date_trunc('QUARTER', ts), + date_trunc('MONTH', ts), + date_trunc('WEEK', ts), + date_trunc('DAY', ts), + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ambiguous +ORDER BY ts diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 2e73ac9c331..a03e749a76e 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -41,6 +41,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH private val crossTimezones = Seq("UTC", "America/Los_Angeles", "Europe/London", "Asia/Tokyo") + /** + * Dates around 2024 stay inside DataFusion `date_trunc`'s TimestampNanosecond range + * (approximately 1678 through 2262). The fuzz generator's default `baseDate` is year 3333, + * which only exercises Comet's wide-range fallback. + */ + private val dataFusionRangeBaseDate: Long = + new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime + private def deepestSparkThrowable(error: Throwable): SparkThrowable with Throwable = causeChain(error) .collect { case e: SparkThrowable with Throwable => e } @@ -80,6 +88,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } test("trunc (TruncDate)") { + checkTruncDate(dataFusionRangeBaseDate) + } + + test("trunc (TruncDate) - dates outside DataFusion nanosecond range") { + checkTruncDate(FuzzDataGenerator.defaultBaseDate) + } + + private def checkTruncDate(baseDate: Long): Unit = { val supportedFormats = CometTruncDate.supportedFormats val unsupportedFormats = Seq("invalid") @@ -88,7 +104,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH Seq( StructField("c0", DataTypes.DateType, true), StructField("c1", DataTypes.StringType, true))) - val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, 1000, DataGenOptions()) + val df = FuzzDataGenerator.generateDataFrame( + r, + spark, + schema, + 1000, + DataGenOptions(baseDate = baseDate)) df.createOrReplaceTempView("tbl") @@ -146,7 +167,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val supportedFormats = CometTruncTimestamp.supportedFormats val unsupportedFormats = Seq("invalid") - createTimestampTestData().createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone // https://github.com/apache/datafusion-comet/issues/2649 @@ -172,7 +193,9 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val unsupportedFormats = Seq("invalid") withTempDir { path => - createTimestampTestData().write.mode(SaveMode.Overwrite).parquet(path.toString) + createTimestampTestData(dataFusionRangeBaseDate).write + .mode(SaveMode.Overwrite) + .parquet(path.toString) spark.read.parquet(path.toString).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone @@ -197,6 +220,15 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } } + test("date_trunc (TruncTimestamp) - timestamps outside DataFusion nanosecond range") { + createTimestampTestData(FuzzDataGenerator.defaultBaseDate).createOrReplaceTempView("tbl") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + for (format <- Seq("year", "week", "hour")) { + checkSparkAnswerAndOperator(s"SELECT c0, date_trunc('$format', c0) from tbl order by c0") + } + } + } + test("date_trunc - non-UTC timezone routes via codegen dispatcher") { // The native date_trunc is Incompatible in non-UTC sessions // (https://github.com/apache/datafusion-comet/issues/2649), so with allowIncompatible=false @@ -241,9 +273,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // for any non-UTC timezone whose DST transitions are 1-hour and whose dates fall within // chrono-tz's precomputed DST horizon (currently ~year 2100). A base date in 2024 keeps the // generator well inside that window. - val baseDate2024 = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime - createTimestampTestData(baseDate2024).createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") val nonUtcTimezones = Seq("America/New_York", "Europe/London", "Asia/Tokyo") for (tz <- nonUtcTimezones) { @@ -327,14 +357,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // Use a reasonable date range (around year 2024) to avoid chrono-tz DST calculation // issues with far-future dates. The default baseDate is year 3333 which is beyond // the range where chrono-tz can reliably calculate DST transitions. - val reasonableBaseDate = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime val ntzDF = FuzzDataGenerator.generateDataFrame( r, spark, ntzSchema, 100, - DataGenOptions(baseDate = reasonableBaseDate)) + DataGenOptions(baseDate = dataFusionRangeBaseDate)) ntzDF.createOrReplaceTempView("ntz_tbl") for (tz <- crossTimezones) { withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz) { From 62c5d44921a25f823c53d8fc58566874fcda3c25 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 16:33:20 +0000 Subject: [PATCH 12/89] perf: skip Spark's encoders for boxed primitive Scala UDF values in the codegen dispatcher (#6714) * perf: skip Spark's encoders for boxed primitive Scala UDF values in the codegen dispatcher A typed Scala UDF converts a boxed primitive parameter (`java.lang.Long`, ...) through its input encoder and a boxed primitive result through its output encoder, each a projection per row. The kernel now drops those encoders when their expression is exactly the boxing or unboxing call, so `ScalaUDF` passes the value through `CatalystTypeConverters` unchanged, as Spark 4 does for a Java UDF. Every other encoder is left alone. Closes #6706. * Compare boxed encoders with the trees Spark's builders produce Check each encoder by building the expected tree with Spark's own `DeserializerBuildHelper` and `SerializerBuildHelper` builders and comparing with `==`, in place of field-by-field matching. Drop parameter encoders with `filterNot` like the result's, and always copy the `ScalaUDF`, since an equal copy leaves its parents untouched. In the tests, register the shared UDFs once, pass vector classes directly, drop two no-op casts, and run the mixed-parameter calls as one query. --- .../codegen/CometBatchKernelCodegen.scala | 86 ++++++++++- .../org/apache/comet/CometCodegenSuite.scala | 133 +++++++++++++++++- 2 files changed, 211 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index e478057ceef..e2edb3c989c 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -24,7 +24,10 @@ import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, Unevaluable} +import org.apache.spark.sql.catalyst.DeserializerBuildHelper.createDeserializerForTypesSupportValueOf +import org.apache.spark.sql.catalyst.SerializerBuildHelper.{createSerializerForBoolean, createSerializerForByte, createSerializerForDouble, createSerializerForFloat, createSerializerForInteger, createSerializerForLong, createSerializerForShort} +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, ScalaUDF, Unevaluable} import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.execution.ExecSubqueryExpression import org.apache.spark.sql.internal.SQLConf @@ -47,7 +50,8 @@ import org.apache.comet.shims.{CometExprTraitShim, CometTypeShim} * Input- and output-side emission live in [[CometBatchKernelCodegenInput]] and * [[CometBatchKernelCodegenOutput]]. This file owns the [[ArrowColumnSpec]] vocabulary, the * [[canHandle]] / [[allocateOutput]] / [[compile]] / [[generateSource]] entry points, and - * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant). + * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant, boxed + * primitive `ScalaUDF` values without their encoders). * * The generated kernel is the `InternalRow` that Spark's `BoundReference.genCode` reads from. See * [[generateSource]] for how the wiring is set up. @@ -263,6 +267,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come inputSchema: Seq[ArrowColumnSpec]): GeneratedSource = { canHandle(boundExpr).foreach(reason => throw new IllegalArgumentException(s"CometBatchKernelCodegen.generateSource: $reason")) + val expr = withoutBoxedPrimitiveEncoders(boundExpr) val ctx = new CodegenContext // `BoundReference.genCode` emits `${ctx.INPUT_ROW}.getUTF8String(ord)`. Aliasing `row` to // `this` at the top of `process` routes those reads to the kernel's typed getters (final @@ -293,19 +298,19 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // `subexprFunctionsCode` is the concatenated helper invocation block, spliced into the // per-row body by `defaultBody`. val ev = if (SQLConf.get.subexpressionEliminationEnabled) { - ctx.generateExpressions(Seq(boundExpr), doSubexpressionElimination = true).head + ctx.generateExpressions(Seq(expr), doSubexpressionElimination = true).head } else { - boundExpr.genCode(ctx) + expr.genCode(ctx) } val subExprsCode = ctx.subexprFunctionsCode val (cls, setup, snippet) = - CometBatchKernelCodegenOutput.emitOutputWriter(boundExpr.dataType, ev.value, ctx) - (cls, setup, defaultBody(boundExpr, inputSchema, ev, snippet, subExprsCode)) + CometBatchKernelCodegenOutput.emitOutputWriter(expr.dataType, ev.value, ctx) + (cls, setup, defaultBody(expr, inputSchema, ev, snippet, subExprsCode)) } val typedFieldDecls = CometBatchKernelCodegenInput.emitInputFieldDecls(inputSchema) val typedInputCasts = CometBatchKernelCodegenInput.emitInputCasts(inputSchema) - val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(boundExpr) + val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(expr) val getters = CometBatchKernelCodegenInput.emitTypedGetters(inputSchema, decimalTypeByOrdinal) val nested = CometBatchKernelCodegenInput.emitNestedClasses(inputSchema) @@ -371,6 +376,73 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come GeneratedSource(code.body, code, ctx.references.toArray) } + /** + * Drops the encoder from each boxed primitive parameter and result of every `ScalaUDF` in the + * tree, so the kernel passes those values to and from the user function directly (#6706). + * + * A typed Scala UDF carries an `ExpressionEncoder` for each parameter and one for its result. + * For a boxed primitive (`(x: java.lang.Long) => ...`) the parameter's deserializer is + * `java.lang.Long.valueOf(input[0])`, which yields null for a null input, and the result's + * serializer is `input[0].longValue()`. `ScalaUDF` runs each one on every row through a + * projection: a `SafeProjection` into a `GenericInternalRow` going in, and an + * `UnsafeProjection` into an `UnsafeRow` coming out. Spark's whole-stage codegen runs the same + * projections. + * + * Without an encoder, `ScalaUDF` converts the value with `CatalystTypeConverters` instead, + * which passes a primitive through in its boxed form: `identity` going in, and coming out an + * `Option` unwrap that a boxed value never takes. The function receives the same value or null, + * and Catalyst gets the same result. Spark 4 builds a Java UDF the same way, with no encoders. + * + * Only an encoder whose expression is exactly the one Spark builds for a boxed primitive is + * dropped. Every other encoder (`String`, `Option`, case classes, collections) stays. + */ + private[comet] def withoutBoxedPrimitiveEncoders(expr: Expression): Expression = + expr.transformUp { case udf: ScalaUDF => + udf.copy( + inputEncoders = udf.inputEncoders.zip(udf.children).map { case (enc, child) => + enc.filterNot(isBoxedPrimitiveDeserializer(_, child.dataType)) + }, + outputEncoder = udf.outputEncoder.filterNot(isBoxedPrimitiveSerializer(_, udf.dataType))) + } + + /** + * Boxed class of each type `CatalystTypeConverters` passes through, and Spark's builder for the + * serializer of that class's encoder. + */ + private def boxedPrimitive(dt: DataType): Option[(Class[_], Expression => Expression)] = + dt match { + case BooleanType => Some((classOf[java.lang.Boolean], createSerializerForBoolean)) + case ByteType => Some((classOf[java.lang.Byte], createSerializerForByte)) + case ShortType => Some((classOf[java.lang.Short], createSerializerForShort)) + case IntegerType => Some((classOf[java.lang.Integer], createSerializerForInteger)) + case LongType => Some((classOf[java.lang.Long], createSerializerForLong)) + case FloatType => Some((classOf[java.lang.Float], createSerializerForFloat)) + case DoubleType => Some((classOf[java.lang.Double], createSerializerForDouble)) + case _ => None + } + + /** + * True iff `enc` deserializes a `dt` value exactly as the encoder of its boxed class does, with + * `valueOf` over a nullable reference. The reference must be nullable because `StaticInvoke` + * propagates null only from a nullable argument: over a non-nullable one, Spark would box the 0 + * that a row holds for a null. + */ + private def isBoxedPrimitiveDeserializer(enc: ExpressionEncoder[_], dt: DataType): Boolean = + boxedPrimitive(dt).exists { case (boxed, _) => + enc.objDeserializer == + createDeserializerForTypesSupportValueOf(BoundReference(0, dt, nullable = true), boxed) + } + + /** + * True iff `enc` serializes a boxed `dt` value exactly as the encoder of the boxed class does, + * by unboxing it. `ScalaUDF` turns a null result into null before it calls the serializer, so + * the serializer only ever sees a boxed value. + */ + private def isBoxedPrimitiveSerializer(enc: ExpressionEncoder[_], dt: DataType): Boolean = + boxedPrimitive(dt).exists { case (boxed, serializerFor) => + enc.objSerializer == serializerFor(BoundReference(0, ObjectType(boxed), nullable = true)) + } + /** * Per-row body. For `NullIntolerant` expressions whose input nulls fully determine a null * output (see [[canShortCircuitNulls]]), prepends a short-circuit on those ordinals so the diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index e9105230d7f..8f6d7c4a8d9 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.util.GenericArrayData @@ -1304,6 +1304,137 @@ class CometCodegenSuite } } + private val boxedPrimitiveTypes = Seq("bool", "byte", "short", "int", "long", "float", "double") + + /** + * Boxed primitive UDFs for the #6706 tests, where the kernel drops Spark's encoders. Each + * `_str` function names the exact value it receives, null included, so a value converted + * differently from Spark's encoder (a null read as 0, a -0.0 read as 0.0) fails the comparison + * with Spark. Each `_id` function returns the boxed value it receives, which covers the + * result side. `long_tag` and `opt_long` pair a boxed parameter with ones that keep their + * encoders. + */ + private def registerBoxedPrimitiveUdfs(): Unit = { + spark.udf.register("bool_str", (x: java.lang.Boolean) => String.valueOf(x)) + spark.udf.register("byte_str", (x: java.lang.Byte) => String.valueOf(x)) + spark.udf.register("short_str", (x: java.lang.Short) => String.valueOf(x)) + spark.udf.register("int_str", (x: java.lang.Integer) => String.valueOf(x)) + spark.udf.register("long_str", (x: java.lang.Long) => String.valueOf(x)) + // Raw bits, so the sign of a zero or a NaN payload would show. + spark.udf.register( + "float_str", + (x: java.lang.Float) => + if (x == null) "null" else java.lang.Float.floatToRawIntBits(x).toString) + spark.udf.register( + "double_str", + (x: java.lang.Double) => + if (x == null) "null" else java.lang.Double.doubleToRawLongBits(x).toString) + spark.udf.register("bool_id", (x: java.lang.Boolean) => x) + spark.udf.register("byte_id", (x: java.lang.Byte) => x) + spark.udf.register("short_id", (x: java.lang.Short) => x) + spark.udf.register("int_id", (x: java.lang.Integer) => x) + spark.udf.register("long_id", (x: java.lang.Long) => x) + spark.udf.register("float_id", (x: java.lang.Float) => x) + spark.udf.register("double_id", (x: java.lang.Double) => x) + spark.udf.register("long_tag", (x: java.lang.Long, s: String) => s"$x:$s") + spark.udf.register("opt_long", (x: Option[Long]) => x.map(_ + 1).getOrElse(-1L)) + } + + private def withBoxedPrimitiveTable(f: => Unit): Unit = { + withTable("t") { + sql( + "CREATE TABLE t (c_bool BOOLEAN, c_byte TINYINT, c_short SMALLINT, c_int INT, " + + "c_long BIGINT, c_float FLOAT, c_double DOUBLE, s STRING) USING parquet") + sql( + "INSERT INTO t VALUES " + + "(true, CAST(-128 AS TINYINT), CAST(-32768 AS SMALLINT), -2147483648, " + + "-9223372036854775808, CAST('-0.0' AS FLOAT), CAST('-0.0' AS DOUBLE), 'a'), " + + "(false, CAST(127 AS TINYINT), CAST(32767 AS SMALLINT), 2147483647, " + + "9223372036854775807, CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), NULL), " + + "(NULL, NULL, NULL, NULL, NULL, NULL, NULL, 'c'), " + + "(true, CAST(0 AS TINYINT), CAST(1000 AS SMALLINT), 1000, 1000, " + + "CAST('Infinity' AS FLOAT), CAST(0.0 AS DOUBLE), 'd'), " + + "(false, CAST(-1 AS TINYINT), CAST(-1000 AS SMALLINT), -1000, -1000, " + + "CAST(1.5 AS FLOAT), CAST('-Infinity' AS DOUBLE), NULL)") + f + } + } + + test("the kernel drops only boxed primitive encoders from a ScalaUDF (#6706)") { + registerBoxedPrimitiveUdfs() + spark.udf.register("str_id", (x: String) => x) + spark.udf.register("long_in_prim_out", (x: java.lang.Long) => if (x == null) -1L else x + 1L) + spark.udf.register("prim_add_one", (x: Long) => x + 1) + withBoxedPrimitiveTable { + def boundUdf(call: String): ScalaUDF = { + val udf = sql(s"SELECT $call FROM t").queryExecution.optimizedPlan.expressions + .flatMap(_.collect { case u: ScalaUDF => u }) + .head + val attrs = udf.collect { case a: AttributeReference => a }.distinct + BindReferences.bindReference(udf, AttributeSeq(attrs)) + } + // Whether the kernel drops the encoder of each parameter and then of the result. + def dropped(call: String): Seq[Boolean] = { + val udf = boundUdf(call) + val kernelUdf = + CometBatchKernelCodegen.withoutBoxedPrimitiveEncoders(udf).asInstanceOf[ScalaUDF] + (udf.inputEncoders :+ udf.outputEncoder) + .zip(kernelUdf.inputEncoders :+ kernelUdf.outputEncoder) + .map { case (before, after) => before.isDefined && after.isEmpty } + } + boxedPrimitiveTypes.foreach { t => + assert(dropped(s"${t}_id(c_$t)") === Seq(true, true), t) + } + assert(dropped("long_in_prim_out(c_long)") === Seq(true, false)) + assert(dropped("long_tag(c_long, s)") === Seq(true, false, false)) + assert(dropped("prim_add_one(c_long)") === Seq(false, false)) + assert(dropped("opt_long(c_long)") === Seq(false, false)) + assert(dropped("str_id(s)") === Seq(false, false)) + + // Spark's encoder path calls the parameter's converter even for a null input. Without the + // encoder, `ScalaUDF` tests for null itself, which shows the kernel compiles the rewrite. + def kernelSource(call: String, vectorClass: Class[_ <: ValueVector]): String = + CometBatchKernelCodegen + .generateSource( + boundUdf(call), + IndexedSeq(ArrowColumnSpec(vectorClass, nullable = true))) + .body + assert(kernelSource("str_id(s)", classOf[VarCharVector]).contains(".apply(null)")) + assert(!kernelSource("long_id(c_long)", classOf[BigIntVector]).contains(".apply(null)")) + } + } + + test("boxed primitive UDF parameters and results match Spark (#6706)") { + registerBoxedPrimitiveUdfs() + withBoxedPrimitiveTable { + val calls = boxedPrimitiveTypes.flatMap(t => Seq(s"${t}_str(c_$t)", s"${t}_id(c_$t)")) + assertCodegenRan { + checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t")) + } + } + } + + test("boxed primitive UDF parameters next to other parameters match Spark (#6706)") { + registerBoxedPrimitiveUdfs() + // Spark guards the call with a null check for the primitive parameter, not the boxed one. + spark.udf.register("long_or", (x: java.lang.Long, y: Long) => if (x == null) y else x + y) + withBoxedPrimitiveTable { + val calls = Seq( + "long_tag(c_long, s)", + "opt_long(c_long)", + "long_or(c_long, c_int)", + // Spark casts the int column to the `java.lang.Long` parameter's type. + "long_str(c_int)", + "long_id(long_id(c_long))", + // Inside a lambda, Spark evaluates the UDF through `eval` rather than generated code, + // with the same converters. + "transform(array(c_int, NULL), x -> int_str(x))") + assertCodegenRan { + checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t")) + } + } + } + test("ScalaUDF returning a different type than its input") { // String -> Int output transition. Identity-loop above keeps input == output. This asserts // the writer can switch types per the UDF's declared return. From f651144b9d3a6ceb7d037907aac6cde6e8a6fa26 Mon Sep 17 00:00:00 2001 From: Bhargava Vadlamani <11091419+coderfender@users.noreply.github.com> Date: Tue, 6 Oct 2026 17:51:01 +0000 Subject: [PATCH 13/89] add benches for remaining comet native kernels (#6648) --- native/spark-expr/Cargo.toml | 16 ++ .../benches/bloom_filter_might_contain.rs | 144 ++++++++++++++ native/spark-expr/benches/cast_nested.rs | 141 ++++++++++++++ native/spark-expr/benches/hll_agg.rs | 184 ++++++++++++++++++ native/spark-expr/benches/hll_scalar.rs | 74 +++++++ 5 files changed, 559 insertions(+) create mode 100644 native/spark-expr/benches/bloom_filter_might_contain.rs create mode 100644 native/spark-expr/benches/cast_nested.rs create mode 100644 native/spark-expr/benches/hll_agg.rs create mode 100644 native/spark-expr/benches/hll_scalar.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 7020f8aab75..965bbd79169 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -83,6 +83,10 @@ harness = false name = "cast_from_string" harness = false +[[bench]] +name = "cast_nested" +harness = false + [[bench]] name = "cast_string_to_date" harness = false @@ -183,6 +187,10 @@ harness = false name = "bloom_filter_agg" harness = false +[[bench]] +name = "bloom_filter_might_contain" +harness = false + [[bench]] name = "collect" harness = false @@ -455,6 +463,14 @@ harness = false name = "aggregate_stats" harness = false +[[bench]] +name = "hll_agg" +harness = false + +[[bench]] +name = "hll_scalar" +harness = false + [[bench]] name = "aggregate_misc" harness = false diff --git a/native/spark-expr/benches/bloom_filter_might_contain.rs b/native/spark-expr/benches/bloom_filter_might_contain.rs new file mode 100644 index 00000000000..c0e926a40ff --- /dev/null +++ b/native/spark-expr/benches/bloom_filter_might_contain.rs @@ -0,0 +1,144 @@ +// 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. + +//! Benchmark for the per-row `might_contain` probe (`BloomFilterMightContain`). The filter it +//! probes against is built up front by running the public `BloomFilterAgg` over a column of +//! longs, so the bench exercises only the probe path (`might_contain_longs`) and not filter +//! construction. + +use arrow::array::{Array, ArrayRef, BinaryArray, Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::common::config::ConfigOptions; +use datafusion::common::ScalarValue; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::logical_expr::{AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::expressions::{lit, Column}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::prelude::SessionContext; +use datafusion_comet_spark_expr::{ + BloomFilterAgg, BloomFilterMightContain, SparkBloomFilterVersion, +}; +use std::hint::black_box; +use std::sync::Arc; +use tokio::runtime::Runtime; + +const DISTINCT_ITEMS: usize = 8192; +const NUM_BITS: i32 = (DISTINCT_ITEMS as i32) * 16; + +/// Build a serialized Spark bloom filter containing `0..DISTINCT_ITEMS` by running the public +/// `BloomFilterAgg` once. Returns the filter's Spark-serialized bytes. +fn build_filter_bytes(rt: &Runtime) -> Vec { + let values: ArrayRef = Arc::new(Int64Array::from( + (0..DISTINCT_ITEMS as i64).collect::>(), + )); + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values]).unwrap(); + + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None).unwrap(), + ))); + let udf = AggregateUDF::new_from_impl(BloomFilterAgg::new( + lit(ScalarValue::Int64(Some(DISTINCT_ITEMS as i64))), + lit(ScalarValue::Int64(Some(NUM_BITS as i64))), + DataType::Binary, + SparkBloomFilterVersion::V1, + )); + let aggr_expr = AggregateExprBuilder::new( + Arc::new(udf), + vec![Arc::new(Column::new("v", 0)) as Arc], + ) + .schema(Arc::clone(&schema)) + .alias("bloom_filter") + .build() + .unwrap(); + + let aggregate: Arc = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new(vec![], vec![], vec![vec![]], false), + vec![aggr_expr.into()], + vec![None], + scan, + schema, + ) + .unwrap(), + ); + + let ctx = SessionContext::new(); + let batches = rt.block_on(collect(aggregate, ctx.task_ctx())).unwrap(); + let col = batches[0].column(0); + let bytes = col.as_any().downcast_ref::().unwrap(); + bytes.value(0).to_vec() +} + +/// Build an Int64 probe column: the first half are hits (`0..rows/2`, present in the filter), +/// the second half are misses (large values not inserted). Every `null_every`-th row is null +/// (`null_every == 0` means no nulls). +fn create_probe_array(rows: usize, null_every: usize) -> ArrayRef { + let arr: Int64Array = (0..rows) + .map(|i| { + if null_every != 0 && i % null_every == 0 { + None + } else if i < rows / 2 { + Some(i as i64) + } else { + Some(i as i64 + 1_000_000_000) + } + }) + .collect(); + Arc::new(arr) +} + +fn run(udf: &BloomFilterMightContain, values: &ArrayRef, rows: usize) { + let args = vec![ColumnarValue::Array(Arc::clone(values))]; + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("result", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap(), + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + let rt = Runtime::new().unwrap(); + let filter_bytes = build_filter_bytes(&rt); + let udf = + BloomFilterMightContain::try_new(lit(ScalarValue::Binary(Some(filter_bytes)))).unwrap(); + + let no_nulls = create_probe_array(rows, 0); + let with_nulls = create_probe_array(rows, 10); + + let mut group = c.benchmark_group("bloom_filter_might_contain"); + group.bench_function("probe no nulls", |b| b.iter(|| run(&udf, &no_nulls, rows))); + group.bench_function("probe with nulls", |b| { + b.iter(|| run(&udf, &with_nulls, rows)) + }); + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/cast_nested.rs b/native/spark-expr/benches/cast_nested.rs new file mode 100644 index 00000000000..5c6efd5cb47 --- /dev/null +++ b/native/spark-expr/benches/cast_nested.rs @@ -0,0 +1,141 @@ +// 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. + +//! Benchmarks for the nested `cast` paths: struct->struct (child retyped), map->map (value +//! retyped), and array->string. Each builds an 8192-row input and casts it through the public +//! `Cast` physical expression, mirroring the scalar cast benches. + +use arrow::array::{ + ArrayRef, Int32Array, ListArray, MapArray, RecordBatch, StringArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; +const ENTRIES_PER_ROW: usize = 2; + +/// Single-column batch named `c` wrapping `array`. +fn batch(array: ArrayRef) -> RecordBatch { + let schema = Schema::new(vec![Field::new("c", array.data_type().clone(), true)]); + RecordBatch::try_new(Arc::new(schema), vec![array]).unwrap() +} + +/// `Struct` with `NUM_ROWS` rows. +fn struct_input() -> ArrayRef { + let a = Arc::new(Int32Array::from((0..NUM_ROWS as i32).collect::>())) as ArrayRef; + let b = Arc::new(StringArray::from( + (0..NUM_ROWS).map(|i| format!("s{i}")).collect::>(), + )) as ArrayRef; + let fields = Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]); + Arc::new(StructArray::new(fields, vec![a, b], None)) +} + +/// `Map` with `NUM_ROWS` rows of `ENTRIES_PER_ROW` entries each. +fn map_input() -> ArrayRef { + let total = NUM_ROWS * ENTRIES_PER_ROW; + let keys = Arc::new(StringArray::from( + (0..total).map(|i| format!("k{i}")).collect::>(), + )) as ArrayRef; + let values = Arc::new(Int32Array::from((0..total as i32).collect::>())) as ArrayRef; + let entries_fields = Fields::from(vec![ + Field::new("key_value_key", DataType::Utf8, false), + Field::new("key_value_value", DataType::Int32, true), + ]); + let entries = StructArray::new(entries_fields, vec![keys, values], None); + let offsets = OffsetBuffer::::new( + (0..=NUM_ROWS) + .map(|i| (i * ENTRIES_PER_ROW) as i32) + .collect::>() + .into(), + ); + let entries_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entries.fields().clone()), + false, + )); + Arc::new(MapArray::new(entries_field, offsets, entries, None, false)) +} + +/// `List` with `NUM_ROWS` rows whose lengths cycle through 0..=3. +fn list_input() -> ArrayRef { + let data = (0..NUM_ROWS).map(|i| { + let len = i % 4; + Some( + (0..len) + .map(|j| Some((i * 4 + j) as i32)) + .collect::>(), + ) + }); + Arc::new(ListArray::from_iter_primitive::(data)) +} + +fn criterion_benchmark(c: &mut Criterion) { + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let col: Arc = Arc::new(Column::new("c", 0)); + + let mut group = c.benchmark_group("cast_nested"); + + // struct -> struct: retypes the first child. + let struct_batch = batch(struct_input()); + let struct_target = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ])); + let cast_struct = Cast::new(col.clone(), struct_target, options.clone(), None, None); + group.bench_function("struct_to_struct", |b| { + b.iter(|| black_box(cast_struct.evaluate(black_box(&struct_batch)).unwrap())) + }); + + // map -> map: retypes the value field. + let map_batch = batch(map_input()); + let map_target = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int64, true), + ])), + false, + )), + false, + ); + let cast_map = Cast::new(col.clone(), map_target, options.clone(), None, None); + group.bench_function("map_to_map", |b| { + b.iter(|| black_box(cast_map.evaluate(black_box(&map_batch)).unwrap())) + }); + + // list -> Utf8: renders each list to its Spark string form. + let list_batch = batch(list_input()); + let cast_list = Cast::new(col, DataType::Utf8, options, None, None); + group.bench_function("array_to_string", |b| { + b.iter(|| black_box(cast_list.evaluate(black_box(&list_batch)).unwrap())) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/hll_agg.rs b/native/spark-expr/benches/hll_agg.rs new file mode 100644 index 00000000000..c933533f7c6 --- /dev/null +++ b/native/spark-expr/benches/hll_agg.rs @@ -0,0 +1,184 @@ +// 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. + +//! Benchmarks for the Comet HyperLogLog aggregates in `agg_funcs`: `hll_sketch_agg` +//! (builds a sketch from raw values) and `hll_union_agg` (merges serialized sketches). +//! Each runs as a grouped `AggregateMode::Partial` over 8192 rows x 10 batches, mirroring +//! the other aggregate benches. The union input is a column of real sketch bytes minted +//! up front with `SparkHllSketch`, since the kernel deserializes and merges each one. + +use arrow::array::{ArrayRef, BinaryBuilder, Int64Builder, RecordBatch, StringBuilder}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::AggregateUDF; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_comet_spark_expr::{HllSketchAgg, HllUnionAgg, SparkHllSketch}; +use futures::StreamExt; +use std::hint::black_box; +use std::sync::Arc; +use tokio::runtime::Runtime; + +/// Spark's default `lgConfigK` for `approx_count_distinct`/HLL sketches. +const LG_CONFIG_K: u8 = 12; +/// Number of distinct sketches minted for the union input; rows cycle through them. +const SKETCH_POOL: usize = 256; + +async fn agg_test( + partitions: &[Vec], + group_col: Arc, + input_exprs: Vec>, + aggregate_udf: Arc, + alias: &str, +) { + let schema = &partitions[0][0].schema(); + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(partitions, Arc::clone(schema), None).unwrap(), + ))); + let aggregate = create_aggregate(scan, group_col, input_exprs, schema, aggregate_udf, alias); + let mut stream = aggregate + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + while let Some(batch) = stream.next().await { + let _batch = batch.unwrap(); + } +} + +fn create_aggregate( + scan: Arc, + group_col: Arc, + input_exprs: Vec>, + schema: &SchemaRef, + aggregate_udf: Arc, + alias: &str, +) -> Arc { + let aggr_expr = AggregateExprBuilder::new(aggregate_udf, input_exprs) + .schema(schema.clone()) + .alias(alias) + .with_ignore_nulls(false) + .with_distinct(false) + .build() + .unwrap(); + + Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![(group_col, "c0".to_string())]), + vec![aggr_expr.into()], + vec![None], // no filter expressions + scan, + Arc::clone(schema), + ) + .unwrap(), + ) +} + +/// A pool of `SKETCH_POOL` distinct non-empty sketches, each covering a disjoint run of i64 +/// values so a union of several estimates to a meaningful cardinality. +fn sketch_pool() -> Vec> { + (0..SKETCH_POOL) + .map(|p| { + let mut sketch = SparkHllSketch::new(LG_CONFIG_K); + let start = (p * 100) as i64; + for v in start..start + 100 { + sketch.update_i64(v); + } + sketch.to_sketch_bytes() + }) + .collect() +} + +/// `c0` is the Utf8 group key (1024 groups); `c1` is a spread-out Int64 input for +/// `hll_sketch_agg`; `c2` is a column of serialized sketches for `hll_union_agg`. +fn create_record_batch(num_rows: usize, pool: &[Vec]) -> RecordBatch { + let mut group_builder = StringBuilder::with_capacity(num_rows, num_rows * 8); + let mut c1_builder = Int64Builder::with_capacity(num_rows); + let mut c2_builder = BinaryBuilder::new(); + for i in 0..num_rows { + group_builder.append_value(format!("group_{}", i % 1024)); + c1_builder.append_value(i as i64); + c2_builder.append_value(&pool[i % pool.len()]); + } + + let fields = vec![ + Field::new("c0", DataType::Utf8, false), + Field::new("c1", DataType::Int64, false), + Field::new("c2", DataType::Binary, false), + ]; + let columns: Vec = vec![ + Arc::new(group_builder.finish()), + Arc::new(c1_builder.finish()), + Arc::new(c2_builder.finish()), + ]; + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let num_rows = 8192; + let pool = sketch_pool(); + let batch = create_record_batch(num_rows, &pool); + let batches = vec![batch; 10]; + let partitions = &[batches]; + + let c0: Arc = Arc::new(Column::new("c0", 0)); + let c1: Arc = Arc::new(Column::new("c1", 1)); + let c2: Arc = Arc::new(Column::new("c2", 2)); + + let rt = Runtime::new().unwrap(); + let mut group = c.benchmark_group("hll_agg"); + + let sketch_udf: Arc = Arc::new(AggregateUDF::new_from_impl(HllSketchAgg::new( + LG_CONFIG_K as i32, + ))); + group.bench_function("hll_sketch_agg", |b| { + b.to_async(&rt).iter(|| { + black_box(agg_test( + partitions, + c0.clone(), + vec![c1.clone()], + sketch_udf.clone(), + "hll_sketch_agg", + )) + }) + }); + + let union_udf: Arc = + Arc::new(AggregateUDF::new_from_impl(HllUnionAgg::new(false))); + group.bench_function("hll_union_agg", |b| { + b.to_async(&rt).iter(|| { + black_box(agg_test( + partitions, + c0.clone(), + vec![c2.clone()], + union_udf.clone(), + "hll_union_agg", + )) + }) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/hll_scalar.rs b/native/spark-expr/benches/hll_scalar.rs new file mode 100644 index 00000000000..d3b4c79070f --- /dev/null +++ b/native/spark-expr/benches/hll_scalar.rs @@ -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. + +//! Benchmarks for the scalar HLL kernels: `hll_sketch_estimate` (Binary sketch -> Long estimate) +//! and `hll_union` (two Binary sketch columns + an `allowDifferentLgConfigK` flag -> Binary). +//! Both deserialize a sketch per row, so inputs are real sketch bytes minted up front with +//! `SparkHllSketch`. + +use arrow::array::{ArrayRef, BinaryArray, BooleanArray}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{spark_hll_sketch_estimate, spark_hll_union, SparkHllSketch}; +use std::hint::black_box; +use std::sync::Arc; + +const LG_CONFIG_K: u8 = 12; + +/// Build a Binary column of `rows` sketches, each covering a distinct run of 100 i64 values +/// offset by `base` so the two columns in a union cover different ranges. +fn sketch_column(rows: usize, base: i64) -> ArrayRef { + let sketches: Vec> = (0..rows) + .map(|i| { + let mut s = SparkHllSketch::new(LG_CONFIG_K); + let start = base + (i as i64) * 100; + for v in start..start + 100 { + s.update_i64(v); + } + s.to_sketch_bytes() + }) + .collect(); + Arc::new(BinaryArray::from_iter_values(sketches)) +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + let left = sketch_column(rows, 0); + let right = sketch_column(rows, 1_000_000); + let allow: ArrayRef = Arc::new(BooleanArray::from(vec![false; rows])); + + let mut group = c.benchmark_group("hll_scalar"); + + group.bench_function("hll_sketch_estimate", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&left))]; + b.iter(|| black_box(spark_hll_sketch_estimate(black_box(&args)).unwrap())) + }); + + group.bench_function("hll_union", |b| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&left)), + ColumnarValue::Array(Arc::clone(&right)), + ColumnarValue::Array(Arc::clone(&allow)), + ]; + b.iter(|| black_box(spark_hll_union(black_box(&args)).unwrap())) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 275d09db8231215a97f6d86154556c074ed691d6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 18:16:39 +0000 Subject: [PATCH 14/89] docs: add 1.1.0 changelog (#6728) * docs: add 1.1.0 changelog (#6282) * docs: [branch-1.1] credit co-authors in the 1.1.0 changelog (#6359) Regenerate the 1.1.0 changelog with the generator from #6358, for the same range as #6282 (1.0.0..021c378fc), and format it with prettier 3.9.9. The generator used for #6282 credited each PR only to the author of its merge commit, so it left out Michael Taranov, whose commits from #6092 are in #6219, and credited the backports only to the person who opened them. The PR lines for #5310, #6219, #6223 and the four backports now name their co-authors, and the credits count each co-authored PR, which adds Michael Taranov and brings the contributor count to 41. * update changelog (#6569) --- docs/source/changelog/1.1.0.md | 503 +++++++++++++++++++++++++++++++++ docs/source/changelog/index.md | 1 + 2 files changed, 504 insertions(+) create mode 100644 docs/source/changelog/1.1.0.md diff --git a/docs/source/changelog/1.1.0.md b/docs/source/changelog/1.1.0.md new file mode 100644 index 00000000000..dfbaf5abda7 --- /dev/null +++ b/docs/source/changelog/1.1.0.md @@ -0,0 +1,503 @@ + + +# DataFusion Comet 1.1.0 Changelog + +This release consists of 416 commits from 41 contributors. See credits at the end of this changelog for more information. + +**Fixed bugs:** + +- fix: skip null slots when checking overflow in unary negation [#5162](https://github.com/apache/datafusion-comet/pull/5162) (Smallfu666) +- fix: surface next_day and make_date ANSI errors as Spark exceptions [#5167](https://github.com/apache/datafusion-comet/pull/5167) (peterxcli) +- fix: normalize nested field nullability in ShuffleScanExec and ExpandExec [#5138](https://github.com/apache/datafusion-comet/pull/5138) (andygrove) +- fix: propagate the Spark task ClassLoader to JVM UDF calls [#5282](https://github.com/apache/datafusion-comet/pull/5282) (andygrove) +- fix: avoid duplicate CheckOverflow evaluation for decimal division [#5225](https://github.com/apache/datafusion-comet/pull/5225) (peterxcli) +- fix: match Spark whitespace trimming in to_time and try_to_time [#5364](https://github.com/apache/datafusion-comet/pull/5364) (sunchao) +- fix: preserve Catalyst nullability and field IDs in native Parquet writes [#5369](https://github.com/apache/datafusion-comet/pull/5369) (sunchao) +- fix: guard against silent fail_on_error loss in scalar wiring (#5074) [#5359](https://github.com/apache/datafusion-comet/pull/5359) (sam-1112) +- fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets [#5234](https://github.com/apache/datafusion-comet/pull/5234) (peterxcli) +- fix: canonicalize NaN in flat arrays_overlap float keys [#5376](https://github.com/apache/datafusion-comet/pull/5376) (sunchao) +- fix: report native shuffle write metrics accurately [#5370](https://github.com/apache/datafusion-comet/pull/5370) (sunchao) +- fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions [#5392](https://github.com/apache/datafusion-comet/pull/5392) (parthchandra) +- fix: format NativeMemoryConsumer id in toString [#5398](https://github.com/apache/datafusion-comet/pull/5398) (ywskycn) +- fix: ignore reader-side parquet.hadoop.vectored.io.enabled in Iceberg native-write detection [#5410](https://github.com/apache/datafusion-comet/pull/5410) (snmvaughan) +- fix: support map casts with NullType elements [#5045](https://github.com/apache/datafusion-comet/pull/5045) (peterxcli) +- fix: distinguish reflection failure from absent accessor in IcebergReflection [#5412](https://github.com/apache/datafusion-comet/pull/5412) (unikdahal) +- fix: use current shuffle config in aggregate test [#5439](https://github.com/apache/datafusion-comet/pull/5439) (peterxcli) +- fix: support wide years in native make_date [#5443](https://github.com/apache/datafusion-comet/pull/5443) (peterxcli) +- fix: report native child spill metrics in shuffle tasks [#5445](https://github.com/apache/datafusion-comet/pull/5445) (sunchao) +- fix: expose native memory usage to Spark [#5408](https://github.com/apache/datafusion-comet/pull/5408) (ywskycn) +- fix: release native shuffle reservation after spill failure [#5461](https://github.com/apache/datafusion-comet/pull/5461) (peterxcli) +- fix: track peak memory before JVM shuffle spill [#5463](https://github.com/apache/datafusion-comet/pull/5463) (peterxcli) +- fix: RAII for memory pool registration [#5464](https://github.com/apache/datafusion-comet/pull/5464) (peterxcli) +- fix: tighten RSS writer visibility and JNI array limits [#5475](https://github.com/apache/datafusion-comet/pull/5475) (pingzh) +- fix: report native operator spill metrics in Spark task metrics for non-shuffle stages [#5497](https://github.com/apache/datafusion-comet/pull/5497) (peterxcli) +- fix: report bounded shuffle allocator memory usage [#5516](https://github.com/apache/datafusion-comet/pull/5516) (ywskycn) +- fix: use Spark type names in ANSI abs overflow errors [#5357](https://github.com/apache/datafusion-comet/pull/5357) (Smallfu666) +- fix: make collect_list/collect_set argument coercion a normalization barrier [#5159](https://github.com/apache/datafusion-comet/pull/5159) (andygrove) +- fix: make task-shared memory pool as ref-counted RAII guard [#5494](https://github.com/apache/datafusion-comet/pull/5494) (peterxcli) +- fix: accept UTC timezone aliases in Python Arrow input [#5556](https://github.com/apache/datafusion-comet/pull/5556) (sunchao) +- fix: normalize scalar float sort and window rank keys [#5469](https://github.com/apache/datafusion-comet/pull/5469) (sunchao) +- fix: make CometExplodeExec respect batch size [#5362](https://github.com/apache/datafusion-comet/pull/5362) (andygrove) +- fix: align Spark 4.2 Python worker configuration [#5561](https://github.com/apache/datafusion-comet/pull/5561) (sunchao) +- fix: prevent memory leak after failed Arrow vector import [#5539](https://github.com/apache/datafusion-comet/pull/5539) (1fanwang) +- fix: preserve Arrow Field metadata across C Data exports [#5552](https://github.com/apache/datafusion-comet/pull/5552) (peterxcli) +- fix: rebase map offsets in mapsort so sliced maps do not overrun entries [#5630](https://github.com/apache/datafusion-comet/pull/5630) (viirya) +- fix: scope Celeborn bootstrap hooks to Comet clients [#5627](https://github.com/apache/datafusion-comet/pull/5627) (pingzh) +- fix: skip codegen dispatcher null short-circuit when a foldable subtree can raise [#5623](https://github.com/apache/datafusion-comet/pull/5623) (andygrove) +- fix: accept dictionary encodings in remote shuffle [#5650](https://github.com/apache/datafusion-comet/pull/5650) (pingzh) +- fix: make CometDiskBlockWriter spill registry per-task instead of executor-global [#5493](https://github.com/apache/datafusion-comet/pull/5493) (peterxcli) +- fix: Bump iceberg-rust so native Iceberg writes URL-escape partition paths [#5651](https://github.com/apache/datafusion-comet/pull/5651) (andygrove) +- fix(celeborn): reject unsafe native push completion tracking [#5665](https://github.com/apache/datafusion-comet/pull/5665) (pingzh) +- fix: match Spark's null short-circuiting in array_join and enable it natively [#5558](https://github.com/apache/datafusion-comet/pull/5558) (Visorgood) +- fix: make native shuffle spill metrics independent of input batching [#5628](https://github.com/apache/datafusion-comet/pull/5628) (sunchao) +- fix: expand object store option references, uniquify constant metadata names, drop dead parquet JNI [#5653](https://github.com/apache/datafusion-comet/pull/5653) (dwsmith1983) +- fix: match Spark's ANSI bound check for float/double to integral casts [#5683](https://github.com/apache/datafusion-comet/pull/5683) (peterxcli) +- fix: return NULL from rpad/lpad when the length column is NULL instead of panicking [#5680](https://github.com/apache/datafusion-comet/pull/5680) (peterxcli) +- fix: fall back for concat_ws with array arguments instead of failing natively [#5679](https://github.com/apache/datafusion-comet/pull/5679) (peterxcli) +- fix: prevent silent overflow when reading Parquet TIMESTAMP_MILLIS values [#5177](https://github.com/apache/datafusion-comet/pull/5177) (peterxcli) +- fix: recover native Celeborn shuffle from oversized rows [#5668](https://github.com/apache/datafusion-comet/pull/5668) (pingzh) +- fix: report native shuffle read metrics [#5554](https://github.com/apache/datafusion-comet/pull/5554) (peterxcli) +- fix: correctly rounded decimal to double/float cast matching BigDecimal.doubleValue/floatValue [#5684](https://github.com/apache/datafusion-comet/pull/5684) (peterxcli) +- fix: restore columnar transitions under the native Iceberg write [#5696](https://github.com/apache/datafusion-comet/pull/5696) (andygrove) +- fix: make columnar-to-row benchmarks exercise Comet [#5718](https://github.com/apache/datafusion-comet/pull/5718) (rich7420) +- fix: distinguish "nothing spilled" from a spill backend with no local path [#5726](https://github.com/apache/datafusion-comet/pull/5726) (andygrove) +- fix: propagate Arrow array copy errors [#5747](https://github.com/apache/datafusion-comet/pull/5747) (rich7420) +- fix: keep the dictionary hash fast path off nested and reseeded buffers [#5757](https://github.com/apache/datafusion-comet/pull/5757) (viirya) +- fix: preserve ANSI errors for rejected TIMESTAMP_NTZ casts [#5752](https://github.com/apache/datafusion-comet/pull/5752) (peterxcli) +- fix: apply the parent struct's null mask before hashing its fields [#5754](https://github.com/apache/datafusion-comet/pull/5754) (viirya) +- fix: read Iceberg tables partitioned by an unknown transform [#5759](https://github.com/apache/datafusion-comet/pull/5759) (andygrove) +- fix: native Iceberg write panics on an evolved partition spec and on a timestamptz partition path [#5729](https://github.com/apache/datafusion-comet/pull/5729) (andygrove) +- fix: let AQE optimize queries over Comet caches [#5733](https://github.com/apache/datafusion-comet/pull/5733) (peterxcli) +- fix: dispatch Iceberg system functions wrapped as ApplyFunctionExpression [#5773](https://github.com/apache/datafusion-comet/pull/5773) (andygrove) +- fix: attach tokio runtime threads to the JVM as daemon threads [#5748](https://github.com/apache/datafusion-comet/pull/5748) (zhangfengcdt) +- fix: check nested TIMESTAMP_MILLIS overflow in unfiltered scans [#5740](https://github.com/apache/datafusion-comet/pull/5740) (peterxcli) +- fix: match iceberg-java's exception for unclustered input to a clustered Iceberg write [#5779](https://github.com/apache/datafusion-comet/pull/5779) (andygrove) +- fix: enable FIRST/LAST partial merge [#5041](https://github.com/apache/datafusion-comet/pull/5041) (peterxcli) +- fix: preserve aggregate result identity during exchange reuse [#5470](https://github.com/apache/datafusion-comet/pull/5470) (sunchao) +- fix: ignore structural tags when lifting expression coverage [#5471](https://github.com/apache/datafusion-comet/pull/5471) (sunchao) +- fix: align string to timestamp parsing with Spark's segment rules [#5682](https://github.com/apache/datafusion-comet/pull/5682) (peterxcli) +- fix: roll native Iceberg data files on iceberg-java's 1000-row grid [#5780](https://github.com/apache/datafusion-comet/pull/5780) (andygrove) +- fix: decide libhdfs routing from the scheme as written [#5825](https://github.com/apache/datafusion-comet/pull/5825) (comphead) +- fix: list a fanout Iceberg write's data files in a stable order [#5810](https://github.com/apache/datafusion-comet/pull/5810) (andygrove) +- fix: isolate object-store registration by backend and configuration [#5503](https://github.com/apache/datafusion-comet/pull/5503) (sunchao) +- fix: Delete completed tasks' data files when an Iceberg write job fails [#5663](https://github.com/apache/datafusion-comet/pull/5663) (andygrove) +- fix: apply Spark's Parquet conversion rules to nested struct/list/map fields [#5681](https://github.com/apache/datafusion-comet/pull/5681) (peterxcli) +- fix: preserve nulls for Boolean/Byte/Short/Integer columns in FuzzDataGenerator [#5855](https://github.com/apache/datafusion-comet/pull/5855) (Smallfu666) +- fix: Accept explicit positive years in timestamp casts [#5858](https://github.com/apache/datafusion-comet/pull/5858) (peterxcli) +- fix: decline structs with duplicate field names before they reach Java Arrow [#5866](https://github.com/apache/datafusion-comet/pull/5866) (dwsmith1983) +- fix: explain ObjectHashAggregate fallback when Comet shuffle is disabled [#5746](https://github.com/apache/datafusion-comet/pull/5746) (0lai0) +- fix: preserve join and generator semantics in plan identity [#5828](https://github.com/apache/datafusion-comet/pull/5828) (ErikBPF) +- fix: propagate Parquet field-name folding failures [#5845](https://github.com/apache/datafusion-comet/pull/5845) (sunchao) +- fix: decode dictionary input for PyArrow UDFs [#5560](https://github.com/apache/datafusion-comet/pull/5560) (sunchao) +- fix: size JVM shuffle pointer array growth from the array, not the data pages [#5907](https://github.com/apache/datafusion-comet/pull/5907) (andygrove) +- fix: render float and double Iceberg partition values like iceberg-java [#5840](https://github.com/apache/datafusion-comet/pull/5840) (andygrove) +- fix: align time parsing and native second extraction with Spark [#5738](https://github.com/apache/datafusion-comet/pull/5738) (peterxcli) +- fix: normalize floating-point values in native collect_set [#5166](https://github.com/apache/datafusion-comet/pull/5166) (peterxcli) +- fix: keep Iceberg complex null checks on native scans [#5732](https://github.com/apache/datafusion-comet/pull/5732) (ErikBPF) +- fix: decode invalid UTF-8 at the JVM to native FFI import boundary [#5310](https://github.com/apache/datafusion-comet/pull/5310) (manuzhang, andygrove) +- fix: fall back when a struct repeats a Parquet field id [#6004](https://github.com/apache/datafusion-comet/pull/6004) (comphead) +- fix: normalize signed zero in nested float array comparisons [#5235](https://github.com/apache/datafusion-comet/pull/5235) (divyankshah) +- fix: fall back to Spark for native Iceberg writes to gs:// through HadoopFileIO [#5935](https://github.com/apache/datafusion-comet/pull/5935) (zhangfengcdt) +- fix: gate the regr_r2 degenerate-case swap on the Spark patch release [#6042](https://github.com/apache/datafusion-comet/pull/6042) (dwsmith1983) +- fix: include ABFS container in object store cache key [#5053](https://github.com/apache/datafusion-comet/pull/5053) (peterxcli) +- fix: preserve Spark row index read errors [#6046](https://github.com/apache/datafusion-comet/pull/6046) (liupoyi-1031) +- fix: remove the ineffective spark.executor.memoryOverhead adjustment from the driver plugin [#6054](https://github.com/apache/datafusion-comet/pull/6054) (andygrove) +- fix: count each memory pool once in analyze_trace [#5991](https://github.com/apache/datafusion-comet/pull/5991) (andygrove) +- fix: revert unsafe partial aggregates after final fallback [#5421](https://github.com/apache/datafusion-comet/pull/5421) (sunchao) +- fix: restore the site's mermaid diagrams and make a dropped one fail CI [#6064](https://github.com/apache/datafusion-comet/pull/6064) (andygrove) +- fix: write cached batches to the schema width, not the batch width [#6090](https://github.com/apache/datafusion-comet/pull/6090) (andygrove) +- fix(iceberg): guard native Iceberg scan driver-metric double-post, add metrics docs and tests [#6085](https://github.com/apache/datafusion-comet/pull/6085) (parthchandra) +- fix: let decimal SUM recover from an intermediate overflow like Spark [#6041](https://github.com/apache/datafusion-comet/pull/6041) (dwsmith1983) +- fix: Nested floating-point IN membership does not match Spark for signed zero [#6073](https://github.com/apache/datafusion-comet/pull/6073) (mizulun) +- fix: let Comet memory pools overcommit on grow instead of panicking [#6128](https://github.com/apache/datafusion-comet/pull/6128) (andygrove) +- fix: correct two nightly test failures on Spark 3.4 and 4.2 [#6156](https://github.com/apache/datafusion-comet/pull/6156) (andygrove) +- fix: support null calendar interval literals [#5133](https://github.com/apache/datafusion-comet/pull/5133) (peterxcli) +- fix: wrap Iceberg split-write failures the way Spark does when abort fails [#6153](https://github.com/apache/datafusion-comet/pull/6153) (andygrove) +- fix: support Utf8/LargeUtf8/Utf8View in native RLike without panicking [#5215](https://github.com/apache/datafusion-comet/pull/5215) (sam-1112) +- fix: normalize noncanonical NaN literals in comparisons [#5472](https://github.com/apache/datafusion-comet/pull/5472) (sunchao) +- fix: preserve map field metadata and honor target sorted flag in cast_map_to_map [#5227](https://github.com/apache/datafusion-comet/pull/5227) (Smallfu666) +- fix: reject duplicate Parquet field names before decoding [#5786](https://github.com/apache/datafusion-comet/pull/5786) (ErikBPF) +- fix: preserve Parquet conversion errors during join filtering [#6067](https://github.com/apache/datafusion-comet/pull/6067) (pingzh) +- fix: plan Iceberg writes with Spark's own operator when Comet is disabled [#6151](https://github.com/apache/datafusion-comet/pull/6151) (andygrove) +- fix: read shuffle write buffer, spill limit and off-heap sizes in bytes [#6191](https://github.com/apache/datafusion-comet/pull/6191) (andygrove) +- fix: bound shuffle schema cache retention and preserve eviction order [#6098](https://github.com/apache/datafusion-comet/pull/6098) (sunchao) +- fix: do not run Comet in on-heap mode without `spark.comet.exec.onHeap.enabled` [#6195](https://github.com/apache/datafusion-comet/pull/6195) (andygrove) +- fix: remove misleading native opt-in for dispatch-only datetime expressions [#6182](https://github.com/apache/datafusion-comet/pull/6182) (LinSimon-901101) +- fix: support CalendarIntervalType hashing [#5135](https://github.com/apache/datafusion-comet/pull/5135) (peterxcli) +- fix: report `Input` column when native Iceberg scan is enabled or native shuffle is enabled [#5265](https://github.com/apache/datafusion-comet/pull/5265) (hsiang-c) +- fix: skip the executor memory overhead warning when a factor is set or in local mode [#6198](https://github.com/apache/datafusion-comet/pull/6198) (andygrove) +- fix: keep operators above a cached relation native after the AQE re-plan [#6208](https://github.com/apache/datafusion-comet/pull/6208) (andygrove) +- fix: preserve current AQE logical-stage links on Comet operators [#5483](https://github.com/apache/datafusion-comet/pull/5483) (sunchao) +- fix: check each fair_unified reservation against its own share [#6205](https://github.com/apache/datafusion-comet/pull/6205) (andygrove) +- fix(iceberg): don't push transform residuals as their source column, fail on residual errors [#6154](https://github.com/apache/datafusion-comet/pull/6154) (andygrove) +- fix: match Spark's duplicate field and field id semantics in parquet field lookup [#5654](https://github.com/apache/datafusion-comet/pull/5654) (dwsmith1983) +- fix: park the native scan loop instead of busy-polling while waiting on native I/O [#6219](https://github.com/apache/datafusion-comet/pull/6219) (andygrove, mixermt) +- fix: refresh S3 policy locations when a location's credential fails [#6223](https://github.com/apache/datafusion-comet/pull/6223) (andygrove, snmvaughan) +- fix: inject Comet session extension rules only once per session [#6230](https://github.com/apache/datafusion-comet/pull/6230) (andygrove) +- fix: [branch-1.1] reject a file without field ids at any depth whether or not id matching is on (#6116) [#6266](https://github.com/apache/datafusion-comet/pull/6266) (andygrove, dwsmith1983) +- fix: [branch-1.1] decline native Iceberg writes with a custom location provider (#6216) [#6305](https://github.com/apache/datafusion-comet/pull/6305) (andygrove, liupoyi-1031) +- fix: [branch-1.1] Native S3 scan on EKS/IRSA turns a transient STS throttle into a hard 403 storm (#6025) [#6323](https://github.com/apache/datafusion-comet/pull/6323) (andygrove, parthchandra) +- fix: [branch-1.1] log partial memory grants at DEBUG and drop the memory usage dump (#6269) [#6346](https://github.com/apache/datafusion-comet/pull/6346) (andygrove) +- fix: [branch-1.1] zero sliced boolean offsets at every level before exporting to the JVM (#6339) [#6449](https://github.com/apache/datafusion-comet/pull/6449) (andygrove) +- fix: [branch-1.1] build with Rust 1.99, which deprecates the legacy f64 constants and fetch_update (#6507) [#6526](https://github.com/apache/datafusion-comet/pull/6526) (andygrove) +- fix: [branch-1.1] make adaptive aggregation skipping opt-in (#6474) [#6488](https://github.com/apache/datafusion-comet/pull/6488) (andygrove) +- fix: [branch-1.1] fall back to Spark for \_metadata.file_block_start and file_block_length (#6510) [#6511](https://github.com/apache/datafusion-comet/pull/6511) (andygrove) +- fix: [branch-1.1] fall back for rank limits over nested float keys (#6468) [#6487](https://github.com/apache/datafusion-comet/pull/6487) (andygrove) +- fix: [branch-1.1] fall back for incompatible regression aggregates (#6451) [#6489](https://github.com/apache/datafusion-comet/pull/6489) (andygrove) +- fix: [branch-1.1] coerce native IF branches to a common type (#6458) [#6491](https://github.com/apache/datafusion-comet/pull/6491) (andygrove) +- fix: [branch-1.1] match pre-epoch Iceberg temporal rounding (#6456) [#6486](https://github.com/apache/datafusion-comet/pull/6486) (andygrove) +- fix: [branch-1.1] defer Parquet conversion errors until a row group is decoded, as Spark does (#6515) [#6540](https://github.com/apache/datafusion-comet/pull/6540) (andygrove) +- fix: [branch-1.1] dispatch StaticInvoke and Invoke only into Spark's own classes (#6542) [#6554](https://github.com/apache/datafusion-comet/pull/6554) (andygrove) +- fix: [branch-1.1] gate array distinct and union signed-zero semantics by Spark version (#5750) [#6561](https://github.com/apache/datafusion-comet/pull/6561) (andygrove, peterxcli) +- fix: [branch-1.1] fall back from the native Iceberg scan when a nested field was added or renamed (#6543) [#6560](https://github.com/apache/datafusion-comet/pull/6560) (andygrove) + +**Performance related:** + +- perf: optimize spark_base64 in spark-expr [#4885](https://github.com/apache/datafusion-comet/pull/4885) (andygrove) +- perf: optimize `spark_floor` (up to 4x faster) [#4911](https://github.com/apache/datafusion-comet/pull/4911) (andygrove) +- perf: cache Iceberg reflection lookups on the planning path [#5222](https://github.com/apache/datafusion-comet/pull/5222) (andygrove) +- perf: vectorize integer-to-decimal cast [#4939](https://github.com/apache/datafusion-comet/pull/4939) (andygrove) +- perf: compute spark_size list lengths with Arrow length kernel [#5233](https://github.com/apache/datafusion-comet/pull/5233) (0lai0) +- perf: make CometShuffleBenchmark completable and fast [#5388](https://github.com/apache/datafusion-comet/pull/5388) (andygrove) +- perf: vectorize Map in spark_size via offset buffer [#5395](https://github.com/apache/datafusion-comet/pull/5395) (0lai0) +- perf: improve ArrowWriter performance for fixed-length vectors [#5046](https://github.com/apache/datafusion-comet/pull/5046) (peterxcli) +- perf: reuse Arrow IPC compression context across shuffle blocks [#5038](https://github.com/apache/datafusion-comet/pull/5038) (peterxcli) +- perf: use Arrow cast for decimal rescale check [#5440](https://github.com/apache/datafusion-comet/pull/5440) (peterxcli) +- perf: bulk copy fixed-width columns in ArrowWriter [#5442](https://github.com/apache/datafusion-comet/pull/5442) (peterxcli) +- perf: serialize Python input directly from Comet Arrow vectors [#5368](https://github.com/apache/datafusion-comet/pull/5368) (sunchao) +- perf: reuse per-partition scratch in the shuffle write path [#5568](https://github.com/apache/datafusion-comet/pull/5568) (dwsmith1983) +- perf: validate shuffle IPC context reuse savings [#5727](https://github.com/apache/datafusion-comet/pull/5727) (peterxcli) +- perf: slice the child instead of gathering it when unnesting [#5667](https://github.com/apache/datafusion-comet/pull/5667) (andygrove) +- perf: evaluate posexplode array expressions once per batch [#5737](https://github.com/apache/datafusion-comet/pull/5737) (rich7420) +- perf: cache expected schemas for remote shuffle decoding [#5722](https://github.com/apache/datafusion-comet/pull/5722) (pingzh) +- perf: use Arrow cast for date to timestamp NTZ [#5735](https://github.com/apache/datafusion-comet/pull/5735) (peterxcli) +- perf: reduce allocations when collecting cache statistics [#5734](https://github.com/apache/datafusion-comet/pull/5734) (peterxcli) +- perf: avoid repeated decimal promotion in expression serialization [#5736](https://github.com/apache/datafusion-comet/pull/5736) (peterxcli) +- perf: give collect_list and collect_set a native GroupsAccumulator [#5803](https://github.com/apache/datafusion-comet/pull/5803) (andygrove) +- perf: vectorize the native map lookup behind element_at and GetMapValue [#5806](https://github.com/apache/datafusion-comet/pull/5806) (andygrove) +- perf: compile user regex patterns once per planned expression [#5612](https://github.com/apache/datafusion-comet/pull/5612) (dwsmith1983) +- perf: batch the nested-element list hash for flat struct elements [#5778](https://github.com/apache/datafusion-comet/pull/5778) (viirya) +- perf: optimize `map_sort` singleton normalization (18x faster) [#5887](https://github.com/apache/datafusion-comet/pull/5887) (viirya) +- perf: optimize `map_sort` for multi-entry string maps (up to 3x faster) [#5901](https://github.com/apache/datafusion-comet/pull/5901) (viirya) +- perf: skip calendar reconstruction in hour/minute/second and dayofweek/weekday [#5771](https://github.com/apache/datafusion-comet/pull/5771) (peterxcli) +- perf: reuse zstd compression contexts across shuffle blocks [#5565](https://github.com/apache/datafusion-comet/pull/5565) (dwsmith1983) +- perf: cache parsed plan data across a stage tasks [#5615](https://github.com/apache/datafusion-comet/pull/5615) (dwsmith1983) +- perf: spill every shuffle partition of a task into one file [#5916](https://github.com/apache/datafusion-comet/pull/5916) (peterxcli) +- perf: decode shuffle blocks against a cached schema instead of re-parsing per block [#5809](https://github.com/apache/datafusion-comet/pull/5809) (peterxcli) +- perf: project cached batches by buffer selection, prune on collated strings [#5543](https://github.com/apache/datafusion-comet/pull/5543) (andygrove) +- perf: make one thread-local access per tracked allocation [#6166](https://github.com/apache/datafusion-comet/pull/6166) (andygrove) +- perf: reuse quantile summary buffers during merge [#4932](https://github.com/apache/datafusion-comet/pull/4932) (peterxcli) +- perf: optimize list_extract without defaults using Arrow take [#5174](https://github.com/apache/datafusion-comet/pull/5174) (peterxcli) +- perf: track decimal overflow without rescanning results [#5044](https://github.com/apache/datafusion-comet/pull/5044) (peterxcli) + +**Implemented enhancements:** + +- feat: support `explode_outer` [#5192](https://github.com/apache/datafusion-comet/pull/5192) (comphead) +- feat: support timestampadd and timestampdiff via codegen dispatch [#5030](https://github.com/apache/datafusion-comet/pull/5030) (andygrove) +- feat: support \_metadata constant columns in native Parquet scan [#5237](https://github.com/apache/datafusion-comet/pull/5237) (mbutrovich) +- feat: add `make_interval` support (codegen dispatch + native) [#5039](https://github.com/apache/datafusion-comet/pull/5039) (peterxcli) +- feat: Optionally split the Iceberg V2 write operator into distinct writer and committer operations [#4658](https://github.com/apache/datafusion-comet/pull/4658) (jordepic) +- feat: detect Iceberg V2 writes and emit fall-back reasons [#5298](https://github.com/apache/datafusion-comet/pull/5298) (jordepic) +- feat: remove native cast from boolean to decimal [#5185](https://github.com/apache/datafusion-comet/pull/5185) (andygrove) +- feat: add micro benchmark runner and EC2 guide [#5374](https://github.com/apache/datafusion-comet/pull/5374) (andygrove) +- feat: build gate + inert wiring for contrib Delta scans [Delta contrib split, part 2] [#4952](https://github.com/apache/datafusion-comet/pull/4952) (schenksj) +- feat: Support native scans with unprojected Spark 4 VARIANT columns [#5377](https://github.com/apache/datafusion-comet/pull/5377) (sunchao) +- feat: expose native aggregate spill and memory metrics [#5423](https://github.com/apache/datafusion-comet/pull/5423) (sunchao) +- feat: support `WindowGroupLimitExec` [#4870](https://github.com/apache/datafusion-comet/pull/4870) (comphead) +- feat: add RSS partition writer and task-owned JNI callback (1/n) [#5473](https://github.com/apache/datafusion-comet/pull/5473) (pingzh) +- feat: native `spark_unbase64` kernel [#5451](https://github.com/apache/datafusion-comet/pull/5451) (0lai0) +- feat: add partition writer destinations to native shuffle plans (2/n) [#5476](https://github.com/apache/datafusion-comet/pull/5476) (pingzh) +- feat: add destination-aware native shuffle execution (3/n) [#5481](https://github.com/apache/datafusion-comet/pull/5481) (pingzh) +- feat: bind task-owned RSS callbacks to native shuffle plans (4/n) [#5491](https://github.com/apache/datafusion-comet/pull/5491) (pingzh) +- feat: add Celeborn shuffle manager and partition pusher (5/n) [#5501](https://github.com/apache/datafusion-comet/pull/5501) (pingzh) +- feat: complete Celeborn map-side shuffle push lifecycle (6/n) [#5513](https://github.com/apache/datafusion-comet/pull/5513) (pingzh) +- feat: add experimental native support for in-memory cache, disabled by default [#5051](https://github.com/apache/datafusion-comet/pull/5051) (andygrove) +- feat: add raw Celeborn native shuffle reader (7/n) [#5531](https://github.com/apache/datafusion-comet/pull/5531) (pingzh) +- feat: support Map for `CreateArray` literal [#5452](https://github.com/apache/datafusion-comet/pull/5452) (comphead) +- feat: route `round` on float/double through the codegen dispatcher [#5600](https://github.com/apache/datafusion-comet/pull/5600) (andygrove) +- feat: enable native-only Celeborn shuffle planning (8/n) [#5537](https://github.com/apache/datafusion-comet/pull/5537) (pingzh) +- feat: support unicode case sensitive field names for reading parquet [#5602](https://github.com/apache/datafusion-comet/pull/5602) (comphead) +- feat: implement native Iceberg V2 writer via iceberg-rust [#5361](https://github.com/apache/datafusion-comet/pull/5361) (jordepic) +- feat: Support Iceberg system functions (bucket, truncate, years/months/days/hours) natively [#5638](https://github.com/apache/datafusion-comet/pull/5638) (andygrove) +- feat: implement regr_slope, regr_intercept, regr_r2, regr_sxx, regr_syy, regr_sxy aggregates [#4775](https://github.com/apache/datafusion-comet/pull/4775) (andygrove) +- feat: carry VariantType identity through schema serialization [#5631](https://github.com/apache/datafusion-comet/pull/5631) (peterxcli) +- feat: route unrecognized StaticInvoke and Invoke through the codegen dispatcher [#5692](https://github.com/apache/datafusion-comet/pull/5692) (andygrove) +- feat: support S3 compliant filesystems [#5314](https://github.com/apache/datafusion-comet/pull/5314) (comphead) +- feat: add native spark_sequence kernel for integral element types [#5614](https://github.com/apache/datafusion-comet/pull/5614) (0lai0) +- feat: support nested types as native shuffle hash partitioning keys [#5567](https://github.com/apache/datafusion-comet/pull/5567) (viirya) +- feat: route next_day and levenshtein collated input through the codegen dispatcher [#5720](https://github.com/apache/datafusion-comet/pull/5720) (0lai0) +- chore: add_benches hash function aggregators [#5730](https://github.com/apache/datafusion-comet/pull/5730) (coderfender) +- test: cover Spark-to-Arrow batch conversion edge cases [#5713](https://github.com/apache/datafusion-comet/pull/5713) (peterxcli) +- feat: normalize marked Variant arrays at the native Parquet boundary [#5715](https://github.com/apache/datafusion-comet/pull/5715) (peterxcli) +- feat: support native concat_ws with string arrays [#5725](https://github.com/apache/datafusion-comet/pull/5725) (peterxcli) +- feat: native dynamic filter pushdown for hash join into Parquet scans [#5699](https://github.com/apache/datafusion-comet/pull/5699) (pingzh) +- feat: enable codegen dispatch for `lpad` and `rpad` [#5764](https://github.com/apache/datafusion-comet/pull/5764) (Satyr09) +- feat: address remaining issues for `CreateArray` [#5766](https://github.com/apache/datafusion-comet/pull/5766) (comphead) +- ci: stop non-gating labels from skipping Preflight [#5784](https://github.com/apache/datafusion-comet/pull/5784) (comphead) +- feat: route `abs` on interval types through the codegen dispatcher [#5622](https://github.com/apache/datafusion-comet/pull/5622) (kazantsev-maksim) +- ci: shard Iceberg Spark tests across four runners [#5459](https://github.com/apache/datafusion-comet/pull/5459) (sunchao) +- test: expand ANSI coverage for round, conv and elt [#5799](https://github.com/apache/datafusion-comet/pull/5799) (rich7420) +- test: strengthen ANSI exception assertions [#5800](https://github.com/apache/datafusion-comet/pull/5800) (rich7420) +- chore: Improve network retry configuration for maven and artifact upload [#5782](https://github.com/apache/datafusion-comet/pull/5782) (comphead) +- ci: gate the Delta contrib build on symbols, not on libcomet size [#5827](https://github.com/apache/datafusion-comet/pull/5827) (andygrove) +- test: add helpers to assert whether an expression ran natively or via codegen dispatch [#5610](https://github.com/apache/datafusion-comet/pull/5610) (andygrove) +- feat: support Spark encode expression via codegen dispatch [#5037](https://github.com/apache/datafusion-comet/pull/5037) (andygrove) +- feat: support native aggregate function `mode` [#4782](https://github.com/apache/datafusion-comet/pull/4782) (andygrove) +- test: name the whole dispatched subtree in the decimal promotion assertion [#5849](https://github.com/apache/datafusion-comet/pull/5849) (andygrove) +- feat: route translate and to_csv through codegen dispatch by default [#5032](https://github.com/apache/datafusion-comet/pull/5032) (andygrove) +- feat: expose native Parquet scan I/O and read-amplification metrics [#5453](https://github.com/apache/datafusion-comet/pull/5453) (sunchao) +- ci: move the job routing policy out of ci.yml expressions and into compute-changes.py [#5850](https://github.com/apache/datafusion-comet/pull/5850) (andygrove) +- feat: support max_by and min_by aggregate expressions [#4817](https://github.com/apache/datafusion-comet/pull/4817) (andygrove) +- ci: add a Required Checks aggregator job so main can have a required status check [#5842](https://github.com/apache/datafusion-comet/pull/5842) (andygrove) +- feat: adapt Parquet storage for Variant projection [#5794](https://github.com/apache/datafusion-comet/pull/5794) (peterxcli) +- ci: move the Spark 3.4/3.5/4.0, Iceberg, macOS and benchmark suites behind a merge queue [#5843](https://github.com/apache/datafusion-comet/pull/5843) (andygrove) +- feat(iceberg): Iceberg table format V3: apply deletion vector on reads [#5853](https://github.com/apache/datafusion-comet/pull/5853) (mbutrovich) +- test: restore ANSI array access error coverage [#5798](https://github.com/apache/datafusion-comet/pull/5798) (rich7420) +- test: cover native memory accounting boundaries [#5856](https://github.com/apache/datafusion-comet/pull/5856) (rich7420) +- test: cover date maps in Parquet temporal fuzz tests [#5877](https://github.com/apache/datafusion-comet/pull/5877) (rich7420) +- feat: Enable adaptive partial aggregation for eligible native shuffle plans [#5723](https://github.com/apache/datafusion-comet/pull/5723) (peterxcli) +- ci: move the Spark 4.1 sql_hive shards behind the merge queue [#5871](https://github.com/apache/datafusion-comet/pull/5871) (andygrove) +- ci: retry the Maven wrapper bootstrap in every job that calls ./mvnw directly [#5852](https://github.com/apache/datafusion-comet/pull/5852) (andygrove) +- test: cover slice over expression-produced non-null element arrays (#… [#5839](https://github.com/apache/datafusion-comet/pull/5839) (sam-1112) +- test: run the libhdfs suite manually instead of in CI [#5892](https://github.com/apache/datafusion-comet/pull/5892) (andygrove) +- feat: Add Lance contrib build gate [#5728](https://github.com/apache/datafusion-comet/pull/5728) (wirybeaver) +- chore: drop support for JDK 11 [#5897](https://github.com/apache/datafusion-comet/pull/5897) (manuzhang) +- test: cover regex expression routing configurations [#5917](https://github.com/apache/datafusion-comet/pull/5917) (rich7420) +- test: cover round expression routing configurations [#5878](https://github.com/apache/datafusion-comet/pull/5878) (rich7420) +- test: cover to_csv routing configurations [#5921](https://github.com/apache/datafusion-comet/pull/5921) (rich7420) +- test: cover interval expression routing configurations [#5920](https://github.com/apache/datafusion-comet/pull/5920) (rich7420) +- test: cover array and map expression routing configurations [#5918](https://github.com/apache/datafusion-comet/pull/5918) (rich7420) +- test: make the cancelled Iceberg abort test yield deterministically [#5919](https://github.com/apache/datafusion-comet/pull/5919) (andygrove) +- chore: deprecate Spark 3.4 rather than removing it in 1.1.0 [#5885](https://github.com/apache/datafusion-comet/pull/5885) (andygrove) +- test: restore the TPC-H suite’s 2 GiB off-heap budget [#5904](https://github.com/apache/datafusion-comet/pull/5904) (ErikBPF) +- test: cover string expression routing and native opt-in [#5915](https://github.com/apache/datafusion-comet/pull/5915) (rich7420) +- chore: run only the cache-writing jobs on push to main [#5930](https://github.com/apache/datafusion-comet/pull/5930) (andygrove) +- ci: publish nightly SNAPSHOT jars to repository.apache.org [#5902](https://github.com/apache/datafusion-comet/pull/5902) (andygrove) +- ci: fold the Delta build gate and PyArrow UDF suite into the merge queue tiers [#5926](https://github.com/apache/datafusion-comet/pull/5926) (andygrove) +- refactor: use Arrow decimal precision validation [#5160](https://github.com/apache/datafusion-comet/pull/5160) (Smallfu666) +- ci: shrink the pull request tier to the Linux build on the default Spark profile [#5939](https://github.com/apache/datafusion-comet/pull/5939) (andygrove) +- ci: run the non-default Spark and Iceberg suites nightly instead of in the merge queue [#5963](https://github.com/apache/datafusion-comet/pull/5963) (andygrove) +- test: cover lpad and rpad routing configurations [#5895](https://github.com/apache/datafusion-comet/pull/5895) (rich7420) +- test: preserve Spark SQL baselines for ordering-sensitive fixtures [#5914](https://github.com/apache/datafusion-comet/pull/5914) (rich7420) +- test: cover json_array_length routing configurations [#5940](https://github.com/apache/datafusion-comet/pull/5940) (peterxcli) +- test: cover SQL fixture metadata and statement parsing [#5941](https://github.com/apache/datafusion-comet/pull/5941) (peterxcli) +- test: cover crc32 on binary inputs [#5942](https://github.com/apache/datafusion-comet/pull/5942) (peterxcli) +- test: cover datetime and timezone routing configurations [#5950](https://github.com/apache/datafusion-comet/pull/5950) (rich7420) +- test: cover collation and predicate routing configurations [#5951](https://github.com/apache/datafusion-comet/pull/5951) (rich7420) +- test: cover array_intersect routing configurations [#5952](https://github.com/apache/datafusion-comet/pull/5952) (rich7420) +- test: guard native plan equality against omitted parameters [#5953](https://github.com/apache/datafusion-comet/pull/5953) (rich7420) +- test: expand replace compatibility regression coverage [#5409](https://github.com/apache/datafusion-comet/pull/5409) (sam-1112) +- feat: add native allocation accounting for memory observability [#5934](https://github.com/apache/datafusion-comet/pull/5934) (andygrove) +- feat: run rlike natively by default for Java-equivalent literal patterns [#5415](https://github.com/apache/datafusion-comet/pull/5415) (sam-1112) +- ci: write large actions/cache entries only on push to main [#5973](https://github.com/apache/datafusion-comet/pull/5973) (andygrove) +- refactor: extract shared runtime filter components [#5937](https://github.com/apache/datafusion-comet/pull/5937) (pingzh) +- test: skip the Spark RocksDB state-store suites when Comet is enabled [#5987](https://github.com/apache/datafusion-comet/pull/5987) (comphead) +- chore: bench nondeterministic and json kernels [#5989](https://github.com/apache/datafusion-comet/pull/5989) (coderfender) +- feat: add BinaryType support for SortMergeJoin [#5928](https://github.com/apache/datafusion-comet/pull/5928) (xhumanoid) +- refactor: replace remaining hand-rolled loops with Arrow kernels [#5367](https://github.com/apache/datafusion-comet/pull/5367) (0lai0) +- feat: support Spark 4 EmptyRelationExec as a native input [#5821](https://github.com/apache/datafusion-comet/pull/5821) (jeffw13) +- chore: Bench agg (welford) stats [#5988](https://github.com/apache/datafusion-comet/pull/5988) (coderfender) +- ci: rename the umbrella workflow from CI to Comet CI [#6003](https://github.com/apache/datafusion-comet/pull/6003) (andygrove) +- ci: bootstrap Maven from the setup actions so every mvnw job is covered [#5881](https://github.com/apache/datafusion-comet/pull/5881) (andygrove) +- ci: add dev/local-ci.sh to run the Spark SQL and Iceberg suites locally [#5974](https://github.com/apache/datafusion-comet/pull/5974) (comphead) +- feat: narrow strict floating-point admission for scalar sort keys [#5981](https://github.com/apache/datafusion-comet/pull/5981) (0lai0) +- feat: route compatible xxhash64 args through SparkXxhash64 [#5960](https://github.com/apache/datafusion-comet/pull/5960) (sam-1112) +- feat: support string arrays in Spark-to-Comet conversion [#5954](https://github.com/apache/datafusion-comet/pull/5954) (rich7420) +- feat: run length on binary input natively [#5874](https://github.com/apache/datafusion-comet/pull/5874) (dwsmith1983) +- chore: drop the redundant width_bucket shim registrations and guard serde uniqueness [#5873](https://github.com/apache/datafusion-comet/pull/5873) (dwsmith1983) +- feat: hook native Parquet writes into Spark's WriteFilesExec seam on Spark 4.0+ [#5763](https://github.com/apache/datafusion-comet/pull/5763) (andygrove) +- feat(iceberg): report native Iceberg scan planning metrics and scan time in the Spark UI [#6027](https://github.com/apache/datafusion-comet/pull/6027) (parthchandra) +- test: guard page skipping in the native Iceberg scan [#6040](https://github.com/apache/datafusion-comet/pull/6040) (dwsmith1983) +- test: stop claiming Spark agrees on signed-zero array literals [#6055](https://github.com/apache/datafusion-comet/pull/6055) (cestercian) +- feat: trace Arrow memory held on the JVM side [#6048](https://github.com/apache/datafusion-comet/pull/6048) (andygrove) +- feat: unix_timestamp codegen dispatch for strings, and fix pre-epoch fractional rounding [#5789](https://github.com/apache/datafusion-comet/pull/5789) (Satyr09) +- chore: remove JaCoCo from the build [#6084](https://github.com/apache/datafusion-comet/pull/6084) (andygrove) +- refactor: compose CometScanRule and CometExecRule into a single CometRule [#6082](https://github.com/apache/datafusion-comet/pull/6082) (andygrove) +- chore: add a native Iceberg write benchmark [#6038](https://github.com/apache/datafusion-comet/pull/6038) (0lai0) +- ci: retry the Spark test pre-compile step on resolution failures [#6079](https://github.com/apache/datafusion-comet/pull/6079) (andygrove) +- ci: focus Miri on unsafe row and hash tests [#6072](https://github.com/apache/datafusion-comet/pull/6072) (rich7420) +- test: run the Iceberg Spark tests with the native Iceberg writer enabled [#5677](https://github.com/apache/datafusion-comet/pull/5677) (andygrove) +- feat: support listagg / string_agg aggregate (Spark 4.0+) [#4816](https://github.com/apache/datafusion-comet/pull/4816) (andygrove) +- ci: run label-triggered CI as a separate workflow [#6161](https://github.com/apache/datafusion-comet/pull/6161) (andygrove) +- feat: always count native allocations and log executor native memory usage [#6162](https://github.com/apache/datafusion-comet/pull/6162) (andygrove) +- feat: use DataFusion unnest_outer instead of Comet ListEmptyToNullExpr [#6132](https://github.com/apache/datafusion-comet/pull/6132) (comphead) +- chore: recommend --force-with-lease when updating PR branches [#6171](https://github.com/apache/datafusion-comet/pull/6171) (mizulun) +- feat: positional round robin shuffle keyed on a row ordinal [#6095](https://github.com/apache/datafusion-comet/pull/6095) (andygrove) +- chore: deprecate spark.comet.exec.memoryPool.fraction [#6163](https://github.com/apache/datafusion-comet/pull/6163) (andygrove) +- feat: add spark.comet.explain.planOnly.enabled [#5394](https://github.com/apache/datafusion-comet/pull/5394) (andygrove) +- feat: support direct Variant projection in native Parquet scans [#5868](https://github.com/apache/datafusion-comet/pull/5868) (peterxcli) +- test: cover JSON and cast expression routing [#6068](https://github.com/apache/datafusion-comet/pull/6068) (rich7420) +- feat: remove memory accounting from Comet's on-heap mode [#6066](https://github.com/apache/datafusion-comet/pull/6066) (andygrove) +- test: accept CometHashAggregateExec in the Spark 4.0 CollationSuite hash agg check [#6220](https://github.com/apache/datafusion-comet/pull/6220) (andygrove) +- feat: per-location credentials for the S3 credential SPI [#6031](https://github.com/apache/datafusion-comet/pull/6031) (snmvaughan) +- refactor: centralize data type support predicates [#5025](https://github.com/apache/datafusion-comet/pull/5025) (peterxcli) +- chore: [branch-1.1] change version from 1.1.0-SNAPSHOT to 1.1.0 [#6236](https://github.com/apache/datafusion-comet/pull/6236) (andygrove) +- feat: [branch-1.1] built-in S3 credential provider adapters for the native Parquet scan (#6023) [#6318](https://github.com/apache/datafusion-comet/pull/6318) (andygrove, parthchandra) +- chore: [branch-1.1] Add 1.1.0 changelog [#6282](https://github.com/apache/datafusion-comet/pull/6282) (andygrove) +- test: [branch-1.1] fix flaky mixed field id directory test on Spark 3.4 and 3.5 (#6312) [#6527](https://github.com/apache/datafusion-comet/pull/6527) (andygrove) +- test: [branch-1.1] cover sliced boolean arrays in explode (#6473) [#6492](https://github.com/apache/datafusion-comet/pull/6492) (andygrove) +- test: [branch-1.1] backport the Iceberg write report and the mid-write retry test (#6155, #6111) [#6307](https://github.com/apache/datafusion-comet/pull/6307) (andygrove, sam-1112) +- feat: [branch-1.1] reuse S3 credentials until shortly before their reported expiry (#6509) [#6545](https://github.com/apache/datafusion-comet/pull/6545) (andygrove, snmvaughan) + +**Documentation updates:** + +- docs: add note about run-iceberg-tests label in CI [#5247](https://github.com/apache/datafusion-comet/pull/5247) (mbutrovich) +- docs: add 1.0.0 TPC-DS benchmark results, remove TPC-H [#5284](https://github.com/apache/datafusion-comet/pull/5284) (mbutrovich) +- docs: add 1.0.0 changelog to main [#5263](https://github.com/apache/datafusion-comet/pull/5263) (andygrove) +- docs: correct Spark 4.2 version and CI test status in installation guide [#5315](https://github.com/apache/datafusion-comet/pull/5315) (andygrove) +- docs: add suggest-native-expression skill for assessing native expression candidates [#5348](https://github.com/apache/datafusion-comet/pull/5348) (andygrove) +- docs: document C2R cost for wide/nested schemas in tuning guide [#5458](https://github.com/apache/datafusion-comet/pull/5458) (DebadityaHait) +- docs: update stale interval multiplication note in expressions.md [#5518](https://github.com/apache/datafusion-comet/pull/5518) (peterxcli) +- docs: add contributing guide link [#5521](https://github.com/apache/datafusion-comet/pull/5521) (Dharan-K) +- doc: fix benchmark examples for MacOS [#5522](https://github.com/apache/datafusion-comet/pull/5522) (xhumanoid) +- docs: add comet meeting link [#5621](https://github.com/apache/datafusion-comet/pull/5621) (coderfender) +- docs: add a contributor guide page for CI and the merge queue [#5863](https://github.com/apache/datafusion-comet/pull/5863) (andygrove) +- docs: add contributor guide page on memory management [#5933](https://github.com/apache/datafusion-comet/pull/5933) (andygrove) +- docs: how to check the scheduled CI runs are actually running [#6000](https://github.com/apache/datafusion-comet/pull/6000) (andygrove) +- docs: explain allocator hazards and diagram where memory is allocated [#6014](https://github.com/apache/datafusion-comet/pull/6014) (andygrove) +- docs: pre-render mermaid diagrams to SVG at build time [#6021](https://github.com/apache/datafusion-comet/pull/6021) (andygrove) +- docs: split the PR review skill by area and correct the shuffle contributor docs [#6018](https://github.com/apache/datafusion-comet/pull/6018) (andygrove) +- docs: correct the stale range-partitioning strict floating-point rule [#6049](https://github.com/apache/datafusion-comet/pull/6049) (andygrove) +- docs: recommend setting spark.executor.memoryOverhead alongside off-heap memory [#6051](https://github.com/apache/datafusion-comet/pull/6051) (andygrove) +- docs: exempt testing-category and internal configs from the versioning policy [#6089](https://github.com/apache/datafusion-comet/pull/6089) (andygrove) +- docs: make doc changes weekly comet sync [#6160](https://github.com/apache/datafusion-comet/pull/6160) (coderfender) +- docs: correct the plugin and shuffle sections of the plugin overview [#6197](https://github.com/apache/datafusion-comet/pull/6197) (andygrove) +- docs: fix stale and missing native Iceberg write details [#6150](https://github.com/apache/datafusion-comet/pull/6150) (andygrove) +- docs: add an Iceberg writes contributor guide and review skill [#6149](https://github.com/apache/datafusion-comet/pull/6149) (andygrove) +- docs: correct which operators run separate native plans in a task [#6194](https://github.com/apache/datafusion-comet/pull/6194) (andygrove) +- docs: update the user guide for the 1.1.0 release [#6168](https://github.com/apache/datafusion-comet/pull/6168) (andygrove) +- docs: [branch-1.1] backport the 1.1.0 upgrade notes and tuning guide updates (#6237, #6244, #6248) [#6265](https://github.com/apache/datafusion-comet/pull/6265) (andygrove) +- ci: [branch-1.1] run every tier on release-branch pull requests, and the full suite before an RC (#6218) [#6284](https://github.com/apache/datafusion-comet/pull/6284) (andygrove) +- docs: [branch-1.1] generate release docs for 1.1.0 [#6344](https://github.com/apache/datafusion-comet/pull/6344) (andygrove) +- docs: [branch-1.1] credit co-authors in the 1.1.0 changelog [#6359](https://github.com/apache/datafusion-comet/pull/6359) (andygrove) + +**Other:** + +- chore: start 1.1.0 development [#5242](https://github.com/apache/datafusion-comet/pull/5242) (andygrove) +- refactor: replace hand-coded rollup of expression fallback reasons onto operators [#5236](https://github.com/apache/datafusion-comet/pull/5236) (andygrove) +- test: rename CometCastSuite to CometNativeCastSuite [#5268](https://github.com/apache/datafusion-comet/pull/5268) (andygrove) +- refactor: use arity helper for Int to Decimal128 reinterpretation [#5193](https://github.com/apache/datafusion-comet/pull/5193) (0lai0) +- test: add guard for Iceberg version on Variant fallback test [#5278](https://github.com/apache/datafusion-comet/pull/5278) (mbutrovich) +- chore: update documentation links for 1.0.0 release [#5290](https://github.com/apache/datafusion-comet/pull/5290) (andygrove) +- chore(deps): bump reqsign-core from 3.2.0 to 3.2.1 in /native in the all-other-cargo-deps group [#5287](https://github.com/apache/datafusion-comet/pull/5287) (dependabot[bot]) +- chore(deps): bump object_store_opendal from 0.57.0 to 0.58.0 in /native [#5289](https://github.com/apache/datafusion-comet/pull/5289) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#5286](https://github.com/apache/datafusion-comet/pull/5286) (dependabot[bot]) +- bug: Revert "chore(deps): bump object_store_opendal from 0.57.0 to 0.58.0 in /native" [#5332](https://github.com/apache/datafusion-comet/pull/5332) (coderfender) +- test: cover the narrowing direction of cast_and_stamp_schema [#5285](https://github.com/apache/datafusion-comet/pull/5285) (andygrove) +- chore(deps): bump opendal from 0.57.0 to 0.58.1 in /native [#5324](https://github.com/apache/datafusion-comet/pull/5324) (manuzhang) +- chore: respect Cargo parallelism settings for native release builds [#5344](https://github.com/apache/datafusion-comet/pull/5344) (pingzh) +- test: make expression benchmark harness fair and reproducible [#5371](https://github.com/apache/datafusion-comet/pull/5371) (andygrove) +- chore(deps): bump the all-other-cargo-deps group in /native with 5 updates [#5360](https://github.com/apache/datafusion-comet/pull/5360) (dependabot[bot]) +- test: fix vacuous signed-zero coverage in SQL file tests [#5393](https://github.com/apache/datafusion-comet/pull/5393) (sam-1112) +- chore: fix clippy warnings for Rust 1.98 [#5400](https://github.com/apache/datafusion-comet/pull/5400) (ywskycn) +- chore(deps): bump the codeql-actions group with 2 updates [#5405](https://github.com/apache/datafusion-comet/pull/5405) (dependabot[bot]) +- test: strengthen signed-zero SQL assertions [#5404](https://github.com/apache/datafusion-comet/pull/5404) (sunchao) +- chore(deps): bump actions/checkout from 6 to 7 [#5406](https://github.com/apache/datafusion-comet/pull/5406) (dependabot[bot]) +- test: add expression fallback-invariance suite [#5329](https://github.com/apache/datafusion-comet/pull/5329) (4ktLuffy) +- refactor: delegate ANSI integer arithmetic to arrow checked kernels [#5280](https://github.com/apache/datafusion-comet/pull/5280) (kazantsev-maksim) +- test: docs and test-coverage hardening for native collect_list / collect_set [#5055](https://github.com/apache/datafusion-comet/pull/5055) (andygrove) +- test: deduplicate Throwable cause-chain traversal [#5441](https://github.com/apache/datafusion-comet/pull/5441) (peterxcli) +- refactor: use Arrow timezone type [#5129](https://github.com/apache/datafusion-comet/pull/5129) (Hashim1999164) +- chore: add md formatting to `make format` [#5460](https://github.com/apache/datafusion-comet/pull/5460) (comphead) +- test: fail unexpected vacuous fallback-invariance checks [#5417](https://github.com/apache/datafusion-comet/pull/5417) (sunchao) +- ci: cache Maven distributions and retry bootstrap downloads [#5422](https://github.com/apache/datafusion-comet/pull/5422) (sunchao) +- chore: add math benches [#5479](https://github.com/apache/datafusion-comet/pull/5479) (coderfender) +- chore: Bench string functions [#5492](https://github.com/apache/datafusion-comet/pull/5492) (coderfender) +- chore(deps): bump actions/setup-java from 5 to 6 [#5524](https://github.com/apache/datafusion-comet/pull/5524) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#5523](https://github.com/apache/datafusion-comet/pull/5523) (dependabot[bot]) +- chore: bench additional math function [#5520](https://github.com/apache/datafusion-comet/pull/5520) (coderfender) +- chore: remove redundant condition and stray println in shuffle code [#5562](https://github.com/apache/datafusion-comet/pull/5562) (viirya) +- test: cover struct data columns in native shuffle [#5564](https://github.com/apache/datafusion-comet/pull/5564) (viirya) +- chore: rename .claude directory to vendor-neutral .ai [#5594](https://github.com/apache/datafusion-comet/pull/5594) (andygrove) +- chore: bench additional scalar functions [#5598](https://github.com/apache/datafusion-comet/pull/5598) (coderfender) +- test: verify Celeborn reflection compatibility [#5604](https://github.com/apache/datafusion-comet/pull/5604) (pingzh) +- test: enable native path in lower/upper_enabled sql fixtures [#5619](https://github.com/apache/datafusion-comet/pull/5619) (cestercian) +- chore: add benches datetime [#5620](https://github.com/apache/datafusion-comet/pull/5620) (coderfender) +- chore: move dead and defensive serde guards out of convert [#5595](https://github.com/apache/datafusion-comet/pull/5595) (andygrove) +- test: exercise the Iceberg write split-operator plan in Iceberg's own suites [#5640](https://github.com/apache/datafusion-comet/pull/5640) (andygrove) +- chore(deps): bump the codeql-actions group with 2 updates [#5669](https://github.com/apache/datafusion-comet/pull/5669) (dependabot[bot]) +- test: add explode operator microbenchmark [#5381](https://github.com/apache/datafusion-comet/pull/5381) (andygrove) +- deps: bump DataFusion 55.0 and Arrow/Parquet 59.2 [#5262](https://github.com/apache/datafusion-comet/pull/5262) (mbutrovich) +- chore: add benches array functions [#5700](https://github.com/apache/datafusion-comet/pull/5700) (coderfender) +- test: restore Comet coverage for recursive HAVING and ORDER BY [#5755](https://github.com/apache/datafusion-comet/pull/5755) (rich7420) +- test: restore Parquet V2 writer and delta encoding coverage [#5760](https://github.com/apache/datafusion-comet/pull/5760) (rich7420) +- chore: Add benches for datetime funcs [#5767](https://github.com/apache/datafusion-comet/pull/5767) (coderfender) +- bench: add a benchmark for the Spark hash kernels [#5765](https://github.com/apache/datafusion-comet/pull/5765) (viirya) +- test: restore Spark 4.1 Variant shredding suites [#5745](https://github.com/apache/datafusion-comet/pull/5745) (rich7420) +- refactor: share one helper for pushing a struct's null mask into its children [#5769](https://github.com/apache/datafusion-comet/pull/5769) (viirya) +- ci: label pull requests by changed paths and title prefix [#5762](https://github.com/apache/datafusion-comet/pull/5762) (dwsmith1983) +- test: cover ambiguous exact nested Parquet field matches [#5751](https://github.com/apache/datafusion-comet/pull/5751) (peterxcli) +- bench: measure nested types as native shuffle hash partitioning keys [#5788](https://github.com/apache/datafusion-comet/pull/5788) (viirya) +- deps: bump to datafusion 55.1.0 [#5865](https://github.com/apache/datafusion-comet/pull/5865) (comphead) +- bench: isolate map normalization and nested key hashing [#5822](https://github.com/apache/datafusion-comet/pull/5822) (viirya) +- bench: add a shuffle read benchmark covering the per-block schema parse [#5805](https://github.com/apache/datafusion-comet/pull/5805) (peterxcli) +- chore(deps): bump actions/setup-java from 4 to 6 [#6012](https://github.com/apache/datafusion-comet/pull/6012) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#6011](https://github.com/apache/datafusion-comet/pull/6011) (dependabot[bot]) +- deps: bump the iceberg-rust pin to bb1e4a4 and document why it is pinned [#6094](https://github.com/apache/datafusion-comet/pull/6094) (andygrove) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. A PR with commits from more than one person counts for each of them. + +``` + 147 Andy Grove + 62 Peter Lee + 26 Chao Sun + 26 KUAN-HAO HUANG + 19 Ping Zhang + 15 Oleks V + 14 dustin + 13 Bhargava Vadlamani + 13 Liang-Chi Hsieh + 11 dependabot[bot] + 10 ChenChen Lai + 8 sam-1112 + 6 Matt Butrovich + 5 Han-Yin Chang + 5 Parth Chandra + 4 Erik Bogado + 4 Steve Vaughan + 4 Wei Yan + 3 Jordan Epstein + 3 Manu Zhang + 2 Alexey + 2 Cestercian + 2 Daipayan Mukherjee + 2 Feng Zhang + 2 Kazantsev Maksim + 2 YuLun Mao + 2 liupoyi-1031 + 1 Dharan-K + 1 Hashim Khan + 1 Henos D + 1 Jeff Wang + 1 LinSimon-901101 + 1 Michael Taranov + 1 Scott Schenkein + 1 Stefan Wang + 1 Unik Dahal + 1 Viacheslav Inozemtsev + 1 Xuanyi Li + 1 debaditya + 1 divyank sameer shah + 1 hsiang-c +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/docs/source/changelog/index.md b/docs/source/changelog/index.md index e79878f7e70..792dfa2df3b 100644 --- a/docs/source/changelog/index.md +++ b/docs/source/changelog/index.md @@ -24,6 +24,7 @@ Per-release change logs for Apache DataFusion Comet. ```{toctree} :maxdepth: 1 +1.1.0 <1.1.0> 1.0.0 <1.0.0> 0.17.1 <0.17.1> 0.17.0 <0.17.0> From f7eb8aa3abd6d74ae74d0fa0d50d98f4ce65002c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 6 Oct 2026 20:01:56 +0000 Subject: [PATCH 15/89] chore: update documentation links for 1.1.0 release (#6736) --- docs/build.sh | 1 + docs/generate-versions.py | 4 ++-- docs/source/user-guide/index.md | 4 ++-- docs/source/user-guide/older-versions.md | 3 ++- 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/build.sh b/docs/build.sh index f4d839d8565..46a2130797f 100755 --- a/docs/build.sh +++ b/docs/build.sh @@ -31,6 +31,7 @@ rm -rf comet-0.15 rm -rf comet-0.16 rm -rf comet-0.17 rm -rf comet-1.0 +rm -rf comet-1.1 python3 generate-versions.py # Generate dynamic content (configs, compatibility matrices) for latest docs diff --git a/docs/generate-versions.py b/docs/generate-versions.py index 3ba6aee12e7..74a422ed375 100644 --- a/docs/generate-versions.py +++ b/docs/generate-versions.py @@ -154,6 +154,6 @@ def generate_docs(snapshot_version: str, latest_released_version: str, previous_ if __name__ == "__main__": print("Generating versioned user guide docs...") snapshot_version = get_version_from_pom() - latest_released_version = "1.0.0" - previous_versions = ["0.13.0", "0.14.0", "0.15.0", "0.16.0", "0.17.0"] + latest_released_version = "1.1.0" + previous_versions = ["0.13.0", "0.14.0", "0.15.0", "0.16.0", "0.17.0", "1.0.0"] generate_docs(snapshot_version, latest_released_version, previous_versions) \ No newline at end of file diff --git a/docs/source/user-guide/index.md b/docs/source/user-guide/index.md index 85b129a4359..c6328493dd9 100644 --- a/docs/source/user-guide/index.md +++ b/docs/source/user-guide/index.md @@ -22,7 +22,7 @@ under the License. The Comet user guide covers installation, configuration, supported data sources, supported operators and expressions, and tuning advice for running Apache Spark with Comet acceleration. -The **current stable release** is 1.0.x. Most users should start there. The **development snapshot** +The **current stable release** is 1.1.x. Most users should start there. The **development snapshot** tracks the upcoming release and may include features and fixes that are not yet generally available. User guides for **older releases** are kept for reference under [Older Versions](older-versions). @@ -31,7 +31,7 @@ User guides for **older releases** are kept for reference under [Older Versions] :caption: User Guide :hidden: -1.0.x (current) <1.0/index> +1.1.x (current) <1.1/index> Development Snapshot Older Versions ``` diff --git a/docs/source/user-guide/older-versions.md b/docs/source/user-guide/older-versions.md index 00d936be54c..ca46ba134ec 100644 --- a/docs/source/user-guide/older-versions.md +++ b/docs/source/user-guide/older-versions.md @@ -20,11 +20,12 @@ under the License. # Older Versions These user guides are kept for reference. They are not maintained: use the -[current stable release](1.0/index) for production guidance. +[current stable release](1.1/index) for production guidance. ```{toctree} :maxdepth: 1 +1.0.x <1.0/index> 0.17.x <0.17/index> 0.16.x <0.16/index> 0.15.x <0.15/index> From 6cd5ea0ed84e10087fac3066baa30ebd65c870ce Mon Sep 17 00:00:00 2001 From: Oleks V Date: Tue, 6 Oct 2026 20:42:54 +0000 Subject: [PATCH 16/89] feat: add a spill compression codec config, defaulting to lz4 (#6698) * feat: allowlist DataFusion configs, starting with spill compression Add `spark.comet.exec.allowedDataFusionConfigs`, a list of `spark.comet.datafusion.*` keys that are passed to native execution even when the testing flag `spark.comet.exec.respectDataFusionConfigs` is false. The default list holds `spark.comet.datafusion.execution.spill_compression`, so compressing DataFusion's spill files no longer needs the testing flag. Spills stay uncompressed unless that key is set. Document spill compression on the memory tuning page and add a nested and wide data section to the tuning guide. Closes #6467. * feat: replace the DataFusion conf allowlist with a spill codec config Replace `spark.comet.exec.allowedDataFusionConfigs` with `spark.comet.exec.spill.compression.codec`, which sets the codec for the files that native sorts, aggregations, sort-merge joins, and nested loop joins spill through DataFusion. It accepts `lz4`, `zstd`, and `none`, and defaults to `lz4`, the codec Spark uses for its own spill files by default, so native spills are now compressed by default. The JVM sends the resolved value, and the native side maps it to DataFusion's `SpillCompression` before the `spark.comet.datafusion.*` testing pass-through, which still wins. Add a test that a spilling native sort writes fewer bytes with the default codec than with `none`, and list nested loop joins in the spill compression docs. --- docs/source/user-guide/latest/tuning.md | 15 +++++++ .../source/user-guide/latest/tuning/memory.md | 10 +++++ native/core/src/execution/jni_api.rs | 25 +++++++++-- native/core/src/execution/spark_config.rs | 2 + .../scala/org/apache/comet/CometConf.scala | 13 ++++++ .../org/apache/comet/CometExecIterator.scala | 1 + .../apache/comet/exec/CometExecSuite.scala | 45 +++++++++++++++++++ 7 files changed, 108 insertions(+), 3 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 158883bb6c1..28f6ce112a7 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -36,6 +36,21 @@ guide is split into the following pages: - [Reducing Row/Columnar Conversion Overhead](tuning/transitions.md): stages in which many operators fall back to Spark. +## Nested and Wide Data + +Rows with many columns, long strings, or nested columns such as arrays of structs make every batch +larger, so queries over such data spill more. For these workloads: + +- Choose the codec for spill files from native sorts, aggregations, and joins, which Comet + compresses with `lz4` by default. See + [Compressing Spill Files](tuning/memory.md#compressing-spill-files). +- Lower the batch size if wide batches cause spilling or out-of-memory errors. See + [Batch Size](tuning/memory.md#batch-size). +- Check which shuffles can run natively when partitioning keys are nested. See + [Native Shuffle](tuning/shuffle.md#native-shuffle). +- Watch for row/columnar conversions, which are expensive for such schemas. See + [Wide or Deeply Nested Schemas](tuning/transitions.md#wide-or-deeply-nested-schemas). + ## Configuring Tokio Runtime Comet uses a global tokio runtime per executor process. By default it starts one worker thread per executor core diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 1f6ea98a2be..2632369df24 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -282,3 +282,13 @@ use more than `N` times this value on shared local disks. If the limit is reache fail and the query errors out. Raise this on workloads with large sort/aggregate/shuffle spills, or lower it to protect executors on shared disks, remembering that the total across an executor is a multiple of this value. + +## Compressing Spill Files + +`spark.comet.exec.spill.compression.codec` sets the codec for the files that native sorts, +aggregations, sort-merge joins, and nested loop joins spill to disk. The default is `lz4`, which +Spark also uses for its own spill files by default. `zstd` writes smaller files, which reduces +spill I/O and how much of `spark.comet.maxTempDirectorySize` they use, but takes more CPU to +compress and decompress them. `none` writes them uncompressed. Comet's shuffle writers use +`spark.comet.shuffle.compression.codec` instead. See +[Shuffle Compression](shuffle.md#shuffle-compression). diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 361d3847da7..98baf952ef3 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -40,6 +40,7 @@ use std::collections::HashSet; use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; +use datafusion::config::SpillCompression; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -129,9 +130,9 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, - COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXPLAIN_NATIVE_ENABLED, - COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, + COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::CometObjectStoreRegistry; @@ -902,6 +903,24 @@ fn prepare_datafusion_session_context( session_config.set_str("datafusion.execution.parquet.reorder_filters", "true"); } + // Translate Comet's spill codec into DataFusion's spill compression, which DataFusion + // otherwise leaves uncompressed. Set before the pass-through below for the same reason, + // so an explicit `spark.comet.datafusion.execution.spill_compression` still wins. + let spill_compression = match spark_config + .get(COMET_EXEC_SPILL_COMPRESSION_CODEC) + .map(String::as_str) + { + None | Some("lz4") => SpillCompression::Lz4Frame, + Some("zstd") => SpillCompression::Zstd, + Some("none") => SpillCompression::Uncompressed, + Some(codec) => { + return Err(CometError::Config(format!( + "Unsupported spill compression codec: {codec}" + ))) + } + }; + session_config = session_config.with_spill_compression(spill_compression); + // Pass through DataFusion configs from Spark. // e.g: spark-shell --conf spark.comet.datafusion.sql_parser.parse_float_as_decimal=true // becomes datafusion.sql_parser.parse_float_as_decimal=true diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index d29bf6f6836..1de88ccbeec 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -26,6 +26,8 @@ pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str = "spark.comet.exec.aggregate.skipPartial.enabled"; +pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = + "spark.comet.exec.spill.compression.codec"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; /// Comet configs read through this trait must be resolved by the JVM first: diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6dee9e3c23e..3d08c72dba8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1265,6 +1265,19 @@ object CometConf extends ShimCometConf { .bytesConf(ByteUnit.BYTE) .createWithDefault(100L * 1024 * 1024 * 1024) // 100 GB + // Used on native side. Check spark_config.rs how the config is used + val COMET_EXEC_SPILL_COMPRESSION_CODEC: ConfigEntry[String] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.spill.compression.codec") + .category(CATEGORY_TUNING) + .doc( + "The codec used to compress the files that native sorts, aggregations, and joins " + + "spill to disk. lz4 and zstd are supported, and none disables compression. zstd " + + "writes smaller files than lz4 but uses more CPU. Comet's shuffle writers use " + + s"${COMET_SHUFFLE_COMPRESSION_CODEC.key} instead. $TUNING_GUIDE.") + .stringConf + .checkValues(Set("lz4", "zstd", "none")) + .createWithDefault("lz4") + val COMET_RESPECT_DATAFUSION_CONFIGS: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.respectDataFusionConfigs") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 545c01fafd4..0f03a6244eb 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -870,6 +870,7 @@ object CometExecIterator extends Logging { CometConf.COMET_DEBUG_ENABLED, CometConf.COMET_DEBUG_MEMORY_ENABLED, CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, + CometConf.COMET_EXEC_SPILL_COMPRESSION_CODEC, CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 5071c40cb61..465da3cdd98 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -112,6 +112,51 @@ class CometExecSuite extends CometTestBase { } } + test("native sort spill files are compressed with the default spill codec") { + val numRows = 20000 + val compressibleValue = "native-sort-spill-compression-" * 8 + withTempPath { path => + // A single input file keeps all rows in one task so its sort outgrows the tiny memory + // pool below and must spill. + spark + .createDataFrame((0 until numRows).map(i => (i, compressibleValue))) + .coalesce(1) + .write + .parquet(path.getAbsolutePath) + + withParquetTable(path.getAbsolutePath, "tbl") { + def sortSpilledBytes(codecConf: (String, String)*): Long = { + var spilledBytes = 0L + withSQLConf( + Seq( + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536") ++ + codecConf: _*) { + val sorted = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc) + assert(sorted.collect().length == numRows) + val plan = sorted.queryExecution.executedPlan + val sorts = collect(plan) { case sort: CometSortExec => sort } + assert(sorts.nonEmpty, s"Expected a native sort:\n$plan") + spilledBytes = sorts.map(_.metrics("spilled_bytes").value).sum + assert(spilledBytes > 0L, "Native sort did not spill") + } + spilledBytes + } + + // `spilled_bytes` counts what DataFusion wrote to disk, so this shows the default codec + // reaching DataFusion's spill writer, not just crossing JNI. + val uncompressed = + sortSpilledBytes(CometConf.COMET_EXEC_SPILL_COMPRESSION_CODEC.key -> "none") + val compressed = sortSpilledBytes() + assert( + compressed < uncompressed, + s"Spilled $compressed bytes with the default codec and $uncompressed with none") + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) From 8d0bb01a5a99924024e07c1722e6eddd9260e167 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Tue, 6 Oct 2026 20:47:30 +0000 Subject: [PATCH 17/89] fix: initialize dispatched kernels with the partition the native plan computes (#6700) * fix: initialize dispatched kernels with the partition the native plan computes The JVM codegen dispatcher seeded kernels from `TaskContext.partitionId()`, which differs from the partition being computed under a union, a coalesce and a cartesian product. It also kept one kernel per task, so the parent partitions of a coalesce shared random state. The native planner now passes its partition index and plan id through the UDF bridge. Nondeterministic kernels are cached per plan and initialized with that index. Deterministic kernels never read it and stay shared across the plans of a task. * fix: drop a closed plan's dispatched kernels and check every node for determinism A coalesce kept the nondeterministic kernel of every parent partition until the task ended. `CometExecIterator.close` now calls `CometUdfBridge.releasePlan`, which hands the closed plan to a new `CometUDF.releasePlan` hook, and the dispatcher drops that plan's kernels there. Kernels shared by every plan stay until the task ends. The dispatcher shared a kernel whenever the root expression reported deterministic, but `Invoke` leaves its target object out of `deterministic`. So `make_valid_utf8` over `spark_partition_id()` still shared one kernel across the plans of a coalesce. It now checks every node. --- native/core/src/execution/planner.rs | 2 + native/core/src/parquet/parquet_exec.rs | 2 + native/jni-bridge/src/comet_udf_bridge.rs | 2 +- native/spark-expr/src/jvm_udf/mod.rs | 13 +++ .../org/apache/comet/udf/CometUdfBridge.java | 33 +++++++- .../org/apache/comet/CometExecIterator.scala | 5 ++ .../scala/org/apache/comet/udf/CometUDF.scala | 19 +++++ .../udf/codegen/CometScalaUDFCodegen.scala | 84 +++++++++++++------ .../comet/CometCodegenSourceSuite.scala | 4 +- .../org/apache/comet/CometCodegenSuite.scala | 63 ++++++++++++-- .../comet/udf/CometUdfBridgeSuite.scala | 4 + 11 files changed, 197 insertions(+), 34 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 99a52fb562f..10f3eddb368 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1045,6 +1045,8 @@ impl PhysicalPlanner { udf.return_nullable, self.task_context.clone(), self.class_loader.clone(), + self.partition, + self.exec_context_id, ))) } ExprStruct::NativeScalarUdf(call) => { diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 8149c3656b7..6f9803d994d 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -417,6 +417,8 @@ mod tests { true, None, None, + 0, + 0, )); let literal: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(1)))); let nested: Arc = diff --git a/native/jni-bridge/src/comet_udf_bridge.rs b/native/jni-bridge/src/comet_udf_bridge.rs index bbd7465d007..6532fe04035 100644 --- a/native/jni-bridge/src/comet_udf_bridge.rs +++ b/native/jni-bridge/src/comet_udf_bridge.rs @@ -42,7 +42,7 @@ impl<'a> CometUdfBridge<'a> { JNIString::new(Self::JVM_CLASS), jni::jni_str!("evaluate"), jni::jni_sig!( - "(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V" + "(Ljava/lang/String;[J[JJJIIJLorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V" ), )?, method_evaluate_ret: ReturnType::Primitive(Primitive::Void), diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 232926744d0..0a37330097b 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -54,9 +54,14 @@ pub struct JvmScalarUdfExpr { /// duration of the call. `None` when no driving Spark task is available (unit tests, direct /// native driver runs); the bridge then installs nothing. class_loader: Option>>>, + /// Index of the partition this native plan computes. See `CometUDF.evaluate`. + partition: i32, + /// Id of this native plan. See `CometUDF.evaluate`. + exec_context_id: i64, } impl JvmScalarUdfExpr { + #[allow(clippy::too_many_arguments)] pub fn new( class_name: String, args: Vec>, @@ -64,6 +69,8 @@ impl JvmScalarUdfExpr { return_nullable: bool, task_context: Option>>>, class_loader: Option>>>, + partition: i32, + exec_context_id: i64, ) -> Self { debug_assert!( !class_name.is_empty(), @@ -76,6 +83,8 @@ impl JvmScalarUdfExpr { return_nullable, task_context, class_loader, + partition, + exec_context_id, } } } @@ -234,6 +243,8 @@ impl PhysicalExpr for JvmScalarUdfExpr { JValue::Long(out_arr_ptr).as_jni(), JValue::Long(out_sch_ptr).as_jni(), JValue::Int(batch.num_rows() as i32).as_jni(), + JValue::Int(self.partition).as_jni(), + JValue::Long(self.exec_context_id).as_jni(), JValue::Object(task_context_ref).as_jni(), JValue::Object(class_loader_ref).as_jni(), ], @@ -276,6 +287,8 @@ impl PhysicalExpr for JvmScalarUdfExpr { self.return_nullable, self.task_context.clone(), self.class_loader.clone(), + self.partition, + self.exec_context_id, ))) } } diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d2881bfe449..3ab34fe0496 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -32,6 +32,8 @@ import org.apache.spark.comet.CometTaskContextShim; import org.apache.spark.util.TaskCompletionListener; +import com.google.common.annotations.VisibleForTesting; + import org.apache.comet.util.ClassLoaders; /** @@ -86,6 +88,8 @@ public class CometUdfBridge { * @param numRows row count of the current batch. Mirrors DataFusion's {@code * ScalarFunctionArgs.number_rows}; the only batch-size signal a zero-input UDF (e.g. a * zero-arg non-deterministic ScalaUDF) ever sees. + * @param partitionIndex index of the partition the calling native plan computes + * @param planId id of the calling native plan * @param taskContext propagated Spark {@link TaskContext} from the driving Spark task thread, or * {@code null} outside a Spark task. Treated as ground truth for the call: installed as the * thread-local on entry, with the prior value (if any) saved and restored in {@code finally}. @@ -113,6 +117,8 @@ public static void evaluate( long outArrayPtr, long outSchemaPtr, int numRows, + int partitionIndex, + long planId, TaskContext taskContext, ClassLoader classLoader) { assert udfClassName != null && !udfClassName.isEmpty() : "udfClassName must be non-empty"; @@ -148,6 +154,8 @@ public static void evaluate( outArrayPtr, outSchemaPtr, numRows, + partitionIndex, + planId, taskContext); } finally { // Unconditional: a no-op when nothing was installed, and it also undoes any change the @@ -170,6 +178,8 @@ private static void evaluateInternal( long outArrayPtr, long outSchemaPtr, int numRows, + int partitionIndex, + long planId, TaskContext taskContext) { long taskAttemptId = (taskContext != null) ? taskContext.taskAttemptId() : NO_TASK_ID; @@ -228,7 +238,7 @@ private static void evaluateInternal( inputs[i] = Data.importVector(importAllocator, inArr, inSch, null); } - result = udf.evaluate(inputs, numRows); + result = udf.evaluate(inputs, numRows, partitionIndex, planId); // Recorded before the checks below, so the invariant holds however this exits. resultIsInput = isOneOf(result, inputs); if (!(result instanceof FieldVector)) { @@ -290,6 +300,27 @@ private static void evaluateInternal( } } + /** + * Called by {@code CometExecIterator} once native plan {@code planId} of task {@code + * taskAttemptId} has closed. Passes the plan to every {@link CometUDF} instance of the task, see + * {@code CometUDF.releasePlan}. + */ + public static void releasePlan(long taskAttemptId, long planId) { + ConcurrentHashMap perTask = INSTANCES.get(taskAttemptId); + if (perTask != null) { + for (CometUDF udf : perTask.values()) { + udf.releasePlan(planId); + } + } + } + + /** The instance of {@code udfClassName} that task {@code taskAttemptId} holds, or null. */ + @VisibleForTesting + public static CometUDF instanceFor(long taskAttemptId, String udfClassName) { + ConcurrentHashMap perTask = INSTANCES.get(taskAttemptId); + return perTask == null ? null : perTask.get(udfClassName); + } + /** Whether the UDF handed back one of the vectors it was given, rather than a new one. */ private static boolean isOneOf(ValueVector result, ValueVector[] inputs) { for (ValueVector input : inputs) { diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 0f03a6244eb..27f880f9132 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -45,6 +45,7 @@ import org.apache.comet.exceptions.CometQueryExecutionException import org.apache.comet.parquet.CometFileKeyUnwrapper import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.shuffle.ShufflePartitionPusher +import org.apache.comet.udf.CometUdfBridge import org.apache.comet.vector.NativeUtil /** @@ -356,6 +357,10 @@ class CometExecIterator( // global refs. attempt(nativeLib.releasePlan(plan)) + // Drops what the task's JVM UDFs keep for this plan, such as its dispatched kernels. `id` is + // the plan id the native plan passed them. + attempt(CometUdfBridge.releasePlan(taskAttemptId, id)) + // Run the diagnostics even when teardown failed: a failed teardown is exactly when the // non-zero memory usage warning below is most informative. attempt { diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala index 6b435c40643..db9178fc32c 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala @@ -45,4 +45,23 @@ import org.apache.arrow.vector.ValueVector */ trait CometUDF { def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector + + /** + * The overload the bridge calls. `partitionIndex` is the index of the partition the calling + * native plan computes, which differs from `TaskContext.partitionId()` under a union, a + * coalesce or a cartesian product. `planId` tells apart the plans one task runs, such as the + * parent partitions of a coalesce. The default ignores both. + */ + def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = + evaluate(inputs, numRows) + + /** + * Called once native plan `planId` has closed, so an implementation that keeps state per plan + * can drop it before the task ends. The default does nothing. + */ + def releasePlan(planId: Long): Unit = () } diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index 1cab882df93..c589fb6fa93 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -29,7 +29,7 @@ import scala.util.control.NonFatal import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.Field -import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.SparkEnv import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.comet.util.Utils @@ -41,8 +41,8 @@ import org.apache.comet.udf.CometUDF /** * Arrow-direct codegen dispatcher. For each `(bound expression, input Arrow schema)` pair, - * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the task's - * partition index, and caches the live instance. + * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the index + * of the partition the calling native plan computes, and caches the live instance. * * Arg 0 is a `VarBinaryVector` scalar carrying the closure-serialized bound `Expression` bytes; * args 1..N are the data columns the `BoundReference`s read in ordinal order. @@ -56,16 +56,20 @@ import org.apache.comet.udf.CometUDF * +----------------------------+ +----------------------------+ +----------------------------+ * | Key: generated Java | | Key: task + UDF class | | Key: bound expression + | * | source | | | | input column shapes | - * | Value: compiled Java class | | Value: dispatcher object | | Value: ready-to-run kernel | - * | Scope: JVM, all queries | | Scope: one Spark task | | with state primed | - * | share it | | | | Scope: one Spark task | - * | Owner: Spark | | Owner: Comet | | (lives inside 2) | + * | Value: compiled Java class | | Value: dispatcher object | | (+ native plan if | + * | Scope: JVM, all queries | | Scope: one Spark task | | nondeterministic) | + * | share it | | | | Value: ready-to-run kernel | + * | Owner: Spark | | Owner: Comet | | with state primed | + * | | | | | Scope: one Spark task | + * | | | | | (lives inside 2), | + * | | | | | or one native plan | + * | | | | | if nondeterministic | * | | | | | Owner: Comet | * +----------------------------+ +----------------------------+ +----------------------------+ * }}} * - * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-task kernel - * across batches. + * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-plan kernel + * across batches. `CometExecIterator.close` drops a plan's kernels through `releasePlan`. * * `evaluate` runs under `this.synchronized` because DataFusion operators like `HashJoinExec` * pipeline build/probe via `OnceAsync` (`tokio::spawn`), so multiple Tokio worker threads can @@ -81,13 +85,34 @@ class CometScalaUDFCodegen extends CometUDF with Logging { * Per-task cache keyed on serialized expression bytes plus per-column specs. The deserialized * `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, `Rand`'s * `XORShiftRandom`) that must not be shared across concurrent tasks running the same query; - * keeping the cache per-task gives each task its own copy. Guarded by `this.synchronized`. + * keeping the cache per-task gives each task its own copy. A nondeterministic kernel is seeded + * from the partition its plan computes, so its key also holds the plan: each plan a task runs, + * such as each parent partition of a coalesce, gets its own, dropped by `releasePlan` when the + * plan closes. Guarded by `this.synchronized`. */ private val kernelCache : mutable.Map[CometScalaUDFCodegen.CacheKey, CometScalaUDFCodegen.CacheEntry] = mutable.HashMap.empty - override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = { + // Kernels shared by every plan stay until the task ends. + override def releasePlan(planId: Long): Unit = this.synchronized { + kernelCache.keys.filter(_.planId == planId).toList.foreach(kernelCache.remove) + } + + /** Plan ids of the cached kernels, `NoPlan` for a shared one. */ + private[comet] def cachedPlanIds: List[Long] = this.synchronized { + kernelCache.keysIterator.map(_.planId).toList.sorted + } + + // Callers that bypass the bridge (unit tests, benchmarks) have no native plan. + override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = + evaluate(inputs, numRows, partitionIndex = 0, planId = CometScalaUDFCodegen.NoPlan) + + override def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = { require( inputs.length >= 1, "CometScalaUDFCodegen requires at least 1 input (serialized expression), " + @@ -115,12 +140,12 @@ class CometScalaUDFCodegen extends CometUDF with Logging { val n = numRows val specsSeq = specs.toIndexedSeq - val key = CometScalaUDFCodegen.CacheKey(ByteBuffer.wrap(bytes), specsSeq) + val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(bytes), specsSeq) // Cache lookup and `process` run under one lock to serialize concurrent Tokio callers that // would otherwise race on the kernel's per-batch instance fields. this.synchronized { - val entry = lookupOrCompile(key, bytes, specsSeq) + val entry = lookupOrCompile(key, bytes, specsSeq, partitionIndex) val out = CometBatchKernelCodegen.allocateOutput( entry.outputField, @@ -144,9 +169,14 @@ class CometScalaUDFCodegen extends CometUDF with Logging { private def lookupOrCompile( key: CometScalaUDFCodegen.CacheKey, bytes: Array[Byte], - specs: IndexedSeq[ArrowColumnSpec]): CometScalaUDFCodegen.CacheEntry = { + specs: IndexedSeq[ArrowColumnSpec], + partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = { assert(Thread.holdsLock(this), "lookupOrCompile must run under this.synchronized") - kernelCache.get(key) match { + // A deterministic kernel never reads the partition index, so one instance under the planless + // key serves every plan in the task. Only a kernel with a nondeterministic node is stored per + // plan. + val sharedKey = key.copy(planId = CometScalaUDFCodegen.NoPlan) + kernelCache.get(sharedKey).orElse(kernelCache.get(key)) match { case Some(entry) => CometScalaUDFCodegen.cacheHitCount.incrementAndGet() entry @@ -168,14 +198,17 @@ class CometScalaUDFCodegen extends CometUDF with Logging { } val compiled = CometBatchKernelCodegen.compile(boundExpr, specs) val kernel = compiled.newInstance() - kernel.init(CometScalaUDFCodegen.currentPartitionIndex()) + kernel.init(partitionIndex) val outputField = CometBatchKernelCodegen.toFfiArrowField( "codegen_result", boundExpr.dataType, boundExpr.nullable) val entry = CometScalaUDFCodegen.CacheEntry(compiled, kernel, boundExpr.dataType, outputField) - kernelCache.put(key, entry) + // Walks the tree, because `deterministic` is not transitive everywhere. `Invoke` skips its + // `targetObject`, which is where Spark 4's `make_valid_utf8` puts its input. + val perPlan = boundExpr.exists(!_.deterministic) + kernelCache.put(if (perPlan) key else sharedKey, entry) CometScalaUDFCodegen.compileCount.incrementAndGet() CometScalaUDFCodegen.recordCompiledSignature(specs, boundExpr.dataType) entry @@ -300,21 +333,22 @@ object CometScalaUDFCodegen { } /** - * Partition index for the kernel's `init`. Expressions whose `doGenCode` calls - * `addPartitionInitializationStatement` (`Rand`, `Randn`, `Uuid`) reseed mutable state from - * this. Falls back to 0 when the dispatcher is exercised outside a Spark task (unit tests). + * Plan id for kernels shared by every plan in a task, and for callers that bypass the bridge. */ - private def currentPartitionIndex(): Int = - Option(TaskContext.get()).map(_.partitionId()).getOrElse(0) + private[comet] val NoPlan = -1L /** - * Cache key: serialized expression bytes plus per-column compile-time invariants. `hashCode` - * walks `bytesKey` per lookup, so for large ScalaUDF closures it scales with closure size. + * Cache key: calling native plan (`NoPlan` for a deterministic expression), serialized + * expression bytes plus per-column compile-time invariants. `hashCode` walks `bytesKey` per + * lookup, so for large ScalaUDF closures it scales with closure size. * * TODO(perf-cache-key): if hot, options are a driver-precomputed hash piggybacked through the * proto, per-instance last-key memoization, or a two-tier cache keyed on the generated source. */ - final case class CacheKey(bytesKey: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec]) + final case class CacheKey( + planId: Long, + bytesKey: ByteBuffer, + specs: IndexedSeq[ArrowColumnSpec]) /** Snapshot of dispatcher cache counters and current size. */ final case class DispatcherStats(compileCount: Long, cacheHitCount: Long, cacheSize: Int) { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 400ec6fcd2f..f1272e151ee 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -1535,8 +1535,8 @@ class CometCodegenSourceSuite extends AnyFunSuite { IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = true)) val nonNullable = IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = false)) - val k1 = CometScalaUDFCodegen.CacheKey(bytes, nullable) - val k2 = CometScalaUDFCodegen.CacheKey(bytes, nonNullable) + val k1 = CometScalaUDFCodegen.CacheKey(0L, bytes, nullable) + val k2 = CometScalaUDFCodegen.CacheKey(0L, bytes, nonNullable) assert( k1 != k2, "expected nullable=true and nullable=false specs to produce distinct cache keys") diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 8f6d7c4a8d9..5a80865da92 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -45,6 +45,7 @@ import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.serde.{CometInvokeTargets, CometScalaUDF, QueryPlanSerde} import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase +import org.apache.comet.udf.CometUdfBridge import org.apache.comet.udf.codegen.CometScalaUDFCodegen import org.apache.comet.vector.CometVector @@ -1621,11 +1622,10 @@ class CometCodegenSuite // `init(int partitionIndex)`, called once per kernel allocation. Spark seeds // `XORShiftRandom(seed + partitionIndex)` per partition, so different partitions produce // different sequences for the same seed. Matching Spark across partitions requires the - // kernel to see the real partition index, which the dispatcher derives from - // `TaskContext.get().partitionId()`, live on this path thanks to the bridge-level - // TaskContext propagation. Composing with a ScalaUDF (identity on Double here) forces the - // tree through codegen dispatch so the Rand evaluation runs inside our kernel's init - // rather than via Spark's normal codegen. + // kernel to see the real partition index, which the native plan passes through the bridge. + // Composing with a ScalaUDF (identity on Double here) forces the tree through codegen + // dispatch so the Rand evaluation runs inside our kernel's init rather than via Spark's + // normal codegen. spark.udf.register("dblId", (d: Double) => d) val df = spark .range(0, 1024, 1, numPartitions = 4) @@ -1645,6 +1645,59 @@ class CometCodegenSuite } } + test("dispatched kernels see the index of the partition the native plan computes") { + // Under a union, a coalesce and a cartesian product the native plan computes a partition + // whose index differs from `TaskContext.partitionId()`, and the coalesce runs two plans in + // one task. See https://github.com/apache/datafusion-comet/issues/6570. + val df = spark + .range(0, 8, 1, numPartitions = 2) + .selectExpr("id", "map(1, spark_partition_id()) AS m", "round(rand(42), 6) AS r") + Seq(df.union(df), df.coalesce(1)).foreach { q => + assertCodegenRan(checkSparkAnswerAndOperator(q)) + } + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + // `CartesianProductExec` stays in Spark, over native children. + assertCodegenRan(checkSparkAnswer(df.crossJoin(spark.range(0, 4, 1, 2)))) + } + if (isSpark40Plus) { + // Spark 4 lowers `make_valid_utf8` to an `Invoke` whose `deterministic` skips its target + // object, so the partition-seeded input sits under a root that reports deterministic. + val invoke = spark + .range(0, 8, 1, numPartitions = 2) + .selectExpr( + "id", + "make_valid_utf8(cast(spark_partition_id() AS STRING)) AS p", + "make_valid_utf8(cast(round(rand(42), 6) AS STRING)) AS r") + assertCodegenRan(checkSparkAnswerAndOperator(invoke.coalesce(1))) + } + } + + test("plans of a coalesce share deterministic kernels and drop their own when they close") { + // The coalesce runs four native plans in one task. The deterministic `round` compiles once + // for all of them and the nondeterministic one once per plan. Each plan's kernel is dropped + // when the plan closes, so once the rows are read only the shared one is left. + val df = spark + .range(0, 16, 1, numPartitions = 4) + .selectExpr("round(id / 3, 2) AS d", "round(rand(42), 6) AS r") + .coalesce(1) + CometScalaUDFCodegen.resetStats() + val cached = df.queryExecution.toRdd + .mapPartitions { rows => + rows.foreach(_ => ()) + val dispatcher = CometUdfBridge.instanceFor( + TaskContext.get().taskAttemptId(), + classOf[CometScalaUDFCodegen].getName) + Iterator(Option(dispatcher).map(_.asInstanceOf[CometScalaUDFCodegen].cachedPlanIds)) + } + .collect() + .toSeq + assert( + cached == Seq(Some(List(CometScalaUDFCodegen.NoPlan))), + s"expected only the shared kernel to outlive the plans, got plan ids $cached") + val stats = CometScalaUDFCodegen.stats() + assert(stats.compileCount == 5, s"expected 1 shared and 4 per-plan compiles, got $stats") + } + test("ScalaUDF composed with reused scalar subquery across projection and filter") { // The same scalar subquery appears in two sites: the projection (which the dispatcher // compiles into a fused kernel) and the filter (a separate operator). Each site holds its diff --git a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala index 7efc5f12cbc..71ec1f856e5 100644 --- a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala @@ -102,6 +102,8 @@ class CometUdfBridgeSuite extends AnyFunSuite { outArrays(0).memoryAddress(), outSchemas(0).memoryAddress(), numRows, + 0, + 0L, null, null) @@ -172,6 +174,8 @@ class CometUdfBridgeSuite extends AnyFunSuite { outArrays(0).memoryAddress(), outSchemas(0).memoryAddress(), numRows, + 0, + 0L, null, null) From 46d91ef02d80f2a42cd89034699550f45fa13394 Mon Sep 17 00:00:00 2001 From: Yu Hsun Hung <75130467+sam-1112@users.noreply.github.com> Date: Tue, 6 Oct 2026 22:32:27 +0000 Subject: [PATCH 18/89] fix: restore Spark write execs when reverting transition-heavy stages (#5957) * fix: restore Spark write execs when reverting transition-heavy stages (#5719) Give Iceberg and parquet native writes a real originalPlan so revertToSpark keeps the commit-message node instead of duplicating or dropping the child. * fix: drop unused string interpolator in transition-revert warning * fix: remove unused import * fix: remove unused import * test: wait for write-plan listener callback before unregister * fix: drop duplicate Parquet write helpers that weaken base-class access * fix: drop duplicate Parquet write helpers that weaken base-class access * fix: keep shuffle-stage transitions when reverting native writes (#6152) Document sparkFallback and remove the CometExecRule unwraps that only matched the old originalPlan = child link. * fix: keep a native stage when transition stripping exposes an exchange root * fix: restore row output after stage reversion * fix: preserve Arrow input for reverted native shuffle stages * refactor: move operator-specific restoration into sparkFallback --------- Co-authored-by: Hung --- .../adding_a_new_operator.md | 28 + .../contributor-guide/iceberg-writes.md | 13 +- .../contributor-guide/native_shuffle.md | 6 + .../apache/comet/rules/CometExecRule.scala | 20 +- ...RevertNativeForTransitionHeavyStages.scala | 133 +++-- .../operator/CometDataWritingCommand.scala | 2 +- .../operator/CometIcebergNativeWrite.scala | 2 + .../comet/CometIcebergNativeScanExec.scala | 10 +- .../sql/comet/CometIcebergWriteExec.scala | 16 +- .../spark/sql/comet/CometLocalTopKExec.scala | 9 + .../spark/sql/comet/CometNativeScanExec.scala | 9 + .../sql/comet/CometNativeWriteExec.scala | 20 +- .../apache/spark/sql/comet/operators.scala | 27 + .../comet/CometIcebergWriteActionSuite.scala | 109 ++++ .../CometIcebergWriteDetectionSuite.scala | 10 +- .../parquet/CometParquetWriterSuite.scala | 219 +++++++- ...tNativeForTransitionHeavyStagesSuite.scala | 514 +++++++++++++++++- 17 files changed, 1057 insertions(+), 90 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 15c4e9169b6..09cac3a1c72 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -731,6 +731,34 @@ Use `QueryPlanSerde.exprToProto` to convert Spark expressions to protobuf: val protoExpr = exprToProto(sparkExpr, inputSchema) ``` +### Restoring the Spark operator (`sparkFallback`) + +`CometExec.originalPlan` is the Spark operator this node replaced. `CometExecRule` copies +`originalPlan.logicalLink` onto the Comet node, which is how AQE finds the node again when it +re-plans a stage. `RevertNativeForTransitionHeavyStages` calls `sparkFallback(newChildren)` to +rebuild that Spark operator with the children of the reverted stage. + +The default implementation is `originalPlan.withNewChildren(newChildren)`. It refuses a null +`originalPlan`, an `originalPlan` that is one of the node's own children, or a different number of +children than the Spark operator has. + +Override `sparkFallback` when conversion changes the plan shape, so the restored node is not that +Spark operator with the same children. `CometNativeWriteExec` replaces a `DataWritingCommandExec` +and drops the `WriteFilesExec` under it; its override puts that wrapper back around the restored +input. `CometIcebergWriteExec` keeps the same shape as `IcebergWriteExec`, so the default is +enough. + +Also override `sparkFallback` when the operator's live state differs from `originalPlan`. +`CometNativeScanExec` restores its current partition and data filters, and +`CometIcebergNativeScanExec` restores its current runtime filters, so AQE's executable DPP +subqueries survive reversion. `CometLocalTopKExec` returns the restored child directly: Comet +inserted that local candidate selection, and only the outer TopK restores Spark's offset and +projection. Rebuilding the original TopK at both nodes would apply it twice. + +Do not point `originalPlan` at a child. If that child is a shuffle or query stage, the copied +logical link puts this node inside the stage's `LogicalQueryStage`. AQE then re-plans a second +copy of the operator around the one that is already there. + ### Handling Fallback Use `withInfo` to tag operators with fallback reasons: diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index a90cff08075..e2547762822 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -107,9 +107,9 @@ serializer requires that tag on Spark 4.1+, leaving stock writers on Spark throu ## From `IcebergWrite` to `CometIcebergWrite` `CometExecRule` converts an `IcebergWriteExec` with the `CometIcebergNativeWrite` operator serde -when `spark.comet.write.iceberg.enabled` is on. Two arms in `CometExecRule` handle it: one unwraps -the double conversion AQE can produce when it re-fires write planning over a sub-tree that already -contains a `CometIcebergWriteExec`, and the other calls `convertToComet`. +when `spark.comet.write.iceberg.enabled` is on. A single arm in `CometExecRule` handles the +conversion by calling `convertToComet`. The converted node keeps the `IcebergWriteExec` as its +`originalPlan`, so AQE re-plans the write from that node. `CometIcebergNativeWrite.requiresNativeChildren` is `true`. The native writer consumes Arrow batches from its child over FFI, so the conversion is declined unless the child is already a Comet @@ -479,8 +479,11 @@ Each of these has caused a bug on this path: ([#5691](https://github.com/apache/datafusion-comet/issues/5691), [#5693](https://github.com/apache/datafusion-comet/issues/5693), [#6141](https://github.com/apache/datafusion-comet/issues/6141)). -- **Plan rewrites must keep the write node.** Rules that restore Spark operators from a Comet - node's `originalPlan` have to handle the write execs, whose `originalPlan` today is their child +- **Plan rewrites must keep the write node.** `CometIcebergWriteExec.originalPlan` is the + `IcebergWriteExec` it replaced. Restoring Spark execution goes through + [`CometExec.sparkFallback`](adding_a_new_operator.md#restoring-the-spark-operator-sparkfallback), + which rebuilds that node around the reverted children. Pointing `originalPlan` at the child + drops the write, and AQE then treats the write as part of that child's stage ([#5719](https://github.com/apache/datafusion-comet/issues/5719)). - **The kill switch must still work.** Code that runs for Iceberg writes has to respect `spark.comet.enabled`, so that disabling Comet restores Spark's own plan diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 61088447882..e9cf7774415 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -39,6 +39,12 @@ Compare this to JVM shuffle's data path: Comet Native (columnar) → ColumnarToRowExec → rows → JVM Shuffle → Arrow IPC → columnar ``` +When `RevertNativeForTransitionHeavyStages` restores the map stage to Spark execution, the +native exchange stays in place. Its input still needs Arrow-backed Comet vectors, even if the +restored Spark operator supports columnar output. The rule adds `CometSparkToColumnarExec` to +convert either Spark rows or Spark columnar batches to Arrow before the native shuffle consumes +them. Spark's `RowToColumnarExec` alone does not satisfy this input contract. + ## When Native Shuffle is Used Native shuffle (`CometExchange`) is selected when all of the following conditions are met: diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index c16d1fc8f8f..e2712b01e34 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -542,23 +542,15 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // DataWritingCommandExec and re-implement the write framework inside CometNativeWriteExec. // This path is retained only for 3.4/3.5 and goes away with them. // - // AQE reoptimization looks for `DataWritingCommandExec` or `WriteFilesExec` - // if there is none it would reinsert write nodes, and since Comet remap those nodes - // to Comet counterparties the write nodes are twice to the plan. - // Checking if AQE inserted another write Command on top of existing write command - case _ @DataWritingCommandExec(_, w: WriteFilesExec) - if !isSpark40Plus && w.child.isInstanceOf[CometNativeWriteExec] => - w.child - + // `originalPlan` is that command. This rule copies `originalPlan.logicalLink` onto the + // Comet node, so AQE re-plans the write with the command rather than with whatever child + // happened to sit under it. No second DataWritingCommandExec is inserted on top. case op: DataWritingCommandExec if !isSpark40Plus => convertToComet(op, CometDataWritingCommand).getOrElse(op) - // AQE re-fires the Iceberg write planning on every stage materialisation, so a - // partitioned write's physical sub-tree may already contain a `CometIcebergWriteExec`. - // Unwrap to avoid a double conversion. - case op: IcebergWriteExec if op.child.isInstanceOf[CometIcebergWriteExec] => - op.child - + // `originalPlan` is this IcebergWriteExec, so AQE re-plans the write as this node. + // A shuffle directly under the native write stays in the child stage and is not wrapped + // again. case op: IcebergWriteExec if CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.get(op.conf) => convertToComet(op, CometIcebergNativeWrite).getOrElse(op) diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index 1181a573bc3..3f3c57f0645 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -23,7 +23,8 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometIcebergNativeScanExec, CometLocalTopKExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike} @@ -62,19 +63,18 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan plan match { case _: BroadcastExchangeLike => plan case exchange: ShuffleExchangeLike => - revertStageIfNeeded(exchange.child, exchange.supportsColumnar) + revertShuffleStageIfNeeded(exchange) .map(reverted => exchange.withNewChildren(Seq(reverted))) .getOrElse(plan) case _ => - // Result stage: its output is collected as rows, so no consumer requires columnar input - // and the reverted stage needs no trailing R2C. + // Result stage: its output is collected as rows. revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan) } } private def applyForNonAQE(plan: SparkPlan): SparkPlan = { val withRevertedStages = plan.transformUp { case exchange: ShuffleExchangeLike => - revertStageIfNeeded(exchange.child, exchange.supportsColumnar) + revertShuffleStageIfNeeded(exchange) .map(reverted => exchange.withNewChildren(Seq(reverted))) .getOrElse(exchange) } @@ -82,12 +82,22 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan .getOrElse(withRevertedStages) } + private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = { + val outputArrow = exchange match { + case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle + case _ => false + } + revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow) + } + /** - * Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar. + * Reverts the stage if C2R count exceeds threshold, restoring the stage's output format when + * the reverted root does not satisfy it. */ private def revertStageIfNeeded( stagePlan: SparkPlan, - outputColumnar: Boolean): Option[SparkPlan] = { + outputColumnar: Boolean, + outputArrow: Boolean = false): Option[SparkPlan] = { val transitionCount = countTransitions(stagePlan) if (transitionCount <= maxTransitions) return None @@ -101,11 +111,27 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan val reason = s"Stage reverted: $transitionCount C2R transitions exceed threshold $maxTransitions" - val reverted = revertToSpark(stagePlan) - val result = if (outputColumnar && !reverted.supportsColumnar) { - RowToColumnarExec(withFallbackReason(reverted, reason)) + val reverted = + try { + revertToSpark(stagePlan) + } catch { + case e: CometExec.InvalidSparkFallbackException => + logWarning( + "Skipping transition-heavy stage reversion because a Comet operator could not " + + s"restore its Spark plan: ${e.getMessage}") + return None + } + val revertedWithReason = withFallbackReason(reverted, reason) + val result = if (outputArrow) { + // Native shuffle consumes Arrow-backed Comet vectors, not arbitrary Spark columnar + // batches. This bridge converts both row-based and vectorized Spark fallback roots. + CometSparkToColumnarExec(revertedWithReason) + } else if (outputColumnar && !reverted.supportsColumnar) { + RowToColumnarExec(revertedWithReason) + } else if (!outputColumnar && reverted.supportsColumnar) { + ColumnarToRowExec(revertedWithReason) } else { - withFallbackReason(reverted, reason) + revertedWithReason } Some(result) } @@ -146,16 +172,31 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan } /** - * Like `transformDown`, never descends stage-boundary children. + * Like `transformDown`, never descends stage-boundary children. If the rule rewrites the + * current node, re-apply it to the result so stacked transitions such as + * `CometSparkToColumnarExec(CometNativeColumnarToRowExec(x))` are fully unwrapped before + * children are visited. Spark's `transformDown` does not do this; leaving the inner C2R in + * place later calls `CometNativeColumnarToRowExec.withNewChildren` with a reverted row-based + * child, which asserts `child.supportsColumnar`. + * + * A rewrite can itself be the stage boundary. Unwrapping a transition that sits directly on a + * shuffle yields that shuffle, and descending into it strips transitions in the next stage. + * `transformStageUp` and `insertTransitions` do not cross the exchange, so those transitions + * would not be restored (#6152). Return the boundary unchanged. */ private def transformStageDown(plan: SparkPlan)( rule: PartialFunction[SparkPlan, SparkPlan]): SparkPlan = { val transformed = rule.applyOrElse(plan, identity[SparkPlan]) - val newChildren = transformed.children.map { child => - if (isStageBoundary(child)) child else transformStageDown(child)(rule) + if (transformed ne plan) { + if (isStageBoundary(transformed)) transformed + else transformStageDown(transformed)(rule) + } else { + val newChildren = transformed.children.map { child => + if (isStageBoundary(child)) child else transformStageDown(child)(rule) + } + if (newChildren == transformed.children) transformed + else transformed.withNewChildren(newChildren) } - if (newChildren == transformed.children) transformed - else transformed.withNewChildren(newChildren) } /** Like `transformUp`, never descends stage-boundary children. */ @@ -184,7 +225,27 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan count } + /** + * Checks for Comet operators whose original Spark plan is also their input. This must run + * before any bottom-up rewrite replaces children. Otherwise a stable `originalPlan` reference + * can keep pointing at the old Comet child after `withNewChildren`, hiding the alias and + * causing fallback to reconstruct that Comet child instead of a Spark operator. + */ + private def validateOriginalPlanAliases(plan: SparkPlan): Unit = plan match { + case _ if isStageBoundary(plan) => () + case cometExec: CometExec => + val sparkPlan = cometExec.originalPlan + if (sparkPlan != null && cometExec.children.exists(_ eq sparkPlan)) { + throw new CometExec.InvalidSparkFallbackException( + s"${cometExec.getClass.getSimpleName} aliases its original Spark plan with a child") + } + cometExec.children.foreach(validateOriginalPlanAliases) + case _ => + plan.children.foreach(validateOriginalPlanAliases) + } + private[rules] def revertToSpark(plan: SparkPlan): SparkPlan = { + validateOriginalPlanAliases(plan) val stripped = transformStageDown(plan) { case CometNativeColumnarToRowExec(child) => child case CometColumnarToRowExec(child) => child @@ -192,40 +253,12 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan case sparkToColumnar: CometSparkToColumnarExec => sparkToColumnar.child case RowToColumnarExec(child) => child } - val reverted = transformStageUp(stripped) { - // Local candidate selection was inserted by Comet. Only the outer TopK owns - // the original Spark operator's offset and projection. - case local: CometLocalTopKExec => local.child - case cometExec: CometExec => - if (cometExec.originalPlan.children.size == cometExec.children.size) { - val originalWithCurrentExpressions = cometExec match { - case scan: CometNativeScanExec => - // AQE's query-stage optimizer rewrites DPP placeholders in the live Comet scan - // before this post-columnar rule runs. The frozen FileSourceScanExec in - // originalPlan still contains SubqueryAdaptiveBroadcastExec, which cannot execute. - // Preserve the rewritten filters when reverting the scan to Spark. - val originalScan = scan.originalPlan.copy( - partitionFilters = scan.partitionFilters, - dataFilters = scan.dataFilters) - scan.originalPlan.logicalLink.foreach(originalScan.setLogicalLink) - originalScan - case scan: CometIcebergNativeScanExec => - // Iceberg's native scan has the same split between live and frozen filters. - // serializedPartitionData rebuilds originalPlan from runtimeFilters before - // execution, but transition reversion skips that path and executes the restored - // BatchScanExec directly. Carry the executable DPP filters across here as well. - val originalScan = scan.originalPlan.copy(runtimeFilters = scan.runtimeFilters) - scan.originalPlan.logicalLink.foreach(originalScan.setLogicalLink) - originalScan - case _ => cometExec.originalPlan - } - originalWithCurrentExpressions.withNewChildren(cometExec.children) - } else { - logWarning( - "Comet plan and original have different child count for " + - s"${cometExec.getClass.getSimpleName}, using originalPlan as-is.") - cometExec.originalPlan - } + if (isStageBoundary(stripped)) { + throw new CometExec.InvalidSparkFallbackException( + "Cannot revert a stage whose stripped root is a stage boundary") + } + val reverted = transformStageUp(stripped) { case cometExec: CometExec => + cometExec.sparkFallback(cometExec.children) } insertTransitions(reverted) } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index 6af2cc18ba6..a56dea862c4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -211,7 +211,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}") } - CometNativeWriteExec(nativeOp, childPlan, outputPath, cmd.mode, committer, jobId) + CometNativeWriteExec(nativeOp, op, childPlan, outputPath, cmd.mode, committer, jobId) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index c6741f3e40c..bb59b924d61 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -1024,7 +1024,9 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { "Native Iceberg write conversion: SparkWrite.outputSpecId reflection failed")) CometIcebergWriteExec( nativeOp, + op, op.child, + op.output, op.batchWrite, table.asInstanceOf[AnyRef], outputSpecId) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index c944d3af8cd..e058a5c0916 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} -import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -65,6 +65,14 @@ case class CometIcebergNativeScanExec( @transient nativeIcebergScanMetadata: CometIcebergNativeScanMetadata) extends CometLeafExec { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + // Native execution rebuilds originalPlan from the live runtimeFilters during partition + // serialization. Reversion skips that path, so carry the executable DPP filters across here. + val restoredScan = originalPlan.copy(runtimeFilters = runtimeFilters) + originalPlan.logicalLink.foreach(restoredScan.setLogicalLink) + restoredScan + } + override val supportsColumnar: Boolean = true override val nodeName: String = "CometIcebergNativeScan" diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 877eb84c1f1..83149ced77c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -24,14 +24,13 @@ import java.util.concurrent.atomic.AtomicReference import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.expressions.UnsafeProjection import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types.BinaryType import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.TaskFailureListener @@ -54,8 +53,12 @@ import org.apache.comet.serde.OperatorOuterClass.Operator * @param nativeOp * Template operator carrying the `IcebergWrite` proto. Per-task `partition_id` / * `task_attempt_id` get stamped on a copy at execution time. + * @param originalPlan + * The JVM Iceberg write operator restored when Comet reverts a transition-heavy stage. * @param child * Comet native child (must be a [[CometNativeExec]] so columnar batches flow through FFI). + * @param output + * Commit-message output attributes copied from the JVM Iceberg write plan. * @param batchWrite * Shared with the outer [[IcebergCommitExec]] -- the same instance the strategy materialised * via `write.toBatch`. Used here only to provide the `dataLocation` / partition spec needed by @@ -66,20 +69,15 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ case class CometIcebergWriteExec( nativeOp: Operator, + @transient override val originalPlan: IcebergWriteExec, child: SparkPlan, + override val output: Seq[Attribute], @transient batchWrite: BatchWrite, @transient table: AnyRef, partitionSpecId: Int) extends CometNativeExec with UnaryExecNode { - override def originalPlan: SparkPlan = child - - // Same output schema as IcebergWriteExec so the outer IcebergCommitExec consumes the - // commit messages identically regardless of which inner exec emitted them. - override def output: Seq[Attribute] = Seq( - AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) - // Native exec emits a single Binary column; the surrounding command framework expects rows, so // the outer commit exec calls executeCollect on us. supportsColumnar = false keeps Spark from // inserting a ColumnarToRow that would clash with our (Nil-output-like) row contract. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala index ec3c85191a3..16ecb208819 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala @@ -91,6 +91,15 @@ case class CometLocalTopKExec( override val serializedPlanOpt: SerializedPlan) extends CometUnaryExec { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = newChildren match { + // Local candidate selection was inserted by Comet. Only the outer TopK owns + // the original Spark operator's offset and projection. + case Seq(restoredChild) => restoredChild + case _ => + throw new CometExec.InvalidSparkFallbackException( + s"CometLocalTopKExec expected one restored child but received ${newChildren.size}") + } + override def outputPartitioning: Partitioning = child.outputPartitioning override def outputOrdering: Seq[SortOrder] = sortOrder diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 8ea48c1002d..37694992211 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -75,6 +75,15 @@ case class CometNativeScanExec( with ShimStreamSourceAwareSparkPlan with CometScanWithPlanData { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + // AQE rewrites DPP placeholders in the live scan before transition reversion. + // The frozen originalPlan can still contain unexecutable SubqueryAdaptiveBroadcastExecs. + val restoredScan = + originalPlan.copy(partitionFilters = partitionFilters, dataFilters = dataFilters) + originalPlan.logicalLink.foreach(restoredScan.setLogicalLink) + restoredScan + } + override lazy val metadata: Map[String, String] = if (originalPlan != null) originalPlan.metadata else Map.empty diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index ced525c1869..bb2c560ba64 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -29,10 +29,13 @@ import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec} import org.apache.spark.rdd.RDD import org.apache.spark.sql.SaveMode import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.Utils @@ -54,6 +57,8 @@ import org.apache.comet.serde.OperatorOuterClass.Operator * * @param nativeOp * The native operator representing the write operation (template, will be modified per task) + * @param originalPlan + * The JVM data-writing command restored when Comet reverts a transition-heavy stage * @param child * The child operator providing the data to write * @param outputPath @@ -69,6 +74,7 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ case class CometNativeWriteExec( nativeOp: Operator, + @transient override val originalPlan: DataWritingCommandExec, child: SparkPlan, outputPath: String, mode: SaveMode, @@ -77,7 +83,19 @@ case class CometNativeWriteExec( extends CometNativeExec with UnaryExecNode { - override def originalPlan: SparkPlan = child + override def output: Seq[Attribute] = child.output + + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = newChildren match { + case Seq(newInput) => + val restoredCommandChild = originalPlan.child match { + case writeFiles: WriteFilesExec => writeFiles.withNewChildren(Seq(newInput)) + case _ => newInput + } + originalPlan.withNewChildren(Seq(restoredCommandChild)) + case _ => + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} expected one reverted input but received ${newChildren.size}") + } // Accumulator to collect TaskCommitMessages from all tasks // Must be eagerly initialized on driver, not lazy diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index fdb07c9acff..bcae614556b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -661,6 +661,30 @@ abstract class CometExec extends CometPlan { /** The original Spark operator from which this Comet operator is converted from */ def originalPlan: SparkPlan + /** + * Rebuilds the Spark operator represented by this Comet operator with reverted children. + * + * Operators whose Comet representation changes the Spark plan shape or whose live state differs + * from the original Spark plan must override this method. + */ + def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + val sparkPlan = originalPlan + if (sparkPlan == null) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} has no original Spark plan") + } + if (newChildren.exists(_ eq sparkPlan)) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} aliases its original Spark plan with a child") + } + if (sparkPlan.children.size != newChildren.size) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} cannot restore ${sparkPlan.getClass.getSimpleName}: " + + s"expected ${sparkPlan.children.size} children but received ${newChildren.size}") + } + sparkPlan.withNewChildren(newChildren) + } + /** Comet always support columnar execution */ override def supportsColumnar: Boolean = true @@ -711,6 +735,9 @@ abstract class CometExec extends CometPlan { } object CometExec { + final class InvalidSparkFallbackException(message: String) + extends IllegalArgumentException(message) + // An unique id for each CometExecIterator, used to identify the native query execution. private val curId = new java.util.concurrent.atomic.AtomicLong() diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index e9ab624c562..0112a025268 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -51,6 +51,7 @@ import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage} import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.Trigger import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} @@ -809,6 +810,114 @@ class CometIcebergWriteActionSuite } } + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback preserves Iceberg writes with AQE=$adaptive") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val suffix = if (adaptive) "aqe" else "no_aqe" + val nativeTable = s"transition_native_$suffix" + val fallbackTable = s"transition_fallback_$suffix" + createTable(warehouseDir, nativeTable, partitionSpec = "") + createTable(warehouseDir, fallbackTable, partitionSpec = "") + val values = "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)" + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + assertNativeWriteEngages(nativeTable, Seq(1, 2, 3)) { + spark.sql(s"INSERT INTO $catalog.$ns.$nativeTable VALUES $values") + } + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + assertNativeWriteDoesNotEngage(fallbackTable, Seq(1, 2, 3)) { + spark.sql(s"INSERT INTO $catalog.$ns.$fallbackTable VALUES $values") + } + } + } + } + } + + // An unpartitioned INSERT never puts an exchange under the write, so it misses #6152: with AQE + // off, unwrapping a transition that sits on a shuffle used to keep walking into the map stage. + // The IN-subquery on a partitioned copy-on-write DELETE is the plan that does, and the AQE-off + // case is the one that fails with `ColumnarBatch cannot be cast to InternalRow`. + // `us-west` keeps a row after deleting id 2. An emptied partition makes the rewrite emit + // nothing, and AQE then replaces the write input with an empty LocalTableScan, so the executed + // plan no longer contains the shuffle this test is checking for. + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback preserves partitioned CoW deletes with AQE=$adaptive") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val suffix = if (adaptive) "aqe" else "no_aqe" + val nativeTable = s"transition_cow_native_$suffix" + val fallbackTable = s"transition_cow_fallback_$suffix" + val props = Some("'write.delete.mode'='copy-on-write'") + val spec = "PARTITIONED BY (region)" + createTable(warehouseDir, nativeTable, spec, props) + createTable(warehouseDir, fallbackTable, spec, props) + val seed = + Seq( + (1, "us-east", 10.0), + (2, "us-west", 20.0), + (3, "eu", 30.0), + (4, "us-east", 40.0), + (5, "us-west", 50.0)) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert(nativeTable, seed) + coalesceInsert(fallbackTable, seed) + } + def delete(table: String): Unit = + spark.sql( + s"DELETE FROM $catalog.$ns.$table WHERE id IN " + + "(SELECT col1 FROM VALUES (2) AS t(col1))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + assertNativeWriteEngages(nativeTable, Seq(1, 3, 4, 5)) { + delete(nativeTable) + } + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val snapshot = withNativeEnabled { + captureWrite(fallbackTable)(delete(fallbackTable)) + } + assertExactlyOneCommit(snapshot) + val nativeExecs = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e } + } + assert( + nativeExecs.isEmpty, + "transition reversion should restore IcebergWriteExec. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + // AdaptiveSparkPlanExec is a leaf, so only the AQE-aware collect reaches the exchange. + val hasExchange = snapshot.plans.exists { plan => + collect(plan) { case _: ShuffleExchangeLike => true }.nonEmpty + } + assert( + hasExchange, + "the delete must keep a shuffle under the write. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + assertRows(fallbackTable, Seq(1, 3, 4, 5)) + } + + val nativeDirs = partitionDirs(warehouseDir, nativeTable) + val fallbackDirs = partitionDirs(warehouseDir, fallbackTable) + assert( + fallbackDirs == nativeDirs, + s"partition layout fallback=$fallbackDirs native=$nativeDirs") + } + } + } + // What Iceberg's Spark writer stamps for `sort_order_id` on appended files changed across // releases: through 1.10 `SparkWrite$WriterFactory` never wires the table sort order (files // get 0 even on a sorted table); 1.11 added `SparkWriteConf.outputSortOrderId` and stamps the diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index da60c14ff0b..49238eee228 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSparkToColumnarExec, IcebergWriteExec} import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, ColumnarToRowExec, CommandExecutionMode, LeafExecNode, SparkPlan} -import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.{BinaryType, IntegerType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} @@ -1851,9 +1851,15 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes * directly. */ private def writeChildAfterTransitionRules(source: SparkPlan): SparkPlan = { + val child = CometSparkToColumnarExec(source) + val output = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + val originalPlan = IcebergWriteExec(null, output, child) val write = CometIcebergWriteExec( Operator.newBuilder().build(), - CometSparkToColumnarExec(source), + originalPlan, + child, + output, batchWrite = null, table = null, partitionSpecId = 0) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index ce981240bd8..5b01a603534 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -35,15 +35,17 @@ import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometWriteFilesExec} -import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} -import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteTaskStats, WriteTaskStatsTracker} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometSparkToColumnarExec, CometWriteFilesExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, FileSourceScanExec, SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteFilesExec, WriteTaskStats, WriteTaskStatsTracker} import org.apache.spark.sql.functions.{array, col, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.rules.RevertNativeForTransitionHeavyStages import org.apache.comet.serde.operator.NativeWriteUtils import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} @@ -126,6 +128,161 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback restores the Spark write plan with AQE=$adaptive") { + // Force a transition below the writer on both paths: Spark 3.x replaces the command, + // whereas Spark 4.x replaces only its WriteFilesExec child. + withTempPath { dir => + withTempPath { inputDir => + val inputPath = createTestData(inputDir) + val nativeOutput = new File(dir, "native-output").getAbsolutePath + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val df = spark.read.parquet(inputPath) + val nativePlan = captureWritePlan(path => df.write.parquet(path), nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + verifyWrittenFile(nativeOutput) + val nativeWrite: SparkPlan = nativePlan + .collectFirst { + case write: CometNativeWriteExec => write + case write: CometWriteFilesExec => write + } + .getOrElse(fail(s"expected a native parquet writer:\n$nativePlan")) + val stageWithTransition = nativeWrite.withNewChildren(Seq( + CometSparkToColumnarExec(CometNativeColumnarToRowExec(nativeWrite.children.head)))) + assert( + stageWithTransition.collect { case _: ColumnarToRowTransition => true }.nonEmpty, + s"test requires a C2R transition:\n$stageWithTransition") + + var fallbackPlan: SparkPlan = null + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val stagePlan = if (isSpark40Plus) { + val command = nativePlan + .collectFirst { case node: DataWritingCommandExec => node } + .getOrElse(fail(s"expected a Spark write command:\n$nativePlan")) + command.withNewChildren(Seq(stageWithTransition)) + } else { + stageWithTransition + } + fallbackPlan = RevertNativeForTransitionHeavyStages(spark)(stagePlan) + } + assertNoCometNativeWriteExec(fallbackPlan) + val commands = fallbackPlan.collect { case command: DataWritingCommandExec => + command + } + assert( + commands.exists(_.child.isInstanceOf[WriteFilesExec]), + s"expected DataWritingCommandExec -> WriteFilesExec after fallback:\n$fallbackPlan") + + deletePath(nativeOutput) + SQLExecution.withNewExecutionId(spark.range(0).queryExecution) { + fallbackPlan.executeCollect() + } + verifyWrittenFile(nativeOutput) + } + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test( + s"transition-heavy fallback restores parquet writes through the query pipeline with AQE=$adaptive") { + withTempPath { dir => + val nativeOutput = new File(dir, "native-output").getAbsolutePath + val fallbackOutput = new File(dir, "fallback-output").getAbsolutePath + // Row source, matching the Iceberg VALUES e2e: native write engages via SparkToColumnar. + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + def rangeWrite(path: String): Unit = + spark.range(0, 1000, 1, numPartitions = 4).toDF("id").write.parquet(path) + + val sparkOutput = new File(dir, "spark-output").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + rangeWrite(sparkOutput) + } + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val nativePlan = captureWritePlan(rangeWrite, nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + assertWrittenRows(nativeOutput, sparkOutput, 1000) + } + + withSQLConf( + (commonConf ++ Seq( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0")): _*) { + val fallbackPlan = captureWritePlan(rangeWrite, fallbackOutput) + assertTransitionHeavyParquetFallback(fallbackPlan) + assertWrittenRows(fallbackOutput, sparkOutput, 1000) + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback restores a union of row sources with AQE=$adaptive") { + withTempPath { dir => + val nativeOutput = new File(dir, "native-union-output").getAbsolutePath + val fallbackOutput = new File(dir, "fallback-union-output").getAbsolutePath + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + def unionedWrite(path: String): Unit = { + val left = spark.range(0, 1000, 1, numPartitions = 2).toDF("id") + val mid = spark.range(1000, 2000, 1, numPartitions = 2).toDF("id") + val right = spark.range(2000, 3000, 1, numPartitions = 2).toDF("id") + left.union(mid).union(right).write.parquet(path) + } + + val sparkOutput = new File(dir, "spark-output").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + unionedWrite(sparkOutput) + } + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val nativePlan = captureWritePlan(unionedWrite, nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + assertWrittenRows(nativeOutput, sparkOutput, 3000) + } + + withSQLConf( + (commonConf ++ Seq( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0")): _*) { + val fallbackPlan = captureWritePlan(unionedWrite, fallbackOutput) + assertTransitionHeavyParquetFallback(fallbackPlan) + assertWrittenRows(fallbackOutput, sparkOutput, 3000) + } + } + } + } + test("basic parquet write with repartition") { withTempPath { dir => // Create test data and write it to a temp parquet file first @@ -1549,6 +1706,58 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + private def deletePath(path: String): Unit = { + def delete(file: File): Unit = { + if (file.isDirectory) { + Option(file.listFiles()).foreach(_.foreach(delete)) + } + file.delete() + } + delete(new File(path)) + } + + /** + * Spark 3.x replaces the whole command with `CometNativeWriteExec`, and Spark then inserts a + * C2R above that columnar node, so `maxTransitions = 0` reverts the stage back to + * `DataWritingCommandExec` -> `WriteFilesExec`. Spark 4.0+ leaves the command in place and does + * not insert a C2R above `CometWriteFilesExec`. These row sources only add `SparkToColumnar`, + * which is an R2C, so the stage is not transition-heavy and the native write stays. + */ + private def assertTransitionHeavyParquetFallback(plan: SparkPlan): Unit = { + if (isSpark40Plus) { + assertHasCometNativeWriteExec(plan) + } else { + assertNoCometNativeWriteExec(plan) + assertRestoredParquetWriteCommand(plan) + } + } + + private def assertRestoredParquetWriteCommand(plan: SparkPlan): Unit = { + assert( + plan + .collect { case command: DataWritingCommandExec => command } + .exists(_.child.isInstanceOf[WriteFilesExec]), + s"expected DataWritingCommandExec -> WriteFilesExec after fallback:\n$plan") + } + + // Compare every row with Spark-written data using the row-based parquet reader, so this + // check is independent of Comet scans and detects duplicates or missing union inputs. + private def assertWrittenRows( + outputPath: String, + sparkOutputPath: String, + expectedRows: Int): Unit = { + val outputDir = new File(outputPath) + val partFiles = outputDir.listFiles().filter(_.getName.startsWith("part-")) + assert(partFiles.length > 1, s"Expected multiple part files under $outputPath") + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") { + val expected = spark.read.parquet(sparkOutputPath).collect().toSeq + assert(expected.size == expectedRows) + checkAnswer(spark.read.parquet(outputPath), expected) + } + } + private def writeWithCometNativeWriteExec( inputPath: String, outputPath: String, @@ -1564,10 +1773,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { Some(plan) } - private def verifyWrittenFile(outputPath: String): Unit = { + private def verifyWrittenFile(outputPath: String, expectedRows: Int = 1000): Unit = { // Verify the data was written correctly val resultDf = spark.read.parquet(outputPath) - assert(resultDf.count() == 1000, "Expected 1000 rows to be written") + assert(resultDf.count() == expectedRows, s"Expected $expectedRows rows to be written") // Verify multiple part files were created val outputDir = new File(outputPath) diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index e350badf8da..d0ca67c60d5 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala @@ -19,19 +19,122 @@ package org.apache.comet.rules -import org.apache.spark.sql.{CometTestBase, Row} +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference + +import org.apache.spark.sql.{CometTestBase, Row, SaveMode} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.BinaryType +import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.serde.OperatorOuterClass.Operator + +private case class AliasingFallbackCometExec( + override val originalPlan: SparkPlan, + child: SparkPlan) + extends CometExec + with UnaryExecNode { + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { + private def cometIcebergWrite(child: SparkPlan): CometIcebergWriteExec = { + val output = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + val originalPlan = IcebergWriteExec(null, output, child) + CometIcebergWriteExec( + Operator.newBuilder().build(), + originalPlan, + child, + output, + batchWrite = null, + table = null, + partitionSpecId = 0) + } + + private def cometFilter(child: SparkPlan): CometFilterExec = { + val condition = Literal.TrueLiteral + val sparkFilter = FilterExec(condition, child) + CometFilterExec( + Operator.newBuilder().build(), + sparkFilter, + sparkFilter.output, + condition, + child, + SerializedPlan(None)) + } + + private def captureDataWritingCommand(path: String): DataWritingCommandExec = { + val captured = new AtomicReference[SparkPlan]() + val callbackCompleted = new CountDownLatch(1) + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + if (funcName == "save" || funcName.contains("command")) { + captured.set(qe.executedPlan) + callbackCompleted.countDown() + } + } + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(1).toDF("id").write.mode("overwrite").parquet(path) + } + assert( + callbackCompleted.await(10, TimeUnit.SECONDS), + "timed out waiting to capture the parquet write plan") + } finally { + spark.listenerManager.unregister(listener) + } + val plan = stripAQEPlan( + Option(captured.get()).getOrElse(fail("expected a captured parquet write plan"))) + plan + .collectFirst { case command: DataWritingCommandExec => command } + .getOrElse(fail(s"expected DataWritingCommandExec:\n$plan")) + } + + private def cometNativeWrite( + child: SparkPlan, + command: DataWritingCommandExec): CometNativeWriteExec = { + CometNativeWriteExec( + Operator.newBuilder().build(), + command, + child, + outputPath = "/tmp/unused-native-write", + mode = SaveMode.Overwrite) + } + + private def assertRestoredParquetWrite(reverted: SparkPlan): WriteFilesExec = { + val command = reverted match { + case node: DataWritingCommandExec => node + case other => fail(s"expected DataWritingCommandExec, got:\n$other") + } + val writeFiles = command.child match { + case node: WriteFilesExec => node + case other => fail(s"expected WriteFilesExec under DataWritingCommandExec, got:\n$other") + } + assert( + command.collect { case _: CometNativeWriteExec => true }.isEmpty, + s"native parquet write should be restored, not erased:\n$reverted") + writeFiles + } + private def createSparkPlan(sql: String): SparkPlan = { var plan: SparkPlan = null withSQLConf(CometConf.COMET_ENABLED.key -> "false") { @@ -59,6 +162,12 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { plan.collect { case _: ColumnarToRowTransition => true }.size } + private def unwrapCodegen(plan: SparkPlan): SparkPlan = plan match { + case wholeStage: WholeStageCodegenExec => unwrapCodegen(wholeStage.child) + case inputAdapter: InputAdapter => unwrapCodegen(inputAdapter.child) + case other => other + } + private def collectCometAggregates(plan: SparkPlan): Seq[CometHashAggregateExec] = { val current = plan match { case aggregate: CometHashAggregateExec => Seq(aggregate) @@ -142,6 +251,207 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } } + test("revertToSpark preserves an Iceberg write with a leaf child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(leaf) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val icebergWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + + assert(icebergWrite.child eq leaf) + assert(icebergWrite.output.map(_.name) == Seq(IcebergWriteExec.CommitMessageColumn)) + assert(icebergWrite.output.map(_.dataType) == Seq(BinaryType)) + } + + test("revertToSpark preserves an Iceberg write over SparkToColumnar of a row source") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(CometSparkToColumnarExec(leaf)) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val icebergWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + + assert( + icebergWrite.child eq leaf, + s"SparkToColumnar should unwrap to the row source:\n$reverted") + assert(icebergWrite.output.map(_.name) == Seq(IcebergWriteExec.CommitMessageColumn)) + assert(icebergWrite.output.map(_.dataType) == Seq(BinaryType)) + assert( + reverted.collect { case _: CometSparkToColumnarExec => true }.isEmpty, + s"SparkToColumnar should be fully unwrapped:\n$reverted") + } + + test("revertToSpark preserves an Iceberg write without duplicating its unary child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1), (2) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(cometFilter(leaf)) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + + assert(reverted.isInstanceOf[IcebergWriteExec], s"expected IcebergWriteExec:\n$reverted") + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + + test("revertToSpark unwraps stacked SparkToColumnar(C2R) under a native write") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val stacked = CometSparkToColumnarExec(CometNativeColumnarToRowExec(cometFilter(leaf))) + val write = cometIcebergWrite(stacked) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + + assert(reverted.isInstanceOf[IcebergWriteExec], s"expected IcebergWriteExec:\n$reverted") + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert( + reverted.collect { case _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec => + true + }.isEmpty, + s"stacked transitions should be fully unwrapped:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + + test("revertToSpark preserves a native parquet write with a leaf child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val write = cometNativeWrite(leaf, captureDataWritingCommand(dir.getAbsolutePath)) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val writeFiles = assertRestoredParquetWrite(reverted) + assert(writeFiles.child eq leaf, s"expected the original leaf child:\n$reverted") + } + } + + test("revertToSpark preserves a native parquet write without duplicating its unary child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1), (2) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val write = + cometNativeWrite(cometFilter(leaf), captureDataWritingCommand(dir.getAbsolutePath)) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + assertRestoredParquetWrite(reverted) + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + } + + test("revertToSpark restores a native parquet write whose command has no WriteFilesExec") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val command = captureDataWritingCommand(dir.getAbsolutePath) + val input = command.child match { + case writeFiles: WriteFilesExec => writeFiles.child + case other => other + } + val commandWithoutWriteFiles = + command.withNewChildren(Seq(input)).asInstanceOf[DataWritingCommandExec] + val write = cometNativeWrite(leaf, commandWithoutWriteFiles) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val restored = reverted match { + case node: DataWritingCommandExec => node + case other => fail(s"expected DataWritingCommandExec, got:\n$other") + } + assert(restored.child eq leaf, s"expected the original leaf child:\n$reverted") + assert( + restored.collect { case _: WriteFilesExec => true }.isEmpty, + s"WriteFilesExec should not be reinserted when the original command lacked it:\n$reverted") + assert( + restored.collect { case _: CometNativeWriteExec => true }.isEmpty, + s"native parquet write should be restored, not erased:\n$reverted") + } + } + + test("invalid original-plan alias skips the entire stage reversion") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val aliasing = AliasingFallbackCometExec(leaf, leaf) + val stagePlan = CometNativeColumnarToRowExec(aliasing) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + val result = rule(stagePlan) + + assert( + result eq stagePlan, + s"invalid fallback must leave the whole stage unchanged:\n$result") + } + } + + test("invalid original-plan alias to a Comet child skips the entire stage reversion") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val cometChild = cometFilter(leaf) + val aliasing = AliasingFallbackCometExec(cometChild, cometChild) + val stagePlan = CometNativeColumnarToRowExec(aliasing) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + val result = rule(stagePlan) + + assert( + result eq stagePlan, + s"invalid nested fallback must leave the whole stage unchanged:\n$result") + } + } + + test("local TopK sparkFallback returns the supplied restored child") { + val child = spark.range(20).queryExecution.sparkPlan + val originalPlan = TakeOrderedAndProjectExec(5, Seq.empty, child.output, child) + val local = CometLocalTopKExec( + Operator.newBuilder().build(), + originalPlan, + child.output, + 5, + Seq.empty, + dynamicFilterEnabled = false, + child, + SerializedPlan(None)) + val restoredChild = spark.range(10).queryExecution.sparkPlan + + assert(local.sparkFallback(Seq(restoredChild)) eq restoredChild) + intercept[CometExec.InvalidSparkFallbackException] { + local.sparkFallback(Seq.empty) + } + } + for (adaptive <- Seq(false, true)) { test(s"transition reversion preserves local TopK: AQE=$adaptive") { withSQLConf( @@ -393,6 +703,206 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } } + test("revertToSpark leaves transitions below a shuffle that a stripped transition sat on") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val df = sql("SELECT _2, count(*) FROM tbl GROUP BY _2") + df.collect() + val cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + val shuffle = cometPlan + .collectFirst { case s: CometShuffleExchangeExec => s } + .getOrElse(fail(s"test requires a native shuffle:\n$cometPlan")) + // The transition that must survive lives in the map stage, under the exchange. + val preserved = + if (shuffle.child.supportsColumnar) CometNativeColumnarToRowExec(shuffle.child) + else CometSparkToColumnarExec(shuffle.child) + val shuffleWithTransition = shuffle.withNewChildren(Seq(preserved)) + // Stacked transitions sit directly on the shuffle, the shape #6152 strips through. + val onExchange = + CometSparkToColumnarExec(CometNativeColumnarToRowExec(shuffleWithTransition)) + val write = cometIcebergWrite(onExchange) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val restoredWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + val restoredShuffle = restoredWrite.child match { + case ColumnarToRowExec(exchange: CometShuffleExchangeExec) => exchange + case exchange: CometShuffleExchangeExec => exchange + case other => + fail(s"expected the shuffle under the restored write, got:\n$other") + } + assert( + restoredShuffle.child eq preserved, + "unwrapping the transition on the shuffle must not strip the stage below it:\n" + + reverted.treeString) + } + } + } + + test("transition-heavy reversion rejects a stripped stage-boundary root") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + var cometPlan: SparkPlan = null + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql("SELECT _1, _2 FROM tbl DISTRIBUTE BY _2") + df.collect() + cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + } + val exchange = cometPlan + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail(s"test requires a native shuffle:\n$cometPlan")) + val stagePlan = CometColumnarToRowExec(exchange) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + var result: SparkPlan = null + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + result = rule(stagePlan) + } + + assert( + result eq stagePlan, + "a transition-heavy stage whose stripped root is an exchange must stay unchanged:\n" + + result.treeString) + } + } + } + + test("transition-heavy revert restores row output for a columnar Spark scan") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val (_, plan) = checkSparkAnswer("SELECT _1, _2 FROM tbl") + val executedPlan = stripAQEPlan(plan) + val resultStageRoot = unwrapCodegen(executedPlan) + + val scan = resultStageRoot match { + case transition: ColumnarToRowExec => + unwrapCodegen(transition.child) match { + case child: FileSourceScanExec => child + case other => + fail(s"expected a vectorized Spark scan under ColumnarToRow, got:\n$other") + } + case other => + fail( + "expected ColumnarToRow over a vectorized Spark scan, got " + + s"${other.getClass.getName}:\n$other") + } + assert(scan.supportsColumnar, s"the reverted scan must use its columnar path:\n$scan") + assert(countCometExecs(executedPlan) == 0, s"the stage must be reverted:\n$executedPlan") + } + } + } + + for (adaptive <- Seq(false, true); columnarRoot <- Seq(false, true)) { + test( + s"transition-heavy map-stage fallback supplies Arrow: AQE=$adaptive, columnar=$columnarRoot") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val rows = (0 until 100).map(i => (i, i % 10)) + withParquetTable(rows, "tbl") { + val df = sql("SELECT _1, _2 FROM tbl DISTRIBUTE BY _2") + df.collect() + val exchange = stripAQEPlan(df.queryExecution.executedPlan) + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail("test requires a native shuffle")) + assert(exchange.child.isInstanceOf[CometNativeScanExec]) + val input = if (columnarRoot) exchange.child else cometFilter(exchange.child) + val stage = exchange.withNewChildren( + Seq(CometSparkToColumnarExec(CometNativeColumnarToRowExec(input)))) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stage.children.head) == 1) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val reverted = rule(stage).asInstanceOf[CometShuffleExchangeExec] + val bridge = reverted.child match { + case node: CometSparkToColumnarExec => node + case other => fail(s"expected an Arrow bridge after map-stage fallback:\n$other") + } + assert(bridge.child.supportsColumnar == columnarRoot) + assert(bridge.child.collect { case _: FileSourceScanExec => true }.nonEmpty) + assert(countCometExecs(bridge.child) == 0) + // Execute the native shuffle, not just its Spark fallback child: Spark batches + // satisfy supportsColumnar but cannot be cast to CometVector by the Arrow stream. + SQLExecution.withNewExecutionId(df.queryExecution) { + val actual = ColumnarToRowExec(reverted) + .executeCollect() + .map(row => (row.getInt(0), row.getInt(1))) + .toSeq + assert(actual.sorted == rows.sorted) + } + } + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy revert preserves native exchange for DISTRIBUTE BY: AQE=$adaptive") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "8", + "spark.sql.adaptive.coalescePartitions.enabled" -> "true", + "spark.sql.adaptive.coalescePartitions.parallelismFirst" -> "false", + "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "67108864") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val query = "SELECT _1, _2 FROM tbl DISTRIBUTE BY _2" + var sparkAnswer: Seq[Row] = Seq.empty + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + sparkAnswer = sql(query).collect().toSeq + } + val df = sql(query) + checkCometAnswer(df, sparkAnswer) + val executedPlan = stripAQEPlan(df.queryExecution.executedPlan) + val resultStageRoot = unwrapCodegen(executedPlan) + + assert( + resultStageRoot.isInstanceOf[ColumnarToRowTransition], + s"the result stage must end with a columnar-to-row transition:\n$executedPlan") + + if (adaptive) { + val read = executedPlan + .collectFirst { case node: AQEShuffleReadExec => node } + .getOrElse(fail(s"expected a coalesced AQE shuffle read:\n$executedPlan")) + assert(read.partitionSpecs.size < 8, s"shuffle must be coalesced:\n$read") + val stage = read.child match { + case node: ShuffleQueryStageExec => node + case other => fail(s"expected a shuffle query stage:\n$other") + } + assert(stage.plan.isInstanceOf[CometShuffleExchangeExec]) + } else { + val exchange = executedPlan + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail(s"expected a native shuffle:\n$executedPlan")) + assert( + exchange.collect { case scan: CometNativeScanExec => scan }.nonEmpty, + s"the map stage must retain its native scan:\n$executedPlan") + assert( + exchange.collect { case scan: FileSourceScanExec => scan }.isEmpty, + s"fallback must not replace the scan below the exchange:\n$executedPlan") + } + } + } + } + } + test("non-AQE apply must not produce an invalid plan when the result stage reverts") { withSQLConf( CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", From c8e6553dad376537c144e36be9dab7e31df2df0e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 7 Oct 2026 00:14:07 +0000 Subject: [PATCH 19/89] fix: count the days partition transform in UTC (#6348) * fix: count the days partition transform in UTC CometDays cast timestamps to dates in the session timezone, while CometHours and Iceberg's days transform count in UTC. Spark cannot evaluate the transform, so there is no Spark result to match. Use UTC so the transforms agree with each other and with Iceberg. Closes #6333. * test: pin UTC flooring of days for a pre-epoch timestamp Check days one microsecond before the epoch in an Asia/Tokyo session. That instant is day -1 in UTC but already 1970-01-01 in Tokyo, so the session-timezone serde returned 0 and fails this check. The existing timestamp literal checks pass either way: they run in the default America/Los_Angeles session, away from a day boundary. * docs: note where days differs from Iceberg for pre-1970 timestamps --- .../scala/org/apache/comet/serde/datetime.scala | 14 +++++++++----- .../comet/CometTemporalExpressionSuite.scala | 14 +++++++++++--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index bf27ce2c552..ed5faa019ec 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import java.util.Locale import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, DivideDTInterval, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -822,8 +821,14 @@ object CometHours extends CometExpressionSerde[Hours] { * For DateType: dates are internally stored as days since epoch, so this is a simple cast to * integer (same as CometUnixDate). * - * For TimestampType: uses a timezone-aware Cast(Timestamp to Date) followed by Cast(Date to Int). - * The first cast respects the session timezone to correctly determine the date boundary. + * For TimestampType: uses Cast(Timestamp to Date) in UTC followed by Cast(Date to Int). Spark + * cannot evaluate `days` itself, so this counts days in UTC like [[CometHours]] and Iceberg's + * `days` transform, rather than in the session timezone. + * + * The cast is a true floor, as [[CometHours]] is. Iceberg's `DateTimeUtil.microsToDays` differs + * in one case: it puts a pre-1970 timestamp exactly 999999 microseconds past midnight, such as + * `1969-01-01 00:00:00.999999`, in the previous day. Comet's native kernel for Iceberg's own + * `days` function matches Iceberg there, and this cast does not. */ object CometDays extends CometExpressionSerde[Days] { @@ -846,9 +851,8 @@ object CometDays extends CometExpressionSerde[Days] { val dateExprOpt = expr.child.dataType match { case DateType => childExpr case TimestampType => - val timezone = SQLConf.get.sessionLocalTimeZone childExpr.flatMap { child => - CometCast.castToProto(expr, Some(timezone), DateType, child, CometEvalMode.LEGACY) + CometCast.castToProto(expr, Some("UTC"), DateType, child, CometEvalMode.LEGACY) } case _ => None } diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index a03e749a76e..e61a57a6418 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -839,7 +839,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { checkDays( tsDF.select(col("ts"), getColumnFromExpression(Days(UnresolvedAttribute("ts")))), - tsDF.selectExpr("ts", "unix_date(cast(ts as date))")) + tsDF.selectExpr("ts", "cast(floor(unix_micros(ts) / 86400000000D) as int)")) } } } @@ -885,8 +885,16 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH java.sql.Timestamp.valueOf("2024-06-15 10:30:00"), DataTypes.TimestampType)))), dummyDF.selectExpr( - "unix_date(cast(TIMESTAMP('1970-01-01 00:00:00') as date))", - "unix_date(cast(TIMESTAMP('2024-06-15 10:30:00') as date))")) + "cast(floor(unix_micros(TIMESTAMP('1970-01-01 00:00:00')) / 86400000000D) as int)", + "cast(floor(unix_micros(TIMESTAMP('2024-06-15 10:30:00')) / 86400000000D) as int)")) + + // One microsecond before the epoch. Days are counted in UTC and floored, so this is day -1 + // even though it is already 1970-01-01 in Asia/Tokyo. + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "Asia/Tokyo") { + checkDays( + dummyDF.select(getColumnFromExpression(Days(Literal(-1L, DataTypes.TimestampType)))), + dummyDF.selectExpr("-1")) + } // Null handling checkDays( From 37ffca54eb19b35416e98cc72f0458ac58571eba Mon Sep 17 00:00:00 2001 From: Unik Dahal <61407386+unikdahal@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:05:44 +0000 Subject: [PATCH 20/89] feat(iceberg): model WriteDelta JVM fallback (#6693) --- .../contributor-guide/iceberg-writes.md | 27 +-- .../user-guide/latest/iceberg-writes.md | 23 +- .../comet/iceberg/DeltaLogicalFields.scala | 46 ++++ .../comet/iceberg/IcebergReflection.scala | 53 ++++- .../comet/iceberg/IcebergWriteDispatch.scala | 28 +++ .../comet/iceberg/IcebergWriteLogical.scala | 2 +- .../comet/iceberg/IcebergWriteStrategy.scala | 124 ++++++++--- .../iceberg/WriteDeltaDispatchInfo.scala | 64 ++++++ .../operator/CometIcebergNativeWrite.scala | 26 ++- .../spark/sql/comet/IcebergCommitExec.scala | 17 +- .../spark/sql/comet/IcebergWriteExec.scala | 54 +++-- .../iceberg/IcebergDeltaLogicalShim.scala | 27 +++ .../iceberg/IcebergDeltaLogicalShim.scala | 47 +++++ .../iceberg/IcebergDeltaWriterShim.scala | 65 ++++++ .../iceberg/IcebergSemanticMetricsShim.scala | 30 +++ .../sql/comet/IcebergWriteSummaryShim.scala | 5 +- .../IcebergDeltaLogicalFieldsShim.scala | 47 +++++ .../iceberg/IcebergSemanticMetricsShim.scala | 30 +++ .../sql/comet/IcebergWriteSummaryShim.scala | 5 +- .../sql/comet/IcebergWriteSummaryShim.scala | 47 +++-- .../IcebergDeltaLogicalFieldsShim.scala | 47 +++++ .../iceberg/IcebergSemanticMetricsShim.scala | 30 +++ .../comet/IcebergDeltaWriteSummaryShim.scala | 34 +++ .../IcebergDeltaLogicalFieldsShim.scala | 47 +++++ .../iceberg/IcebergSemanticMetricsShim.scala | 45 ++++ .../comet/IcebergDeltaWriteSummaryShim.scala | 58 ++++++ .../IcebergTransactionalCommitExec.scala | 10 +- .../iceberg/IcebergDeltaLogicalShim.scala | 34 +++ .../iceberg/IcebergDeltaWriterShim.scala | 70 +++++++ .../comet/CometIcebergWriteActionSuite.scala | 196 +++++++++++++++--- .../iceberg/IcebergReflectionSuite.scala | 34 +++ .../IcebergTransactionalCommitTestShim.scala | 8 +- 32 files changed, 1232 insertions(+), 148 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala create mode 100644 spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala create mode 100644 spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala create mode 100644 spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala create mode 100644 spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala create mode 100644 spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala create mode 100644 spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala create mode 100644 spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala create mode 100644 spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala create mode 100644 spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala create mode 100644 spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala create mode 100644 spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala create mode 100644 spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala create mode 100644 spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala create mode 100644 spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index e2547762822..1ef7539329a 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -70,13 +70,13 @@ IcebergCommit driver: collect task commit messages, BatchWrite.c +- scans, projects, exchanges, sorts; inside AQE with or without the split ``` -| Component | Location | Role | -| ------------------------------------------------------------------------------- | -------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `IcebergWriteStrategy` | `spark/src/main/scala/org/apache/comet/iceberg/` | Planner strategy. Matches `AppendData`, `OverwriteByExpression`, `OverwritePartitionsDynamic`, `ReplaceData` (and Iceberg's own `ReplaceIcebergData`, which Iceberg 1.5.2 plans on Spark 3.4) whose `Write` is an Iceberg `SparkWrite`. | -| `IcebergWriteLogical` | same | Logical anchor for the writer, so AQE re-plans re-emit only the writer and not a second committer. | -| `IcebergWriteExec` | `spark/src/main/scala/org/apache/spark/sql/comet/` | JVM writer. Runs iceberg-java's `DataWriter` per task and returns the serialized `WriterCommitMessage` as one binary row. | -| `IcebergCommitExec` | same | Driver committer. A `V2CommandExec`, so `run()` is memoized and the commit happens once. | -| `IcebergReplaceDataShim`, `IcebergRefreshCacheShim`, `IcebergDriverMetricsShim` | `spark/src/main/spark-*/org/apache/comet/iceberg/` | Version differences: Spark 4.x operation-coded `ReplaceData` rows, cache refresh by name on 4.1+, driver metric reporting. | +| Component | Location | Role | +| ------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `IcebergWriteStrategy` | `spark/src/main/scala/org/apache/comet/iceberg/` | Planner strategy. Matches `AppendData`, `OverwriteByExpression`, `OverwritePartitionsDynamic`, `ReplaceData` (and Iceberg's own `ReplaceIcebergData`, which Iceberg 1.5.2 plans on Spark 3.4), plus Spark 3.5+ Iceberg `WriteDelta`. | +| `IcebergWriteLogical` | same | Logical anchor for the writer, so AQE re-plans re-emit only the writer and not a second committer. | +| `IcebergWriteExec` | `spark/src/main/scala/org/apache/spark/sql/comet/` | JVM writer. Runs iceberg-java's `DataWriter` or, for `WriteDelta`, `DeltaWriter` per task and returns the serialized `WriterCommitMessage` as one binary row. | +| `IcebergCommitExec` | same | Driver committer. A `V2CommandExec`, so `run()` is memoized and the commit happens once. | +| `IcebergReplaceDataShim`, `IcebergDeltaLogicalShim`, `IcebergDeltaWriterShim`, `IcebergRefreshCacheShim`, `IcebergDriverMetricsShim` | `spark/src/main/spark-*/org/apache/comet/iceberg/` | Version differences: operation-coded `ReplaceData` / `WriteDelta` rows, cache refresh by name on 4.1+, driver metric reporting. | Things to know before changing this layer: @@ -90,16 +90,19 @@ Things to know before changing this layer: - **Messages are collected per task as tasks finish** (`sparkContext.runJob` with a result handler), not with `executeCollect`. That is how a failed job still knows which tasks completed, so it can delete their files. -- **Commit-coordinator writes are not intercepted.** Iceberg's `SparkWrite` never asks for one; the - check in `buildTwoOp` is defensive. -- **What is not intercepted:** merge-on-read (`WriteDelta`), streaming writes, and CTAS/RTAS on - Spark 3.4. Those keep Spark's plan. +- **Commit-coordinator writes are not intercepted.** Iceberg's current `SparkWrite` and + position-delta write do not ask for one; the checks in `buildTwoOp` and `buildDeltaTwoOp` + are defensive. +- **WriteDelta stays JVM-backed.** Spark 3.5+ merge-on-read commands are intercepted by the split + plan, but `IcebergWriteExec` delegates their rows to Iceberg's JVM `DeltaWriter`; the native + Iceberg writer is explicitly declined. Spark 3.4 `WriteDelta` keeps Spark's plan. +- **What is not intercepted:** streaming writes and CTAS/RTAS on Spark 3.4. Those keep Spark's plan. On Spark 4.1+, `IcebergWriteSummaryShim` finds either Spark's `MergeRowsExec` or `CometMergeRowsExec` in the executed query and forwards its eight action counters to `BatchWrite.commit(messages, summary)`. Native instructions carry the context of each `Keep`; `Discard` counts a deletion and `Split` counts one update. Spark 4.2 uses last-attempt -accumulators, read through `MergeRowsMetricsShim`, so summary values come from the tasks that +accumulators, read through `IcebergSemanticMetricsShim`, so summary values come from the tasks that produced the committed output. Stock V2 writers still require the concrete Spark node, so `IcebergWriteStrategy` tags a copy of its logical MERGE query before planning and AQE. The serializer requires that tag on Spark 4.1+, leaving stock writers on Spark throughout replanning. diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index 831dcc20386..a04f6be61cf 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -119,10 +119,15 @@ For an unpartitioned copy-on-write `MERGE`, the native Iceberg writer is reachab On Spark 4.1+, the native MergeRows path also preserves the semantic counters required by the summary-aware writer commit contract. -The mechanism behind row-level DML differs by Spark version: on Spark 4.0+ the analyzer emits -operation-coded rows that Comet's writer dispatches through `ReplaceData`'s projections, while -on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The supported set of -operations is the same either way. +For copy-on-write row-level DML, the mechanism differs by Spark version: on Spark 4.0+ the +analyzer emits operation-coded rows that Comet's writer dispatches through `ReplaceData`'s +projections, while on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The +supported set of operations is the same either way. + +On Spark 3.5+, merge-on-read uses Spark's `WriteDelta`. The split plan intercepts that command so +Comet can keep the same driver commit and reporting path, but task-side row-level writes stay on +Iceberg's JVM `DeltaWriter`; `CometIcebergWriteExec` is never used for position-delta rows. +Spark 3.4 leaves `WriteDelta` on Spark's stock write plan. On Spark 4.1+ the split plan matches two further stock-Spark behaviours: MERGE metrics are forwarded to the writer's commit (Iceberg 1.11+ records them in the snapshot summary), and @@ -134,8 +139,9 @@ changes. The rewrite is skipped — and the write runs through Spark's stock combined operator — when: - `spark.comet.write.iceberg.splitOperator.enabled` is `false` (the default); -- the write is not an Iceberg `SparkWrite` (any other V2 data source); -- the table uses merge-on-read: delta writes (Iceberg `WriteDelta`) are not intercepted; +- the write is neither an Iceberg `SparkWrite` nor a supported Iceberg position-delta write; +- the table uses merge-on-read on Spark 3.4; Spark 3.5+ `WriteDelta` is intercepted but remains + on Iceberg's JVM `DeltaWriter`; - the statement is CTAS / RTAS on Spark 3.4, where the staged exec writes inline; on Spark 3.5+ those statements re-plan their inner append, which is intercepted normally; - the write requires Spark's commit coordinator, which Comet's per-task commit protocol does @@ -273,8 +279,9 @@ Other `write.*` properties are intentionally not gated because they cannot make writer produce different data files: distribution and ordering settings shape the Spark plan identically on both paths, WAP / branch / snapshot properties act on the JVM committer, `write.avro.*` / `write.orc.*` apply only to formats already excluded, and merge-on-read -settings route the write through `WriteDelta`, which the split plan never intercepts. Every -rule is pinned by `CometIcebergWriteDetectionSuite`. +settings route the write through `WriteDelta`, whose task writer remains iceberg-java rather +than the native writer. Every native eligibility rule is pinned by +`CometIcebergWriteDetectionSuite`. Manifest `DataFile` metrics are assembled on the JVM before commit: each written file's metrics are re-derived from its parquet footer through the version-matched diff --git a/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala b/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala new file mode 100644 index 00000000000..feebad3a100 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala @@ -0,0 +1,46 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.util.WriteDeltaProjections +import org.apache.spark.sql.connector.write.Write + +sealed trait DeltaCommand extends Product with Serializable +case object DeltaDelete extends DeltaCommand +case object DeltaUpdate extends DeltaCommand +case object DeltaMerge extends DeltaCommand + +case class DeltaLogicalFields( + query: LogicalPlan, + originalTable: NamedRelation, + projections: WriteDeltaProjections, + write: Option[Write], + command: Option[DeltaCommand] = None) + +private[iceberg] trait IcebergDeltaLogicalShimApi { + def extract(plan: LogicalPlan): Option[DeltaLogicalFields] +} + +/** Version-specific extraction of Spark's public WriteDelta logical node. */ +private[iceberg] trait IcebergDeltaLogicalFieldsShimApi { + def extract(plan: LogicalPlan): Option[DeltaLogicalFields] +} diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 9a0dfee85fb..96626631393 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -205,6 +205,20 @@ object IcebergReflection extends Logging { batchWrite.getClass.getName.startsWith(ClassNames.SPARK_WRITE + "$") } + private val positionDeltaWriteClassName = + "org.apache.iceberg.spark.source.SparkPositionDeltaWrite" + private val positionDeltaBatchWriteClassPrefix = + "org.apache.iceberg.spark.source.SparkPositionDeltaWrite$PositionDeltaBatchWrite" + + /** True only for Iceberg's JVM DeltaWrite implementation. */ + def isIcebergPositionDeltaWrite(write: Any): Boolean = + write != null && tryLoadClass(positionDeltaWriteClassName).exists(_.isInstance(write)) + + /** True only for the BatchWrite enclosed by Iceberg's SparkPositionDeltaWrite. */ + def isIcebergPositionDeltaBatchWrite(batchWrite: Any): Boolean = + batchWrite != null && batchWrite.getClass.getName.startsWith( + positionDeltaBatchWriteClassPrefix) + def getOuterSparkWrite(batchWrite: Any): Option[Any] = { if (batchWrite == null) None else { @@ -1423,6 +1437,14 @@ object IcebergReflection extends Logging { def getTableFromSparkWrite(sparkWrite: Any): Option[Any] = getSparkWriteField(sparkWrite, "table") + /** Table owned by either Iceberg's regular SparkWrite or its position-delta write. */ + def getTableFromIcebergWrite(write: Any): Option[Any] = + if (isIcebergPositionDeltaWrite(write)) { + reflectField(write, "table") + } else { + getTableFromSparkWrite(write) + } + def getWritePropertiesFromSparkWrite(sparkWrite: Any): Option[Map[String, String]] = { import scala.jdk.CollectionConverters._ getSparkWriteField(sparkWrite, "writeProperties") @@ -2010,19 +2032,36 @@ object IcebergReflection extends Logging { } /** - * The locations of the data files carried by a `SparkWrite$TaskCommit` message (its - * package-private `files()`), or empty when `message` is not one. Used to clean up after a - * write job that failed before any commit was attempted. - */ - def taskCommitFileLocations(message: AnyRef): Seq[String] = + * The newly written file locations carried by an Iceberg task commit message. Plain + * `SparkWrite$TaskCommit` exposes `files()`; position-delta commits expose separate + * `dataFiles()` and `deleteFiles()` arrays. Do not include `rewrittenDeleteFiles()`: those + * files pre-date this job and are only candidates for removal after a successful delta commit. + * + * Used to clean completed tasks after a write job fails before any commit is attempted. Both + * Iceberg write implementations keep their own cleanup disabled in that failure mode. + */ + def taskCommitFileLocations(message: AnyRef): Seq[String] = { + def locations(methodName: String): Seq[String] = + findMethodInHierarchy(message.getClass, methodName).toSeq + .flatMap { method => + method.invoke(message) match { + case array: Array[_] => + array.toSeq.flatMap(file => extractFileLocation(file.getClass, file)) + case _ => Seq.empty + } + } + findMethodInHierarchy(message.getClass, "files") match { case Some(files) => files.invoke(message) match { - case array: Array[_] => array.toSeq.flatMap(f => extractFileLocation(f)) + case array: Array[_] => + array.toSeq.flatMap(file => extractFileLocation(file.getClass, file)) case _ => Seq.empty } - case None => Seq.empty + case None => + locations("dataFiles") ++ locations("deleteFiles") } + } /** The table's `FileIO` (`table.io()`). Iceberg requires `FileIO` to be `Serializable`. */ def getTableIO(table: Any): Option[AnyRef] = diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala new file mode 100644 index 00000000000..970179753b4 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala @@ -0,0 +1,28 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +sealed trait IcebergWriteDispatch extends Product with Serializable + +case object PlainIcebergWrite extends IcebergWriteDispatch + +case class ReplaceDataWrite(info: ReplaceDataDispatchInfo) extends IcebergWriteDispatch + +case class PositionDeltaWrite(info: WriteDeltaDispatchInfo) extends IcebergWriteDispatch diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala index 6ceb526d4c9..c79e32365a1 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala @@ -30,7 +30,7 @@ case class IcebergWriteLogical( child: LogicalPlan, // Driver-side only: AQE re-planning is driver-local and write commands aren't cached. @transient batchWrite: BatchWrite, - replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + dispatch: IcebergWriteDispatch = PlainIcebergWrite) extends UnaryNode { // Owns the commit-message attribute so the physical writer keeps the same exprId across diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala index 85fbcc8ce3a..dbb19fba797 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala @@ -31,8 +31,8 @@ import org.apache.comet.CometSparkSessionExtensions.isCometLoaded import org.apache.comet.shims.ShimCometMergeRows /** - * Spark Strategy that intercepts Iceberg V2 copy-on-write logical writes and emits Comet's - * two-operator physical tree. + * Spark strategy for Comet's split Iceberg V2 writer. WriteDelta is modeled and dispatched here, + * but position-delta rows are still executed by Iceberg's JVM DeltaWriter. */ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { @@ -44,38 +44,68 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { return Nil } // Planner strategies run before CometRule, so plan-only mode needs its own guard here. - if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) { + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf)) { return Nil } plan match { case ad: AppendData => - matchedSparkWrite(ad.table, ad.write, ad.query, replaceDataDispatch = None).toList + matchedSparkWrite(ad.table, ad.write, ad.query, PlainIcebergWrite).toList case obe: OverwriteByExpression => - matchedSparkWrite(obe.table, obe.write, obe.query, replaceDataDispatch = None).toList + matchedSparkWrite(obe.table, obe.write, obe.query, PlainIcebergWrite).toList case opd: OverwritePartitionsDynamic => - matchedSparkWrite(opd.table, opd.write, opd.query, replaceDataDispatch = None).toList + matchedSparkWrite(opd.table, opd.write, opd.query, PlainIcebergWrite).toList case rd: ReplaceData => matchedSparkWrite( rd.originalTable, rd.write, rd.query, - replaceDataDispatch = IcebergReplaceDataShim.extractProjections(rd)).toList - case plan if IcebergReflection.isReplaceIcebergData(plan) => + IcebergReplaceDataShim + .extractProjections(rd) + .map(ReplaceDataWrite) + .getOrElse(PlainIcebergWrite)).toList + case replace if IcebergReflection.isReplaceIcebergData(replace) => IcebergReflection - .extractReplaceIcebergDataFields(plan) + .extractReplaceIcebergDataFields(replace) .flatMap { case (_, query, originalTable, write) => matchedSparkWrite( originalTable.asInstanceOf[org.apache.spark.sql.catalyst.analysis.NamedRelation], write.asInstanceOf[Option[Write]], query.asInstanceOf[LogicalPlan], - replaceDataDispatch = None) + PlainIcebergWrite) } .toList // Hit by AQE. - case l @ IcebergWriteLogical(child, batchWrite, replaceDataDispatch) => - Seq(IcebergWriteExec(batchWrite, l.output, planLater(child), replaceDataDispatch)) - case _ => Nil + case l @ IcebergWriteLogical(child, batchWrite, dispatch) => + Seq(IcebergWriteExec(batchWrite, l.output, planLater(child), dispatch)) + case delta => + IcebergDeltaLogicalShim + .extract(delta) + .flatMap { fields => + fields.command.flatMap { command => + fields.write.flatMap { deltaWrite => + if (!IcebergReflection.isIcebergPositionDeltaWrite(deltaWrite)) { + None + } else { + WriteDeltaDispatchInfo + .build( + fields.projections, + fields.query.output, + IcebergDeltaWriterShim.OperationCodes, + Some(command)) + .flatMap { info => + buildDeltaTwoOp( + deltaWrite, + fields.originalTable, + fields.query, + PositionDeltaWrite(info), + Some(command)) + } + } + } + } + } + .toList } } @@ -83,12 +113,12 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { table: org.apache.spark.sql.catalyst.analysis.NamedRelation, write: Option[Write], query: LogicalPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + dispatch: IcebergWriteDispatch): Option[SparkPlan] = { table match { case rel: DataSourceV2Relation => write.flatMap { w => if (IcebergReflection.isIcebergSparkWrite(w)) { - buildTwoOp(w, rel, query, replaceDataDispatch) + buildTwoOp(w, rel, query, dispatch) } else { None } @@ -104,32 +134,66 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { * writer's child is wrapped in [[IcebergWriteLogical]] so AQE re-emits only the data-writing * operator on each re-plan as opposed to multiple new commit operators. * - * Iceberg's `SparkWrite` never asks for Spark's commit coordinator, so the - * `useCommitCoordinator` fallback below is defensive coverage in case a future Iceberg version - * changes that; the split writer's per-task commit protocol does not use it. + * Iceberg's current JVM writers do not ask for Spark's commit coordinator, so the + * `useCommitCoordinator` fallback is defensive coverage in case that changes; the split + * writer's per-task commit protocol does not use it. */ private def buildTwoOp( write: Write, rel: DataSourceV2Relation, query: LogicalPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + dispatch: IcebergWriteDispatch): Option[SparkPlan] = { val batchWrite = write.toBatch if (batchWrite.useCommitCoordinator()) { return None } - // To mirror Spark ReplaceData semantics we invalidate our cache of the state of - // `originalTable`. + // To mirror Spark's V2 write semantics, invalidate the cached state of `originalTable`. val refresh: () => Unit = () => IcebergRefreshCacheShim.refreshCache(session, rel) Some( - IcebergCommitPlanShim.wrap(IcebergCommitExec( - batchWrite, - write, - refresh, - // `replaceDataDispatch` may project the data into the format the writer expects. - planLater( - IcebergWriteLogical( - ShimCometMergeRows.withNativeMergeSummary(query), + IcebergCommitPlanShim.wrap( + IcebergCommitExec( + batchWrite, + write, + refresh, + planLater( + IcebergWriteLogical( + ShimCometMergeRows.withNativeMergeSummary(query), + batchWrite, + dispatch))))) + } + + /** + * Builds the same split commit/write tree for Iceberg's position-delta writer. The task-side + * writer remains Iceberg's JVM `DeltaWriter`; this only lets Comet own the split-plan + * lifecycle, AQE re-planning, reporting, and driver commit. + */ + private def buildDeltaTwoOp( + write: Write, + table: org.apache.spark.sql.catalyst.analysis.NamedRelation, + query: LogicalPlan, + dispatch: IcebergWriteDispatch, + command: Option[DeltaCommand]): Option[SparkPlan] = { + table match { + case rel: DataSourceV2Relation => + val batchWrite = write.toBatch + if (!IcebergReflection.isIcebergPositionDeltaBatchWrite(batchWrite) || + batchWrite.useCommitCoordinator()) { + None + } else { + val refresh: () => Unit = () => IcebergRefreshCacheShim.refreshCache(session, rel) + val commit = IcebergCommitExec( batchWrite, - replaceDataDispatch))))) + write, + refresh, + planLater( + IcebergWriteLogical( + ShimCometMergeRows.withNativeMergeSummary(query), + batchWrite, + dispatch)), + command) + Some(IcebergCommitPlanShim.wrap(commit)) + } + case _ => None + } } } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala b/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala new file mode 100644 index 00000000000..4aaf43aebec --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala @@ -0,0 +1,64 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.ProjectingInternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.util.{RowDeltaUtils, WriteDeltaProjections} +import org.apache.spark.sql.types.IntegerType + +case class DeltaOperationCodes(delete: Int, update: Int, insert: Int, reinsert: Option[Int]) + +/** + * Version-neutral task contract for Spark's operation-coded WriteDelta row stream. + * + * This intentionally contains only state consumed by Iceberg's JVM DeltaWriter path. Native + * position-delete schemas and metadata layouts belong to the native WriteDelta implementation. + */ +case class WriteDeltaDispatchInfo( + operationOrdinal: Int, + operationCodes: DeltaOperationCodes, + rowProjection: Option[ProjectingInternalRow], + rowIdProjection: ProjectingInternalRow, + metadataProjection: Option[ProjectingInternalRow], + command: Option[DeltaCommand] = None) + +private[iceberg] object WriteDeltaDispatchInfo { + + def build( + projections: WriteDeltaProjections, + childOutput: Seq[Attribute], + operationCodes: DeltaOperationCodes, + command: Option[DeltaCommand] = None): Option[WriteDeltaDispatchInfo] = { + childOutput.headOption.collect { + case attribute + if attribute.name == RowDeltaUtils.OPERATION_COLUMN && + attribute.dataType == IntegerType && + !attribute.nullable => + WriteDeltaDispatchInfo( + operationOrdinal = 0, + operationCodes = operationCodes, + rowProjection = projections.rowProjection, + rowIdProjection = projections.rowIdProjection, + metadataProjection = projections.metadataProjection, + command = command) + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index bb59b924d61..e55014812be 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, Icebe import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason -import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.iceberg.{IcebergReflection, PositionDeltaWrite, ReplaceDataWrite} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.Operator @@ -292,6 +292,12 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { } private def checkTriggers(op: IcebergWriteExec): Option[String] = { + op.dispatch match { + case PositionDeltaWrite(_) => + return Some("Iceberg WriteDelta executes through the JVM DeltaWriter") + case _ => + } + val batchWrite = op.batchWrite if (!IcebergReflection.isIcebergBatchWrite(batchWrite)) { return Some(s"not an Iceberg SparkWrite: ${batchWrite.getClass.getName}") @@ -935,13 +941,12 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { * giving it the wider row (e.g. 6 columns when the schema has 3) and * `decorate_batch_with_field_ids` rejects the batch. * - * For Spark 3.4 / 3.5 the strategy shim returns `None` for `replaceDataDispatch` and the - * upstream plan already projects to the data columns -- no extra projection needed. For 4.x we - * splice a `Projection` proto between our `IcebergWrite` op and the FFI `Scan`, selecting the - * upstream attributes whose names match the Iceberg schema's columns. The JVM-side child stays - * at the original wide output, so its `executeColumnar()` still emits the wide batches the FFI - * scan declares; the projection then strips them inside the native runtime before the writer - * sees the data. + * Plain writes already present only data columns, so no extra projection is needed. For 4.x + * ReplaceData we splice a `Projection` proto between our `IcebergWrite` op and the FFI `Scan`, + * selecting the upstream attributes whose names match the Iceberg schema's columns. The + * JVM-side child stays at the original wide output, so its `executeColumnar()` still emits the + * wide batches the FFI scan declares; the projection then strips them inside the native runtime + * before the writer sees the data. */ private def dropNonDataColumns( op: IcebergWriteExec, @@ -950,7 +955,10 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // write schema has no row lineage columns: when it has, Iceberg's writer reads their values // from the metadata columns (`ExtractRowLineage`), which this projection discards. // `requireNoMetadataColumns` declines those writes. - if (op.replaceDataDispatch.isEmpty) return Some(scan) + op.dispatch match { + case ReplaceDataWrite(_) => + case _ => return Some(scan) + } val sparkWrite = IcebergReflection.getOuterSparkWrite(op.batchWrite).getOrElse { withFallbackReason(op, "Could not unwrap outer SparkWrite for ReplaceData projection") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala index a2323e0acce..e066cd30ab3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.execution.{SparkPlan, SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.datasources.v2.V2CommandExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.comet.iceberg.{IcebergDriverMetricsShim, IcebergReflection} +import org.apache.comet.iceberg.{DeltaCommand, IcebergDriverMetricsShim, IcebergReflection} /** * Driver-side committer for Comet's split-operator Iceberg V2 write. @@ -38,7 +38,8 @@ case class IcebergCommitExec( @transient batchWrite: BatchWrite, @transient write: Write, @transient refreshCache: IcebergCommitExec.RefreshCache, - child: SparkPlan) + child: SparkPlan, + command: Option[DeltaCommand] = None) extends V2CommandExec with UnaryExecNode with Logging { @@ -95,9 +96,9 @@ case class IcebergCommitExec( cause) // The BatchWrite contract expects a job-level abort. Iceberg's own `SparkWrite.abort` // only deletes files after a cleanable *commit* failure and skips cleanup otherwise, so - // the data files of tasks that completed before the job failed would stay behind even - // though nothing can reference them (no commit was attempted). Delete them here; the - // failed task's own files are cleaned up by the task itself. + // the files of tasks that completed before the job failed would stay behind even though + // nothing can reference them (no commit was attempted). Delete newly written data/delete + // files here; the failed task's own files are cleaned up by the task itself. val failure = abortAfter(completed, cause) try deleteCompletedTaskFiles(completed) catch { @@ -110,7 +111,7 @@ case class IcebergCommitExec( try { messages.foreach(batchWrite.onDataWriterCommit) - IcebergWriteSummaryShim.commit(batchWrite, messages, child) + IcebergWriteSummaryShim.commit(batchWrite, messages, child, command) logInfo(s"Iceberg commit succeeded with ${messages.length} task message(s)") } catch { case cause: Throwable => @@ -143,7 +144,7 @@ case class IcebergCommitExec( if (locations.nonEmpty) { val io = IcebergReflection .getOuterSparkWrite(batchWrite) - .flatMap(IcebergReflection.getTableFromSparkWrite) + .flatMap(IcebergReflection.getTableFromIcebergWrite) .flatMap(IcebergReflection.getTableIO) io match { case Some(fileIO) => @@ -153,7 +154,7 @@ case class IcebergCommitExec( s"job abort, ${completed.length} completed task(s)") case None => logWarning( - s"Could not resolve the table FileIO; leaving ${locations.size} data file(s) from " + + s"Could not resolve the table FileIO; leaving ${locations.size} written file(s) from " + "completed tasks of a failed write job for remove_orphan_files") } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala index 90bfdaab048..e18ff4ac5dd 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.types.{BinaryType, StructField, StructType} import org.apache.spark.util.Utils -import org.apache.comet.iceberg.ReplaceDataDispatchInfo +import org.apache.comet.iceberg.{IcebergDeltaWriterShim, IcebergSemanticMetricsShim, IcebergWriteDispatch, PlainIcebergWrite, PositionDeltaWrite, ReplaceDataDispatchInfo, ReplaceDataWrite} /** * Executor-side file writer for Comet's split-operator Iceberg V2 write. @@ -41,7 +41,7 @@ case class IcebergWriteExec( @transient batchWrite: BatchWrite, override val output: Seq[Attribute], child: SparkPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + dispatch: IcebergWriteDispatch = PlainIcebergWrite) extends UnaryExecNode { // Spark already adds a distribution for the V2 write; adding another here is redundant. @@ -51,7 +51,12 @@ case class IcebergWriteExec( // through the driver-side commit report, and its DataWriters expose no task-level values. // Registering them here too would double-count them in the SQL UI. override lazy val metrics: Map[String, SQLMetric] = - Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) ++ + (dispatch match { + case PositionDeltaWrite(info) => + IcebergSemanticMetricsShim.deltaMetrics(sparkContext, info.command) + case _ => Map.empty[String, SQLMetric] + }) override protected def doExecute(): RDD[InternalRow] = { val rdd = { @@ -72,7 +77,8 @@ case class IcebergWriteExec( val rowsMetric = longMetric("numOutputRows") val schemaTypes = output.map(_.dataType).toArray - val capturedReplaceDataDispatch = replaceDataDispatch + val capturedDispatch = dispatch + val semanticMetrics = metrics -- Set("numOutputRows") rdd.mapPartitionsInternal { iter => val partId = TaskContext.getPartitionId() val taskId = TaskContext.get().taskAttemptId() @@ -83,7 +89,8 @@ case class IcebergWriteExec( iter, rowsMetric, projection, - capturedReplaceDataDispatch) + capturedDispatch, + semanticMetrics) } } @@ -106,17 +113,40 @@ object IcebergWriteExec { iter: Iterator[InternalRow], rowsMetric: SQLMetric, projection: UnsafeProjection, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + dispatch: IcebergWriteDispatch, + semanticMetrics: Map[String, SQLMetric] = Map.empty): Iterator[InternalRow] = { val iterWithMetrics = new IteratorWithMetrics(iter, rowsMetric) // Serialization happens inside the guarded block: if the commit message cannot be // serialised the task must abort so the already-finalised data files get deleted. val messageBytes = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { - if (replaceDataDispatch.isDefined) { - runReplaceDataWriter(writer, iterWithMetrics, replaceDataDispatch.get) - } else { - while (iterWithMetrics.hasNext) { - writer.write(iterWithMetrics.next()) - } + dispatch match { + case PlainIcebergWrite => + while (iterWithMetrics.hasNext) writer.write(iterWithMetrics.next()) + case ReplaceDataWrite(info) => + runReplaceDataWriter(writer, iterWithMetrics, info) + case PositionDeltaWrite(info) => + val deltaWriter = writer match { + case delta: org.apache.spark.sql.connector.write.DeltaWriter[_] => + delta.asInstanceOf[org.apache.spark.sql.connector.write.DeltaWriter[InternalRow]] + case other => + throw new IllegalStateException( + s"Expected DeltaWriter for position delta write, got ${other.getClass.getName}") + } + var numUpdatedRows = 0L + var numDeletedRows = 0L + while (iterWithMetrics.hasNext) { + val row = iterWithMetrics.next() + val operation = row.getInt(info.operationOrdinal) + if (operation == info.operationCodes.delete) { + numDeletedRows += 1L + } else if (operation == info.operationCodes.update || + info.operationCodes.reinsert.contains(operation)) { + numUpdatedRows += 1L + } + IcebergDeltaWriterShim.writeOperation(deltaWriter, operation, row, info) + } + semanticMetrics.get("numUpdatedRows").foreach(_.add(numUpdatedRows)) + semanticMetrics.get("numDeletedRows").foreach(_.add(numDeletedRows)) } serializeMessage(writer.commit()) })( diff --git a/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..afbf7d60b70 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -0,0 +1,27 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + +/** WriteDelta interception starts with Spark 3.5. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = None +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..d85d0c19fd8 --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -0,0 +1,47 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 3.5 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala new file mode 100644 index 00000000000..c652eb3558a --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala @@ -0,0 +1,65 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{DELETE_OPERATION, INSERT_OPERATION, UPDATE_OPERATION} +import org.apache.spark.sql.connector.write.DeltaWriter + +/** Spark 3.x supports DELETE, UPDATE, and INSERT operation codes in WriteDelta. */ +object IcebergDeltaWriterShim { + val OperationCodes: DeltaOperationCodes = + DeltaOperationCodes( + delete = DELETE_OPERATION, + update = UPDATE_OPERATION, + insert = INSERT_OPERATION, + reinsert = None) + + def writeOperation( + writer: DeltaWriter[InternalRow], + operation: Int, + row: InternalRow, + info: WriteDeltaDispatchInfo): Unit = { + def metadata(): InternalRow = info.metadataProjection.map { projection => + projection.project(row) + projection + }.orNull + + operation match { + case code if code == info.operationCodes.delete => + info.rowIdProjection.project(row) + writer.delete(metadata(), info.rowIdProjection) + case code if code == info.operationCodes.update => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("UPDATE operation has no row projection")) + data.project(row) + info.rowIdProjection.project(row) + writer.update(metadata(), info.rowIdProjection, data) + case code if code == info.operationCodes.insert => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("INSERT operation has no row projection")) + data.project(row) + writer.insert(data) + case _ => + throw new SparkException(s"Unexpected operation ID: $operation") + } + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index d6a78ca2ea4..135384f5d9f 100644 --- a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -22,12 +22,15 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.SparkPlan +import org.apache.comet.iceberg.DeltaCommand + /** The `BatchWrite.commit(messages, summary)` overload only exists on Spark 4.1+. */ private[comet] object IcebergWriteSummaryShim { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { batchWrite.commit(messages) } } diff --git a/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..ea2ceba1823 --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.0 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index d6a78ca2ea4..135384f5d9f 100644 --- a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -22,12 +22,15 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.SparkPlan +import org.apache.comet.iceberg.DeltaCommand + /** The `BatchWrite.commit(messages, summary)` overload only exists on Spark 4.1+. */ private[comet] object IcebergWriteSummaryShim { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { batchWrite.commit(messages) } } diff --git a/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index f3ea21e3171..af6c834bde1 100644 --- a/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec -import org.apache.comet.shims.MergeRowsMetricsShim +import org.apache.comet.iceberg.{DeltaCommand, IcebergSemanticMetricsShim} /** * Spark 4.1 collects MERGE metrics from the executed plan and hands them to @@ -34,27 +34,30 @@ private[comet] object IcebergWriteSummaryShim extends AdaptiveSparkPlanHelper { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { - collectFirst(query) { - case m: MergeRowsExec => m.metrics - case m: CometMergeRowsExec => m.metrics - } match { - case Some(metrics) => - def metricValue(name: String): Long = - metrics.get(name).map(MergeRowsMetricsShim.value).getOrElse(-1L) - batchWrite.commit( - messages, - MergeSummaryImpl( - metricValue("numTargetRowsCopied"), - metricValue("numTargetRowsDeleted"), - metricValue("numTargetRowsUpdated"), - metricValue("numTargetRowsInserted"), - metricValue("numTargetRowsMatchedUpdated"), - metricValue("numTargetRowsMatchedDeleted"), - metricValue("numTargetRowsNotMatchedBySourceUpdated"), - metricValue("numTargetRowsNotMatchedBySourceDeleted"))) - case None => - batchWrite.commit(messages) + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { + if (!IcebergDeltaWriteSummaryShim.commit(batchWrite, messages, query, command)) { + collectFirst(query) { + case m: MergeRowsExec => m.metrics + case m: CometMergeRowsExec => m.metrics + } match { + case Some(metrics) => + def metricValue(name: String): Long = + metrics.get(name).map(IcebergSemanticMetricsShim.value).getOrElse(-1L) + batchWrite.commit( + messages, + MergeSummaryImpl( + metricValue("numTargetRowsCopied"), + metricValue("numTargetRowsDeleted"), + metricValue("numTargetRowsUpdated"), + metricValue("numTargetRowsInserted"), + metricValue("numTargetRowsMatchedUpdated"), + metricValue("numTargetRowsMatchedDeleted"), + metricValue("numTargetRowsNotMatchedBySourceUpdated"), + metricValue("numTargetRowsNotMatchedBySourceDeleted"))) + case None => + batchWrite.commit(messages) + } } } } diff --git a/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..8aab1f72b57 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.1 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala new file mode 100644 index 00000000000..36d70e65670 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala @@ -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. + */ + +package org.apache.spark.sql.comet + +import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} +import org.apache.spark.sql.execution.SparkPlan + +import org.apache.comet.iceberg.DeltaCommand + +/** Spark 4.1 has MERGE summaries; UPDATE and DELETE summaries arrive with Spark 4.2. */ +private[comet] object IcebergDeltaWriteSummaryShim { + def commit( + batchWrite: BatchWrite, + messages: Array[WriterCommitMessage], + query: SparkPlan, + command: Option[DeltaCommand]): Boolean = false +} diff --git a/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..8964df79ab2 --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.2 exposes the row-level fields directly; keep this extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..4808c4ebfb2 --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,45 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.{SQLLastAttemptMetric, SQLLastAttemptMetrics, SQLMetric} + +/** Spark 4.2 uses last-attempt accumulators for row-level write summaries. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + command match { + case Some(DeltaUpdate) => + Map( + "numUpdatedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of updated rows"), + "numCopiedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of copied rows")) + case Some(DeltaDelete) => + Map( + "numDeletedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of deleted rows"), + "numCopiedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of copied rows")) + case _ => Map.empty + } + + def value(metric: SQLMetric): Long = metric match { + case lastAttempt: SQLLastAttemptMetric => + lastAttempt.lastAttemptValueForHighestRDDId().getOrElse(lastAttempt.value) + case other => other.value + } +} diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala new file mode 100644 index 00000000000..72879e35ebc --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala @@ -0,0 +1,58 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import org.apache.spark.sql.connector.write.{BatchWrite, DeleteSummaryImpl, UpdateSummaryImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.metric.SQLMetric + +import org.apache.comet.iceberg.{DeltaCommand, DeltaDelete, DeltaUpdate, IcebergSemanticMetricsShim} + +/** Spark 4.2 DELETE/UPDATE summary handling for the JVM PositionDelta writer. */ +private[comet] object IcebergDeltaWriteSummaryShim extends AdaptiveSparkPlanHelper { + def commit( + batchWrite: BatchWrite, + messages: Array[WriterCommitMessage], + query: SparkPlan, + command: Option[DeltaCommand]): Boolean = command match { + case Some(DeltaUpdate) => + deltaWriterMetrics(query).exists { metrics => + batchWrite.commit( + messages, + UpdateSummaryImpl(value(metrics, "numUpdatedRows"), value(metrics, "numCopiedRows"))) + true + } + case Some(DeltaDelete) => + deltaWriterMetrics(query).exists { metrics => + batchWrite.commit( + messages, + DeleteSummaryImpl(value(metrics, "numDeletedRows"), value(metrics, "numCopiedRows"))) + true + } + case _ => false + } + + private def deltaWriterMetrics(query: SparkPlan): Option[Map[String, SQLMetric]] = + collectFirst(query) { case writer: IcebergWriteExec => writer.metrics } + + private def value(metrics: Map[String, SQLMetric], name: String): Long = + metrics.get(name).map(IcebergSemanticMetricsShim.value).getOrElse(-1L) +} diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala index 7dec0be1d23..a749217393d 100644 --- a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala +++ b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala @@ -33,7 +33,8 @@ private[org] final class IcebergTransactionalCommitExec private ( delegate.batchWrite, delegate.write, delegate.refreshCache, - delegate.child) + delegate.child, + delegate.command) with TransactionalExec { override def withTransaction(txn: Option[Transaction]): SparkPlan = copyWith(transaction = txn) @@ -62,13 +63,14 @@ private[org] final class IcebergTransactionalCommitExec private ( } override def makeCopy(newArgs: Array[AnyRef]): IcebergTransactionalCommitExec = { - require(newArgs.length == 4, s"expected 4 IcebergCommitExec arguments, got ${newArgs.length}") + require(newArgs.length == 5, s"expected 5 IcebergCommitExec arguments, got ${newArgs.length}") val copied = new IcebergTransactionalCommitExec( IcebergCommitExec( newArgs(0).asInstanceOf[org.apache.spark.sql.connector.write.BatchWrite], newArgs(1).asInstanceOf[org.apache.spark.sql.connector.write.Write], newArgs(2).asInstanceOf[IcebergCommitExec.RefreshCache], - newArgs(3).asInstanceOf[SparkPlan]), + newArgs(3).asInstanceOf[SparkPlan], + newArgs(4).asInstanceOf[Option[org.apache.comet.iceberg.DeltaCommand]]), transaction) copied.copyTagsFrom(this) copied @@ -85,7 +87,7 @@ private[org] final class IcebergTransactionalCommitExec private ( child: SparkPlan = this.child, transaction: Option[Transaction] = this.transaction): IcebergTransactionalCommitExec = new IcebergTransactionalCommitExec( - IcebergCommitExec(batchWrite, write, refreshCache, child), + IcebergCommitExec(batchWrite, write, refreshCache, child, command), transaction) } diff --git a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..43354cbec61 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + +/** Spark 4.x uses the version-neutral WriteDelta descriptor for Iceberg row-level writes. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + private val writeDeltaClass = "org.apache.spark.sql.catalyst.plans.logical.WriteDelta" + + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = + if (plan.getClass.getName == writeDeltaClass) { + IcebergDeltaLogicalFieldsShim.extract(plan) + } else { + None + } +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala new file mode 100644 index 00000000000..ceb6ff18460 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala @@ -0,0 +1,70 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{DELETE_OPERATION, INSERT_OPERATION, REINSERT_OPERATION, UPDATE_OPERATION} +import org.apache.spark.sql.connector.write.DeltaWriter + +/** Spark 4.x adds REINSERT to the WriteDelta operation stream. */ +object IcebergDeltaWriterShim { + val OperationCodes: DeltaOperationCodes = + DeltaOperationCodes( + delete = DELETE_OPERATION, + update = UPDATE_OPERATION, + insert = INSERT_OPERATION, + reinsert = Some(REINSERT_OPERATION)) + + def writeOperation( + writer: DeltaWriter[InternalRow], + operation: Int, + row: InternalRow, + info: WriteDeltaDispatchInfo): Unit = { + def metadata(): InternalRow = info.metadataProjection.map { projection => + projection.project(row) + projection + }.orNull + + operation match { + case code if code == info.operationCodes.delete => + info.rowIdProjection.project(row) + writer.delete(metadata(), info.rowIdProjection) + case code if code == info.operationCodes.update => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("UPDATE operation has no row projection")) + data.project(row) + info.rowIdProjection.project(row) + writer.update(metadata(), info.rowIdProjection, data) + case code if code == info.operationCodes.insert => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("INSERT operation has no row projection")) + data.project(row) + writer.insert(data) + case code if info.operationCodes.reinsert.contains(code) => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("REINSERT operation has no row projection")) + data.project(row) + writer.reinsert(metadata(), data) + case _ => + throw new SparkException(s"Unexpected operation ID: $operation") + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 0112a025268..45a4fe597fd 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -56,8 +56,8 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.Trigger import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} -import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} -import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener, PositionDeltaWrite} import org.apache.comet.serde.Unsupported import org.apache.comet.serde.operator.CometIcebergNativeWrite @@ -744,6 +744,134 @@ class CometIcebergWriteActionSuite } } + Seq("DELETE", "UPDATE", "MERGE").foreach { command => + test(s"WriteDelta $command is recognized and stays on Iceberg JVM DeltaWriter") { + assume(icebergAvailable, "Iceberg not available in classpath") + assume(isSpark35Plus, "WriteDelta interception starts with Spark 3.5") + withIcebergCatalog { warehouseDir => + val suffix = command.toLowerCase(java.util.Locale.ROOT) + val cometTable = s"write_delta_${suffix}_comet" + val jvmTable = s"write_delta_${suffix}_jvm" + val properties = + "'format-version'='2', 'write.delete.mode'='merge-on-read', " + + "'write.update.mode'='merge-on-read', 'write.merge.mode'='merge-on-read'" + + Seq(cometTable, jvmTable).foreach { table => + createTable( + warehouseDir, + table, + partitionSpec = "PARTITIONED BY (region)", + properties = Some(properties)) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + // Keep ids 1 and 2 in the same Iceberg data file so DELETE id = 1 cannot be + // satisfied as a metadata/file delete and must exercise the WriteDelta path. + coalesceInsert(table, Seq((1, "a", 10.0), (2, "a", 20.0), (3, "c", 30.0))) + } + } + + def run(table: String): Unit = command match { + case "DELETE" => + spark.sql(s"DELETE FROM $catalog.$ns.$table WHERE id = 1") + case "UPDATE" => + spark.sql( + s"UPDATE $catalog.$ns.$table SET region = 'updated', amount = amount + 5 " + + "WHERE id = 2") + case "MERGE" => + spark.sql(s""" + |MERGE INTO $catalog.$ns.$table t + |USING ( + | SELECT 2 AS id, 'updated' AS region, 200.0 AS amount + | UNION ALL + | SELECT 4 AS id, 'new' AS region, 40.0 AS amount + |) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.region = s.region, t.amount = s.amount + |WHEN NOT MATCHED THEN + | INSERT (id, region, amount) VALUES (s.id, s.region, s.amount) + |""".stripMargin) + } + + val jvmBefore = countSnapshots(jvmTable) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + run(jvmTable) + } + assert(countSnapshots(jvmTable) - jvmBefore == 1L) + + val snapshot = withNativeEnabled { + captureWrite(cometTable) { + run(cometTable) + } + } + assert(snapshot.snapshotDelta == 1L, s"unexpected snapshot count: $snapshot") + + val deltaWrites = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { + case writer: IcebergWriteExec if writer.dispatch.isInstanceOf[PositionDeltaWrite] => + writer + } + } + assert( + deltaWrites.nonEmpty, + s"expected PositionDeltaWrite dispatch. Plans:\n${snapshot.plans.mkString("\n--\n")}") + assert( + deltaWrites.forall { writer => + IcebergReflection + .getOuterSparkWrite(writer.batchWrite) + .flatMap(IcebergReflection.getTableFromIcebergWrite) + .isDefined + }, + "expected PositionDelta batch writes to expose their Iceberg table") + + val nativeWrites = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case writer: CometIcebergWriteExec => writer } + } + assert( + nativeWrites.isEmpty, + "WriteDelta must stay on Iceberg's JVM DeltaWriter. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + + def rows(table: String): Seq[Row] = + spark + .sql(s"SELECT id, region, amount FROM $catalog.$ns.$table ORDER BY id") + .collect() + .toSeq + + assert( + rows(cometTable) == rows(jvmTable), + s"$command result differs from the stock Iceberg JVM path") + + val summaryPrefix = command match { + case "DELETE" => "spark.delete." + case "UPDATE" => "spark.update." + case "MERGE" => "spark.merge-into." + } + def rowLevelSummary(table: String): Map[String, String] = + spark + .sql(s"SELECT summary FROM $catalog.$ns.$table.snapshots " + + "ORDER BY committed_at DESC LIMIT 1") + .collect()(0) + .getMap[String, String](0) + .filter { case (key, _) => key.startsWith(summaryPrefix) } + .toMap + + val cometSummary = rowLevelSummary(cometTable) + val jvmSummary = rowLevelSummary(jvmTable) + assert( + cometSummary == jvmSummary, + s"$command snapshot summary differs from the stock Iceberg JVM path: " + + s"Comet=$cometSummary JVM=$jvmSummary") + val summaryExpected = + (command == "MERGE" && isSpark41Plus) || + ((command == "DELETE" || command == "UPDATE") && isSpark42Plus) + if (summaryExpected) { + assert( + cometSummary.nonEmpty, + s"expected Spark row-level summary metrics for $command, got $cometSummary") + } + } + } + } + test("sanity check: Spark's default DELETE path works against a Hadoop catalog") { assume(icebergAvailable, "Iceberg not available in classpath") withIcebergCatalog { warehouseDir => @@ -1490,39 +1618,30 @@ class CometIcebergWriteActionSuite val commits = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { case c: IcebergCommitExec => c } } - if (mode == "copy-on-write") { - assert( - commits.nonEmpty, - s"expected >= 1 IcebergCommitExec for $mode, got 0. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } else { - assert( - commits.isEmpty, - s"merge-on-read should stay on Iceberg WriteDelta, got ${commits.size} IcebergCommitExec. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } - val mergeExecs = writeSnapshot.plans.flatMap { plan => - collectWithSubqueries(plan) { case e: CometMergeRowsExec => e } - } - if (mode == "merge-on-read" && isSpark41Plus) { - val sparkMergeRows = writeSnapshot.plans.flatMap { plan => + assert( + commits.nonEmpty, + s"expected >= 1 IcebergCommitExec for $mode, got 0. Plans:\n" + + writeSnapshot.plans.mkString("\n--\n")) + if (mode == "merge-on-read") { + val deltaWrites = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { - case e if e.getClass.getSimpleName == "MergeRowsExec" => e + case writer: IcebergWriteExec + if writer.dispatch.isInstanceOf[PositionDeltaWrite] => + writer } } assert( - mergeExecs.isEmpty, - "Spark 4.1+ merge-on-read uses the stock V2 writer and must retain MergeRowsExec") - assert( - sparkMergeRows.nonEmpty, - s"expected Spark MergeRowsExec for Spark 4.1+ $mode. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } else { - assert( - mergeExecs.nonEmpty, - s"expected CometMergeRowsExec for $mode. Plans:\n" + + deltaWrites.nonEmpty, + s"expected PositionDeltaWrite dispatch for $mode. Plans:\n" + writeSnapshot.plans.mkString("\n--\n")) } + val mergeExecs = writeSnapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case e: CometMergeRowsExec => e } + } + assert( + mergeExecs.nonEmpty, + s"expected CometMergeRowsExec for $mode. Plans:\n" + + writeSnapshot.plans.mkString("\n--\n")) val nativeWrites = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e } @@ -4087,6 +4206,16 @@ class CometIcebergWriteActionSuite "report_orc", partitionSpec = "", properties = Some("'write.format.default'='orc'")) + if (isSpark35Plus) { + createTable( + warehouseDir, + "report_delta", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'format-version'='2', 'write.delete.mode'='merge-on-read'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert("report_delta", Seq((1, "us-east", 1.5), (2, "us-east", 2.5))) + } + } val writes = reportedWrites { withNativeEnabled { @@ -4096,14 +4225,19 @@ class CometIcebergWriteActionSuite .sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (1, 'us-east', 1.5)") .collect() spark.sql(s"INSERT INTO $catalog.$ns.report_orc VALUES (2, 'eu', 2.5)") + if (isSpark35Plus) { + spark.sql(s"DELETE FROM $catalog.$ns.report_delta WHERE id = 1") + } } withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { spark.sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (3, 'eu', 3.5)") } } - assert(writes.map(_.writer) == Seq("native", "jvm", "spark"), writes.mkString("\n")) + val expectedWriters = + Seq("native", "jvm") ++ (if (isSpark35Plus) Seq("jvm") else Nil) ++ Seq("spark") + assert(writes.map(_.writer) == expectedWriters, writes.mkString("\n")) assert(writes(1).reasons.exists(_.contains("only parquet"))) - assert(writes(2).node == "AppendData") + assert(writes.last.node == "AppendData") assert(writes.forall(!_.failed)) } } diff --git a/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala b/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala index e7bff634db4..f5f3a75f493 100644 --- a/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala @@ -138,6 +138,27 @@ class IcebergReflectionSuite extends AnyFunSuite { assert(ex.getCause.getMessage == "boom") } + test("taskCommitFileLocations reads plain SparkWrite task commits") { + val commit = new PlainTaskCommit( + Array( + new LocationFile("s3://bucket/data/a.parquet"), + new LocationFile("s3://bucket/data/b.parquet"))) + assert( + IcebergReflection.taskCommitFileLocations(commit) == + Seq("s3://bucket/data/a.parquet", "s3://bucket/data/b.parquet")) + } + + test("taskCommitFileLocations reads new files from position-delta task commits") { + val commit = new DeltaTaskCommit( + Array(new LocationFile("s3://bucket/data/a.parquet")), + Array(new LocationFile("s3://bucket/delete/d.parquet")), + Array(new LocationFile("s3://bucket/delete/old.parquet"))) + + assert( + IcebergReflection.taskCommitFileLocations(commit) == + Seq("s3://bucket/data/a.parquet", "s3://bucket/delete/d.parquet")) + } + test("getFileFormat reads format() when declared") { val file = new FormatFile("PARQUET") assert(IcebergReflection.getFileFormat(classOf[FormatFile], file) == Some("PARQUET")) @@ -309,6 +330,19 @@ class IcebergReflectionSuite extends AnyFunSuite { assert(spec.partitionType().fields().get(0).`type`().toString == "string") } + class PlainTaskCommit(taskFiles: Array[LocationFile]) { + def files(): Array[LocationFile] = taskFiles + } + + class DeltaTaskCommit( + data: Array[LocationFile], + deletes: Array[LocationFile], + rewrittenDeletes: Array[LocationFile]) { + def dataFiles(): Array[LocationFile] = data + def deleteFiles(): Array[LocationFile] = deletes + def rewrittenDeleteFiles(): Array[LocationFile] = rewrittenDeletes + } + /** Mimics a newer Iceberg ContentFile, which exposes location(). */ class LocationFile(loc: String) { def location(): String = loc diff --git a/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala b/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala index adb0b05f8c5..64d78504933 100644 --- a/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala +++ b/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala @@ -42,7 +42,8 @@ object IcebergTransactionalCommitTestShim { null.asInstanceOf[BatchWrite], null.asInstanceOf[Write], () => (), - LocalTableScanExec(Nil, Nil, None)) + LocalTableScanExec(Nil, Nil, None), + command = Some(DeltaUpdate)) val original = IcebergCommitPlanShim .wrap(commit) .asInstanceOf[IcebergTransactionalCommitExec] @@ -58,6 +59,9 @@ object IcebergTransactionalCommitTestShim { val attached = transformed.asInstanceOf[IcebergTransactionalCommitExec] assert(attached.transaction.contains(transaction)) - assert(attached.clone().transaction.contains(transaction)) + assert(attached.command.contains(DeltaUpdate)) + val cloned = attached.clone() + assert(cloned.transaction.contains(transaction)) + assert(cloned.command.contains(DeltaUpdate)) } } From e038cf05fa665dc61fdecbdde1c07b8b6a1ce3cb Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 7 Oct 2026 10:17:45 +0000 Subject: [PATCH 21/89] feat: use native shuffle for row input by converting it to Arrow (#6607) * feat: use native shuffle for row input by converting it to Arrow When a shuffle's child is a Spark row-based operator, Comet uses its JVM columnar shuffle. With the new spark.comet.convert.shuffleInput.enabled (off by default), CometExecRule instead puts CometSparkToColumnarExec over the child and uses native shuffle, where native shuffle supports the partitioning and the conversion supports the columns. A shuffle that hashes a decimal wider than 18 digits stays on the JVM columnar shuffle, because native shuffle does not hash those as Spark does. As for the typed Dataset conversion, the child's subtree gets its columnar transitions from Spark's own rule first, since Spark inserts none below a RowToColumnarTransition. The revert of a shuffle between two Spark aggregates covers the new shape too. * test: run CometShuffleInputConversionSuite in the PR builds * fix: keep a shuffle that hashes a string on the JVM columnar shuffle Spark's partitioner hashes a string's bytes as they are. Once CometSparkToColumnarExec has converted the rows, the import into native replaces invalid UTF-8 before native shuffle hashes the string, so such a key can go to a different partition than Spark's partitioner sends it to. A join whose other input stays on the JVM columnar shuffle then loses the matches for that key. The conversion now leaves a shuffle that hashes a string on the JVM columnar shuffle, as it does one that hashes a wide decimal. * fix: have native shuffle read a converted child's Arrow stream A native shuffle over a CometSparkToColumnarExec wrapped the child's executeColumnar batches in a ColumnarBatchArrowReader, which closes each batch once native has it. Those batches share the vectors that the conversion reuses for the next batch, and closing a struct vector drops its children, so the next batch failed to import. Native shuffle now reads such a child's Arrow stream, as a native operator does, which also fixes the same failure with spark.comet.sparkToColumnar.enabled. * fix: leave calendar intervals out of the shuffle input conversion Arrow holds the time part of an interval in nanoseconds, so IntervalMonthDayNanoWriter overflows on a calendar interval with more microseconds than that can hold, which Spark accepts. The JVM columnar shuffle leaves calendar intervals to Spark's shuffle, and the conversion now does too. * test: turn off RDD scan conversion with its own config after the merge of main Since #6602, CometTestBase turns on the conversion of RDD scans with spark.comet.convert.rdd.enabled rather than spark.comet.sparkToColumnar.enabled, so setting the old key to false no longer kept the RDD scans in CometShuffleInputConversionSuite on Spark. The scans were converted before the shuffle input conversion saw them, and six tests failed. The suite now turns off the conversions CometTestBase turns on, and the struct test and the benchmark use the RDD scan config. * fix: keep a shuffle whose key is computed from a string on the JVM columnar shuffle The string guard looked only at the type of each hash key, so a key that is a number computed from a string, such as hash(s), passed it. Native shuffle evaluates such a key after the import into native has replaced invalid UTF-8 in the string, so it can still send a row to a different partition than Spark's partitioner does, and a join whose other input stays on the JVM columnar shuffle loses the matches. The guard now looks for a string anywhere in the key expression. A wide decimal still matters only as the type of the key itself, because Comet's hash of a wide decimal already falls back to Spark. The string join test also joins on the hash of the strings, and uses two keys instead of four. Both shuffles write the strings with invalid UTF-8 replaced, so the keys stay apart only while each key expression sends them to different partitions, and two of the four hashes share one. * fix: put a converted shuffle back on the JVM columnar shuffle when its stage is reverted RevertNativeForTransitionHeavyStages strips CometSparkToColumnarExec with the rest of a reverted stage, which left a native shuffle reading the batches of Spark's RowToColumnarExec, and the task failed casting them to CometVector. The shuffle now goes back to the JVM columnar shuffle that the conversion replaced, which reads the stage's rows. * refactor: convert only the input of a shuffle that would use the JVM columnar shuffle convertsInputForNativeShuffle now starts from shuffleSupported, whose checks it restated, and which also rules out Celeborn and calendar intervals. With spark.comet.shuffle.mode=native, a Spark child keeps Spark's shuffle, as the config doc says. * perf: sample a conversion's input rows for range partitioning The range partitioner reads only the sort keys, so it samples the rows the conversion reads rather than converting all of their columns to Arrow first. The benchmark gets a range-partitioned case. * test: check where native shuffle puts each key type the conversion admits Checks the partition of every row and a join against the JVM columnar shuffle for each hash key type the conversion admits, with NULL and boundary values. Adds array and map columns with NULLs to the test over several batches, and plan assertions to the test of aggregates whose partial aggregate runs on Spark. * docs: describe the shuffle input conversion in the contributor guide Also corrects the revert's log message in the tuning guide. * test: count a range-partitioned conversion's rows once, and range-partition every column in the benchmark Before the sampling change, the range partitioner's sampling job ran through the conversion and added its rows to the conversion's metrics, so a range-partitioned shuffle reported twice the rows it wrote. * docs: say that the conversion gains with hash partitioning and comes out even with range partitioning Also rewords the contributor guide's rule for when native shuffle takes a Spark child. * test: correct why the shuffle between the two sort aggregates stays native Since #4565, the revert of a shuffle between two Spark aggregates takes sort aggregates too, so the reason the test gave, that the revert covers hash aggregates only, no longer held. The shuffle stays native because the final sort aggregate reads it through a sort. * docs: name sort aggregates in the tuning guide's shuffle revert section Since #4565, the revert of a shuffle between two Spark aggregates takes SortAggregateExec too, but only where the final aggregate reads the shuffle directly. One with grouping keys reads it through a sort, so that shuffle is not reverted. The last sentence also no longer calls the shuffle that disabling the revert keeps a columnar shuffle, since with the shuffle input conversion it can be a native one. * docs: say that the shuffle revert takes any Spark aggregate Since #4565, the revert of a shuffle between two Spark aggregates matches any BaseAggregateExec, not only hash aggregates, so the config doc and the scaladoc of revertRedundantColumnarShuffle no longer name hash aggregates alone. The scaladoc also says why the shuffle below a sort aggregate with grouping keys stays. * test: check the Celeborn exclusion of the shuffle input conversion in every shuffle mode In native mode the JVM columnar shuffle is off, so the conversion stays off with or without the Celeborn checks. Run the test in auto and jvm too, where Celeborn is what keeps the Spark child on Spark's shuffle, and check that the fallback reason names Celeborn, which in native mode only the Celeborn check in shuffleSupported gives. * fix: keep native shuffle over a stacked bridge when reverting a transition-heavy stage --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + docs/source/contributor-guide/jvm_shuffle.md | 6 +- .../contributor-guide/native_shuffle.md | 21 +- .../user-guide/latest/tuning/shuffle.md | 25 +- .../scala/org/apache/comet/CometConf.scala | 30 +- .../apache/comet/rules/CometExecRule.scala | 82 +++- ...RevertNativeForTransitionHeavyStages.scala | 49 +- .../shuffle/CometShuffleExchangeExec.scala | 105 +++-- .../CometShuffleInputConversionSuite.scala | 428 ++++++++++++++++++ ...CometShuffleInputConversionBenchmark.scala | 192 ++++++++ .../CometCelebornShufflePlanningSuite.scala | 20 + 12 files changed, 877 insertions(+), 83 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index d6d487eb146..69f0fa89612 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -486,6 +486,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometShuffleInputConversionSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index bf6a86b8f42..0b525419943 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -190,6 +190,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometShuffleInputConversionSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 11014f34909..dca7144a01a 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -43,8 +43,10 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE 1. **Shuffle mode is explicitly set to "jvm"**: When `spark.comet.shuffle.mode` is set to `jvm`. 2. **Child plan is not a Comet native operator**: When the child plan is a Spark row-based operator - (not a `CometPlan`), JVM shuffle is the only option since native shuffle requires columnar input - from Comet operators. + (not a `CometPlan`), JVM shuffle is used, since native shuffle requires columnar input from + Comet operators. The exception is `spark.comet.convert.shuffleInput.enabled`, which converts + the child's rows to Arrow with `CometSparkToColumnarExec` so that native shuffle can take the + shuffle instead. See [When Native Shuffle is Used](native_shuffle.md#when-native-shuffle-is-used). 3. **Unsupported partition key types**: `RangePartitioning` keys must be primitive, so a complex range key always falls back here. `HashPartitioning` keys must be primitive only by default: diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index e9cf7774415..704fa2a0ea7 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -52,7 +52,13 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 1. **Shuffle mode allows native**: `spark.comet.shuffle.mode` is `native` or `auto`. 2. **Child plan is a Comet native operator**: The child must be a `CometPlan` that produces - columnar output. Row-based Spark operators require JVM shuffle. + columnar output. Row-based Spark operators require JVM shuffle, unless + `spark.comet.convert.shuffleInput.enabled` is set. Then a shuffle that JVM shuffle would take + uses native shuffle instead, provided native shuffle supports the partitioning and the + columns, and a `CometSparkToColumnarExec` converts the child's rows to Arrow for it. A shuffle + that hashes a string, a value computed from a string, or a decimal wider than 18 digits stays + on JVM shuffle, because native shuffle would not put every row in the partition Spark's + partitioner does (`CometShuffleExchangeExec.convertsInputForNativeShuffle`). 3. **Supported partitioning type**: Native shuffle supports: - `HashPartitioning` @@ -166,11 +172,14 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet- - The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a `CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator` as the writer, with no JVM-to-native batch boundary between them. - - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the dep was built via the - convenience `prepareShuffleDependency(rdd, ...)` overload (used by - `CometCollectLimitExec` and `CometTakeOrderedAndProjectExec`, or when the - exchange's child is a non-native `CometPlan` such as `CometSparkToColumnarExec`). Native - code reads `ColumnarBatch`es from the JVM input iterator via Arrow C Stream Interface. + - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the exchange's child is a + non-native `CometPlan`, and for the dependencies that `CometCollectLimitExec` and + `CometTakeOrderedAndProjectExec` build. Native code reads the input through the Arrow C + Stream Interface. A `CometNativeArrowSource` child, such as `CometSparkToColumnarExec`, + exports its own Arrow stream from `doExecuteAsArrowStream()`, which + `prepareArrowStreamShuffleDependency` hands to native. Any other input is an + `RDD[ColumnarBatch]`, which the convenience `prepareShuffleDependency(rdd, ...)` overload + wraps in a stream. 2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan. diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index 71c9c431ab0..b7ad3d2f090 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -66,23 +66,34 @@ write side. To restrict columnar shuffle to cases where the child is already a C as native Spark shuffles, which avoids the row to columnar conversion but means the downstream stage will also start on Spark. +Setting `spark.comet.convert.shuffleInput.enabled=true` makes such a shuffle use native shuffle instead. Comet +converts the child's rows to Arrow, as the `spark.comet.convert.*` settings do for Spark scans, and shuffles the +batches natively. With hash partitioning that is usually faster than columnar shuffle, while with range partitioning +the two come out about even. This applies only where native shuffle supports the partitioning and Comet can convert +all of the child's columns, which rules out calendar intervals, and arrays and maps other than `array` and +`map`. A shuffle that hashes a decimal wider than 18 digits, or a string or a value computed from a +string, stays on columnar shuffle, so that it partitions rows as Spark does. The converted batches are allocated on +the JVM outside Comet's memory pools. The setting is disabled by default. + ### Automatic Revert to Spark Shuffle -When a Comet columnar shuffle ends up between a partial and a final aggregate that Comet could not convert (both -remain Spark `HashAggregateExec` or `ObjectHashAggregateExec` operators), Comet reverts it to Spark's built-in shuffle. +When a Comet columnar shuffle, or a native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted, +ends up between a partial and a final aggregate that Comet could not convert (both remain Spark `HashAggregateExec`, +`ObjectHashAggregateExec` or `SortAggregateExec` operators), Comet reverts it to Spark's built-in shuffle. Keeping columnar shuffle between the two row-based aggregates would add `row -> Arrow -> shuffle -> Arrow -> row` -conversions with no Comet consumer on either side to benefit from columnar output. Other shuffles between non-Comet -operators are not reverted. +conversions with no Comet consumer on either side to benefit from columnar output. The final aggregate has to read the +shuffle directly, which a `SortAggregateExec` with grouping keys does not: it reads the shuffle through a sort, so that +shuffle is not reverted. Other shuffles between non-Comet operators are not reverted either. This shifts the affected shuffles from Comet's off-heap memory pool back to the JVM execution memory pool. Clusters tuned for a small JVM heap may see `ExternalSorter` spills on queries where this revert fires. Shuffle I/O may also grow marginally because Spark's row-based serializer generally compresses less well than Comet's Arrow IPC format. -Each revert is logged at `INFO` level on the driver as `Reverting Comet columnar shuffle to Spark shuffle between - and `, which lets you correlate any unexpected behavior with this optimization. +Each revert is logged at `INFO` level on the driver as `Reverting Comet shuffle to Spark shuffle between +and `, which lets you correlate any unexpected behavior with this optimization. This optimization is enabled by default and can be disabled by setting -`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep the columnar shuffle +`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep its shuffle even when both of those aggregates run on Spark. ## Shuffle Compression diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 3d08c72dba8..15889d53277 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -271,6 +271,20 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.convert.shuffleInput.enabled") + .category(CATEGORY_EXEC) + .doc( + "When enabled, a shuffle that would use Comet's JVM columnar shuffle because its input " + + "comes from a Spark operator converts that input to Arrow format and uses Comet " + + "native shuffle instead. This applies only where native shuffle supports the " + + "partitioning and Comet can convert all of the input's columns, which rules out " + + "calendar intervals. A shuffle that hashes a decimal wider than 18 digits, or a " + + "string or a value computed from a string, stays on the JVM columnar shuffle, so " + + "that it partitions rows as Spark does.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.enabled") .category(CATEGORY_EXEC) .doc( @@ -726,14 +740,14 @@ object CometConf extends ShimCometConf { conf("spark.comet.shuffle.revertRedundantColumnar.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.revertRedundantColumnar.enabled") .category(CATEGORY_SHUFFLE) - .doc( - "When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " + - "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " + - "hash aggregate operators. This avoids a redundant " + - "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " + - "side can consume columnar output. Disable to keep Comet columnar shuffle even in " + - "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " + - "the extra conversion.") + .doc("When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " + + "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " + + "aggregate operators. The same applies to a native shuffle whose input " + + s"`${COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key}` converted. This avoids a redundant " + + "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " + + "side can consume columnar output. Disable to keep the Comet shuffle even in " + + "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " + + "the extra conversion.") .booleanConf .createWithDefault(true) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index e2712b01e34..d339f27965d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -205,11 +205,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) /** * Revert any `CometShuffleExchangeExec` with `CometColumnarShuffle` whose parent and child are - * both non-Comet `HashAggregateExec` / `ObjectHashAggregateExec` operators back to the original - * Spark `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't + * both non-Comet aggregate operators (any `BaseAggregateExec`) back to the original Spark + * `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't * convert either aggregate; keeping a columnar shuffle between them only adds * row->arrow->shuffle->arrow->row conversion overhead with no Comet consumer on either side. - * See https://github.com/apache/datafusion-comet/issues/4004. + * See https://github.com/apache/datafusion-comet/issues/4004. A native shuffle over rows that + * `convertShuffleInput` converted is the same pattern, and is reverted the same way. A + * `SortAggregateExec` with grouping keys reads the shuffle through the `SortExec` that + * `EnsureRequirements` adds, so it is not the shuffle's parent and the shuffle stays. * * The match is intentionally narrow (both sides must be row-based aggregates that remained JVM * after the main transform pass). Running the revert post-transform means we only fire when the @@ -230,26 +233,35 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = { def isAggregate(p: SparkPlan): Boolean = p.isInstanceOf[BaseAggregateExec] - def isRedundantShuffle(child: SparkPlan): Boolean = child match { - case s: CometShuffleExchangeExec => - s.shuffleType == CometColumnarShuffle && isAggregate(s.child) - case _ => false + // The Spark operator whose rows a Comet shuffle writes: the child of a JVM columnar shuffle, + // or the input of the conversion that `convertShuffleInput` put under a native shuffle. + def rowInput(s: CometShuffleExchangeExec): Option[SparkPlan] = + (s.shuffleType, s.child) match { + case (CometColumnarShuffle, child) => Some(child) + case (CometNativeShuffle, conversion: CometSparkToColumnarExec) => Some(conversion.child) + case _ => None + } + + def redundantShuffleInput(child: SparkPlan): Option[SparkPlan] = child match { + case s: CometShuffleExchangeExec => rowInput(s).filter(isAggregate) + case _ => None } plan.transform { - case op if isAggregate(op) && op.children.exists(isRedundantShuffle) => - val newChildren = op.children.map { - case s: CometShuffleExchangeExec - if s.shuffleType == CometColumnarShuffle && isAggregate(s.child) => - val reverted = - s.originalPlan.withNewChildren(Seq(s.child)).asInstanceOf[ShuffleExchangeExec] - reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ()) - logInfo( - "Reverting Comet columnar shuffle to Spark shuffle between " + - s"${op.getClass.getSimpleName} and ${s.child.getClass.getSimpleName} " + - "(no Comet operator on either side to consume columnar output)") - reverted - case other => other + case op if isAggregate(op) && op.children.exists(redundantShuffleInput(_).isDefined) => + val newChildren = op.children.map { child => + (child, redundantShuffleInput(child)) match { + case (s: CometShuffleExchangeExec, Some(input)) => + val reverted = + s.originalPlan.withNewChildren(Seq(input)).asInstanceOf[ShuffleExchangeExec] + reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ()) + logInfo( + "Reverting Comet shuffle to Spark shuffle between " + + s"${op.getClass.getSimpleName} and ${input.getClass.getSimpleName} " + + "(no Comet operator on either side to consume columnar output)") + reverted + case _ => child + } } op.withNewChildren(newChildren) } @@ -608,6 +620,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case s: ShuffleExchangeExec if shouldSkipCometShuffle(s) => preserveSparkAggregateBuffers(s) + case s: ShuffleExchangeExec if CometShuffleExchangeExec.convertsInputForNativeShuffle(s) => + convertShuffleInput(s) + case s: ShuffleExchangeExec => convertToComet(s, CometShuffleExchangeExec) .getOrElse(preserveSparkAggregateBuffers(s)) @@ -635,7 +650,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) op case _: ColumnarToRowTransition => // A transition does no work of its own. This rule only meets one that - // `convertTypedDatasetOutput` inserted on an earlier pass over the same plan. + // `convertTypedDatasetOutput` or `convertShuffleInput` inserted on an earlier pass + // over the same plan. op case _: WriteFilesExec => // The write is converted at the enclosing DataWritingCommandExec above: on Spark 3.x @@ -1397,6 +1413,30 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } } + /** + * Converts the rows a Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as + * native shuffle instead of the JVM columnar shuffle. See + * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and + * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies. + * + * Spark inserts the columnar transitions after this rule, but it does not look below a + * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark + * operator in the child's subtree that reads a Comet operator would do so through + * `CometExec.doExecute`, Spark's interpreted columnar-to-row path. So the subtree gets its + * transitions now, from Spark's own rule, and `EliminateRedundantTransitions` later replaces + * each one over a Comet child with Comet's own. Spark's rule leaves existing transitions alone, + * which matters because this rule runs over the same plan twice under AQE. + */ + private def convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = { + val child = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child) + convertToComet(child, CometSparkToColumnarExec) + .flatMap(converted => + convertToComet(s.withNewChildren(Seq(converted)), CometShuffleExchangeExec)) + .getOrElse( + convertToComet(s, CometShuffleExchangeExec).getOrElse(preserveSparkAggregateBuffers(s))) + } + private def shouldApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = { // Only consider converting leaf nodes to columnar currently, so that all the following // operators can have a chance to be converted to columnar. Leaf operators that output diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index 3f3c57f0645..9ad2d003841 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec} -import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike} @@ -63,9 +63,7 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan plan match { case _: BroadcastExchangeLike => plan case exchange: ShuffleExchangeLike => - revertShuffleStageIfNeeded(exchange) - .map(reverted => exchange.withNewChildren(Seq(reverted))) - .getOrElse(plan) + revertShuffleStageIfNeeded(exchange).getOrElse(plan) case _ => // Result stage: its output is collected as rows. revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan) @@ -74,20 +72,47 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan private def applyForNonAQE(plan: SparkPlan): SparkPlan = { val withRevertedStages = plan.transformUp { case exchange: ShuffleExchangeLike => - revertShuffleStageIfNeeded(exchange) - .map(reverted => exchange.withNewChildren(Seq(reverted))) - .getOrElse(exchange) + revertShuffleStageIfNeeded(exchange).getOrElse(exchange) } revertStageIfNeeded(withRevertedStages, outputColumnar = false) .getOrElse(withRevertedStages) } - private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = { - val outputArrow = exchange match { - case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle - case _ => false + /** + * Reverts the stage below a shuffle if needed and returns the shuffle over the reverted stage. + * A native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted loses + * the conversion with the rest of the stage, so it goes back to the JVM columnar shuffle that + * the conversion replaced, which reads the stage's rows. Any other native shuffle consumes + * Arrow-backed Comet vectors, so the reverted stage is bridged back to them. + */ + private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = + exchange match { + case s: CometShuffleExchangeExec + if s.shuffleType == CometNativeShuffle && convertsSparkRows(s.child) => + revertStageIfNeeded(s.child, s.supportsColumnar).map { reverted => + val columnar = s.copy(child = reverted, shuffleType = CometColumnarShuffle) + columnar.copyTagsFrom(s) + columnar + } + case _ => + val outputArrow = exchange match { + case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle + case _ => false + } + revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow) + .map(reverted => exchange.withNewChildren(Seq(reverted))) } - revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow) + + /** + * Whether `plan` is the conversion that `spark.comet.convert.shuffleInput.enabled` puts over a + * Spark operator. A Comet transition directly under a `CometSparkToColumnarExec` is a stacked + * bridge that the revert strips and then restores for a native shuffle, not that conversion. + */ + private def convertsSparkRows(plan: SparkPlan): Boolean = plan match { + case conversion: CometSparkToColumnarExec => + !conversion.child.isInstanceOf[CometNativeColumnarToRowExec] && + !conversion.child.isInstanceOf[CometColumnarToRowExec] + case _ => false } /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 3216b6a4301..c5e5c7641e5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -21,6 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import java.util.function.Supplier +import scala.collection.mutable.ListBuffer import scala.concurrent.Future import scala.jdk.CollectionConverters._ @@ -214,17 +215,18 @@ case class CometShuffleExchangeExec( val nativeChild = child.asInstanceOf[CometNativeExec] // RangePartitioner needs real rows for sampling. Reuse the precomputed context so we // don't re-walk the SparkPlan tree or re-broadcast the encryption Hadoop conf. - val samplingRDD: Option[RDD[ColumnarBatch]] = outputPartitioning match { + val samplingRows: Option[RDD[InternalRow]] = outputPartitioning match { case _: RangePartitioning => Some( - nativeChild.executeColumnarWithContext( - ctx, - nativeChildMetricNode.withoutAggregateMetrics(nativeChild))) + CometShuffleExchangeExec.rowsOf( + nativeChild.executeColumnarWithContext( + ctx, + nativeChildMetricNode.withoutAggregateMetrics(nativeChild)))) case _ => None } CometShuffleExchangeExec.prepareNativeShuffleDependency( inputRDD.asInstanceOf[CometNativeShuffleInputRDD], - samplingRDD, + samplingRows, child.output, outputPartitioning, serializer, @@ -236,12 +238,19 @@ case class CometShuffleExchangeExec( positionalRoundRobin)) case None => child match { - case _: CometNativeArrowSource => + case source: CometNativeArrowSource => CometShuffleExchangeExec.prepareArrowStreamShuffleDependency( inputRDD.asInstanceOf[RDD[ArrowArrayStream]], - // The range partitioner samples rows, so it needs them as batches. outputPartitioning match { - case _: RangePartitioning => Some(child.executeColumnar()) + // The range partitioner samples only the sort keys, so a conversion of rows + // gives it the rows it reads, rather than all of their columns in Arrow. + case _: RangePartitioning => + Some(source match { + case conversion: CometSparkToColumnarExec + if !conversion.child.supportsColumnar => + conversion.child.execute() + case _ => CometShuffleExchangeExec.rowsOf(child.executeColumnar()) + }) case _ => None }, child.output, @@ -547,6 +556,48 @@ object CometShuffleExchangeExec None } + /** + * Whether a shuffle that would use the JVM columnar shuffle, because its child is a Spark + * row-based plan, can convert the child's rows to Arrow with `CometSparkToColumnarExec` and use + * native shuffle instead. See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The + * native checks are those that [[shuffleSupported]] would make for the converted plan, so + * `CometExecRule` only converts a shuffle that will become native. Like [[shuffleSupported]], + * tags the node when no Comet shuffle can take it. + */ + def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(s.conf) && + !isCometPlan(s.child) && + // This also rules out Celeborn, which has no JVM columnar shuffle, and calendar intervals. + // Arrow holds the time part of an interval in nanoseconds, so the conversion would + // overflow on one with more microseconds than that can hold. + shuffleSupported(s).contains(CometColumnarShuffle) && + CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) && + nativeShuffleFailureReasons(s).isEmpty && + !hashesDifferentlyFromSpark(s) + + /** + * Whether the shuffle hashes into more than one partition a key that native shuffle, reading + * converted rows, may put in a different partition from Spark's partitioner. Native shuffle + * must not take over such a shuffle from the JVM columnar shuffle: a join with an input that is + * still partitioned by Spark would put matching keys in different partitions. The keys are: + * + * - A decimal wider than 18 digits, which native shuffle hashes differently (#5994). + * - A string, or any value computed from one, such as `hash(s)`. Spark reads the string's + * bytes as they are, but the import of the converted batch into native replaces invalid + * UTF-8 before native shuffle evaluates the key. + * + * TODO: allow wide decimals once native hashing matches Spark for them. + */ + private def hashesDifferentlyFromSpark(s: ShuffleExchangeExec): Boolean = + s.outputPartitioning match { + case HashPartitioning(expressions, numPartitions) => + numPartitions > 1 && expressions.exists { key => + key.dataType.existsRecursively(DecimalType.isByteArrayDecimalType) || + key.exists(_.dataType.existsRecursively(_.isInstanceOf[StringType])) + } + case _ => false + } + /** * Whether the stage feeding a shuffle starts at a typed Dataset conversion (see * [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]). `CometExecRule` decides the shuffle @@ -862,6 +913,10 @@ object CometShuffleExchangeExec } } + /** The rows of `batches`, for the range partitioner to sample. */ + private def rowsOf(batches: RDD[ColumnarBatch]): RDD[InternalRow] = + batches.mapPartitionsInternal(_.flatMap(_.rowIterator().asScala)) + /** * Build a Comet native shuffle dependency around an existing `RDD[ColumnarBatch]` of real * batches. Used by [[org.apache.spark.sql.comet.CometCollectLimitExec]] and @@ -891,7 +946,7 @@ object CometShuffleExchangeExec "ShuffleWriterInput") prepareArrowStreamShuffleDependency( streamRDD, - Some(rdd), + Some(rowsOf(rdd)), outputAttributes, outputPartitioning, serializer, @@ -900,12 +955,12 @@ object CometShuffleExchangeExec /** * [[prepareShuffleDependency]] for input that native reads as one Arrow stream per partition, - * such as a [[CometNativeArrowSource]]'s. `samplingRDD` gives the same rows as batches, and is - * only required for [[RangePartitioning]]. + * such as a [[CometNativeArrowSource]]'s. `samplingRows` gives the same rows, and is only + * required for [[RangePartitioning]]. */ private def prepareArrowStreamShuffleDependency( streamRDD: RDD[ArrowArrayStream], - samplingRDD: Option[RDD[ColumnarBatch]], + samplingRows: Option[RDD[InternalRow]], outputAttributes: Seq[Attribute], outputPartitioning: Partitioning, serializer: Serializer, @@ -946,7 +1001,7 @@ object CometShuffleExchangeExec // is `shuffleWriterMetrics` at the root with one empty leaf for the Scan child. prepareNativeShuffleDependency( thinRDD, - samplingRDD, + samplingRows, outputAttributes, outputPartitioning, serializer, @@ -965,13 +1020,14 @@ object CometShuffleExchangeExec * @param thinRDD * scheduling-anchor RDD whose `compute` returns a [[CometNativeShuffleInputIterator]]; * produces no batches itself. - * @param samplingRDD - * regular columnar execution of the child, only required for [[RangePartitioning]] (sampling - * needs real rows). `None` for hash / single / round-robin. + * @param samplingRows + * the child's rows, such as those of a regular columnar execution of the child, only required + * for [[RangePartitioning]] (sampling needs real rows). `None` for hash / single / + * round-robin. */ def prepareNativeShuffleDependency( thinRDD: CometNativeShuffleInputRDD, - samplingRDD: Option[RDD[ColumnarBatch]], + samplingRows: Option[RDD[InternalRow]], outputAttributes: Seq[Attribute], outputPartitioning: Partitioning, serializer: Serializer, @@ -1011,11 +1067,11 @@ object CometShuffleExchangeExec // ShuffleExchangeExec::prepareShuffleDependency val (partitioner, rangePartitionBounds) = outputPartitioning match { case rangePartitioning: RangePartitioning => - // Sampling needs real rows; use the dedicated samplingRDD (a regular columnar execution - // of the child). The thin RDD itself yields nothing. - val samplingInput = samplingRDD.getOrElse( + // Sampling needs real rows; use the dedicated samplingRows (the child's rows). The thin + // RDD itself yields nothing. + val samplingInput = samplingRows.getOrElse( throw new IllegalStateException( - "RangePartitioning requires a samplingRDD on the native-shuffle path")) + "RangePartitioning requires samplingRows on the native-shuffle path")) // Extract only fields used for sorting to avoid collecting large fields that does not // affect sorting result when deciding partition bounds in RangePartitioner val rddForSampling = samplingInput.mapPartitionsInternal { iter => @@ -1025,12 +1081,7 @@ object CometShuffleExchangeExec // Internally, RangePartitioner runs a job on the RDD that samples keys to compute // partition bounds. To get accurate samples, we need to copy the mutable keys. - iter.flatMap { batch => - val rowIter = batch.rowIterator().asScala - rowIter.map { row => - mutablePair.update(projection(row).copy(), null) - } - } + iter.map(row => mutablePair.update(projection(row).copy(), null)) } // Construct ordering on extracted sort key. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala new file mode 100644 index 00000000000..03181136005 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala @@ -0,0 +1,428 @@ +/* + * 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. + */ + +package org.apache.comet.exec + +import java.sql.{Date, Timestamp} +import java.time.LocalDateTime + +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan} +import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, spark_partition_id, sum} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.unsafe.types.CalendarInterval + +import org.apache.comet.{CometConf, ExtendedExplainInfo} + +// Top-level, so the encoder needs no outer pointer. +case class ShuffleInputRec(a: Int, b: String) + +/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */ +class CometShuffleInputConversionSuite extends CometTestBase { + + import testImplicits._ + + /** + * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would + * convert the scans below the shuffles here before this conversion sees them. It is off by + * default. + */ + private def withoutLeafConversion(f: => Unit): Unit = + withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f) + + /** Defines a test that runs with the shuffle input conversion enabled. */ + private def convertTest(name: String)(f: => Unit): Unit = + test(name) { + withoutLeafConversion { + withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f) + } + } + + private val rowSchema = new StructType() + .add("k", IntegerType) + .add("l", LongType) + .add("s", DataTypes.StringType) + .add("d", DataTypes.DoubleType) + .add("m", DecimalType(18, 2)) + + /** + * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of + * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default. + */ + private def rowsDf(rows: Int = 1000): DataFrame = { + val data = (0 until rows).map { i => + Row( + i % 23, + i.toLong, + if (i % 11 == 0) null else s"value-$i", + i + 0.5d, + java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2)) + } + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema) + } + + /** Rows with an `array` column, which `CometSparkToColumnarExec` does not convert. */ + private def arraysDf(rows: Int = 100): DataFrame = { + val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType)) + val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1))) + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + } + + private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] = + collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c } + + private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s } + + /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */ + private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + cometShuffles(plan).filter { s => + s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec] + } + + /** + * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds + * them for the operators under the conversion. Without one, an operator reads its Comet child + * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the + * right answer slowly, so only the plan shows it. + */ + private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = { + val converted = conversions(plan) + assert(converted.nonEmpty, plan) + converted.foreach { conversion => + // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`. + val bare = conversion.child.collect { + case p + if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] && + p.children.exists(_.supportsColumnar) => + p + } + assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan") + assert( + conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty, + s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan") + } + // The transitions are added on the first of the rule's passes under AQE. A later pass must + // not report them as operators Comet failed to convert. + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(!reasons.exists(_.contains("ColumnarToRow")), reasons) + } + + /** Runs `f` with each join planned as a sort-merge join of two shuffles into 10 partitions. */ + private def withShuffledJoins(f: => Unit): Unit = + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10")(f) + + /** + * Values of each type the conversion admits as a hash key, with boundary values. The floating + * point values include NaN and both zeros, which a join normalizes. + */ + private val hashKeys: Seq[(DataType, Seq[Any])] = Seq( + BooleanType -> Seq(true, false), + ByteType -> Seq(0.toByte, -1.toByte, Byte.MinValue, Byte.MaxValue), + ShortType -> Seq(0.toShort, -1.toShort, Short.MinValue, Short.MaxValue), + IntegerType -> Seq(0, -1, Int.MinValue, Int.MaxValue), + LongType -> Seq(0L, -1L, Long.MinValue, Long.MaxValue), + FloatType -> Seq( + 0.0f, + -0.0f, + 1.5f, + Float.NaN, + Float.MinPositiveValue, + Float.MaxValue, + Float.NegativeInfinity), + DoubleType -> Seq( + 0.0d, + -0.0d, + 1.5d, + Double.NaN, + Double.MinPositiveValue, + Double.MaxValue, + Double.NegativeInfinity), + DecimalType(18, 2) -> Seq("0.00", "-0.01", "9999999999999999.99", "-9999999999999999.99") + .map(new java.math.BigDecimal(_)), + DateType -> Seq("1970-01-01", "1969-12-31", "0001-01-01", "9999-12-31").map(Date.valueOf), + TimestampType -> Seq( + new Timestamp(0L), + new Timestamp(-1L), + Timestamp.valueOf("0001-01-01 00:00:00"), + Timestamp.valueOf("9999-12-31 23:59:59.999999")), + TimestampNTZType -> Seq( + LocalDateTime.of(1970, 1, 1, 0, 0), + LocalDateTime.of(1969, 12, 31, 23, 59, 59, 999999000), + LocalDateTime.of(1, 1, 1, 0, 0), + LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999000)), + BinaryType -> Seq(Array.emptyByteArray, Array[Byte](0), Array[Byte](-1, -128, 127))) + + /** Four rows of each value and of a NULL, as column `k`, over four input partitions. */ + private def keysDf(keyType: DataType, values: Seq[Any]): DataFrame = { + val keys = values :+ null + val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong)) + val schema = new StructType().add("k", keyType).add("v", LongType) + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + } + + convertTest("the shuffle of a Spark operator's rows runs as native shuffle") { + Seq("true", "false").foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val (_, plan) = checkSparkAnswer( + rowsDf() + .repartition(10, col("k")) + .groupBy("k") + .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m"))) + assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan") + assert( + !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle), + s"AQE $aqe:\n$plan") + } + } + } + + convertTest("Spark operators over Comet operators read them through transitions") { + withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") { + Seq("true", "false").foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec] + // The shuffle reads the typed operation's SerializeFromObject. + val (_, repartitioned) = checkSparkAnswer( + ds.map(r => ShuffleInputRec(r.a % 13, r.b)) + .repartition(7, col("a")) + .groupBy("a") + .agg(count("b"))) + assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned") + assertRowOperatorsReadThroughTransitions(repartitioned) + // The shuffle reads a Spark partial aggregate over the typed operation. The final + // aggregate stays on Spark too, so the shuffle would go back to Spark's own. + withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") { + val (_, aggregated) = + checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count()) + assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated") + assertRowOperatorsReadThroughTransitions(aggregated) + } + } + } + } + } + + convertTest("aggregates whose partial aggregate runs on Spark") { + // The partial aggregate reads Spark rows, so it runs on Spark, and native shuffle carries its + // buffers to the final aggregate. With these aggregates the final one runs natively. + val (_, split) = checkSparkAnswer( + rowsDf() + .groupBy((col("k") % 7).as("g")) + .agg(sum("l"), avg("d"), max("l"), min("l"))) + assert(convertedShuffles(split).length == 1, split) + assert(collect(split) { case a: HashAggregateExec => a }.length == 1, split) + assert(collect(split) { case a: CometHashAggregateExec => a }.length == 1, split) + // The min of a string makes Spark plan sort aggregates, which stay on Spark, so the final + // aggregate reads the buffers back as rows. The shuffle stays native: the final aggregate + // reads it through a sort, and the revert in the next test only takes a shuffle that an + // aggregate reads directly. + val (_, sorted) = checkSparkAnswer( + rowsDf() + .groupBy((col("k") % 7).as("g")) + .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s"))) + assert(convertedShuffles(sorted).length == 1, sorted) + assert(collect(sorted) { case a: SortAggregateExec => a }.length == 2, sorted) + } + + convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") { + withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") { + val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l"))) + assert(cometShuffles(plan).isEmpty, plan) + withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") { + val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l"))) + assert(convertedShuffles(kept).length == 1, kept) + } + } + } + + convertTest("a stage that the transition revert puts back on Spark keeps the JVM shuffle") { + // The typed operation reads the Comet scan through a transition. With no transitions allowed, + // the revert puts the stage below the shuffle back on Spark, the conversion included, and the + // shuffle goes back to the JVM columnar shuffle, which reads the stage's rows. + withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") { + Seq("true", "false").foreach { aqe => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec] + val (_, plan) = checkSparkAnswer( + ds.map(r => ShuffleInputRec(r.a % 13, r.b)) + .repartition(7, col("a")) + .groupBy("a") + .agg(max("b"))) + assert(conversions(plan).isEmpty, s"AQE $aqe:\n$plan") + val reverted = cometShuffles(plan).filter(_.shuffleType == CometColumnarShuffle) + assert(reverted.length == 1, s"AQE $aqe:\n$plan") + assert(!reverted.head.child.exists(_.isInstanceOf[CometPlan]), s"AQE $aqe:\n$plan") + } + } + } + } + + hashKeys.foreach { case (keyType, values) => + convertTest(s"native shuffle puts ${keyType.simpleString} keys where Spark does") { + withShuffledJoins { + val keys = keysDf(keyType, values) + // The partition of each row, against Spark's partitioner, NULL keys included. + val (_, placed) = checkSparkAnswer( + keys.repartition(10, col("k")).select(col("k"), col("v"), spark_partition_id())) + assert(convertedShuffles(placed).length == 1, placed) + // The left input converts and the right one, with its array column, does not, so + // matching keys meet only if both shuffles put them in the same partition. + val left = keys.select(col("k").as("lk"), col("v")) + val right = keys.select(col("k").as("rk"), array(col("v")).as("xs")) + val (_, joined) = checkSparkAnswer(left.join(right, col("lk") === col("rk"))) + assert(convertedShuffles(joined).length == 1, joined) + assert(cometShuffles(joined).count(_.shuffleType == CometColumnarShuffle) == 1, joined) + } + } + } + + convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") { + // Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner + // (#5994). The right input, with its array column, stays on the JVM columnar shuffle, + // so the left one has to as well, or matching keys land in different partitions. + withShuffledJoins { + val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l")) + val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs")) + val df = left.join(right, col("lk") === col("rk")) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, plan) + checkCometExchange(df, 2, native = false) + } + } + + convertTest("a shuffle that hashes a string stays on the JVM columnar shuffle") { + // Spark's partitioner hashes a string's bytes as they are, but native shuffle hashes them + // after the import into native has replaced invalid UTF-8. The right input, with its + // array column, stays on the JVM columnar shuffle, so the left one has to as well, or + // matching keys with invalid UTF-8 land in different partitions. A key computed from a + // string, such as its hash, is no different: native shuffle would compute it from the + // replaced string. Both shuffles write the rows with invalid UTF-8 replaced, and the join + // compares what they wrote, so the two keys stay apart only because each key expression puts + // them in different partitions. + withShuffledJoins { + val schema = new StructType().add("b", BinaryType).add("v", LongType) + val data = (0 until 40).map { i => + Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong) + } + val keys = spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + val left = keys.select(col("b").cast(StringType).as("lk"), col("v")) + val right = keys.select( + col("b").cast(StringType).as("rk"), + array(col("v").cast(IntegerType)).as("xs")) + Seq(col("lk") === col("rk"), hash(col("lk")) === hash(col("rk"))).foreach { condition => + val df = left.join(right, condition) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, s"$condition:\n$plan") + checkCometExchange(df, 2, native = false) + } + } + } + + convertTest("nested columns over several batches") { + // The rows of each input partition fill several batches, and the conversion reuses its + // vectors from one batch to the next. With the conversion of RDD scans on, native shuffle + // reads the converted scan the same way. + val schema = new StructType() + .add("k", IntegerType) + .add("payload", new StructType().add("v", LongType).add("s", StringType)) + .add("xs", ArrayType(StringType)) + .add("m", MapType(StringType, StringType)) + val data = (0 until 24).map { i => + Row( + i, + if (i % 5 == 4) null else Row(i.toLong, if (i % 3 == 0) null else s"s$i"), + if (i % 6 == 5) null else Seq.tabulate(i % 4)(j => if (j == 1) null else s"x$i-$j"), + if (i % 7 == 6) null + else Seq.tabulate(i % 3)(j => s"k$j" -> (if (j == 1) null else s"v$i-$j")).toMap) + } + Seq("false", "true").foreach { leafConversion => + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "7", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> leafConversion) { + val df = spark.createDataFrame(spark.sparkContext.parallelize(data, 1), schema) + Seq(df.repartition(3, col("k")), df.repartitionByRange(3, col("k"))).foreach { shuffled => + val (_, plan) = checkSparkAnswer(shuffled) + assert(convertedShuffles(plan).length == 1, s"leaf conversion $leafConversion:\n$plan") + } + } + } + } + + convertTest("calendar intervals keep Spark's shuffle") { + // Arrow holds an interval's time part in nanoseconds, which can't hold every number of + // microseconds Spark can. The JVM columnar shuffle doesn't take calendar intervals either. + val schema = new StructType().add("k", IntegerType).add("i", CalendarIntervalType) + val data = Seq(Row(1, new CalendarInterval(0, 0, 10800000000000000L))) + val df = spark + .createDataFrame(spark.sparkContext.parallelize(data, 1), schema) + .repartition(2, col("k")) + val (_, plan) = checkSparkAnswer(df) + assert(cometShuffles(plan).isEmpty, plan) + } + + convertTest("columns the conversion does not support keep the JVM columnar shuffle") { + val df = arraysDf().repartition(5, col("k")).groupBy("k").agg(sum(size(col("xs")))) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, plan) + assert(cometShuffles(plan).forall(_.shuffleType == CometColumnarShuffle), plan) + } + + convertTest("range partitioning") { + val (_, plan) = checkSparkAnswer(rowsDf().orderBy(col("l").desc)) + assert(convertedShuffles(plan).length == 1, plan) + // The range partitioner samples the rows that the conversion reads, without converting them, + // so the conversion only counts the rows that the shuffle writes. + assert(conversions(plan).head.metrics("numInputRows").value == 1000, plan) + } + + convertTest("the conversion only takes over from the JVM columnar shuffle") { + // In `jvm` mode the shuffle keeps the JVM columnar shuffle. In `native` mode, where Comet has + // no shuffle for a Spark operator's rows, it keeps Spark's shuffle. + Seq("jvm" -> Seq(CometColumnarShuffle), "native" -> Seq.empty).foreach { + case (mode, shuffleTypes) => + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val (_, plan) = + checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count()) + assert(conversions(plan).isEmpty, s"$mode:\n$plan") + assert( + cometShuffles(plan).map(_.shuffleType).distinct == shuffleTypes, + s"$mode:\n$plan") + } + } + } + + test("off by default") { + withoutLeafConversion { + val (_, plan) = checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count()) + assert(conversions(plan).isEmpty, plan) + assert(cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan) + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala new file mode 100644 index 00000000000..3942f1859ef --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala @@ -0,0 +1,192 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.comet.{CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.functions.{col, count, length, lit, sum} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataTypes, DecimalType, IntegerType, LongType, StructType} + +import org.apache.comet.CometConf + +// Top-level, so the encoder needs no outer pointer. +case class ShuffleInputBenchRec(a: Long, b: String) + +/** + * Compares three ways to run a shuffle whose input comes from a Spark operator: + * + * - Spark: Comet disabled. + * - Comet: the default. Comet's JVM columnar shuffle buffers the rows and converts them to + * Arrow natively, one partition at a time. + * - Comet, converted: `spark.comet.convert.shuffleInput.enabled`, which converts the rows with + * `CometSparkToColumnarExec` and uses native shuffle. + * + * Each case partitions its rows into 200 partitions and aggregates them. Most hash-partition on a + * number, because the conversion leaves a shuffle that hashes a string on the JVM columnar + * shuffle, and one range-partitions. Every arm's result and plan are checked before it is timed, + * and the Comet arm runs again at the end of each case to show the noise. To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometShuffleInputConversionBenchmark + * }}} + * Results will be written to + * "spark/benchmarks/CometShuffleInputConversionBenchmark-**results.txt". + */ +object CometShuffleInputConversionBenchmark extends CometBenchmarkBase { + + private val numRows = 4L * 1024 * 1024 + private val numPartitions = 200 + + import spark.implicits._ + + private case class Arm(name: String, confs: Seq[(String, String)], check: SparkPlan => Unit) + + // No AQE, so every arm plans the same shape on every iteration. + private val planConfs = Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") + + private def cometConfs(convert: Boolean) = planConfs ++ Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + // Spark scans the RDDs, as it does by default, so the shuffle reads a Spark operator. + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> convert.toString) + + private def shuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + collect(plan) { case s: CometShuffleExchangeExec => s } + + private val arms = Seq( + Arm( + "Spark", + planConfs :+ (CometConf.COMET_ENABLED.key -> "false"), + plan => assert(collect(plan) { case c: CometPlan => c }.isEmpty, plan)), + Arm( + "Comet", + cometConfs(convert = false), + plan => assert(shuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan)), + Arm( + "Comet, converted", + cometConfs(convert = true), + plan => + assert( + shuffles(plan).exists(s => + s.shuffleType == CometNativeShuffle && s.child + .isInstanceOf[CometSparkToColumnarExec]), + plan))) + + private val rowSchema = new StructType() + .add("k", IntegerType) + .add("l", LongType) + .add("s", DataTypes.StringType) + .add("d", DataTypes.DoubleType) + .add("m", DecimalType(18, 2)) + + /** `numRows` rows from an RDD, generated as the query runs. */ + private def rddRows(): DataFrame = { + val rows = spark.sparkContext.range(0, numRows, 1, 1).map { id => + val i = id.toInt + Row( + i % 1000, + id, + s"value-${i * 31 % 100003}", + i + 0.5d, + java.math.BigDecimal.valueOf(i * 7919L % 1000000000L, 2)) + } + spark.createDataFrame(rows, rowSchema) + } + + /** Totals the per-group results in `s`, so each case returns one row. */ + private def total(grouped: DataFrame): DataFrame = grouped.agg(sum("s"), count(lit(1))) + + private val cases: Seq[(String, () => DataFrame)] = Seq( + "RDD rows: int, long, double" -> (() => + total( + rddRows() + .select("k", "l", "d") + .repartition(numPartitions, col("k")) + .groupBy("k") + .agg((sum("l") + sum("d")).as("s")))), + "RDD rows: int, long, string, double, decimal(18,2)" -> (() => + total( + rddRows() + .repartition(numPartitions, col("k")) + .groupBy("k") + .agg((sum("l") + sum(length(col("s"))) + sum("d") + sum("m")).as("s")))), + "RDD rows: int, long, string, double, decimal(18,2), range partitioned" -> (() => + total( + rddRows() + .repartitionByRange(numPartitions, col("l")) + .select((col("k") + col("l") + length(col("s")) + col("d").cast("long") + col("m")) + .as("s")))), + "map over a Parquet scan" -> (() => + total( + spark + .table("parquetV1Table") + .select(col("id").as("a"), col("key").as("b")) + .as[ShuffleInputBenchRec] + .map(r => ShuffleInputBenchRec(r.a % 1000, r.b)) + .repartition(numPartitions, col("a")) + .groupBy("a") + .agg(sum(length(col("b"))).as("s"))))) + + private def runArm(arm: Arm, query: () => DataFrame): (Seq[Row], SparkPlan) = { + var result: (Seq[Row], SparkPlan) = null + withSQLConf(arm.confs: _*) { + val df = query() + result = (df.collect().toSeq, df.queryExecution.executedPlan) + } + result + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.range(numRows).selectExpr("id", "CAST(id % 1000 AS STRING) AS key")) + + cases.foreach { case (name, query) => + runBenchmark(name) { + // Check every arm's plan and result before timing anything. + val results = arms.map(runArm(_, query)) + arms.zip(results).foreach { case (arm, (rows, plan)) => + arm.check(plan) + assert( + rows == results.head._1, + s"${arm.name} returned $rows, expected ${results.head._1}") + } + + val benchmark = new Benchmark(name, numRows, output = output) + (arms :+ arms(1).copy(name = "Comet (repeat)")).foreach { arm => + benchmark.addCase(arm.name) { _ => + withSQLConf(arm.confs: _*) { + query().collect() + } + } + } + benchmark.run() + } + } + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index a63560f3376..1d217bf6a25 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -265,6 +265,26 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { } } + test("the shuffle input conversion leaves a Spark child to Spark's shuffle") { + // The conversion only takes over from the JVM columnar shuffle, which Celeborn does not + // have, and Celeborn's native shuffle needs a native child. Without Celeborn, auto and jvm + // would convert this child. Native mode has no JVM columnar shuffle either, so the reason + // is what shows that Celeborn kept Spark's shuffle there. + Seq("auto", "jvm", "native").foreach { mode => + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> mode, + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true") { + val leaf = sparkLeaf() + val exchange = ShuffleExchangeExec(HashPartitioning(leaf.output, 2), leaf) + assert(!CometShuffleExchangeExec.convertsInputForNativeShuffle(exchange), mode) + assert(reasons(exchange).exists(_.contains("Celeborn")), s"$mode: ${reasons(exchange)}") + val planned = CometExecRule(spark)(exchange) + assert(planned.isInstanceOf[ShuffleExchangeExec], planned) + assert(collect(planned) { case c: CometSparkToColumnarExec => c }.isEmpty, planned) + } + } + } + test("fallback is sticky after AQE reshapes the child or native mode becomes available") { val child = nativeChild() val exchange = ShuffleExchangeExec(SinglePartition, child) From b56349697b786ff2ad1c1bcf6ecf45b809af5f30 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 7 Oct 2026 14:13:40 +0000 Subject: [PATCH 22/89] fix: normalize float operands in the native comparison builder (#6447) * fix: normalize float operands in the native comparison builder Part of #6385. CometExecRule.normalize wrapped comparison operands in NormalizeNaNAndZero only inside ProjectExec and FilterExec, so a comparison in an aggregate argument, a FILTER clause, a join condition, a sort key or a generator compared raw Arrow values, where a NaN with the sign bit set sorts below every other value. Arrays and structs had the same gap for <=>, <, <=, > and >= in every operator. spark_comparison, which builds every native comparison, now normalizes float operands for all eight comparison operators: scalars through NormalizeNaNAndZero, arrays and structs through NormalizeNestedFloats, and literals while planning. Nested = and <> keep spark_equality. A scan's data filters are built with FloatOperands::Raw so that Parquet pruning still sees a column compared with a literal; the Filter above the scan evaluates them again. CometExecRule.normalize is removed, including its Divide and Remainder divisor normalization, which native division and remainder do not need. NormalizeNaNAndZero keeps a scalar child a scalar. TPC-DS golden files count one fewer native expression where the rewrite had added one. * fix: match Spark's float ordering in min, max, greatest and least Spark orders floats with SQLOrderingUtil.compareDoubles, in which NaN is larger than every other value and -0.0 equals 0.0, and its max, min, greatest and least keep the first of equal values. DataFusion's versions order floats by IEEE 754 total order, so a NaN with the sign bit set is the smallest value; its grouped max and min start each group at f64::MIN and f64::MAX, so a group of only -Infinity (or Infinity) returns that finite value; and its greatest and least let a later argument win a tie. Add SparkMinMax for float min and max in aggregates and windows, with ungrouped, grouped and sliding-frame accumulators, and SparkGreatestLeast for float greatest and least, including arrays and structs with a float leaf. Other types keep DataFusion's implementations. Strict floating-point mode no longer makes min and max fall back. Part of #6385. * test: add min/max and greatest/least benchmarks and FLOAT sliding-window coverage Add criterion benchmarks comparing SparkMinMax and SparkGreatestLeast with DataFusion's versions, which check that both return the same values before timing. Cover SlidingMinMaxAccumulator over FLOAT in the unit test and the SQL fixture, share the edge values and the first-wins oracle between the two test modules through float_semantics, and route window min/max through min_max_udaf so the float routing lives in one place. * fix: normalize computed operands in scan data filters With spark.comet.parquet.rowFilterPushdown.enabled the Parquet reader evaluates a scan's data filters on each row, and a row it drops never reaches the Filter above the scan. Data filters left every float operand raw, so a computed operand such as -d, whose NaN has the sign bit set, dropped the NaN row that Spark keeps. Leave only a float column compared with a literal raw, the shape Parquet pruning reads, and normalize every other comparison. Run float_comparisons.sql with row-level pushdown off and on, add filters on stored edge values, add a benchmark of Spark's float comparisons against DataFusion's, and share one helper for nested types with a float leaf. * test: avoid scan-order assumptions in float min/max checks * test: drop the sweep suite's known gaps that normalized comparisons close With comparisons normalized in the native comparison builder, the cases for comparisons outside Project and Filter, Parquet scan filters against a sign-bit NaN literal, and nested array and struct comparisons (#6157) match Spark, so the suite now requires them to. * test: retire fixed float min/max and greatest/least gaps * fix: harden float extrema planner routing * test: share float_semantics' edge values with the array function tests array_remove and sort_array, from #6518, kept their own copy of the edge values that min/max and greatest/least share through float_semantics. The two arrays were identical. * test: check signed-zero ties in aggregate min and max again eb5c7af4f stopped depending on scan order by dropping the aggregate tie queries from the SQL fixture, and by testing each zero on its own in CometAggregateSuite, where any tie rule returns the same value. Write both tables as a single file instead, so that Spark and Comet read the rows in the same order in one task, and check ties through the ungrouped and grouped accumulators again. The ungrouped checks run once per group, because over groups 2 and 3 together the first and the last zero have the same sign. * refactor: drop the unreachable missing-result branch from greatest and least invoke_with_args checks for at least two arguments before it starts, so the result can no longer be missing after the loop. Start from the first argument instead of an Option, as return_type already does. * test: cover IS DISTINCT FROM, two-level nesting and a NaN row group float_comparisons.sql now runs IS DISTINCT FROM, which Comet plans as one native operator, on scalars, constants, arrays and structs, and compares ARRAY> and STRUCT> columns under every operator. The float pruning test puts a NaN in a row group whose other values are all below 500, and checks that `d > 500.0D` still returns it. * test: hash quotients of a NaN divisor with the sign bit set The removed Project and Filter rewrite normalized the divisor of Divide and Remainder, which also canonicalized a NaN divisor. Without it the quotient keeps the sign bit, and hash and xxhash64 must still hash it as Spark does, in both ANSI modes. * test: order quotients of a NaN divisor in greatest, least, max and min Without the removed divisor rewrite, least and greatest of 1.0D / (-d) see a NaN with the sign bit set, and so do max and min of a quotient computed in a projection. The merged #6457 orders floats the Spark way in all four, so cover them in nan_divisor.sql with ANSI off and on. * fix: keep normalizing float divisors, and normalize data filters that drop rows With spark.comet.parquet.rowFilterPushdown.enabled, the Parquet reader drops the rows a scan's data filter rejects. FloatOperands::Raw left a float column raw against a normalized literal, so a stored NaN with other bits, such as one with the sign bit set, failed `d = -double('NaN')` and every ordering comparison that Spark matches. Data filters now leave the column raw only when row-level pushdown is off, where the reader uses them for pruning alone and the Filter above the scan applies Spark's semantics to each row. With it on, they normalize both sides and float comparisons give up statistics pruning. Removing normalizePlan also removed the normalization of Divide and Remainder divisors. Comparisons no longer need it, but it canonicalized a NaN quotient, and percentile_approx still orders a NaN with the sign bit set below every other value (#6519). normalizeDivisors keeps that part, in Projects and Filters as before, until #6519 is fixed. The TPC-DS plans count NormalizeNaNAndZero again, so their expression counts go back to main's. * fix: match either signed zero when a float data filter probes a bloom filter A data filter that only prunes compared the raw float column with a normalized literal, and a bloom filter probe hashes the literal's bits, so `d = -0.0D` probed for `0.0` and pruned a row group holding `-0.0` but no `0.0`. `=` against either zero now probes for both. Other comparisons keep the literal as it is, since DataFusion's comparisons, and with them statistics pruning, already treat the two zeros as equal. A NaN literal normalizes the column too, as the IN path does, because no literal's bits match every NaN a file can hold. Also document that float comparisons in data filters do not prune with spark.comet.parquet.rowFilterPushdown.enabled=true (#6702), in the floating-point guide, the config doc and the scan tuning page. * docs: add a contributor guide page on Spark's floating-point semantics Adds a Floating-Point Semantics page next to Timezone Handling. It covers which Spark rule a function follows for -0.0 and NaN, the releases in which a rule changed, how Arrow and DataFusion differ, where Comet applies each rule (the float_semantics helpers, key normalization in the planner, the comparison builder and the Scala version policies), and how to test for it. The expression guide's testing tips link to it. Part of #6385. * docs: correct the collect_set and OpenHashSet rules, and note where Spark's reader skips the other zero Before Spark 4.2, collect_set keys its buffer with Scala's `==`, which merges the two zeros and keeps every NaN apart, so the page now lists that rule separately. OpenHashSet and OpenHashMap follow Double.equals only from Spark 3.5.2 and 4.0.0 (SPARK-45599), which also qualifies mode. Spark's own Parquet reader keeps the two zeros apart in its dictionary and bloom filters, so it can skip a row group that Comet reads. The compatibility guide says so, and the contributor guide's data filter paragraph now describes the bloom filter expansion and the NaN literal. --- .../adding_a_new_expression.md | 1 + .../contributor-guide/floating_point.md | 207 +++++++ docs/source/contributor-guide/index.md | 1 + .../latest/compatibility/floating-point.md | 37 +- docs/source/user-guide/latest/tuning/scans.md | 3 +- native/core/src/execution/planner.rs | 181 +++++- native/core/src/execution/planner/macros.rs | 1 + native/core/src/parquet/parquet_exec.rs | 73 +++ native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/float_comparison.rs | 114 ++++ .../spark-expr/benches/nested_comparison.rs | 6 +- native/spark-expr/src/array_funcs/mod.rs | 2 +- .../src/array_funcs/nested_comparison.rs | 536 +++++++++++++++++- native/spark-expr/src/float_semantics/mod.rs | 14 +- .../src/float_semantics/normalize.rs | 182 +++++- .../scala/org/apache/comet/CometConf.scala | 10 +- .../apache/comet/rules/CometExecRule.scala | 91 ++- .../conditional/float_comparisons.sql | 207 +++++++ .../sql-tests/expressions/math/arithmetic.sql | 4 + .../expressions/math/nan_divisor.sql | 52 ++ .../comet/CometFloatSemanticsSuite.scala | 20 - .../comet/exec/CometNativeReaderSuite.scala | 86 +++ 22 files changed, 1691 insertions(+), 141 deletions(-) create mode 100644 docs/source/contributor-guide/floating_point.md create mode 100644 native/spark-expr/benches/float_comparison.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index d9e7707bab4..a246ea70bb6 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -408,6 +408,7 @@ For full documentation on the test file format, including directives like `Confi - **Cover both column references and literals.** Comet often uses different code paths for each. The Comet SQL Tests suite automatically disables constant folding, so all-literal queries are evaluated natively. - **Include edge cases** such as `NULL`, empty strings, boundary values, `NaN`, and multibyte UTF-8 characters. +- **For `FLOAT` and `DOUBLE` input, include `-0.0` and a NaN with the sign bit set.** [Floating-Point Semantics](floating_point.md) explains which Spark rule an expression follows and how to test it. - **Keep one file per expression** to make failures easy to locate. - **Pin the mechanism when it depends on the input.** If your serde routes some input types natively and others through the JVM codegen dispatcher, say so with `expect_native(...)` and `expect_dispatch(...)`. A plain `query` cannot tell the two apart, so without these a later change that swaps one for the other passes silently. See [Comet SQL Tests](sql-file-tests.md). diff --git a/docs/source/contributor-guide/floating_point.md b/docs/source/contributor-guide/floating_point.md new file mode 100644 index 00000000000..c113b86b8d5 --- /dev/null +++ b/docs/source/contributor-guide/floating_point.md @@ -0,0 +1,207 @@ + + +# Floating-Point Semantics + +This page describes how Comet matches Spark's handling of `-0.0` and NaN in `FLOAT` and `DOUBLE` +values. It is aimed at contributors working on an expression or operator that compares, orders, +hashes, deduplicates or groups floating-point values, including values nested in arrays, structs +and maps. For user-facing differences from Spark, see +[Floating-point Number Comparison](../user-guide/latest/compatibility/floating-point.md). The work +to apply these rules systematically is tracked in +[#6385](https://github.com/apache/datafusion-comet/issues/6385). + +The short version: Spark has no single rule for floating-point equality. Each function inherits one +from the Java API that its implementation calls. Arrow and DataFusion follow IEEE 754 total order, +which matches none of them. A native implementation must follow the rule of the Spark function it +replaces, using the shared helpers in `native/spark-expr/src/float_semantics/`, and its tests must +include a NaN with the sign bit set. + +## Spark's rules + +| Rule | `-0.0` and `0.0` | NaN | Used by | +| ----------------------------------------------------------- | ------------------------- | ---------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| SQL ordering (`SQLOrderingUtil.compareDoubles`, `genEqual`) | Equal | All NaNs are equal, and NaN sorts above all values | Comparisons, `IN`, sorting and ranking, `min`/`max`, `greatest`/`least`, `sort_array`, `array_min`/`array_max`, `array_contains`, `array_position`, `array_remove`, and comparisons of arrays and structs | +| `NormalizeNaNAndZero`, inserted by Spark's optimizer | `-0.0` becomes `0.0` | Every NaN becomes the canonical NaN | Grouping keys, join keys and window partition keys | +| `Murmur3Hash` and `XxHash64` | Same hash | Hashed through `doubleToLongBits`, which canonicalizes NaN | `hash`, `xxhash64` and hash partitioning | +| `java.lang.Double.equals`, for boxed values | Distinct | All NaNs are equal | Java hash sets and maps of boxed values, and Spark's `OpenHashSet` and `OpenHashMap` from Spark 3.5.2 and 4.0.0 | +| Scala `==`, for boxed values (`BoxesRunTime.equals`) | Equal | A NaN equals nothing, not even another NaN | Scala sets and maps of `Any`, such as `collect_set`'s buffer before Spark 4.2 | +| `java.lang.Double.compare` | `-0.0` sorts before `0.0` | As in SQL ordering | An ascending `sort_array` of elements that cannot be null, in Spark's generated code | + +Some functions changed rules in a Spark release, so a native path has to follow the Spark version +it runs against: + +- `collect_set` keys its buffer with Scala's `==` before Spark 4.2, so `-0.0` and `0.0` are one + value and every NaN is a value of its own. Spark 4.2 normalizes NaN and `-0.0` first and keys the + buffer by the normalized bits, so all NaNs are one value too (SPARK-57298). +- Spark's `OpenHashSet`, and the `OpenHashMap` built on it, follow `Double.equals` only from Spark + 3.5.2 and 4.0.0 (SPARK-45599). In every 3.4 release and in 3.5.0 and 3.5.1 they match a key with + `==` but hash its bits (`doubleToLongBits`), so two NaNs never match, and `-0.0` and `0.0` match + only when probing from one reaches the other. `mode` and `percentile` count values in an + `OpenHashMap`. `array_distinct`, `array_union`, `array_intersect` and `array_except` look up the + elements of a flat array in an `OpenHashSet`, but treat all NaNs as one value. +- `mode` follows `Double.equals` from Spark 3.5.2 and 4.0.0 until Spark 4.2, which folds `-0.0` + into `0.0` first (SPARK-57329). +- `array_distinct` and `array_union` keep signed zeros apart in a flat array before Spark 4.2.0 + (before 3.5.2, unless probing merges them). Spark 4.2.0 normalizes their arguments in the plan + (SPARK-54918). From 4.0.5, 4.1.4 and 4.2.1 they normalize while they evaluate instead + (SPARK-59602). +- Map construction (`ArrayBasedMapBuilder`) finds duplicate keys by `Double.equals` in Spark 3.4 + and 3.5. From Spark 4.0 it normalizes each key first, unless + `spark.sql.legacy.disableMapKeyNormalization` is set + ([#6549](https://github.com/apache/datafusion-comet/issues/6549)). + +## How Arrow and DataFusion differ + +Arrow compares floats by IEEE 754 total order. `-0.0` sorts below `0.0`, NaNs compare by their bit +patterns, and a NaN with the sign bit set sorts below `-Infinity`. Arrow's sort, row format and +hash kernels all work from the same bits. DataFusion 55 folds `-0.0` into `0.0` in some kernels but +does not canonicalize NaN, and that has changed between DataFusion releases. Don't rely on it: +normalize the values, or compare them with the helpers below. + +NaNs with the sign bit set are the normal case on x86-64, not a corner case. Every NaN that +arithmetic produces at run time, such as `sqrt(-1)` or `Infinity - Infinity`, is +`0xfff8000000000000` in both Rust and the JVM on x86-64, while aarch64 produces +`0x7ff8000000000000`. Spark hides the difference through `doubleToLongBits`. A native path that +compares raw values puts those NaNs below every other value, so a query that passes on an Apple +Silicon laptop can fail on a Linux x86 cluster. + +## Where Comet applies the rules + +### The `float_semantics` module + +`native/spark-expr/src/float_semantics/` holds the per-value rules and the kernels built on them. +Its module documentation is the reference. In short: + +| Helper | Rule | Use it to | +| ---------------------------------------------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------- | +| `compare_floats`, `float_lt`, `float_gt` | SQL ordering | Compare values in a kernel that returns the original bits, as `array_min` and `greatest` do | +| `spark_comparator`, `spark_equality` | SQL ordering, at any depth | Compare arrays and structs in place | +| `normalize_float`, `normalize_floats`, `normalize_nested_floats` | `NormalizeNaNAndZero` | Normalize values before an Arrow kernel sorts, row-encodes, hashes or compares them | +| `NormalizeNaNAndZero`, `NormalizeNestedFloats` | `NormalizeNaNAndZero` | Wrap a key or operand expression. `wrap_if_needed` skips other types and keys that Spark already normalized | +| `canonicalize_nan` | `java.lang.Double.equals` | Key a hash set or map the way boxed Java values do | +| `compare_floats_java` | `java.lang.Double.compare` | Sort the way `java.util.Arrays.sort` sorts a primitive array | +| `hash_input` | `Murmur3Hash` and `XxHash64` | Get the value to hash in place of a float | + +Once `-0.0` is folded and every NaN is canonical, Arrow's total order agrees with Spark's SQL +ordering. That is why normalizing the inputs of an Arrow kernel works, as long as only the keys are +normalized and the output keeps the original values. + +### Keys, sorting and comparisons + +- Spark's optimizer wraps grouping, join and window partition keys in `NormalizeNaNAndZero` + (`NormalizeFloatingNumbers`). Comet serializes it like any other expression. +- `create_normalized_key_expr` in `native/core/src/execution/planner.rs` normalizes sort keys, + window order and partition keys, and the partition keys of `WindowGroupLimit`, so that Arrow's + sort orders them the Spark way. +- `spark_comparison` in `native/spark-expr/src/array_funcs/nested_comparison.rs` builds every native + `=`, `<>`, `<=>`, `<`, `<=`, `>`, `>=` and `IS DISTINCT FROM`, whichever operator evaluates it. It + normalizes float operands, and folds a literal while the plan is built so that it stays a + literal. Nested `=` and `<>` compare in place with `spark_equality`. +- A scan's pushed-down data filters leave a float column compared with a literal other than NaN + unwrapped (`FloatOperands::Raw`) so that Parquet pruning still recognizes it, but only while the + reader prunes with them without filtering rows. A bloom filter probe hashes the literal's bits, so + `=` against either zero becomes `= -0.0 OR = 0.0`. With + `spark.comet.parquet.rowFilterPushdown.enabled=true` they normalize both sides, because the + reader drops the rows a filter rejects, and a raw column would reject a stored NaN that Spark + matches ([#6702](https://github.com/apache/datafusion-comet/issues/6702)). Spark's own reader + keeps the two zeros apart in its dictionary and bloom filters, so it can skip a row group that + Comet reads. A test of this pruning writes out the expected rows instead of comparing them with + Spark's. +- `IN` normalizes its operands in the serde (`normalizeInOperand` in `predicates.scala`). +- `hash`, `xxhash64`, the native shuffle's hash partitioner and `approx_count_distinct` hash floats + through `hash_input`. + +### Spark version differences + +Rules that depend on the Spark version are decided in Scala. The serde either chooses the native +path or passes a flag to native code, as `Mode`'s `normalize_neg_zero` does. Where a patch release +changed the rule, check the runtime version down to the patch with +`Utils.majorMinorPatchVersion(SPARK_VERSION)`, as `ArraySetSupport` in `arrays.scala` does. + +Until an expression has a native implementation that follows its rule, report it `Incompatible` +for floating-point input so that Spark evaluates it. + +`spark.comet.exec.strictFloatingPoint=true` makes Comet fall back for floating-point operations +that can still differ from Spark. To gate such an operation, use +`SupportLevel.strictFloatingPointReason(dataType, what)`. It returns a reason only in strict mode, +and only for a type that contains a `FLOAT` or `DOUBLE`. + +## Choosing the rule for an expression + +1. Read Spark's implementation, both `eval` and `doGenCode`, in every Spark version that Comet + supports, at the latest patch release of each. What it calls decides the rule. `ordering.compare`, + `ordering.equiv`, `ctx.genComp` and `ctx.genEqual` mean SQL ordering. A `java.util.HashMap` or + another Java collection of boxed values means `Double.equals`, and so does an `OpenHashSet` or + `OpenHashMap` from Spark 3.5.2 and 4.0.0. A Scala collection of `Any`, such as + `mutable.HashSet[Any]`, means Scala's `==`. `java.util.Arrays.sort` on a primitive array means + `Double.compare`. +2. Check whether Spark's optimizer normalizes the input in the plan. If it does, Comet receives + normalized values on that version only. +3. Apply the rule at every depth. A float inside an array, struct or map follows the same rule, and + some DataFusion kernels fold `-0.0` only in a flat array. +4. Record what you found in the expression's [audit](expression-audits/index.md) and, for a + difference users can see, in the compatibility guide. + +## Guidelines + +- Don't compare floating-point values with Arrow's comparison, sort or hash kernels, or with + `total_cmp` or `partial_cmp`, where Spark uses SQL ordering, unless you normalize them first. +- Normalize keys and operands, not results. Spark returns the original bits: `greatest(-0.0D, 0.0D)` + returns `-0.0`, `max` returns whichever zero it reads first, and `array_min` returns the first of + equal elements. +- Don't write a local copy of a rule. If a helper is missing, add it to `float_semantics`. +- In a hot loop, use the `float_lt` and `float_gt` predicates. Comparing `compare_floats(a, b)` with + an ordering known only at run time was several times slower in a scan for a minimum. +- When equality on nested values can stop early, keep that exit. `spark_equality` rejects lists of + different lengths without comparing their elements, which an equality built on an ordering + comparator loses. + +## Testing + +- Use the edge values `0.0`, `-0.0`, a canonical NaN, a NaN with the sign bit set, `1.0`, `-1.0`, + `Infinity`, `-Infinity` and `NULL`, for both `FLOAT` and `DOUBLE`. +- Make the sign-bit NaN at query time. Spark's Parquet writer canonicalizes NaN, so a stored NaN + reads back canonical. Negate it in the query (`-d`), which flips the sign bit on every platform. + Arithmetic such as `sqrt(-1)` yields a sign-bit NaN only on x86-64, so a test that relies on it + passes on Apple Silicon without exercising anything. The writer keeps the sign of zero, so `-0.0` + can be stored directly. +- Read the values from a Parquet table, and also compare them with literals on both sides. The + [Comet SQL Tests](sql-file-tests.md) turn constant folding off, so `-0.0D` stays a literal with the + sign bit set, while `double('NaN')` is a cast. +- Cover every place the rule applies. For a comparison, that is `Project`, `Filter`, aggregate + arguments, `FILTER` clauses, join conditions, sort keys, generators and windows. +- Nest the values in arrays and structs, two levels deep, and make the nested NaN a sign-bit one. +- When arithmetic produces the NaN, run with ANSI mode both on and off + (`-- ConfigMatrix: spark.sql.ansi.enabled=false,true`), because division and remainder take + different native paths. +- Where the rule changed in a Spark release, gate the fixtures with `MinSparkVersion` or + `MaxSparkVersion`, and add the `run-all-spark-profiles` label to the pull request. CI pins one + patch release per Spark line, so when a patch release changed the rule, also test the routing in + Scala. +- Add the expression to `CometFloatSemanticsSuite`, which crosses the edge values with operator + contexts and with every expression and aggregate that accepts floating-point input. New + expressions go in its `expressions` or `aggregates` list. If one does not match Spark yet, add a + `KnownGap` that links the issue. Known gaps are strict: a case that starts matching Spark fails, + so the fix also removes the entry. + +Fixtures to start from include `expressions/conditional/float_comparisons.sql`, +`expressions/math/nan_divisor.sql`, `expressions/math/greatest_least_floating_point.sql`, +`expressions/aggregate/min_max_floating_point.sql`, `expressions/array/sort_array_floating_point.sql` +and `windows/nested_float_order_keys.sql`, all under `spark/src/test/resources/sql-tests/`. diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 475bb77358a..1abd8891d20 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -47,6 +47,7 @@ Development Guide Comet Plugin Overview Arrow FFI Timezone Handling +Floating-Point Semantics JVM Shuffle Native Shuffle Memory Management diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index 455f58c3b63..f9e257dba3f 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -23,20 +23,43 @@ Spark normalizes NaN and zero for floating point numbers for several cases. See However, one exception is comparison. Spark does not normalize NaN and zero when comparing values because they are handled well in Spark (e.g., `SQLOrderingUtil.compareFloats`). But the comparison functions of arrow-rs used by DataFusion do not normalize NaN and zero (e.g., [arrow::compute::kernels::cmp::eq](https://docs.rs/arrow/latest/arrow/compute/kernels/cmp/fn.eq.html#)). -For top-level `FLOAT` and `DOUBLE` comparisons, Comet normalizes both operands before native -execution, including noncanonical NaN literals. Top-level `IN`, `InSet`, and `NOT IN` membership -also normalize dynamic candidates and lists containing NaN. When every candidate is a non-NaN -literal, Comet keeps DataFusion's static filter and pruning path, enumerating both signed-zero -forms when a list contains zero. +For `FLOAT` and `DOUBLE` comparisons (`=`, `<>`, `<=>`, `<`, `<=`, `>` and `>=`), Comet +normalizes both operands before native execution, including noncanonical NaN literals. This +applies wherever a comparison appears: projections, filters, aggregate arguments and `FILTER` +clauses, join conditions, sort keys, and generator arguments. + +A native Parquet scan prunes row groups and pages with its data filters. So that this pruning +still applies, a `FLOAT` or `DOUBLE` column compared with a constant other than NaN in a data +filter is compared without normalizing the column, and the filter above the scan evaluates the +comparison again with Spark's semantics. Every other comparison in a data filter is normalized. +Bloom filters store `-0.0` and `0.0` separately, so `=` against either zero checks them for both. +With `spark.comet.parquet.rowFilterPushdown.enabled=true` the scan also drops the rows its data +filters reject, so every comparison in a data filter is normalized, and a `FLOAT` or `DOUBLE` +comparison in a data filter does not prune row groups or pages +([#6702](https://github.com/apache/datafusion-comet/issues/6702)). + +Spark's own Parquet reader keeps `-0.0` and `0.0` apart when it prunes row groups with dictionaries +and bloom filters, so it can skip a row group holding rows that its filter matches. For example, +`WHERE d = -0.0D` can skip a row group that holds `0.0` but not `-0.0`, and so can `<=>` or `IN` +against `-0.0D`. `WHERE d >= 0.0D` can skip a dictionary-encoded row group whose largest value is +`-0.0`, and `WHERE d <= -0.0D` one whose smallest value is `0.0`. Comet reads those row groups, so +it can return rows that Spark does not. Comet's result is the one that Spark's comparison semantics +call for, and Spark returns the same rows with `spark.sql.parquet.filterPushdown=false`. + +Top-level `IN`, `InSet`, and `NOT IN` membership also normalize dynamic candidates and lists +containing NaN. When every candidate is a non-NaN literal, Comet keeps DataFusion's static filter +and pruning path, enumerating both signed-zero forms when a list contains zero. This scalar membership handling does not yet recurse into floating-point leaves nested in arrays or structs; see [#6019](https://github.com/apache/datafusion-comet/issues/6019). -## Nested equality and membership +## Nested comparisons and membership For arrays and structs containing `FLOAT` or `DOUBLE`, native `=`, `<>`, `IN`, and `NOT IN` compare signed zeros as equal and all NaN representations as equal, matching Spark. This also -covers single-candidate membership that Spark rewrites into equality. +covers single-candidate membership that Spark rewrites into equality. `<=>`, `<`, `<=`, `>`, and +`>=` normalize the floating-point values inside both operands first, so they follow Spark's order +too, in which NaN sorts above every other value. Equality and dynamic membership compare nested elements directly and stop at the first mismatch. Constant membership sets use normalized comparison values for static lookup. These operations diff --git a/docs/source/user-guide/latest/tuning/scans.md b/docs/source/user-guide/latest/tuning/scans.md index f23d32f41c7..40bc82634ec 100644 --- a/docs/source/user-guide/latest/tuning/scans.md +++ b/docs/source/user-guide/latest/tuning/scans.md @@ -28,7 +28,8 @@ decode step and lazily materializes projected columns for surviving rows. This c CPU and memory when the filter is highly selective on a small subset of columns. It is disabled by default because it can hurt when the filter is not selective or when most columns must be read anyway. Row-group, page-index, and bloom-filter pruning happen regardless of this flag whenever Spark's -`spark.sql.parquet.filterPushdown` is on. +`spark.sql.parquet.filterPushdown` is on, except for filters that compare `FLOAT` or `DOUBLE` values, +which do not prune while this flag is on (see [Floating-point Number Comparison](../compatibility/floating-point.md)). ### Parquet Native Scans diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 10f3eddb368..221d6b05360 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -144,10 +144,10 @@ use datafusion_comet_proto::{ use datafusion_comet_spark_expr::{ create_case_when, create_if_expr, jvm_udf::JvmScalarUdfExpr, spark_in_list, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, - DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, - HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, NormalizeNestedFloats, - Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal, ToJson, UnboundColumn, - Variance, WideDecimalBinaryExpr, WideDecimalOp, + DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields, GetStructField, HllPlusPlus, + HllSketchAgg, HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, + NormalizeNestedFloats, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal, + ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -325,6 +325,7 @@ pub struct BinaryExprOptions { pub const TEST_EXEC_CONTEXT_ID: i64 = -1; /// The query planner for converting Spark query plans to DataFusion query plans. +#[derive(Clone)] pub struct PhysicalPlanner { // The execution context id of this planner. exec_context_id: i64, @@ -346,6 +347,9 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + /// How comparisons treat floating-point operands. `Raw` only while planning a scan's data + /// filters; see [`Self::create_data_filter`]. + float_operands: FloatOperands, } impl Default for PhysicalPlanner { @@ -365,6 +369,7 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + float_operands: FloatOperands::Normalize, } } @@ -477,6 +482,11 @@ impl PhysicalPlanner { self } + /// How comparisons treat floating-point operands. + pub fn float_operands(&self) -> FloatOperands { + self.float_operands + } + /// Return session context of this planner. pub fn session_ctx(&self) -> &Arc { &self.session_ctx @@ -1129,6 +1139,45 @@ impl PhysicalPlanner { } } + /// Create a data filter that a scan pushes into the Parquet reader, with float operands + /// treated as [`Self::data_filter_float_operands`] decides. + fn create_data_filter( + &self, + spark_expr: &Expr, + input_schema: SchemaRef, + float_operands: FloatOperands, + ) -> Result, ExecutionError> { + let planner = Self { + float_operands, + ..self.clone() + }; + planner.create_expr(spark_expr, input_schema) + } + + /// How a scan's data filters treat float operands. The Parquet reader prunes row groups and + /// pages with them, and with row-level pushdown (`pushdown_filters`) it also drops the rows + /// they reject, which Spark's Filter above the scan then never sees. + /// + /// Without row-level pushdown the filters only prune, and pruning only recognizes a column + /// compared with a literal, so that shape keeps the raw column ([`FloatOperands::Raw`]). With + /// it, every operand is normalized and float comparisons give up pruning: a raw column would + /// drop a stored NaN that Spark matches, such as one with the sign bit set, which Arrow orders + /// below every other value. + fn data_filter_float_operands(&self) -> FloatOperands { + if self + .session_ctx + .copied_config() + .options() + .execution + .parquet + .pushdown_filters + { + FloatOperands::Normalize + } else { + FloatOperands::Raw + } + } + /// Normalize floating-point comparison keys without changing output values: a `FLOAT` or /// `DOUBLE` key, and an array or struct key with a float at any depth, whose order Arrow /// otherwise takes from the raw bits. Sort, Window, and WindowGroupLimit must use identical @@ -1910,10 +1959,17 @@ impl PhysicalPlanner { .cloned() .collect::>(), )); + let float_operands = self.data_filter_float_operands(); common .data_filters .iter() - .map(|expr| self.create_expr(expr, Arc::clone(&filter_schema))) + .map(|expr| { + self.create_data_filter( + expr, + Arc::clone(&filter_schema), + float_operands, + ) + }) .collect() }; @@ -5826,6 +5882,121 @@ mod tests { assert_eq!(0, filter_exec.additional_native_plans.len()); } + /// Comparisons normalize float operands, except in the data filters that a scan pushes into + /// the Parquet reader, where pruning has to see the column itself. + #[test] + fn scan_data_filters_compare_float_columns_directly() { + use datafusion::physical_expr::expressions::BinaryExpr; + use datafusion_comet_spark_expr::NormalizeNaNAndZero; + let double = spark_expression::DataType { + type_id: 6, + type_info: None, + }; + let operand = |expr_struct| Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + }; + let expr = operand(Gt(Box::new(spark_expression::BinaryExpr { + left: Some(Box::new(operand(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(double.clone()), + })))), + right: Some(Box::new(operand(Literal(spark_expression::Literal { + value: Some(literal::Value::DoubleVal(500.0)), + datatype: Some(double), + is_null: false, + })))), + }))); + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, true)])); + let left_operand = |expr: Arc| { + let comparison = expr.downcast_ref::().expect("a comparison"); + Arc::clone(comparison.left()) + }; + let planner = PhysicalPlanner::default(); + let comparison = planner.create_expr(&expr, Arc::clone(&schema)).unwrap(); + assert!(left_operand(comparison) + .downcast_ref::() + .is_some()); + let data_filter = planner + .create_data_filter(&expr, schema, planner.data_filter_float_operands()) + .unwrap(); + assert!(left_operand(data_filter).downcast_ref::().is_some()); + } + + /// With row-level pushdown the Parquet reader drops the rows a data filter rejects, so the + /// filter normalizes the column as well as the literal. A raw column would drop a stored NaN + /// whose bits differ from the normalized literal, and a stored NaN with the sign bit set + /// under any ordering comparison, both of which Spark matches. + #[test] + fn scan_data_filters_normalize_float_columns_with_row_level_pushdown() { + use arrow::array::{AsArray, BooleanArray}; + use datafusion::physical_expr::expressions::BinaryExpr; + use datafusion_comet_spark_expr::{FloatOperands, NormalizeNaNAndZero}; + let double = spark_expression::DataType { + type_id: 6, + type_info: None, + }; + let operand = |expr_struct| Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + }; + let column_and = |value: f64| { + Box::new(spark_expression::BinaryExpr { + left: Some(Box::new(operand(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(double.clone()), + })))), + right: Some(Box::new(operand(Literal(spark_expression::Literal { + value: Some(literal::Value::DoubleVal(value)), + datatype: Some(double.clone()), + is_null: false, + })))), + }) + }; + // Spark folds `-double('NaN')` into a literal with the sign bit set. + let negative_nan = f64::from_bits(0xfff8_0000_0000_0000); + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float64Array::from(vec![ + negative_nan, + f64::NAN, + 1.0, + ]))], + ) + .unwrap(); + let config = + SessionConfig::new().set_bool("datafusion.execution.parquet.pushdown_filters", true); + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(config)), 0); + assert_eq!( + planner.data_filter_float_operands(), + FloatOperands::Normalize + ); + for (expr, expected) in [ + (operand(Eq(column_and(negative_nan))), [true, true, false]), + (operand(Gt(column_and(0.0))), [true, true, true]), + ] { + let data_filter = planner + .create_data_filter( + &expr, + Arc::clone(&schema), + planner.data_filter_float_operands(), + ) + .unwrap(); + let comparison = data_filter + .downcast_ref::() + .expect("a comparison"); + assert!(comparison + .left() + .downcast_ref::() + .is_some()); + let matched = data_filter.evaluate(&batch).unwrap().into_array(3).unwrap(); + assert_eq!(matched.as_boolean(), &BooleanArray::from(expected.to_vec())); + } + } + #[test] fn spark_plan_metrics_hash_join() { let op_scan = create_scan(); diff --git a/native/core/src/execution/planner/macros.rs b/native/core/src/execution/planner/macros.rs index d146f6633bb..5c11770f09d 100644 --- a/native/core/src/execution/planner/macros.rs +++ b/native/core/src/execution/planner/macros.rs @@ -98,6 +98,7 @@ macro_rules! binary_expr_builder { $operator, right, input_schema.as_ref(), + planner.float_operands(), ) .map_err(Into::into) } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 6f9803d994d..93cffbca67a 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -1348,6 +1348,79 @@ mod tests { assert!(scan_metric(&scan, "scan_io_metadata_bytes") >= bloom_bytes); } + /// A bloom filter holds the bits of each value, so `-0.0` and `0.0` are separate entries in it, + /// while Spark's `=` matches either zero. A data filter that only prunes (no row-level + /// pushdown) probes the filter for both zeros, so the row group holding `-0.0` and no `0.0` + /// is read for `d = 0.0` and for `d = -0.0`, and a value the filter does not hold still + /// prunes it. The column statistics cover all three values, so only the bloom filter prunes. + #[tokio::test] + async fn bloom_filter_pruning_matches_either_signed_zero() { + use arrow::array::Float64Array; + use datafusion::physical_plan::metrics::MetricValue; + use datafusion_comet_spark_expr::{spark_comparison, FloatOperands}; + + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float64Array::from(vec![-1.0, -0.0, 1.0]))], + ) + .unwrap(); + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let props = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + for (value, read) in [(-0.0, true), (0.0, true), (0.5, false)] { + let filter = spark_comparison( + Arc::new(Column::new("d", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Float64(Some(value)))), + &schema, + FloatOperands::Raw, + ) + .unwrap(); + let scan = init_test_scan( + Arc::clone(&schema), + Arc::clone(&schema), + PartitionedFile::from_path(filename.clone()).unwrap(), + None, + Some(vec![filter]), + &session_ctx, + ); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + } + let bloom_filter = scan + .metrics() + .unwrap() + .sum_by_name("row_groups_pruned_bloom_filter"); + let Some(MetricValue::PruningMetrics { + pruning_metrics, .. + }) = bloom_filter + else { + panic!("missing bloom filter pruning metrics: {bloom_filter:?}"); + }; + let expected = if read { (0, 1, 3) } else { (1, 0, 0) }; + assert_eq!( + (pruning_metrics.pruned(), pruning_metrics.matched(), rows), + expected, + "d = {value:?}" + ); + } + } + #[tokio::test] async fn reports_projected_and_predicate_pruned_data_io() { let (filename, schema) = write_scan_io_fixture(); diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 965bbd79169..3011f1c5c6f 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -491,6 +491,10 @@ harness = false name = "nested_comparison" harness = false +[[bench]] +name = "float_comparison" +harness = false + [[bench]] name = "float_arrays" harness = false diff --git a/native/spark-expr/benches/float_comparison.rs b/native/spark-expr/benches/float_comparison.rs new file mode 100644 index 00000000000..7a411e05b59 --- /dev/null +++ b/native/spark-expr/benches/float_comparison.rs @@ -0,0 +1,114 @@ +// 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. + +//! Compare Spark's float comparisons, which normalize each float operand first, with DataFusion's +//! on the same data, and time `normalize_floats` on its own. Both comparisons return the same +//! answer except on a NaN with the sign bit set, which is checked before timing; special-value +//! semantics belong in tests. + +use arrow::array::{Array, ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::{normalize_floats, spark_comparison, FloatOperands}; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +const ROWS: usize = 8192; + +/// A NaN with the sign bit set, which arithmetic produces on x86-64. +const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); + +/// Two columns of doubles, with `special` in every tenth row of each if given. +fn batch(special: Option) -> RecordBatch { + let column = |seed: usize| -> ArrayRef { + Arc::new(Float64Array::from_iter_values((0..ROWS).map( + |i| match special { + Some(value) if i % 10 == seed => value, + _ => ((i * 7 + seed) % 1000) as f64 * 1.5, + }, + ))) + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Float64, true), + ])); + RecordBatch::try_new(schema, vec![column(1), column(3)]).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("float_comparison"); + group.sample_size(20); + group.warm_up_time(Duration::from_millis(250)); + group.measurement_time(Duration::from_secs(1)); + for (data, special) in [ + ("no_special", None), + ("neg_zero_10pct", Some(-0.0)), + ("nan_10pct", Some(f64::NAN)), + ("sign_bit_nan_10pct", Some(NEGATIVE_NAN)), + ] { + let batch = batch(special); + let schema = batch.schema(); + let a: Arc = Arc::new(Column::new("a", 0)); + let b: Arc = Arc::new(Column::new("b", 1)); + let literal: Arc = + Arc::new(Literal::new(ScalarValue::Float64(Some(500.0)))); + for (shape, right) in [("column_column", b), ("column_literal", literal)] { + let datafusion: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&a), + Operator::Lt, + Arc::clone(&right), + )); + let spark = spark_comparison( + Arc::clone(&a), + Operator::Lt, + right, + &schema, + FloatOperands::Normalize, + ) + .unwrap(); + let evaluate = |expr: &Arc| { + expr.evaluate(&batch) + .unwrap() + .into_array(ROWS) + .unwrap() + .to_data() + }; + if special.map(f64::to_bits) != Some(NEGATIVE_NAN.to_bits()) { + assert_eq!(evaluate(&spark), evaluate(&datafusion)); + } + for (engine, expr) in [("spark", &spark), ("datafusion", &datafusion)] { + group.bench_function(BenchmarkId::new(format!("{shape}_{engine}"), data), |b| { + b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) + }); + } + } + let column = Arc::clone(batch.column(0)); + group.bench_function(BenchmarkId::new("normalize_floats", data), |b| { + b.iter(|| black_box(normalize_floats(black_box(&column)))) + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/nested_comparison.rs b/native/spark-expr/benches/nested_comparison.rs index cf3fd44a077..85e1905c19d 100644 --- a/native/spark-expr/benches/nested_comparison.rs +++ b/native/spark-expr/benches/nested_comparison.rs @@ -27,7 +27,9 @@ use datafusion::common::ScalarValue; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, Literal}; use datafusion::physical_expr::PhysicalExpr; -use datafusion_comet_spark_expr::{spark_comparison, spark_in_list, NormalizeNestedFloats}; +use datafusion_comet_spark_expr::{ + spark_comparison, spark_in_list, FloatOperands, NormalizeNestedFloats, +}; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -65,7 +67,7 @@ fn expression(version: &str, mode: &str, batch: &RecordBatch) -> Arc = Arc::new(Column::new("b", 1)); if mode == "eq" { return if version == "new" { - spark_comparison(a, Operator::Eq, b, &schema).unwrap() + spark_comparison(a, Operator::Eq, b, &schema, FloatOperands::Normalize).unwrap() } else { Arc::new(BinaryExpr::new(a, Operator::Eq, b)) }; diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 01c19c418f9..e79d8f18a2e 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -42,7 +42,7 @@ pub use flatten::SparkFlatten; pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; pub use list_positions::ListPositionsExpr; -pub use nested_comparison::{spark_comparison, spark_in_list}; +pub use nested_comparison::{spark_comparison, spark_in_list, FloatOperands}; pub use sequence::spark_sequence; pub use size::{spark_size, SparkSizeFunc}; pub use sort_array::SparkSortArray; diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index 395a5ddc3f5..e945eb8f269 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -15,10 +15,12 @@ // specific language governing permissions and limitations // under the License. -//! Spark equality for nested floating-point values without materializing normalized columns. +//! Spark comparisons for floating-point values, and Spark equality for nested floating-point +//! values without materializing normalized columns. use crate::float_semantics::{ - has_float_leaf, normalize_nested_floats, spark_equality, NormalizeNestedFloats, + is_nested_with_float_leaf, normalize_comparison_operand, normalize_nested_floats, + spark_equality, NormalizeNestedFloats, }; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::buffer::{BooleanBuffer, NullBuffer}; @@ -27,17 +29,13 @@ use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; use datafusion::common::{internal_err, DFSchema, Result, ScalarValue}; use datafusion::logical_expr::{ColumnarValue, Operator}; -use datafusion::physical_expr::expressions::{in_list, BinaryExpr, InListExpr}; +use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, InListExpr, Literal}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::physical_expr::is_volatile; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; -fn needs_spark_equality(dt: &DataType) -> bool { - dt.is_nested() && has_float_leaf(dt) -} - struct Operand { array: ArrayRef, scalar: bool, @@ -219,31 +217,133 @@ impl PhysicalExpr for NestedPredicate { } } -/// Build equality after the planner has reconciled nested operand nullability. +/// How [`spark_comparison`] treats floating-point operands. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FloatOperands { + /// Normalize them, so that the comparison follows Spark's SQL ordering. + Normalize, + /// Leave a Float32 or Float64 column compared with a literal other than NaN as it is, and + /// normalize every other operand. Only a scan's pushed-down data filters use this, and only + /// when the Parquet reader prunes with them but does not filter rows: Parquet pruning + /// recognizes a column compared with a literal but not a normalized column, and Spark's + /// Filter above the scan applies Spark's semantics to every row. A NaN literal is normalized + /// with the column, giving up pruning, because a stored NaN can have other bits than the + /// literal and a writer can leave NaNs out of the column statistics. With row-level pushdown + /// the reader would drop the rows such a comparison rejects, including a stored NaN with the + /// sign bit set, which Arrow orders below every other value, so the data filters use + /// [`FloatOperands::Normalize`] there instead. + Raw, +} + +/// Builds a comparison with Spark's SQL ordering for floats, in which `-0.0` equals `0.0`, all +/// NaNs are equal and NaN sorts above every other value, at any depth of a list or struct. +/// +/// Arrow compares floats by IEEE 754 total order instead, so each float operand is normalized +/// first with [`normalize_comparison_operand`], after which the two orders agree. Nested `=` and +/// `<>` compare with `spark_equality` instead, without building normalized copies of the nested +/// values. Any other operator, such as `AND`, builds a plain [`BinaryExpr`]. +/// +/// The planner reconciles the nullability of nested operands before calling this. pub fn spark_comparison( left: Arc, op: Operator, right: Arc, schema: &Schema, + float_operands: FloatOperands, ) -> Result> { + use Operator::*; + if !matches!( + op, + Eq | NotEq | Lt | LtEq | Gt | GtEq | IsDistinctFrom | IsNotDistinctFrom + ) { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + } // An operand whose type does not resolve against this schema falls back to the plain // comparison, the way `reconcile_nested_comparison_types` already leaves such operands alone. - let nested = matches!(op, Operator::Eq | Operator::NotEq) - && match (left.data_type(schema), right.data_type(schema)) { - (Ok(lt), Ok(_)) => needs_spark_equality(<), - _ => false, - }; - if nested { + let (Ok(left_type), Ok(_)) = (left.data_type(schema), right.data_type(schema)) else { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + }; + if matches!(op, Eq | NotEq) && is_nested_with_float_leaf(&left_type) { validate_types(&left, std::slice::from_ref(&right), schema)?; - Ok(Arc::new(NestedPredicate { + return Ok(Arc::new(NestedPredicate { value: left, candidates: vec![right], - negated: op == Operator::NotEq, + negated: op == NotEq, membership: false, - })) - } else { - Ok(Arc::new(BinaryExpr::new(left, op, right))) + })); + } + if float_operands == FloatOperands::Raw { + if let Some(comparison) = raw_float_comparison(&left, op, &right, schema) { + return Ok(comparison); + } } + Ok(Arc::new(BinaryExpr::new( + normalize_comparison_operand(left, schema)?, + op, + normalize_comparison_operand(right, schema)?, + ))) +} + +/// The comparison that [`FloatOperands::Raw`] builds for a Float32 or Float64 column compared with +/// a literal other than NaN: the column and the literal as they are. `None` for any other operands. +/// +/// Statistics pruning compares `-0.0` and `0.0` as equal, but a bloom filter probe hashes the +/// literal's bits, so `=` against a zero literal becomes `=` against either zero. +fn raw_float_comparison( + left: &Arc, + op: Operator, + right: &Arc, + schema: &Schema, +) -> Option> { + let (literal, literal_on_left) = match ( + left.downcast_ref::(), + right.downcast_ref::(), + ) { + (None, Some(literal)) if is_float_column(left, schema) => (literal, false), + (Some(literal), None) if is_float_column(right, schema) => (literal, true), + _ => return None, + }; + let (negative_zero, positive_zero) = match literal.value() { + ScalarValue::Float32(Some(v)) if v.is_nan() => return None, + ScalarValue::Float64(Some(v)) if v.is_nan() => return None, + ScalarValue::Float32(Some(v)) if op == Operator::Eq && *v == 0.0 => ( + ScalarValue::Float32(Some(-0.0)), + ScalarValue::Float32(Some(0.0)), + ), + ScalarValue::Float64(Some(v)) if op == Operator::Eq && *v == 0.0 => ( + ScalarValue::Float64(Some(-0.0)), + ScalarValue::Float64(Some(0.0)), + ), + _ => { + return Some(Arc::new(BinaryExpr::new( + Arc::clone(left), + op, + Arc::clone(right), + ))) + } + }; + let equal_to = |zero: ScalarValue| -> Arc { + let zero: Arc = Arc::new(Literal::new(zero)); + if literal_on_left { + Arc::new(BinaryExpr::new(zero, op, Arc::clone(right))) + } else { + Arc::new(BinaryExpr::new(Arc::clone(left), op, zero)) + } + }; + Some(Arc::new(BinaryExpr::new( + equal_to(negative_zero), + Operator::Or, + equal_to(positive_zero), + ))) +} + +/// Whether `expr` is a Float32 or Float64 column, the operand that Parquet pruning reads. +fn is_float_column(expr: &Arc, schema: &Schema) -> bool { + expr.downcast_ref::().is_some() + && matches!( + expr.data_type(schema), + Ok(DataType::Float32 | DataType::Float64) + ) } fn validate_types( @@ -268,7 +368,7 @@ pub fn spark_in_list( negated: bool, schema: &Schema, ) -> Result> { - if !needs_spark_equality(&value.data_type(schema)?) || candidates.is_empty() { + if !is_nested_with_float_leaf(&value.data_type(schema)?) || candidates.is_empty() { return in_list(value, candidates, &negated, schema); } validate_types(&value, &candidates, schema)?; @@ -441,6 +541,7 @@ mod tests { }, Arc::clone(&b), batch.schema().as_ref(), + FloatOperands::Normalize, )?; let inside = spark_in_list( Arc::clone(&a), @@ -458,6 +559,7 @@ mod tests { Operator::Eq, Arc::clone(&lit), batch.schema().as_ref(), + FloatOperands::Normalize, )?; let inside = spark_in_list( Arc::clone(&a), @@ -471,6 +573,7 @@ mod tests { Operator::Eq, Arc::clone(&a), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert_eq!(results(&eq, &batch), results(&reversed, &batch)); } @@ -558,6 +661,7 @@ mod tests { Operator::Eq, Arc::clone(&zero), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert!(matches!( scalar.evaluate(&RecordBatch::new_empty(batch.schema()))?, @@ -568,6 +672,7 @@ mod tests { Operator::Eq, Arc::clone(&b), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert!(eq.nullable(batch.schema().as_ref())?); assert_eq!(eq.data_type(batch.schema().as_ref())?, DataType::Boolean); @@ -623,7 +728,8 @@ mod tests { Arc::clone(&primitive), Operator::Eq, primitive, - batch.schema().as_ref() + batch.schema().as_ref(), + FloatOperands::Normalize, )? .as_ref() .is::()); @@ -635,11 +741,15 @@ mod tests { (Arc::clone(&unresolved), Arc::clone(&a)), (Arc::clone(&a), Arc::clone(&unresolved)), ] { - assert!( - spark_comparison(l, Operator::Eq, r, batch.schema().as_ref())? - .as_ref() - .is::() - ); + assert!(spark_comparison( + l, + Operator::Eq, + r, + batch.schema().as_ref(), + FloatOperands::Normalize + )? + .as_ref() + .is::()); } Ok(()) } @@ -762,4 +872,378 @@ mod tests { assert!(expr.evaluate(&batch).is_err()); Ok(()) } + + const COMPARISONS: [Operator; 8] = [ + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + Operator::IsDistinctFrom, + Operator::IsNotDistinctFrom, + ]; + + /// Spark's answer for `left op right`, given the ordering of two non-null values. + fn spark_answer(op: Operator, ordering: Option>) -> Option { + use Operator::*; + match ordering { + // Both non-null. + Some(Some(ord)) => Some(match op { + Eq | IsNotDistinctFrom => ord.is_eq(), + NotEq | IsDistinctFrom => ord.is_ne(), + Lt => ord.is_lt(), + LtEq => ord.is_le(), + Gt => ord.is_gt(), + GtEq => ord.is_ge(), + _ => unreachable!(), + }), + // Both null. + Some(None) => match op { + IsNotDistinctFrom => Some(true), + IsDistinctFrom => Some(false), + _ => None, + }, + // One null. + None => match op { + IsNotDistinctFrom => Some(false), + IsDistinctFrom => Some(true), + _ => None, + }, + } + } + + fn float_ordering(l: Option, r: Option) -> Option> { + match (l, r) { + (Some(l), Some(r)) => Some(Some(crate::float_semantics::compare_floats(l, r))), + (None, None) => Some(None), + _ => None, + } + } + + /// Every pair of edge values, under every operator, as two columns and as a column and a + /// literal on either side, for Float64 and Float32. + #[test] + fn float_operands_follow_spark_ordering() -> Result<()> { + let values = [ + Some(f64::NEG_INFINITY), + Some(-1.0), + Some(-0.0), + Some(0.0), + Some(f64::MIN_POSITIVE), + Some(f64::INFINITY), + Some(f64::NAN), + // A NaN with the sign bit set, as arithmetic produces on x86-64. + Some(f64::from_bits(0xfff8_0000_0000_0000)), + // A NaN with a payload. + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + None, + ]; + // The Float32 pass computes Spark's answers from the values as Float32 holds them. + let to_f32 = |v: Option| v.map(|v| v as f32 as f64); + for float32 in [false, true] { + let values: Vec<_> = if float32 { + values.iter().map(|&v| to_f32(v)).collect() + } else { + values.to_vec() + }; + let pairs: Vec<_> = values + .iter() + .flat_map(|&l| values.iter().map(move |&r| (l, r))) + .collect(); + let left: Vec<_> = pairs.iter().map(|p| p.0).collect(); + let right: Vec<_> = pairs.iter().map(|p| p.1).collect(); + let array = |v: &[Option]| -> ArrayRef { + if float32 { + Arc::new(arrow::array::Float32Array::from( + v.iter().map(|v| v.map(|v| v as f32)).collect::>(), + )) + } else { + Arc::new(Float64Array::from(v.to_vec())) + } + }; + let batch = batch(array(&left), array(&right)); + let schema = batch.schema(); + let (a, b) = columns(); + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, r)| spark_answer(op, float_ordering(l, r))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} b"); + } + // The first rows pair `values[0]` with each value in turn, so row `i` of column `b` + // holds `values[i]`. + for (i, &value) in values.iter().enumerate() { + let lit = literal(batch.column(1), i); + for op in COMPARISONS { + let column_first = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = left + .iter() + .map(|&l| spark_answer(op, float_ordering(l, value))) + .collect(); + assert_eq!(results(&column_first, &batch), expected, "a {op} {lit}"); + let literal_first = spark_comparison( + Arc::clone(&lit), + op, + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = left + .iter() + .map(|&l| spark_answer(op, float_ordering(value, l))) + .collect(); + assert_eq!(results(&literal_first, &batch), expected, "{lit} {op} a"); + } + } + } + Ok(()) + } + + /// A literal operand is normalized while planning, so the comparison stays `column op + /// literal`. With `FloatOperands::Raw`, as for a scan's data filters, the column and the + /// literal of that shape are left as they are. + #[test] + fn float_operand_shapes() -> Result<()> { + use crate::float_semantics::NormalizeNaNAndZero; + let batch = batch( + Arc::new(Float64Array::from(vec![1.0])), + Arc::new(Float64Array::from(vec![1.0])), + ); + let schema = batch.schema(); + let negative_zero: Arc = + Arc::new(Literal::new(ScalarValue::Float64(Some(-0.0)))); + let expr = spark_comparison( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::clone(&negative_zero), + &schema, + FloatOperands::Normalize, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(binary + .left() + .downcast_ref::() + .is_some()); + let folded = binary.right().downcast_ref::().unwrap(); + assert!(matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == 0)); + + // With `FloatOperands::Raw`, a float column compared with a literal keeps both as they + // are, on either side. + let expr = spark_comparison( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::clone(&negative_zero), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(binary.left().downcast_ref::().is_some()); + assert!(Arc::ptr_eq(binary.right(), &negative_zero)); + let expr = spark_comparison( + Arc::clone(&negative_zero), + Operator::Lt, + Arc::new(Column::new("a", 0)), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &negative_zero)); + assert!(binary.right().downcast_ref::().is_some()); + + // Any other shape is normalized even with `FloatOperands::Raw`, because a reader with + // row-level pushdown drops the rows it rejects. + let (a, b) = columns(); + let expr = spark_comparison( + Arc::clone(&a), + Operator::Lt, + Arc::clone(&b), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(binary + .left() + .downcast_ref::() + .is_some()); + assert!(binary + .right() + .downcast_ref::() + .is_some()); + + // Other operators are not comparisons and keep their operands. + let expr = spark_comparison( + Arc::clone(&a), + Operator::Plus, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &a)); + assert!(Arc::ptr_eq(binary.right(), &b)); + Ok(()) + } + + /// A bloom filter holds the bits of each value, so with `FloatOperands::Raw` an `=` against + /// either zero probes the filter for both zeros, while other operators keep the literal. A + /// NaN literal is normalized together with the column instead, because no single literal + /// stands for every NaN a file can hold. + #[test] + fn raw_float_comparison_literals() -> Result<()> { + use crate::float_semantics::NormalizeNaNAndZero; + use datafusion::physical_expr::utils::{Guarantee, LiteralGuarantee}; + use std::collections::HashSet; + let schema = Schema::new(vec![ + Field::new("f", DataType::Float32, true), + Field::new("d", DataType::Float64, true), + ]); + let scalar = |value: f64, data_type: &DataType| match data_type { + DataType::Float32 => ScalarValue::Float32(Some(value as f32)), + _ => ScalarValue::Float64(Some(value)), + }; + for (index, name, data_type) in [(0, "f", DataType::Float32), (1, "d", DataType::Float64)] { + let column: Arc = Arc::new(Column::new(name, index)); + for zero in [-0.0, 0.0] { + let literal: Arc = + Arc::new(Literal::new(scalar(zero, &data_type))); + for (left, right) in [(&column, &literal), (&literal, &column)] { + let expr = spark_comparison( + Arc::clone(left), + Operator::Eq, + Arc::clone(right), + &schema, + FloatOperands::Raw, + )?; + let guarantees = LiteralGuarantee::analyze(&expr); + assert_eq!(guarantees.len(), 1, "{expr}"); + assert_eq!(guarantees[0].guarantee, Guarantee::In, "{expr}"); + assert_eq!(guarantees[0].column.name(), name, "{expr}"); + assert_eq!( + guarantees[0].literals, + HashSet::from([scalar(-0.0, &data_type), scalar(0.0, &data_type)]), + "{expr}" + ); + } + let expr = spark_comparison( + Arc::clone(&column), + Operator::GtEq, + Arc::clone(&literal), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &column)); + assert!(Arc::ptr_eq(binary.right(), &literal)); + } + } + + let column: Arc = Arc::new(Column::new("d", 1)); + for nan in [f64::NAN, f64::from_bits(0xfff8_0000_0000_0000)] { + let expr = spark_comparison( + Arc::clone(&column), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Float64(Some(nan)))), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(binary + .left() + .downcast_ref::() + .is_some()); + let folded = binary.right().downcast_ref::().unwrap(); + assert!( + matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == f64::NAN.to_bits()) + ); + } + Ok(()) + } + + /// Ordering and null-safe comparisons of lists and structs match `spark_comparator`, which + /// compares float leaves in Spark's order without normalizing them. + #[test] + fn nested_float_operands_follow_spark_ordering() -> Result<()> { + use crate::float_semantics::spark_comparator; + let leaves = [ + Some(-0.0), + Some(0.0), + Some(f64::NAN), + Some(f64::from_bits(0xfff8_0000_0000_0000)), + Some(f64::INFINITY), + Some(1.0), + None, + ]; + // Lists of every pair of leaves, including lists of different lengths. + let mut left_lists = vec![]; + let mut right_lists = vec![]; + for &l in &leaves { + for &r in &leaves { + left_lists.push(Some(vec![l])); + right_lists.push(Some(vec![r])); + left_lists.push(Some(vec![Some(1.0), l])); + right_lists.push(Some(vec![Some(1.0), r, Some(2.0)])); + } + } + left_lists.push(None); + right_lists.push(Some(vec![])); + left_lists.push(None); + right_lists.push(None); + let list = |lists: &[Option>>]| -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::( + lists.to_vec(), + )) + }; + let (left, right) = (list(&left_lists), list(&right_lists)); + let as_struct = |array: &ArrayRef| -> ArrayRef { + Arc::new(StructArray::new( + vec![Arc::new(Field::new("v", array.data_type().clone(), true))].into(), + vec![Arc::clone(array)], + None, + )) + }; + for (l, r) in [ + (Arc::clone(&left), Arc::clone(&right)), + (as_struct(&left), as_struct(&right)), + ] { + let compare = spark_comparator(l.as_ref(), r.as_ref())?; + let batch = batch(Arc::clone(&l), Arc::clone(&r)); + let (a, b) = columns(); + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + batch.schema().as_ref(), + FloatOperands::Normalize, + )?; + let expected: Vec<_> = (0..l.len()) + .map(|i| { + let ordering = match (l.is_null(i), r.is_null(i)) { + (false, false) => Some(Some(compare(i, i))), + (true, true) => Some(None), + _ => None, + }; + spark_answer(op, ordering) + }) + .collect(); + assert_eq!(results(&expr, &batch), expected, "{} {op}", l.data_type()); + } + } + Ok(()) + } } diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs index 60f8614c285..b295f3ea545 100644 --- a/native/spark-expr/src/float_semantics/mod.rs +++ b/native/spark-expr/src/float_semantics/mod.rs @@ -37,10 +37,11 @@ //! [`hash_input`]. //! //! A native expression must follow the rule of the Spark function it replaces, so build it from -//! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes or hashes the -//! values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's total order -//! agrees with `compareDoubles`. Where Comet compares values itself, or has to return the original -//! bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], +//! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes, hashes or +//! compares the values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's +//! total order agrees with `compareDoubles`. Comparison operands go through +//! [`normalize_comparison_operand`]. Where Comet compares values itself, or has to return the +//! original bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], //! [`spark_comparator`] or [`spark_equality`]. Non-canonical NaNs are not a corner case: on x86-64 //! every NaN that arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. @@ -48,9 +49,10 @@ mod compare; mod normalize; pub use compare::{spark_comparator, spark_equality}; +pub(crate) use normalize::is_nested_with_float_leaf; pub use normalize::{ - has_float_leaf, normalize_floats, normalize_nested_floats, NormalizeNaNAndZero, - NormalizeNestedFloats, + has_float_leaf, normalize_comparison_operand, normalize_floats, normalize_nested_floats, + NormalizeNaNAndZero, NormalizeNestedFloats, }; use num::Float; diff --git a/native/spark-expr/src/float_semantics/normalize.rs b/native/spark-expr/src/float_semantics/normalize.rs index b518e826f04..ea575aa0454 100644 --- a/native/spark-expr/src/float_semantics/normalize.rs +++ b/native/spark-expr/src/float_semantics/normalize.rs @@ -24,6 +24,7 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion::common::{Result, ScalarValue}; use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -92,15 +93,17 @@ impl PhysicalExpr for NormalizeNaNAndZero { } fn evaluate(&self, batch: &RecordBatch) -> Result { - let cv = self.child.evaluate(batch)?; - let array = cv.into_array(batch.num_rows())?; - match &self.data_type { - DataType::Float32 | DataType::Float64 => { - Ok(ColumnarValue::Array(normalize_floats(&array))) - } + DataType::Float32 | DataType::Float64 => {} dt => panic!("Unexpected data type {dt:?}"), } + // A scalar stays a scalar, so a constant operand is not expanded into a column. + match self.child.evaluate(batch)? { + ColumnarValue::Array(array) => Ok(ColumnarValue::Array(normalize_floats(&array))), + ColumnarValue::Scalar(value) => { + Ok(ColumnarValue::Scalar(normalize_float_scalar(value))) + } + } } fn children(&self) -> Vec<&Arc> { @@ -124,7 +127,8 @@ impl Display for NormalizeNaNAndZero { } } -/// Normalizes nested IN operands, preserving constants for static membership lookup. +/// Applies [`normalize_nested_floats`] to a list or struct operand of `IN` or a comparison. A +/// constant operand stays a scalar, which keeps static membership lookup available. #[derive(Debug, Eq)] pub struct NormalizeNestedFloats { child: Arc, @@ -132,19 +136,13 @@ pub struct NormalizeNestedFloats { impl NormalizeNestedFloats { /// Wrap nested floating-point operands only; scalar floats keep their existing semantics. + /// An operand that is already wrapped is returned as is. pub fn wrap_if_needed( child: Arc, schema: &Schema, ) -> Result> { let dt = child.data_type(schema)?; - if matches!( - dt, - DataType::List(_) - | DataType::LargeList(_) - | DataType::FixedSizeList(_, _) - | DataType::Struct(_) - ) && has_float_leaf(&dt) - { + if is_nested_with_float_leaf(&dt) && child.downcast_ref::().is_none() { Ok(Arc::new(Self { child })) } else { Ok(child) @@ -214,6 +212,32 @@ impl PhysicalExpr for NormalizeNestedFloats { } } +/// Normalizes an operand of a comparison, so that Arrow's comparison kernels order it the way +/// Spark's SQL ordering does. Arrow orders floats by IEEE 754 total order, which agrees with +/// Spark's once `-0.0` is folded and every NaN canonicalized. +/// +/// A Float32 or Float64 operand is wrapped in [`NormalizeNaNAndZero`], and a list or struct with +/// a float leaf in [`NormalizeNestedFloats`]. A literal is normalized now instead, so that it stays +/// a literal. Other operands, and operands that are already wrapped, are returned as is. +pub fn normalize_comparison_operand( + operand: Arc, + schema: &Schema, +) -> Result> { + if let Some(literal) = operand.downcast_ref::() { + let value = literal.value(); + let normalized = match value.data_type() { + DataType::Float32 | DataType::Float64 => normalize_float_scalar(value.clone()), + dt if is_nested_with_float_leaf(&dt) => { + ScalarValue::try_from_array(&normalize_nested_floats(&value.to_array()?), 0)? + } + _ => return Ok(operand), + }; + return Ok(Arc::new(Literal::new(normalized))); + } + let operand = NormalizeNaNAndZero::wrap_if_needed(operand, schema)?; + NormalizeNestedFloats::wrap_if_needed(operand, schema) +} + /// Applies [`normalize_float`] to a Float32 or Float64 array. Any other array is returned as is, /// including a nested one: [`normalize_nested_floats`] reaches floats inside lists and structs. pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { @@ -232,6 +256,15 @@ pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { } } +/// Applies [`normalize_float`] to a Float32 or Float64 scalar. Other scalars are returned as is. +fn normalize_float_scalar(value: ScalarValue) -> ScalarValue { + match value { + ScalarValue::Float32(v) => ScalarValue::Float32(v.map(normalize_float)), + ScalarValue::Float64(v) => ScalarValue::Float64(v.map(normalize_float)), + other => other, + } +} + /// Whether a Float32 or Float64 field occurs at any depth of a list or struct type. The keys and /// values of a map are not searched. pub fn has_float_leaf(dt: &DataType) -> bool { @@ -245,6 +278,11 @@ pub fn has_float_leaf(dt: &DataType) -> bool { } } +/// Whether `dt` is a list or struct with a Float32 or Float64 field at any depth. +pub(crate) fn is_nested_with_float_leaf(dt: &DataType) -> bool { + dt.is_nested() && has_float_leaf(dt) +} + /// Recursively rebuilds nested arrays with `-0.0` normalized to `0.0` and NaN canonicalized /// in any Float32/Float64 leaves. /// @@ -526,6 +564,120 @@ mod tests { assert!(Arc::ptr_eq(&normalize_floats(&nested), &nested)); } + #[test] + fn test_normalize_nan_and_zero_keeps_scalars() -> Result<()> { + use datafusion::physical_expr::expressions::Literal; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + for (value, expected) in [ + ( + ScalarValue::Float64(Some(-0.0)), + ScalarValue::Float64(Some(0.0)), + ), + ( + ScalarValue::Float64(Some(f64::from_bits(0xfff8_0000_0000_0000))), + ScalarValue::Float64(Some(f64::NAN)), + ), + ( + ScalarValue::Float32(Some(-0.0)), + ScalarValue::Float32(Some(0.0)), + ), + (ScalarValue::Float64(None), ScalarValue::Float64(None)), + ] { + let expr = + NormalizeNaNAndZero::new(value.data_type(), Arc::new(Literal::new(value.clone()))); + // Compare the bits, since every NaN prints the same. + let bits = |v: &ScalarValue| match v { + ScalarValue::Float64(v) => v.map(f64::to_bits), + ScalarValue::Float32(v) => v.map(|v| u64::from(v.to_bits())), + other => panic!("unexpected {other:?}"), + }; + match expr.evaluate(&batch)? { + ColumnarValue::Scalar(actual) => { + assert_eq!(actual.data_type(), expected.data_type()); + assert_eq!(bits(&actual), bits(&expected), "{value:?}"); + } + ColumnarValue::Array(_) => panic!("{value:?} was expanded into an array"), + } + } + Ok(()) + } + + #[test] + fn test_normalize_comparison_operand() -> Result<()> { + use datafusion::physical_expr::expressions::{Column, Literal}; + let list = DataType::List(Arc::new(Field::new("item", DataType::Float64, true))); + let schema = Schema::new(vec![ + Field::new("f64", DataType::Float64, true), + Field::new("list", list, true), + Field::new("i32", DataType::Int32, true), + ]); + let column = |i: usize| -> Arc { + Arc::new(Column::new(schema.field(i).name(), i)) + }; + + // Columns are wrapped once, and only when they hold floats. + let wrapped = normalize_comparison_operand(column(0), &schema)?; + assert!(wrapped.downcast_ref::().is_some()); + let again = normalize_comparison_operand(Arc::clone(&wrapped), &schema)?; + assert!(Arc::ptr_eq(&again, &wrapped)); + let wrapped = normalize_comparison_operand(column(1), &schema)?; + assert!(wrapped.downcast_ref::().is_some()); + let again = normalize_comparison_operand(Arc::clone(&wrapped), &schema)?; + assert!(Arc::ptr_eq(&again, &wrapped)); + let int = column(2); + assert!(Arc::ptr_eq( + &normalize_comparison_operand(Arc::clone(&int), &schema)?, + &int + )); + + // Literals are normalized now and stay literals. + let literal = + |value: ScalarValue| -> Arc { Arc::new(Literal::new(value)) }; + let value = |expr: Arc| -> ScalarValue { + expr.downcast_ref::() + .expect("a literal operand stays a literal") + .value() + .clone() + }; + let bits = |expr: Arc| match value(expr) { + ScalarValue::Float64(Some(v)) => v.to_bits(), + ScalarValue::Float32(Some(v)) => u64::from(v.to_bits()), + other => panic!("unexpected {other:?}"), + }; + let folded = + normalize_comparison_operand(literal(ScalarValue::Float64(Some(-0.0))), &schema)?; + assert_eq!(bits(folded), 0.0f64.to_bits()); + let folded = normalize_comparison_operand( + literal(ScalarValue::Float32(Some(f32::from_bits(0xffc0_0000)))), + &schema, + )?; + assert_eq!(bits(folded), u64::from(f32::NAN.to_bits())); + let nested = ScalarValue::List(Arc::new( + ListArray::from_iter_primitive::([Some(vec![ + Some(-0.0), + Some(f64::from_bits(0xfff8_0000_0000_0000)), + ])]), + )); + let folded = value(normalize_comparison_operand(literal(nested), &schema)?); + let ScalarValue::List(folded) = folded else { + panic!("expected a list literal, got {folded:?}"); + }; + let folded = folded.value(0); + let folded = folded.as_primitive::(); + assert_eq!(folded.value(0).to_bits(), 0.0f64.to_bits()); + assert_eq!(folded.value(1).to_bits(), f64::NAN.to_bits()); + let int = literal(ScalarValue::Int32(Some(1))); + assert!(Arc::ptr_eq( + &normalize_comparison_operand(Arc::clone(&int), &schema)?, + &int + )); + Ok(()) + } + #[test] fn test_wrap_float_keys_once() -> Result<()> { use datafusion::physical_expr::expressions::Column; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 15889d53277..248772fb6c4 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -178,10 +178,12 @@ object CometConf extends ShimCometConf { "When enabled, the native Parquet reader evaluates pushed filters during decode " + "and lazily materializes projected columns for surviving rows (DataFusion's " + "pushdown_filters / late-materialization). Format-level pruning (row-group " + - "statistics, page index, bloom filters) is independent of this flag and runs " + - "whenever Spark's spark.sql.parquet.filterPushdown is enabled. Disabling this " + - "flag still lets format-level pruning work; the per-row eval falls back to " + - "the CometFilter operator above the scan.") + "statistics, page index, bloom filters) runs whenever Spark's " + + "spark.sql.parquet.filterPushdown is enabled, except that while this flag is " + + "enabled, a pushed filter that compares FLOAT or DOUBLE values does not prune, " + + "because the reader must evaluate it with Spark's NaN and signed-zero semantics. " + + "Disabling this flag still lets format-level pruning work; the per-row eval falls " + + "back to the CometFilter operator above the scan.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d339f27965d..4327256dbb5 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,7 +25,7 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, Expression, FloatLiteral, KnownFloatingPointNormalized, NamedExpression, Remainder, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule @@ -895,58 +895,6 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } } - private def normalizePlan(plan: SparkPlan): SparkPlan = { - plan.transformUp { - case p: ProjectExec => - val newProjectList = p.projectList.map(normalize(_).asInstanceOf[NamedExpression]) - ProjectExec(newProjectList, p.child) - case f: FilterExec => - val newCondition = normalize(f.condition) - FilterExec(newCondition, f.child) - } - } - - // Spark will normalize NaN and zero for floating point numbers for several cases. - // See `NormalizeFloatingNumbers` optimization rule in Spark. - // However, one exception is for comparison operators. Spark does not normalize NaN and zero - // because they are handled well in Spark (e.g., `SQLOrderingUtil.compareFloats`). But the - // comparison functions in arrow-rs do not normalize NaN and zero. So we need to normalize NaN - // and zero for comparison operators in Comet. - private def normalize(expr: Expression): Expression = { - expr.transformUp { - case EqualTo(left, right) => - EqualTo(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case EqualNullSafe(left, right) => - EqualNullSafe(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case GreaterThan(left, right) => - GreaterThan(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case GreaterThanOrEqual(left, right) => - GreaterThanOrEqual(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case LessThan(left, right) => - LessThan(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case LessThanOrEqual(left, right) => - LessThanOrEqual(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case Divide(left, right, evalMode) => - Divide(left, normalizeNaNAndZero(right), evalMode) - case Remainder(left, right, evalMode) => - Remainder(left, normalizeNaNAndZero(right), evalMode) - } - } - - private def normalizeNaNAndZero(expr: Expression): Expression = { - expr match { - case _: KnownFloatingPointNormalized => expr - case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr - case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr - case _ => - expr.dataType match { - case _: FloatType | _: DoubleType => - KnownFloatingPointNormalized(NormalizeNaNAndZero(expr)) - case _ => expr - } - } - } - /** * A relation keeps the cache format it was stored in, since `spark.sql.cache.serializer` is * static, so a plan that runs without Comet's native execution still reads relations cached in @@ -977,6 +925,41 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) newPlan } + /** + * Wraps the divisor of a floating-point `Divide` or `Remainder` in a Project or Filter in + * `NormalizeNaNAndZero`. Native comparisons already follow Spark's ordering, including the + * check for a zero divisor, so the quotient does not need this. It also canonicalizes a NaN + * divisor, though, and with it the NaN quotient: `1.0D / (-d)` of a NaN `d` is otherwise a NaN + * with the sign bit set, which `percentile_approx` still orders below every other value + * (https://github.com/apache/datafusion-comet/issues/6519). Remove this once that is fixed. + */ + private def normalizeDivisors(plan: SparkPlan): SparkPlan = { + def normalize(expr: Expression): Expression = expr.transformUp { + case Divide(left, right, evalMode) => Divide(left, normalizeNaNAndZero(right), evalMode) + case Remainder(left, right, evalMode) => + Remainder(left, normalizeNaNAndZero(right), evalMode) + } + plan.transformUp { + case p: ProjectExec => + ProjectExec(p.projectList.map(normalize(_).asInstanceOf[NamedExpression]), p.child) + case f: FilterExec => FilterExec(normalize(f.condition), f.child) + } + } + + private def normalizeNaNAndZero(expr: Expression): Expression = { + expr match { + case _: KnownFloatingPointNormalized => expr + case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr + case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr + case _ => + expr.dataType match { + case _: FloatType | _: DoubleType => + KnownFloatingPointNormalized(NormalizeNaNAndZero(expr)) + case _ => expr + } + } + } + private def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) { @@ -997,7 +980,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { - val normalizedPlan = normalizePlan(plan) + val normalizedPlan = normalizeDivisors(plan) val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { normalizedPlan.transformUp { case p => diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql b/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql new file mode 100644 index 00000000000..61b4b36f0e9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql @@ -0,0 +1,207 @@ +-- 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. + +-- Comparisons follow Spark's SQL ordering for floats (SQLOrderingUtil.compareDoubles): -0.0 +-- equals 0.0, every NaN equals every other NaN, and NaN sorts above every other value. This +-- must hold in every operator that evaluates a comparison, not only in Project and Filter. +-- +-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every +-- platform. That is the NaN that arithmetic produces on x86-64, and Arrow's total order sorts +-- it below -Infinity. +-- +-- A scan pushes its filters into the Parquet reader, which with row-level pushdown drops rows +-- itself, before Spark's Filter above the scan sees them. +-- ConfigMatrix: spark.comet.parquet.rowFilterPushdown.enabled=false,true + +statement +CREATE TABLE float_cmp(id INT, d DOUBLE, f FLOAT) USING parquet + +statement +INSERT INTO float_cmp VALUES + (1, 0.0D, float('0.0')), + (2, double('-0.0'), float('-0.0')), + (3, double('NaN'), float('NaN')), + (4, 1.0D, float('1.0')), + (5, -1.0D, float('-1.0')), + (6, double('Infinity'), float('Infinity')), + (7, NULL, NULL) + +-- Project. `IS DISTINCT FROM` is `NOT (a <=> b)`, which Comet plans as one native operator. +query +SELECT id, -d = d, -d <=> d, -d < d, -d <= d, -d > d, -d >= d, -d != d, -d IS DISTINCT FROM d +FROM float_cmp + +query +SELECT id, -f = f, -f <=> f, -f < f, -f <= f, -f > f, -f >= f, -f != f, -f IS DISTINCT FROM f +FROM float_cmp + +-- Comparisons against a constant on either side. `-0.0D` and `-0.0F` are literals with the sign +-- bit set, while `double('NaN')` is a cast, because the suite turns off constant folding. +query +SELECT id, -d > 0.0D, -d >= -0.0D, -0.0D = d, -d = double('NaN'), double('NaN') <= -d, + -d < double('-0.0'), -d IS DISTINCT FROM double('NaN'), -0.0D IS DISTINCT FROM d, + -f > 0.0F, -0.0F <=> f, -f = float('NaN'), float('-0.0') >= -f +FROM float_cmp + +-- Filter +query +SELECT id FROM float_cmp WHERE -d > 0.0D + +query +SELECT id FROM float_cmp WHERE -f >= f + +-- A stored column compared with a constant, the filter shape that Parquet pruning reads. The +-- scan still returns the -0.0 and NaN rows that Spark's ordering keeps. +query +SELECT id FROM float_cmp WHERE d >= 0.0D + +query +SELECT id FROM float_cmp WHERE f <= -0.0F + +query +SELECT id FROM float_cmp WHERE d = double('NaN') + +query +SELECT id FROM float_cmp WHERE d > double('Infinity') + +-- Aggregate argument +query +SELECT sum(if(-d > 0.0D, 1, 0)), sum(if(-d = d, 1, 0)), sum(if(-f < f, 1, 0)) FROM float_cmp + +-- Aggregate FILTER clause +query +SELECT count(*) FILTER (WHERE -d >= d), count(*) FILTER (WHERE -f <=> f), + count(*) FILTER (WHERE -d IS DISTINCT FROM d) +FROM float_cmp + +-- Grouping by a comparison +query +SELECT -d > 0.0D AS k, count(*) FROM float_cmp GROUP BY -d > 0.0D + +-- Join conditions, for each join strategy +query +SELECT /*+ BROADCAST(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.d >= b.d + +query +SELECT /*+ SHUFFLE_HASH(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.d >= b.d + +query +SELECT /*+ MERGE(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.f >= b.f + +query +SELECT /*+ BROADCAST(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON -a.d > b.d WHERE a.id = 3 + +-- Sort key +query +SELECT id FROM float_cmp ORDER BY -d > 0.0D, id + +-- Generate +query +SELECT id, x FROM float_cmp LATERAL VIEW explode(array(-d > 0.0D, -d = d)) v AS x + +-- Window function argument, which Spark evaluates in a Project below the Window +query +SELECT id, sum(if(-d > 0.0D, 1, 0)) OVER (ORDER BY id) FROM float_cmp + +-- Arrays and structs compare their float leaves the same way, for every operator +statement +CREATE TABLE float_cmp_nested( + id INT, + a ARRAY, b ARRAY, + s STRUCT, u STRUCT, + x ARRAY, y ARRAY) USING parquet + +statement +INSERT INTO float_cmp_nested VALUES + (1, array(double('-0.0')), array(0.0D), + named_struct('v', double('-0.0')), named_struct('v', 0.0D), + array(float('-0.0')), array(float('0.0'))), + (2, array(double('NaN')), array(double('Infinity')), + named_struct('v', double('NaN')), named_struct('v', double('Infinity')), + array(float('NaN')), array(float('Infinity'))), + (3, array(1.0D, NULL), array(1.0D, 2.0D), + named_struct('v', NULL), named_struct('v', 1.0D), + array(float('1.0'), NULL), array(float('1.0'), float('2.0'))), + (4, array(1.0D), array(1.0D, double('-0.0')), + named_struct('v', 1.0D), named_struct('v', 1.0D), + array(), array(float('-0.0'))), + (5, NULL, array(0.0D), NULL, named_struct('v', 0.0D), NULL, array(float('0.0'))) + +query +SELECT id, a = b, a <=> b, a < b, a <= b, a > b, a >= b, a IS DISTINCT FROM b +FROM float_cmp_nested + +query +SELECT id, s = u, s <=> u, s < u, s <= u, s > u, s >= u, s IS DISTINCT FROM u +FROM float_cmp_nested + +query +SELECT id, x = y, x <=> y, x < y, x <= y, x > y, x >= y, x IS DISTINCT FROM y +FROM float_cmp_nested + +-- Negating the elements gives the NaN in row 2 the sign bit +query +SELECT id, transform(a, e -> -e) < b, transform(a, e -> -e) <=> a, + transform(a, e -> -e) IS DISTINCT FROM a +FROM float_cmp_nested + +-- Nested comparisons outside Project and Filter +query +SELECT sum(if(a < b, 1, 0)), count(*) FILTER (WHERE s <=> u) FROM float_cmp_nested + +-- Floats two levels deep, in an array of structs and in a struct of arrays +statement +CREATE TABLE float_cmp_deep( + id INT, + a ARRAY>, b ARRAY>, + s STRUCT>, u STRUCT>) USING parquet + +statement +INSERT INTO float_cmp_deep VALUES + (1, array(named_struct('v', double('-0.0'))), array(named_struct('v', 0.0D)), + named_struct('a', array(double('-0.0'))), named_struct('a', array(0.0D))), + (2, array(named_struct('v', double('NaN'))), array(named_struct('v', double('Infinity'))), + named_struct('a', array(double('NaN'))), named_struct('a', array(double('Infinity')))), + (3, array(named_struct('v', 1.0D), named_struct('v', NULL)), + array(named_struct('v', 1.0D), named_struct('v', 2.0D)), + named_struct('a', array(1.0D, NULL)), named_struct('a', array(1.0D, 2.0D))), + (4, array(named_struct('v', 1.0D)), + array(named_struct('v', 1.0D), named_struct('v', double('-0.0'))), + named_struct('a', NULL), named_struct('a', array(double('-0.0')))), + (5, NULL, array(named_struct('v', 0.0D)), NULL, named_struct('a', array(0.0D))) + +query +SELECT id, a = b, a <=> b, a < b, a <= b, a > b, a >= b, a IS DISTINCT FROM b +FROM float_cmp_deep + +query +SELECT id, s = u, s <=> u, s < u, s <= u, s > u, s >= u, s IS DISTINCT FROM u +FROM float_cmp_deep + +-- Negating the leaves gives the NaN in row 2 the sign bit +query +SELECT id, transform(a, e -> named_struct('v', -e.v)) < b, + transform(a, e -> named_struct('v', -e.v)) <=> a, + named_struct('a', transform(s.a, e -> -e)) <=> s +FROM float_cmp_deep + +query +SELECT sum(if(a < b, 1, 0)), count(*) FILTER (WHERE s <=> u) FROM float_cmp_deep diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql index 969ddf18511..a2cb88d3433 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql @@ -79,3 +79,7 @@ INSERT INTO test_div_zero VALUES (1, 0, 0.0), (0, 0, 0.0) query SELECT a / b, a % b, d / 0.0 FROM test_div_zero + +-- `-0.0D` is a double literal with the sign bit set, and is a zero divisor too +query +SELECT d / -0.0D, d % -0.0D, d / double('-0.0'), d % double('-0.0') FROM test_div_zero diff --git a/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql new file mode 100644 index 00000000000..51135545d44 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql @@ -0,0 +1,52 @@ +-- 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. + +-- Division and remainder pass a NaN divisor through, so `1.0D / (-d)` of a NaN is a NaN with the +-- sign bit set, the NaN that arithmetic produces on x86-64. Spark hashes every NaN alike, and +-- orders every NaN above every other value. ANSI mode divides through a different native path. +-- ConfigMatrix: spark.sql.ansi.enabled=false,true + +statement +CREATE TABLE nan_divisor(d double, f float) USING parquet + +statement +INSERT INTO nan_divisor VALUES (double('NaN'), float('NaN')), (1.5, float(1.5)), (NULL, NULL) + +query +SELECT hash(1.0D / (-d)), xxhash64(1.0D / (-d)), hash(1.0D % (-d)), xxhash64(1.0D % (-d)) +FROM nan_divisor + +query +SELECT hash(1.0F % (-f)), xxhash64(1.0F % (-f)) FROM nan_divisor + +query +SELECT least(1.0D / (-d), 0.0D), greatest(1.0D / (-d), 0.0D), least(1.0D % (-d), 0.0D), + greatest(1.0F % (-f), 0.0F) +FROM nan_divisor + +-- The union keeps the quotient in a projection below the aggregate. +query +SELECT max(q), min(q) +FROM (SELECT 1.0D / (-d) AS q FROM nan_divisor UNION ALL SELECT 0.0D AS q FROM nan_divisor) + +-- percentile_approx orders a NaN with the sign bit set below every other value +-- (https://github.com/apache/datafusion-comet/issues/6519). Until that is fixed, Comet keeps +-- normalizing the divisor, which makes the quotient a canonical NaN. +query +SELECT percentile_approx(q, 1.0D), percentile_approx(q, 0.0D) +FROM (SELECT /*+ COALESCE(1) */ q + FROM (SELECT 1.0D / (-d) AS q FROM nan_divisor UNION ALL SELECT 0.0D AS q FROM nan_divisor)) diff --git a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala index 46683611d50..442cb7c41a7 100644 --- a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala @@ -333,26 +333,6 @@ object CometFloatSemanticsSuite { c.group == group && contexts.contains(c.context) val knownGaps: Seq[KnownGap] = Seq( - KnownGap( - issue(6385), - "A comparison outside Project and Filter compares raw Arrow values, where a sign-bit NaN " + - "sorts below every other value. Equi-join conditions already match.", - c => - in("comparison", "aggregate argument", "aggregate FILTER", "sort key", "Generate")(c) || - (c.group == "comparison" && c.context.endsWith("join condition") && - !Set("=", "<=>").contains(c.variant))), - KnownGap( - issue(6385), - "A data filter pushed into the Parquet reader compares a stored column with a sign-bit " + - "NaN literal raw, and drops rows that Spark keeps.", - c => in("literal comparison", "scan filter")(c) && c.variant != "!="), - KnownGap( - issue(6157), - "Ordering and null-safe comparisons of arrays and structs compare float leaves raw.", - c => - c.group == "comparison" && - (c.context.startsWith("array operands") || c.context.startsWith("struct operands")) && - !Set("=", "!=").contains(c.variant)), KnownGap( issue(5312), "collect_set before Spark 4.2 treats -0.0 and 0.0 as one value and NaNs as distinct.", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala index 63a7c3f0420..d5dff80cb2f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala @@ -1581,6 +1581,92 @@ class CometNativeReaderSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("row-group statistics pruning fires for a floating-point comparison") { + // Native comparisons normalize float operands to follow Spark's ordering, which pruning + // cannot see through. A scan's data filters are built without that normalization, so they + // still prune, and the Filter above the scan applies Spark's semantics to the rows. Spark + // orders NaN above every other value, so `d > 500.0D` matches the NaN in the first row + // group even though every other value there is smaller: pruning must keep that row group. + withTempPath { dir => + withSQLConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") { + spark + .range(0, 1000) + .selectExpr("IF(id = 10, double('NaN'), CAST(id AS DOUBLE)) AS d") + .repartition(1) + .write + .option("parquet.block.size", "1024") + .format("parquet") + .save(dir.toString) + + val parquetFile = dir + .listFiles() + .find(_.getName.endsWith(".parquet")) + .getOrElse(fail("No parquet file was written")) + val reader = ParquetFileReader.open( + org.apache.parquet.hadoop.util.HadoopInputFile + .fromPath(new Path(parquetFile.getAbsolutePath), spark.sessionState.newHadoopConf())) + val numRowGroups = + try reader.getRowGroups.size() + finally reader.close() + assert(numRowGroups > 1, s"Test setup needs >1 row groups, got $numRowGroups") + + val df = spark.read.parquet(dir.toString).where("d > 500.0D") + val (_, cometPlan) = checkSparkAnswerAndOperator(df) + val nativeScans = cometPlan.collect { case n: CometNativeScanExec => n } + assert(nativeScans.nonEmpty, "Expected a CometNativeScanExec") + val metrics = nativeScans.head.metrics + val pruned = metrics("row_groups_pruned_statistics").value + val matched = metrics("row_groups_matched_statistics").value + assert( + pruned > 0 && pruned + matched == numRowGroups, + "Row-group statistics pruning did not fire " + + s"(pruned=$pruned, matched=$matched of $numRowGroups total)") + assert(df.collect().exists(_.getDouble(0).isNaN), "The NaN row was filtered out") + } + } + } + + test("bloom filter pruning matches either signed zero for a floating-point equality") { + // A bloom filter holds the bits of each value, so `-0.0` and `0.0` are separate entries, + // while Spark's `=` matches either zero. The scan's data filter probes for both, so the file + // holding `-0.0` and the file holding `0.0` are both read for either zero. Spark's own reader + // probes with the literal's bits and skips the file holding the other zero, so the expected + // rows are written out instead of compared with Spark's. + import testImplicits._ + withTempPath { dir => + val path = dir.getCanonicalPath + Seq(Seq(-1.0, -0.0, 1.0), Seq(-2.0, 0.0, 2.0)).foreach { values => + values + .toDF("d") + .coalesce(1) + .write + .mode("append") + .option("parquet.bloom.filter.enabled#d", "true") + .option("parquet.enable.dictionary", "false") + .parquet(path) + } + + def bloomFilterPruning(df: DataFrame): (Long, Long) = { + val scans = collect(df.queryExecution.executedPlan) { case s: CometNativeScanExec => s } + assert(scans.size == 1, s"Expected one CometNativeScanExec:\n${df.queryExecution}") + val metrics = scans.head.metrics + ( + metrics("row_groups_pruned_bloom_filter").value, + metrics("row_groups_matched_bloom_filter").value) + } + + Seq("-0.0D", "0.0D").foreach { zero => + val df = spark.read.parquet(path).where(s"d = $zero") + checkAnswer(df, Seq(Row(-0.0), Row(0.0))) + assert(bloomFilterPruning(df) == ((0L, 2L)), s"d = $zero") + } + // The statistics of both files cover 0.5, so only their bloom filters prune them. + val df = spark.read.parquet(path).where("d = 0.5D") + checkAnswer(df, Seq.empty[Row]) + assert(bloomFilterPruning(df) == ((2L, 0L)), "d = 0.5D") + } + } + test("datafusion escape hatch pruning=false disables row-group statistics pruning") { // Regression test pinning the `pruning` field of the newly-plumbed session // `ParquetOptions` at the level users care about: an explicit From 6ce39140f87a862c4cdebd94958f455600034ee9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 7 Oct 2026 15:57:39 +0000 Subject: [PATCH 23/89] perf: key the codegen dispatcher's kernel cache on a digest of the expression (#6709) * perf: key the codegen dispatcher's kernel cache on a digest of the expression Every batch copied the closure-serialized expression out of the dispatcher's first argument, about 7 KB for a Scala UDF, and hashed the copy byte by byte to find the kernel. The serde now computes a SHA-256 digest of the bytes once on the driver and sends it ahead of them. The dispatcher looks the kernel up by the 32-byte digest and reads the serialized expression only to compile on a cache miss. Closes #6705. * perf: build a JVM UDF's literal arguments once instead of for every batch JvmScalarUdfExpr evaluated each literal argument for every batch, which cloned its value and copied the clone into a one-row array. For the codegen dispatcher that is two copies of the serialized expression per batch. The arrays are now built when the expression is created and shared by every batch. * test: add a benchmark of a Scala UDF's per-batch cost on the codegen dispatcher Runs SELECT max(f(c)) end to end at three batch sizes, and calls the dispatcher directly next to the generated kernel it runs, so a cost the dispatcher pays per batch shows up apart from the cost of the UDF itself. * docs: say that a CometUDF must not modify its inputs A literal argument is now the same vector on every batch of a task, so a CometUDF that wrote into one would change what later batches see. A column argument can already share its buffers with the rest of the plan. CometUDF's scaladoc now says inputs are read-only. * test: check that new children rebuild a JVM UDF's literal arrays JvmScalarUdfExpr derives its literal arrays from its children. If with_new_children kept the old arrays, a position that no longer held a literal would go on sending the old literal's value. * docs: say that a CometUDF's literal argument shares buffers, not a vector CometUdfBridge.evaluate imports a new vector for every input on every call and closes it when the call returns. What stays the same is the native array behind a literal argument, which JvmScalarUdfExpr builds once and sends with every batch, so its buffers last as long as the expression rather than the task. * test: check that the serde ships a digest of the serialized expression --- native/spark-expr/src/jvm_udf/mod.rs | 70 +++++- .../apache/comet/serde/CometScalaUDF.scala | 21 +- .../scala/org/apache/comet/udf/CometUDF.scala | 3 + .../udf/codegen/CometScalaUDFCodegen.scala | 89 +++++--- .../comet/CometCodegenSourceSuite.scala | 12 +- .../org/apache/comet/CometCodegenSuite.scala | 80 ++++++- .../CometScalaUDFDispatchBenchmark.scala | 204 ++++++++++++++++++ .../benchmark/CometTimeExtractBenchmark.scala | 7 +- 8 files changed, 432 insertions(+), 54 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 0a37330097b..6355b1a9da6 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -26,6 +26,7 @@ use arrow::record_batch::RecordBatch; use datafusion::common::Result as DFResult; use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::PhysicalExpr; use datafusion_comet_common::{decode_string_arrays, zero_offsets}; @@ -39,6 +40,12 @@ use jni::objects::{Global, JObject, JValue}; pub struct JvmScalarUdfExpr { class_name: String, args: Vec>, + /// The length-1 array sent for each argument that is a literal. Built once here, because + /// building it for every batch copies the value twice: `Literal::evaluate` clones it and + /// `to_array_of_size` copies the clone. The codegen dispatcher's serialized expression is + /// such a literal, several KB for a Scala UDF. `None` for any other argument, and for a + /// literal whose array cannot be built, which `evaluate` then reports. + literal_arrays: Vec>, return_type: DataType, return_nullable: bool, /// Captured at `createPlan` time and threaded here by the planner. Passed through the @@ -76,9 +83,17 @@ impl JvmScalarUdfExpr { !class_name.is_empty(), "JvmScalarUdfExpr requires a non-empty class name" ); + let literal_arrays = args + .iter() + .map(|arg| { + arg.downcast_ref::() + .and_then(|literal| literal.value().to_array_of_size(1).ok()) + }) + .collect(); Self { class_name, args, + literal_arrays, return_type, return_nullable, task_context, @@ -143,9 +158,13 @@ impl PhysicalExpr for JvmScalarUdfExpr { let arrays: Vec = self .args .iter() - .map(|e| match e.evaluate(batch)? { - ColumnarValue::Array(a) => Ok(a), - ColumnarValue::Scalar(s) => s.to_array_of_size(1), + .zip(&self.literal_arrays) + .map(|(e, literal_array)| match literal_array { + Some(a) => Ok(Arc::clone(a)), + None => match e.evaluate(batch)? { + ColumnarValue::Array(a) => Ok(a), + ColumnarValue::Scalar(s) => s.to_array_of_size(1), + }, }) .collect::>()?; @@ -292,3 +311,48 @@ impl PhysicalExpr for JvmScalarUdfExpr { ))) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, BinaryArray}; + use datafusion::common::ScalarValue; + use datafusion::physical_expr::expressions::Column; + + #[test] + fn builds_literal_arguments_once() { + let bytes = vec![7u8; 4096]; + let udf = JvmScalarUdfExpr::new( + "org.example.Udf".to_string(), + vec![ + Arc::new(Literal::new(ScalarValue::Binary(Some(bytes.clone())))), + Arc::new(Column::new("a", 0)), + ], + DataType::Int64, + true, + None, + None, + 0, + 0, + ); + let literal = udf.literal_arrays[0] + .as_ref() + .expect("a literal argument's array is built up front"); + let literal = literal.as_any().downcast_ref::().unwrap(); + assert_eq!(literal.len(), 1); + assert_eq!(literal.value(0), bytes.as_slice()); + assert!(udf.literal_arrays[1].is_none()); + + // New children rebuild the arrays, so a position that no longer holds a literal does not + // keep sending the old literal's array. + let swapped = Arc::new(udf) + .with_new_children(vec![ + Arc::new(Column::new("a", 0)), + Arc::new(Literal::new(ScalarValue::Binary(Some(bytes)))), + ]) + .unwrap(); + let swapped = swapped.downcast_ref::().unwrap(); + assert!(swapped.literal_arrays[0].is_none()); + assert!(swapped.literal_arrays[1].is_some()); + } +} diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 58caedcff14..c79b4839bb3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -117,7 +117,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { } // Serialize via Spark's closure serializer: respects the task context classloader (so user - // UDF jars are visible) and matches Spark's wire format. The bytes become arg 0 of the + // UDF jars are visible) and matches Spark's wire format. The bytes become arg 1 of the // JvmScalarUdf proto and self-describe the expression so this works in cluster mode without // executor-side driver registry state. // @@ -145,12 +145,17 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { s"(${e.getClass.getSimpleName}: ${e.getMessage})") return None } - val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse { - withFallbackReason( - expr, - s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + - "expression payload") - return None + // Arg 0 is a digest of the bytes. The dispatcher finds each batch's kernel by it, so it reads + // the bytes only to compile on a cache miss instead of copying and hashing them for every + // batch. + val exprArgs = Seq(CometScalaUDFCodegen.digest(bytes), bytes).map { payload => + exprToProtoInternal(Literal(payload, BinaryType), inputs, binding).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + + "expression payload") + return None + } } val dataArgs = attrs.map { a => @@ -169,7 +174,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val udfBuilder = ExprOuterClass.JvmScalarUdf .newBuilder() .setClassName(classOf[CometScalaUDFCodegen].getName) - .addArgs(exprArg) + exprArgs.foreach(udfBuilder.addArgs) dataArgs.foreach(udfBuilder.addArgs) udfBuilder .setReturnType(returnTypeProto) diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala index db9178fc32c..a1d72939914 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala @@ -28,6 +28,9 @@ import org.apache.arrow.vector.ValueVector * - Vector arguments arrive at the row count of the current batch. * - Scalar (literal-folded) arguments arrive as length-1 vectors and must be read at index 0. * - The returned vector's length must match `numRows`. + * - Inputs belong to native execution and must not be modified. Each call receives new vectors, + * but a scalar argument's buffers are the same on every call for the life of the expression, + * and a vector argument can share its buffers with the rest of the plan. * * `numRows` mirrors DataFusion's `ScalarFunctionArgs.number_rows` and is the batch row count. * UDFs that always have at least one batch-length input can read length from it and ignore diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index c589fb6fa93..11b2594fdde 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -20,6 +20,7 @@ package org.apache.comet.udf.codegen import java.nio.ByteBuffer +import java.security.MessageDigest import java.util.Collections import java.util.concurrent.atomic.AtomicLong @@ -44,8 +45,11 @@ import org.apache.comet.udf.CometUDF * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the index * of the partition the calling native plan computes, and caches the live instance. * - * Arg 0 is a `VarBinaryVector` scalar carrying the closure-serialized bound `Expression` bytes; - * args 1..N are the data columns the `BoundReference`s read in ordinal order. + * Arg 0 is a `VarBinaryVector` scalar carrying the [[CometScalaUDFCodegen.digest]] of the + * closure-serialized bound `Expression`; arg 1 is a `VarBinaryVector` scalar carrying the + * serialized bytes themselves; args 2..N are the data columns the `BoundReference`s read in + * ordinal order. A batch finds its kernel by the digest, so the serialized expression, several KB + * for a Scala UDF's closure, is read only to compile a kernel on a cache miss. * * Caching hierarchy, broadest scope on the left: * {{{ @@ -82,13 +86,13 @@ import org.apache.comet.udf.CometUDF class CometScalaUDFCodegen extends CometUDF with Logging { /** - * Per-task cache keyed on serialized expression bytes plus per-column specs. The deserialized - * `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, `Rand`'s - * `XORShiftRandom`) that must not be shared across concurrent tasks running the same query; - * keeping the cache per-task gives each task its own copy. A nondeterministic kernel is seeded - * from the partition its plan computes, so its key also holds the plan: each plan a task runs, - * such as each parent partition of a coalesce, gets its own, dropped by `releasePlan` when the - * plan closes. Guarded by `this.synchronized`. + * Per-task cache keyed on the serialized expression's digest plus per-column specs. The + * deserialized `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, + * `Rand`'s `XORShiftRandom`) that must not be shared across concurrent tasks running the same + * query; keeping the cache per-task gives each task its own copy. A nondeterministic kernel is + * seeded from the partition its plan computes, so its key also holds the plan: each plan a task + * runs, such as each parent partition of a coalesce, gets its own, dropped by `releasePlan` + * when the plan closes. Guarded by `this.synchronized`. */ private val kernelCache : mutable.Map[CometScalaUDFCodegen.CacheKey, CometScalaUDFCodegen.CacheEntry] = @@ -114,25 +118,25 @@ class CometScalaUDFCodegen extends CometUDF with Logging { partitionIndex: Int, planId: Long): ValueVector = { require( - inputs.length >= 1, - "CometScalaUDFCodegen requires at least 1 input (serialized expression), " + - s"got ${inputs.length}") - val exprVec = inputs(0).asInstanceOf[VarBinaryVector] + inputs.length >= 2, + "CometScalaUDFCodegen requires at least 2 inputs (expression digest and serialized " + + s"expression), got ${inputs.length}") + val digest = binaryScalar(inputs(0), "expression digest at arg 0") require( - exprVec.getValueCount >= 1 && !exprVec.isNull(0), - "CometScalaUDFCodegen requires non-null serialized expression bytes at arg 0") - val bytes = exprVec.get(0) + digest.length == CometScalaUDFCodegen.DigestLength, + s"CometScalaUDFCodegen requires a ${CometScalaUDFCodegen.DigestLength}-byte expression " + + s"digest at arg 0, got ${digest.length} bytes") // TODO(dict-encoded): kernels assume materialized inputs. Dict-encoded vectors would fail the // cast in `specFor` below. Fix is to materialize at the dispatcher (via // `CDataDictionaryProvider`) or widen `emitTypedGetters` with a dict-index + lookup path. - val numDataCols = inputs.length - 1 + val numDataCols = inputs.length - 2 val dataCols = new Array[ValueVector](numDataCols) val specs = new Array[ArrowColumnSpec](numDataCols) var di = 0 while (di < numDataCols) { - val v = inputs(di + 1) + val v = inputs(di + 2) dataCols(di) = v specs(di) = specFor(v) di += 1 @@ -140,12 +144,12 @@ class CometScalaUDFCodegen extends CometUDF with Logging { val n = numRows val specsSeq = specs.toIndexedSeq - val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(bytes), specsSeq) + val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(digest), specsSeq) // Cache lookup and `process` run under one lock to serialize concurrent Tokio callers that // would otherwise race on the kernel's per-batch instance fields. this.synchronized { - val entry = lookupOrCompile(key, bytes, specsSeq, partitionIndex) + val entry = lookupOrCompile(key, inputs(1), specsSeq, partitionIndex) val out = CometBatchKernelCodegen.allocateOutput( entry.outputField, @@ -168,7 +172,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging { private def lookupOrCompile( key: CometScalaUDFCodegen.CacheKey, - bytes: Array[Byte], + exprVec: ValueVector, specs: IndexedSeq[ArrowColumnSpec], partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = { assert(Thread.holdsLock(this), "lookupOrCompile must run under this.synchronized") @@ -181,6 +185,7 @@ class CometScalaUDFCodegen extends CometUDF with Logging { CometScalaUDFCodegen.cacheHitCount.incrementAndGet() entry case None => + val bytes = binaryScalar(exprVec, "serialized expression at arg 1") val loader = Option(Thread.currentThread().getContextClassLoader) .getOrElse(classOf[Expression].getClassLoader) val boundExpr = @@ -215,6 +220,15 @@ class CometScalaUDFCodegen extends CometUDF with Logging { } } + /** The value of a binary scalar argument, which arrives as a length-1 vector. */ + private def binaryScalar(v: ValueVector, what: String): Array[Byte] = { + val vec = v.asInstanceOf[VarBinaryVector] + require( + vec.getValueCount >= 1 && !vec.isNull(0), + s"CometScalaUDFCodegen requires a non-null $what") + vec.get(0) + } + /** * Build the compile-time spec for one input Arrow vector. Recurses on complex types. * @@ -225,9 +239,10 @@ class CometScalaUDFCodegen extends CometUDF with Logging { * * Top-level `nullable=true` is hardcoded: the cache key does not specialize on per-batch null * density. Schema-declared nullability still reaches the kernel via `BoundReference.nullable` - * embedded in `bytesKey`, so `BoundReference.doGenCode` elides its own `isNullAt` probe on - * non-null columns. `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata, - * which is a schema property and therefore stable across batches. + * embedded in the serialized expression, which the key's digest covers, so + * `BoundReference.doGenCode` elides its own `isNullAt` probe on non-null columns. + * `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata, which is a schema + * property and therefore stable across batches. */ private def specFor(v: ValueVector): ArrowColumnSpec = v match { case map: MapVector => @@ -337,18 +352,24 @@ object CometScalaUDFCodegen { */ private[comet] val NoPlan = -1L + /** Length in bytes of a [[digest]]. */ + val DigestLength = 32 + /** - * Cache key: calling native plan (`NoPlan` for a deterministic expression), serialized - * expression bytes plus per-column compile-time invariants. `hashCode` walks `bytesKey` per - * lookup, so for large ScalaUDF closures it scales with closure size. - * - * TODO(perf-cache-key): if hot, options are a driver-precomputed hash piggybacked through the - * proto, per-instance last-key memoization, or a two-tier cache keyed on the generated source. + * Identifies a closure-serialized bound expression in the kernel cache. The serde computes it + * once on the driver and ships it as arg 0, next to the bytes, so the per-batch lookup hashes + * 32 bytes instead of the whole serialized expression. SHA-256 rather than a cheaper hash, + * because a hit is trusted without comparing the bytes: two expressions whose digests collided + * would share a kernel. + */ + def digest(serialized: Array[Byte]): Array[Byte] = + MessageDigest.getInstance("SHA-256").digest(serialized) + + /** + * Cache key: calling native plan (`NoPlan` for a deterministic expression), the bound + * expression's [[digest]] plus per-column compile-time invariants. */ - final case class CacheKey( - planId: Long, - bytesKey: ByteBuffer, - specs: IndexedSeq[ArrowColumnSpec]) + final case class CacheKey(planId: Long, digest: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec]) /** Snapshot of dispatcher cache counters and current size. */ final case class DispatcherStats(compileCount: Long, cacheHitCount: Long, cacheSize: Int) { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index f1272e151ee..54b59dec9dd 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -1436,8 +1436,8 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("closure-serialized bytes diverge for failOnError / roundOff / timeZoneId variants") { - // The dispatcher caches kernels by the closure-serialized bytes of the bound expression - // (`CacheKey.bytesKey`). For expression classes that carry a runtime-dependent boolean + // The dispatcher caches kernels by a digest of the closure-serialized bytes of the bound + // expression (`CacheKey.digest`). For expression classes that carry a runtime-dependent boolean // (failOnError, roundOff) or string (timeZoneId), two plan instances differing only in that // field must serialize to distinct byte sequences so they receive distinct cache entries. // A collision would let a kernel compiled for one variant (e.g. ANSI throw site) silently @@ -1523,20 +1523,20 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("CacheKey discriminates on ArrowColumnSpec.nullable") { - // Structural regression: same expression bytes and same Arrow vector class with different + // Structural regression: same expression digest and same Arrow vector class with different // `nullable` must produce non-equal cache keys. The dispatcher today hardcodes `nullable=true` // for top-level specs, so the two variants don't both arise from runtime data, but the case // class equality contract still has to discriminate so that any future tiered cache or test // construction can rely on it. The non-nullable variant's generated source emits a literal // `false` from `isNullAt`, distinct codegen output that we never want to silently share with // the nullable variant. - val bytes = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3)) + val digest = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3)) val nullable = IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = true)) val nonNullable = IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = false)) - val k1 = CometScalaUDFCodegen.CacheKey(0L, bytes, nullable) - val k2 = CometScalaUDFCodegen.CacheKey(0L, bytes, nonNullable) + val k1 = CometScalaUDFCodegen.CacheKey(0L, digest, nullable) + val k2 = CometScalaUDFCodegen.CacheKey(0L, digest, nonNullable) assert( k1 != k2, "expected nullable=true and nullable=false specs to produce distinct cache keys") diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5a80865da92..7af83a5494c 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -624,6 +624,30 @@ class CometCodegenSuite .contains(Set("hypot", "cast", "checkoverflow", "add"))) } + test("the serde ships a digest of the serialized expression at arg 0 (#6705)") { + // The dispatcher trusts a cache hit on the digest without comparing the bytes, so the digest + // the serde ships must be the digest of the bytes it ships next to it. Calling the dispatcher + // directly cannot catch a mismatch, because those tests build their own digest. + val x = AttributeReference("x", DoubleType, nullable = false)() + def payload(e: Expression): (Array[Byte], Array[Byte]) = { + val proto = QueryPlanSerde.exprToProto(e, Seq(x)).get + assert(proto.hasJvmScalarUdf) + val args = proto.getJvmScalarUdf.getArgsList + ( + args.get(0).getLiteral.getBytesVal.toByteArray, + args.get(1).getLiteral.getBytesVal.toByteArray) + } + + val (hypotDigest, hypotBytes) = payload(Hypot(x, Literal(4.0d))) + val (otherDigest, otherBytes) = payload(Hypot(x, Literal(5.0d))) + assert(hypotDigest.sameElements(CometScalaUDFCodegen.digest(hypotBytes))) + assert(otherDigest.sameElements(CometScalaUDFCodegen.digest(otherBytes))) + + // Expressions that differ only in a literal must not share a kernel. + assert(!hypotBytes.sameElements(otherBytes)) + assert(!hypotDigest.sameElements(otherDigest)) + } + test("tags copied onto the shared TrueLiteral do not leak into unrelated plans") { // Catalyst copies a rewritten node's tags onto its replacement, so a tagged expression that an // earlier query rewrote into `Literal.TrueLiteral` brands that process-wide singleton for the @@ -866,7 +890,7 @@ class CometCodegenSuite "same UDF over nullable and non-nullable columns gets distinct kernels with independent state") { // Two columns, same type, different schema-declared nullability. Same UDF applied to each // alongside a per-projection MonotonicallyIncreasingID. Each projection has its own MII - // child (different bytesKey), so each kernel must have its own counter advancing 0..N-1. + // child (a different digest), so each kernel must have its own counter advancing 0..N-1. // If the dispatcher collapses them onto one kernel or shares state somehow, the counters // would interleave and the output would diverge from Spark. spark.udf.register("withId", (s: String, id: Long) => s"${s}_${id}") @@ -2329,6 +2353,7 @@ class CometCodegenSuite // because Spark 4.1 still rejects TIME columns in file-based data sources, so no SQL query // can produce a TIME input today. val timeVec = new TimeNanoVector("tm", CometArrowAllocator) + val digestVec = new VarBinaryVector("digest", CometArrowAllocator) val exprVec = new VarBinaryVector("expr", CometArrowAllocator) var out: ValueVector = null try { @@ -2342,21 +2367,72 @@ class CometCodegenSuite val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr) val bytes = new Array[Byte](serialized.remaining()) serialized.get(bytes) + digestVec.allocateNew() + digestVec.setSafe(0, CometScalaUDFCodegen.digest(bytes)) + digestVec.setValueCount(1) exprVec.allocateNew() exprVec.setSafe(0, bytes) exprVec.setValueCount(1) - out = new CometScalaUDFCodegen().evaluate(Array(exprVec, timeVec), 2) + out = new CometScalaUDFCodegen().evaluate(Array(digestVec, exprVec, timeVec), 2) val comet = CometVector.getVector(out.asInstanceOf[FieldVector], null) assert(comet.getLong(0) === 45296000000000L) assert(comet.isNullAt(1)) } finally { if (out != null) out.close() + digestVec.close() exprVec.close() timeVec.close() } } + test("dispatcher finds a compiled kernel by the expression digest alone (#6705)") { + // The serde ships a digest of the serialized expression at arg 0 and the bytes at arg 1, and + // the dispatcher reads the bytes only to compile on a cache miss. The second call passes a + // null at arg 1, so it succeeds only if the digest finds the kernel the first call compiled. + val expr = Add(BoundReference(0, LongType, nullable = true), Literal(1L)) + val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr) + val bytes = new Array[Byte](serialized.remaining()) + serialized.get(bytes) + + def binaryScalar(name: String, value: Array[Byte]): VarBinaryVector = { + val v = new VarBinaryVector(name, CometArrowAllocator) + v.allocateNew() + if (value == null) v.setNull(0) else v.setSafe(0, value) + v.setValueCount(1) + v + } + val digestVec = binaryScalar("digest", CometScalaUDFCodegen.digest(bytes)) + val exprVec = binaryScalar("expr", bytes) + val nullExprVec = binaryScalar("expr", null) + val input = new BigIntVector("x", CometArrowAllocator) + try { + input.allocateNew(2) + input.set(0, 41L) + input.setNull(1) + input.setValueCount(2) + + val dispatcher = new CometScalaUDFCodegen() + Seq(exprVec, nullExprVec).foreach { arg1 => + val out = + dispatcher.evaluate(Array(digestVec, arg1, input), 2).asInstanceOf[BigIntVector] + try { + assert(out.get(0) === 42L) + assert(out.isNull(1)) + } finally out.close() + } + + // The layout before the digest, with the serialized expression at arg 0, is refused + // rather than taken for a digest. + val e = intercept[IllegalArgumentException] { + dispatcher.evaluate(Array(exprVec, input), 2) + } + assert(e.getMessage.contains("expression digest"), e.getMessage) + } finally { + Seq(digestVec, exprVec, nullExprVec, input).foreach(_.close()) + } + } + // Runtime coverage for nullable nested `getStruct` / `getArray` / `getMap` element reads is // exercised through HOFs in `CometCodegenHOFSuite`. Static emitter assertions live in // `CometCodegenSourceSuite`. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala new file mode 100644 index 00000000000..91a9337f26a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala @@ -0,0 +1,204 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.arrow.vector.{BigIntVector, ValueVector, VarBinaryVector} +import org.apache.spark.SparkEnv +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSeq, BindReferences, Expression, ScalaUDF} +import org.apache.spark.sql.execution.adaptive.QueryStageExec + +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * What a row-level Scala UDF costs on the codegen dispatcher, split into the cost paid per batch + * and the cost paid per row. + * + * Part 1 runs `SELECT max(f(c))` over a Parquet `bigint` column end to end at several batch + * sizes. The total of a cost paid once per batch grows as batches shrink, and the total of a cost + * paid per row does not. A tenth of the rows are null. Spark wraps a UDF with a primitive + * parameter in a null check, `if(isnull(c), null, f(knownnotnull(c)))`, which runs natively + * around the dispatched call; a UDF with a boxed parameter gets no such wrapper and converts its + * input through Spark's encoder instead. + * + * Part 2 calls the dispatcher directly on the same in-memory batch over and over, with no Spark + * query and no native code, and does the same with the generated kernel it runs. The difference + * between the two is the dispatcher's own work for each batch: reading its arguments and finding + * the kernel. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometScalaUDFDispatchBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometScalaUDFDispatchBenchmark-**results.txt". + */ +object CometScalaUDFDispatchBenchmark extends CometBenchmarkBase { + + private val Rows = 4L * 1024 * 1024 + private val BatchSizes = Seq(1024, 8192, 65536) + private val LayerBatchRows = Seq(1024, 8192) + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + spark.udf.register("prim_add_one", (x: Long) => x + 1) + spark.udf.register( + "boxed_add_one", + (x: java.lang.Long) => if (x == null) null else java.lang.Long.valueOf(x + 1)) + + withTempPath { dir => + withTempTable(tbl, "parquetV1Table") { + spark.range(Rows).createOrReplaceTempView(tbl) + prepareTable(dir, spark.sql(s"SELECT IF(id % 10 = 0, NULL, id * 7919) AS c FROM $tbl")) + + checkPlans() + BatchSizes.foreach(endToEnd) + LayerBatchRows.foreach(jvmLayers) + } + } + } + + /** `dispatch` turns the dispatcher on; the cases without a UDF leave it off. */ + private case class EndToEndCase(label: String, f: String, dispatch: Boolean) + + private val endToEndCases = Seq( + EndToEndCase("dispatch, primitive param", "prim_add_one(c)", dispatch = true), + EndToEndCase("dispatch, boxed param", "boxed_add_one(c)", dispatch = true), + EndToEndCase("native c + 1", "c + 1", dispatch = false), + EndToEndCase("no function, max(c)", "c", dispatch = false)) + + private def configs(dispatch: Boolean, batchSize: Int): Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> batchSize.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) + + private def query(f: String): String = s"SELECT max($f) FROM parquetV1Table" + + /** Fails unless Comet runs every operator of every case. */ + private def checkPlans(): Unit = endToEndCases.foreach { + case EndToEndCase(label, f, dispatch) => + withSQLConf(configs(dispatch, BatchSizes.head): _*) { + val df = spark.sql(query(f)) + df.noop() + val plan = stripAQEPlan(df.queryExecution.executedPlan) + // A query stage is a leaf, so check the plan inside each one separately. + val stagePlans = collect(plan) { case s: QueryStageExec => s.plan } + (plan +: stagePlans) + .flatMap(findFirstNonCometOperator(_, classOf[QueryStageExec])) + .headOption + .foreach { op => + throw new IllegalStateException(s"$label ran ${op.nodeName} in Spark:\n$plan") + } + } + } + + private def endToEnd(batchSize: Int): Unit = { + val benchmark = + new Benchmark(s"max(f(c)), batch size $batchSize", Rows, minNumIters = 5, output = output) + endToEndCases.foreach { case EndToEndCase(label, f, dispatch) => + benchmark.addCase(label) { _ => + withSQLConf(configs(dispatch, batchSize): _*) { + spark.sql(query(f)).noop() + } + } + } + benchmark.run() + } + + /** The bound UDF subtree the serde would dispatch, and its closure-serialized bytes. */ + private def boundUdf(f: String): (Expression, Array[Byte]) = { + val plan = spark.sql(s"SELECT $f FROM parquetV1Table").queryExecution.optimizedPlan + val udf = plan.expressions.flatMap(_.collect { case u: ScalaUDF => u }).head + val attrs = udf.collect { case a: AttributeReference => a }.distinct + val bound = BindReferences.bindReference(udf: Expression, AttributeSeq(attrs)) + val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(bound) + val bytes = new Array[Byte](buffer.remaining()) + buffer.get(bytes) + (bound, bytes) + } + + private def binaryScalar(value: Array[Byte]): VarBinaryVector = { + val v = new VarBinaryVector("arg", CometArrowAllocator) + v.allocateNew() + v.setSafe(0, value) + v.setValueCount(1) + v + } + + /** The arguments the serde ships ahead of the data columns. */ + private def dispatcherArgs(bytes: Array[Byte]): Seq[VarBinaryVector] = + Seq(CometScalaUDFCodegen.digest(bytes), bytes).map(binaryScalar) + + private def jvmLayers(batchRows: Int): Unit = { + val in = new BigIntVector("c", CometArrowAllocator) + in.allocateNew(batchRows) + var i = 0 + while (i < batchRows) { + if (i % 10 == 0) in.setNull(i) else in.set(i, i * 7919L) + i += 1 + } + in.setValueCount(batchRows) + val inputs: Array[ValueVector] = Array(in) + val batchesPerIter = (Rows / batchRows).toInt + + val benchmark = + new Benchmark(s"JVM layers, $batchRows-row batches", Rows, minNumIters = 5, output = output) + + def addBatchCase(label: String)(evaluate: () => ValueVector): Unit = + benchmark.addCase(label) { _ => + var b = 0 + while (b < batchesPerIter) { + evaluate().close() + b += 1 + } + } + + val toClose = Seq("prim_add_one(c)" -> "primitive param", "boxed_add_one(c)" -> "boxed param") + .flatMap { case (f, param) => + val (bound, bytes) = boundUdf(f) + val specs = IndexedSeq(ArrowColumnSpec(classOf[BigIntVector], nullable = true)) + val kernel = CometBatchKernelCodegen.compile(bound, specs).newInstance() + kernel.init(0) + val field = CometBatchKernelCodegen.toFfiArrowField("r", bound.dataType, bound.nullable) + addBatchCase(s"generated kernel only, $param") { () => + val out = CometBatchKernelCodegen.allocateOutput(field, batchRows, 0) + kernel.process(inputs, out, batchRows) + out.setValueCount(batchRows) + out + } + + val leading = dispatcherArgs(bytes) + val dispatcher = new CometScalaUDFCodegen + val dispatchInputs: Array[ValueVector] = (leading :+ in).toArray + addBatchCase(s"dispatcher evaluate, $param (${bytes.length}-byte expression)") { () => + dispatcher.evaluate(dispatchInputs, batchRows) + } + leading + } + + try benchmark.run() + finally { + toClose.foreach(_.close()) + in.close() + } + } +} diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala index 025fe43d459..38d4ac8b549 100644 --- a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala +++ b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala @@ -85,6 +85,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { private def runKernels(precision: Int, nulls: Boolean): Unit = { val size = 10000 val time = new TimeNanoVector("t", CometArrowAllocator) + val digest = new VarBinaryVector("digest", CometArrowAllocator) val serialized = new VarBinaryVector("expr", CometArrowAllocator) try { time.allocateNew(size) @@ -114,11 +115,14 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(expr) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) + digest.allocateNew() + digest.setSafe(0, CometScalaUDFCodegen.digest(bytes)) + digest.setValueCount(1) serialized.allocateNew() serialized.setSafe(0, bytes) serialized.setValueCount(1) val dispatcher = new CometScalaUDFCodegen() - val inputs = Array[org.apache.arrow.vector.ValueVector](serialized, time) + val inputs = Array[org.apache.arrow.vector.ValueVector](digest, serialized, time) val result = dispatcher.evaluate(inputs, size) try { val vector = @@ -170,6 +174,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { } } } finally { + digest.close() serialized.close() time.close() } From 4b543b90d9a7a05219fa810098b6f1a5bd0f273a Mon Sep 17 00:00:00 2001 From: KAZUYUKI TANIMURA Date: Wed, 7 Oct 2026 16:46:56 +0000 Subject: [PATCH 24/89] fix: keep a Spark scan unconverted when the plan reads input_file_name (#6703) * fix: keep a Spark scan unconverted when the plan reads input_file_name input_file_name, input_file_block_start and input_file_block_length read InputFileBlockHolder, a thread-local that the reader producing the rows sets as it moves from file to file. CometScanRule already keeps Spark's scan when a plan uses them, but with spark.comet.convert.parquet.enabled CometExecRule then wrapped that scan in CometSparkToColumnarExec. The conversion fills a batch and Comet operators above it pull more before they emit, so by the time the Spark project evaluated the expressions the reader had moved to a later file or unset the holder. Rows reported another file's values or "", -1, -1. Move CometScanRule's check into a shared helper and refuse to convert any leaf to Arrow when the plan reads these expressions, recording a fallback reason. This covers the Parquet, CSV and JSON conversions for V1 and V2 scans, and other sparkToColumnar leaves such as RDDScan, whose HadoopRDD sets the same holder. Closes #6573. * fix: keep a typed Dataset's output unconverted when the plan reads input_file_name The typed Dataset conversion from upstream converts the rows a SerializeFromObjectExec produces to Arrow. Like the conversion of a leaf, it fills a batch before the operators above it see a row, so it reads ahead of the file that a Spark project evaluating input_file_name expects. Apply the same fallback to it, sharing the reason with the leaf conversion. Also switch the RDD test to spark.comet.convert.rdd.enabled, which replaced the deprecated spark.comet.sparkToColumnar setting for RDDScan, and mention typed Dataset output in the user guide note. * docs: note the input_file_name fallback for typed Dataset output in the operator guide --- .../adding_a_new_operator.md | 5 +- docs/source/user-guide/latest/datasources.md | 5 ++ .../apache/comet/rules/CometExecRule.scala | 40 ++++++++++++++-- .../apache/comet/rules/CometScanRule.scala | 26 +++++++---- .../apache/comet/exec/CometExecSuite.scala | 46 +++++++++++++++++++ .../comet/exec/CometTypedDatasetSuite.scala | 20 ++++++++ 6 files changed, 129 insertions(+), 13 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 09cac3a1c72..85946f44fd9 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -148,7 +148,10 @@ own operators itself. Spark computes a typed operation's rows one at a time, as the conversion fills a whole Arrow batch first. So the rule leaves the output unconverted where a limit, a `mapPartitions` function, or code reading `Dataset.rdd` could stop reading it early, unless an operator that reads all of its input first, such as an exchange, a sort, or a hash aggregate, -sits in between. Fusing the deserializer, the `Invoke` that calls the user function, and the +sits in between. For the same reason, it leaves the output unconverted when the plan uses +`input_file_name()`, `input_file_block_start()` or `input_file_block_length()`: filling the batch +moves the scan's reader past the file that these report when Spark evaluates them above the +conversion. Fusing the deserializer, the `Invoke` that calls the user function, and the serializer of `Dataset.map` into one projection in the JVM codegen dispatcher was tried in [#5714](https://github.com/apache/datafusion-comet/pull/5714) and dropped. The dispatcher only calls into Spark's own classes, and the conversion gets nearly the same speedup for `map` while also diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index eb56d1e8970..8682633d71c 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -87,6 +87,11 @@ The same types apply to the output of typed `Dataset` operations, such as `map`, converts when `spark.comet.convert.typedDataset.enabled=true`. A column of any other type keeps the operators above the typed operation on Spark. +Comet does not convert a source, or the output of a typed `Dataset` operation, when the query uses +`input_file_name()`, `input_file_block_start()` or `input_file_block_length()`, so the operators +above it run in Spark. These functions report the file that Spark's reader is on when Spark evaluates +them, and the conversion and the Comet operators above it would read ahead of that. + ## Data Catalogs ### Apache Iceberg diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 4327256dbb5..fdf37b4e2bf 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -186,6 +186,11 @@ object CometExecRule { */ private val TYPED_DATASET_PARTIAL_READER: TreeNodeTag[String] = TreeNodeTag[String]("comet.typedDatasetPartialReader") + + /** Why Comet does not convert rows to Arrow in a plan that reads `InputFileBlockHolder`. */ + private val INPUT_FILE_BLOCK_FALLBACK_REASON: String = + "Spark to Arrow conversion is not compatible with input_file_name, " + + "input_file_block_start, or input_file_block_length" } /** @@ -409,6 +414,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) */ // spotless:on private def transform(plan: SparkPlan): SparkPlan = { + // Walks the whole plan, so it is lazy: only consulted once a node could be converted. + lazy val readsInputFileBlock = CometScanRule.readsInputFileBlock(plan) + if (CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf)) { tagPartiallyReadTypedDatasetOutputs(plan) } @@ -484,7 +492,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.") } - if (shouldApplySparkToColumnar(conf, scan)) { + if (shouldApplySparkToColumnar(conf, scan, readsInputFileBlock)) { convertToComet(scan, CometSparkToColumnarExec).getOrElse(scan) } else { scan @@ -506,7 +514,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case op: LeafExecNode if hasEnabledHandler(op) => convertToComet(op, allExecs(op.getClass)) .orElse { - if (shouldApplySparkToColumnar(conf, op)) { + if (shouldApplySparkToColumnar(conf, op, readsInputFileBlock)) { convertToComet(op, CometSparkToColumnarExec) } else { None @@ -514,7 +522,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } .getOrElse(op) - case op if shouldApplySparkToColumnar(conf, op) => + case op if shouldApplySparkToColumnar(conf, op, readsInputFileBlock) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) // Typed Dataset operations (`map`, `flatMap`, `mapPartitions`, `mapGroups`, ...) pass JVM @@ -530,6 +538,10 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) "Comet does not convert the output of a typed Dataset operation when " + s"$reader can stop reading it early, because filling an Arrow batch would " + "run the user function on rows that Spark never reaches") + // The conversion reads ahead of input_file_name and friends, as it does over a leaf. + // See shouldApplySparkToColumnar. + case None if readsInputFileBlock => + withFallbackReason(op, CometExecRule.INPUT_FILE_BLOCK_FALLBACK_REASON) case None => convertTypedDatasetOutput(op) } @@ -1420,7 +1432,27 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) convertToComet(s, CometShuffleExchangeExec).getOrElse(preserveSparkAggregateBuffers(s))) } - private def shouldApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = { + private def shouldApplySparkToColumnar( + conf: SQLConf, + op: SparkPlan, + readsInputFileBlock: => Boolean): Boolean = { + // A converted leaf reads ahead of the Spark operator that evaluates input_file_name and + // friends: the conversion fills a whole batch, and Comet operators above it may pull more + // before they emit. By the time Spark evaluates them, the leaf's reader may have moved on to + // a later file or unset InputFileBlockHolder at the end of its input, so rows would report + // another file's values or the unset defaults. Leave the leaf on Spark so that the plan above + // it stays on Spark too. + if (!canApplySparkToColumnar(conf, op)) { + false + } else if (readsInputFileBlock) { + withFallbackReason(op, CometExecRule.INPUT_FILE_BLOCK_FALLBACK_REASON) + false + } else { + true + } + } + + private def canApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = { // Only consider converting leaf nodes to columnar currently, so that all the following // operators can have a chance to be converted to columnar. Leaf operators that output // columnar batches, such as Spark's vectorized readers, will also be converted to native diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index b109fa158ed..b80136314b7 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -354,14 +354,10 @@ case class CometScanRule(session: SparkSession) withFallbackReason(scanExec, "Native Parquet Variant scans do not support encryption") return None } - // input_file_name, input_file_block_start, and input_file_block_length read from - // InputFileBlockHolder, a thread-local set by Spark's FileScanRDD. The native DataFusion - // scan does not use FileScanRDD, so these expressions would return empty/default values. - if (plan.exists(node => - node.expressions.exists(_.exists { - case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength => true - case _ => false - }))) { + // input_file_name, input_file_block_start, and input_file_block_length read values that + // Spark's FileScanRDD sets. The native DataFusion scan does not use FileScanRDD, so these + // expressions would return empty/default values. + if (CometScanRule.readsInputFileBlock(plan)) { withFallbackReason( scanExec, "Native Parquet scan is not compatible with input_file_name, " + @@ -1185,6 +1181,20 @@ case class CometScanTypeChecker() extends DataTypeSupport with CometTypeShim { object CometScanRule extends Logging { + /** + * Whether any node in `plan` evaluates `input_file_name`, `input_file_block_start` or + * `input_file_block_length`. These read `InputFileBlockHolder`, a thread-local that the reader + * producing the rows (`FileScanRDD`, the V2 file readers, `HadoopRDD`, `NewHadoopRDD`, + * connector readers such as Iceberg's) sets as it moves from file to file, so they only return + * each row's values when Spark evaluates them as that reader produces the row. + */ + def readsInputFileBlock(plan: SparkPlan): Boolean = + plan.exists(node => + node.expressions.exists(_.exists { + case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength => true + case _ => false + })) + // Memo of `NativeBase.isObjectStoreSchemeSupported`, keyed by the probe URL rather than the // scheme: object_store's parser keys on (scheme, host-presence), so an authorityless URL would // otherwise poison the authority-bearing form of the same scheme. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 465da3cdd98..478c62f4eda 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -4713,6 +4713,52 @@ class CometExecSuite extends CometTestBase { } } + // https://github.com/apache/datafusion-comet/issues/6573 + private val inputFileColumns = + Seq("input_file_name()", "input_file_block_start()", "input_file_block_length()", "id") + private val inputFileFallbackReason = + "Spark to Arrow conversion is not compatible with input_file_name" + + test("input_file_name above a converted Spark file scan keeps the scan on Spark") { + for (format <- Seq("parquet", "csv", "json"); v1List <- Seq("", format)) { + withTempPath { dir => + spark.range(9000).repartition(3).write.format(format).save(dir.toString) + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> v1List, + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_CSV_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "true") { + // The filter would run in Comet above the conversion. Without a Comet operator there, + // Comet drops the conversion and Spark's project reads the scan directly. + val df = spark.read + .schema("id LONG") + .format(format) + .load(dir.toString) + .where("id >= 0") + .selectExpr(inputFileColumns: _*) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(df, inputFileFallbackReason) + assert(collect(cometPlan) { case c: CometSparkToColumnarExec => c }.isEmpty) + } + } + } + } + + test("input_file_name above a converted RDD scan keeps the scan on Spark") { + withTempPath { dir => + spark.range(9000).repartition(3).selectExpr("CAST(id AS STRING)").write.text(dir.toString) + withSQLConf(CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + // HadoopRDD sets the file Spark evaluates input_file_name against, as FileScanRDD does. + val rows = spark.sparkContext.textFile(dir.toString).map(line => Row(line.toLong)) + val df = spark + .createDataFrame(rows, new StructType().add("id", LongType)) + .where("id >= 0") + .selectExpr(inputFileColumns: _*) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(df, inputFileFallbackReason) + assert(collect(cometPlan) { case c: CometSparkToColumnarExec => c }.isEmpty) + } + } + } + test("read CSV file") { Seq("", "csv").foreach { v1List => withSQLConf( diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala index 9b809e24e8d..c9e1a349e3d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala @@ -340,6 +340,26 @@ class CometTypedDatasetSuite extends CometTestBase { } } + // https://github.com/apache/datafusion-comet/issues/6573 + convertTest("input_file_name above a map keeps the map's output on Spark") { + withTempPath { dir => + spark.range(9000).repartition(3).write.parquet(dir.toString) + // The filter would run in Comet above the conversion, which would read ahead of the files + // Spark's reader is on when Spark evaluates input_file_name in the project. + val df = spark.read + .parquet(dir.toString) + .as[Long] + .map(_ + 1) + .toDF("id") + .where("id >= 0") + .selectExpr("input_file_name()", "input_file_block_start()", "input_file_block_length()") + val (_, plan) = checkSparkAnswerAndFallbackReason( + df, + "Spark to Arrow conversion is not compatible with input_file_name") + assert(conversions(plan).isEmpty, plan) + } + } + test("off by default") { withRecs() { ds => val (_, plan) = checkSparkAnswer(ds.map(r => TypedDsRec(r.a + 1, r.b)).groupBy("b").count()) From e76c35f1e782b78525d447ba09d6f580f9bf6dec Mon Sep 17 00:00:00 2001 From: dustin Date: Wed, 7 Oct 2026 18:30:51 +0000 Subject: [PATCH 25/89] fix: load the Iceberg storage scheme lists from the native factory (#6065) * feat: publish the native Iceberg storage scheme list over JNI The native storage factory now exposes the schemes it can open per access mode, and a JNI entry returns them, so the JVM read and write gates can load the list instead of mirroring it by hand. Tests keep the list and the factory's match arms in step in both directions. * fix: load the Iceberg gate scheme lists from the native storage factory The JVM Iceberg scan and write gates hardcoded the schemes the native storage factory can open and had already drifted from it. Load both lists over JNI through a new IcebergStorageSchemes object, keeping fallback constants only for when the library cannot load, and pin those constants to what native publishes in CometScanSchemeFallbackSuite. A wasb-backed Iceberg table now falls back to Spark with a reason naming the scheme instead of failing natively at execution. * fix: gate the Iceberg scheme lists on the native list and match schemes verbatim Make builtin_storage_schemes the single point of change by rejecting any scheme it does not list before the factory match, drop memory from the read list since an OpenDAL memory backend is a fresh empty store, and load the JVM sets lazily behind the library-loaded check so a disabled Comet never touches the native library. The JVM gates now match the built-in set verbatim, as OpenDAL strips the scheme prefix case-sensitively at open time. Tests round-trip the JNI lists directly and cover the wasb write gate. * fix: drop the fallback scheme lists and match Iceberg aliases verbatim The JVM scheme sets are only read after the native library has loaded, so the hand-maintained fallback lists never did work. `load` now returns the empty set when the library is not loaded, and a blank native list parses to the empty set, each with a warning. iceberg-rust builds the S3 prefix from the parsed scheme, which is lowercase, and compares it to the raw path, so a mixed-case alias location passed the gates and failed at open. Aliases are now matched as written on the Iceberg path, on the JVM gate and in the native factory, like the built-in schemes. The data sources guide says an Iceberg location must spell the alias in lowercase. --- docs/source/user-guide/latest/datasources.md | 5 +- .../src/execution/operators/iceberg_common.rs | 197 +++++++++++++++--- native/core/src/lib.rs | 25 +++ .../parquet/objectstore/s3_blob_fs_support.rs | 6 +- .../java/org/apache/comet/NativeBase.java | 12 ++ .../comet/iceberg/IcebergStorageSchemes.scala | 77 +++++++ .../apache/comet/rules/CometScanRule.scala | 44 ++-- .../operator/CometIcebergNativeWrite.scala | 13 +- .../hadoop/fs/FakeWasbSchemeFileSystem.java | 51 +++++ .../CometIcebergWriteDetectionSuite.scala | 14 ++ .../rules/CometScanSchemeFallbackSuite.scala | 160 +++++++++++++- 11 files changed, 538 insertions(+), 66 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala create mode 100644 spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index 8682633d71c..f351aff5b6f 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -294,7 +294,10 @@ URLs. This is opt-in and disabled by default. Enable it by listing the schemes to treat as S3-compliant aliases in `spark.hadoop.fs.comet.s3Compliant.schemes` (Hadoop key `fs.comet.s3Compliant.schemes`), a comma-separated, case-insensitive list. This mirrors the -existing `fs.comet.libhdfs.schemes` config. +existing `fs.comet.libhdfs.schemes` config. The list entries are case-insensitive, but an Iceberg +table location must write the alias in lowercase (`blob://`, not `BLOB://`): the native Iceberg +reader opens each recorded location as written, and its S3 backend accepts only a lowercase +scheme prefix, so Comet declines such a location up front and leaves that scan to Spark. ```shell --conf spark.hadoop.fs.comet.s3Compliant.schemes=blob diff --git a/native/core/src/execution/operators/iceberg_common.rs b/native/core/src/execution/operators/iceberg_common.rs index 262325c30cb..a17eb076e0d 100644 --- a/native/core/src/execution/operators/iceberg_common.rs +++ b/native/core/src/execution/operators/iceberg_common.rs @@ -48,16 +48,12 @@ const ICEBERG_PROVIDER_CLASS_PROPERTY: &str = "s3.comet.credential.provider.clas /// iceberg-storage-opendal's own settings, such as `opendal.io-timeout-ms`. const STORAGE_PROPERTY_PREFIXES: &[&str] = &["s3.", "gcs.", "adls.", "client.", "opendal."]; -/// Pick an OpenDAL storage backend from a URI's scheme. `file` (or no scheme) falls through to -/// the local file system. `memory` is used by the write path to assemble manifest bytes that -/// stay entirely in-process. For S3, the Comet credential bridge is wired in when a provider -/// class is configured; `access_mode` is forwarded to the JVM SPI so the read and write paths can -/// be granted different (e.g. read-only vs read-write) credentials. +/// Pick an OpenDAL storage backend for a URI whose scheme `builtin_storage_schemes` lists for +/// `access_mode`, or that is an opted-in S3-compliant alias written in lowercase; anything else is +/// rejected before the match, so the list alone decides what each mode admits. For S3, the Comet +/// credential bridge is wired in when a provider class is configured and `access_mode` is +/// forwarded to the JVM SPI. /// -/// The JVM planner mirrors the schemes admitted here by hand, so that it can decline a scan -/// cleanly instead of failing at execution. Changing the arms below means updating -/// `CometScanRule.icebergReadableSchemes` (reads) and -/// `CometIcebergNativeWrite.SupportedStorageSchemes` (writes). /// The bool is false when a read fell back to opendal's default chain because the configured S3 /// access provider failed to initialise; such a `FileIO` must not be cached, so the next task /// retries. @@ -67,24 +63,22 @@ pub(crate) fn storage_factory_for( catalog_name: &str, access_mode: AccessMode, ) -> Result<(Arc, bool), DataFusionError> { + // Verbatim match: OpenDAL strips the scheme prefix from every path case-sensitively at open + // time, so admitting `S3://` here would only defer the failure. Aliases are held to the same + // rule (`is_iceberg_alias_scheme`). The JVM gates match verbatim. let scheme = scheme_of(path); + if !builtin_storage_schemes(access_mode).contains(&scheme) + && !is_s3_family_scheme(scheme, catalog_properties) + { + return Err(DataFusionError::Execution(format!( + "Unsupported storage scheme: {scheme}" + ))); + } match scheme { "file" => Ok((Arc::new(OpenDalStorageFactory::Fs), true)), "memory" => Ok((Arc::new(OpenDalStorageFactory::Memory), true)), "gs" => Ok((Arc::new(OpenDalStorageFactory::Gcs), true)), - // Reads keep the OSS backend they have always had (CometScanRule admits `oss` scan - // locations through HadoopFileIO). Writes fail closed: Comet does not forward `oss.*` - // properties into the FileIO and no test covers the write path, so OSS-specific - // endpoint/credential configuration could silently be dropped. The JVM write gate - // already declines `oss` locations; this is the native-side backstop. - "oss" => match access_mode { - AccessMode::Read => Ok((Arc::new(OpenDalStorageFactory::Oss), true)), - AccessMode::Write => Err(DataFusionError::Execution( - "OSS is not supported for native Iceberg writes (oss.* properties are not \ - forwarded to the native FileIO)" - .to_string(), - )), - }, + "oss" => Ok((Arc::new(OpenDalStorageFactory::Oss), true)), // s3, s3a, and any opted-in s3-compliant alias (e.g. blob) route to the S3 backend. Listed // last so the built-in backends above stay authoritative even if one of their schemes is // also named in `fs.comet.s3Compliant.schemes`. An alias additionally gets a wrapper that @@ -93,7 +87,7 @@ pub(crate) fn storage_factory_for( // location-scoped provider gets a storage that serves each file with the credential of // its policy location -- see iceberg_location_scoped. s if is_s3_family_scheme(s, catalog_properties) => { - let alias = is_s3_compliant_alias_scheme(s, catalog_properties); + let alias = is_iceberg_alias_scheme(s, catalog_properties); let (access, cacheable) = build_s3_access(path, catalog_properties, catalog_name, access_mode)?; let factory: Arc = match access { @@ -111,6 +105,7 @@ pub(crate) fn storage_factory_for( }; Ok((factory, cacheable)) } + // Only a listed scheme without an arm reaches here; the tests below fail on that. _ => Err(DataFusionError::Execution(format!( "Unsupported storage scheme: {scheme}" ))), @@ -303,6 +298,17 @@ fn cached_file_io( Ok(file_io) } +/// The single point of change for the built-in schemes `storage_factory_for` admits per access +/// mode; the JVM read and write gates load these lists over JNI. `memory` is write-only: an OpenDAL +/// memory backend is a fresh empty in-process store the write path assembles manifests in, so a +/// read finds nothing. `oss` is read-only: no `oss.*` property is forwarded and no test covers it. +pub(crate) fn builtin_storage_schemes(access_mode: AccessMode) -> &'static [&'static str] { + match access_mode { + AccessMode::Read => &["file", "gs", "oss", "s3", "s3a"], + AccessMode::Write => &["file", "memory", "gs", "s3", "s3a"], + } +} + pub(crate) fn load_file_io( catalog_properties: &HashMap, reference_path: &str, @@ -587,7 +593,19 @@ pub(crate) fn scheme_of(path: &str) -> &str { /// iceberg-rust's storage factory has no `s3n` backend and the Scala Iceberg scheme gate /// (`isIcebergReadableScheme`) rejects `s3n` before a path ever reaches this operator. fn is_s3_family_scheme(scheme: &str, catalog_properties: &HashMap) -> bool { - matches!(scheme, "s3" | "s3a") || is_s3_compliant_alias_scheme(scheme, catalog_properties) + matches!(scheme, "s3" | "s3a") || is_iceberg_alias_scheme(scheme, catalog_properties) +} + +/// True if `scheme` is an opted-in S3-compliant alias written exactly as the lowercase form of its +/// list entry. The shared `is_s3_compliant_alias_scheme` is case-insensitive, which is safe for +/// the Parquet path because it rewrites an alias URL to `s3://` before anything opens it. The +/// Iceberg path opens the recorded location as written, and OpenDAL's S3 backend checks it against +/// a `scheme://bucket/` prefix whose scheme comes from `Url::parse` and so is lowercase, so +/// `BLOB://bucket/key` would pass a case-insensitive gate here and fail at open time. The JVM +/// Iceberg gate matches the same way. +fn is_iceberg_alias_scheme(scheme: &str, catalog_properties: &HashMap) -> bool { + !scheme.bytes().any(|b| b.is_ascii_uppercase()) + && is_s3_compliant_alias_scheme(scheme, catalog_properties) } #[cfg(test)] @@ -823,7 +841,25 @@ mod tests { // oss.* property forwarding exists and is tested. assert!(factory_result("oss://bucket/db/table", AccessMode::Read).is_ok()); let err = factory_result("oss://bucket/db/table", AccessMode::Write).unwrap_err(); - assert!(err.contains("OSS"), "unexpected error: {err}"); + assert!( + err.contains("Unsupported storage scheme: oss"), + "unexpected error: {err}" + ); + } + + #[test] + fn memory_scheme_is_writable_but_not_readable() { + // The write path assembles manifests in a fresh in-process memory store. A read against + // a new memory store can never find data, so the factory must decline it up front. + assert!(factory_result("memory:manifest.avro", AccessMode::Write).is_ok()); + assert!(factory_result("memory:///manifest.avro", AccessMode::Write).is_ok()); + for path in ["memory:manifest.avro", "memory:///key.parquet"] { + let err = factory_result(path, AccessMode::Read).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: memory"), + "unexpected error for {path}: {err}" + ); + } } #[test] @@ -831,13 +867,85 @@ mod tests { for mode in [AccessMode::Read, AccessMode::Write] { assert!(factory_result("file:///tmp/x", mode).is_ok()); assert!(factory_result("/tmp/no-scheme", mode).is_ok()); - assert!(factory_result("memory:manifest.avro", mode).is_ok()); // No credential provider configured: the default chain applies in both modes. assert!(factory_result("s3://bucket/db/table", mode).is_ok()); assert!(factory_result("gs://bucket/db/table", mode).is_ok()); } } + #[test] + fn scheme_list_pre_check_is_load_bearing() { + // The oss and memory arms build a backend for either mode; only their absence from the + // list for the other mode rejects them. Both facts are asserted so that adding an + // access-mode branch back into an arm, or listing the scheme, breaks this test. + assert!(factory_result("oss://bucket/path", AccessMode::Read).is_ok()); + assert!(!builtin_storage_schemes(AccessMode::Write).contains(&"oss")); + let err = factory_result("oss://bucket/path", AccessMode::Write).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: oss"), + "unexpected error: {err}" + ); + assert!(factory_result("memory:///path", AccessMode::Write).is_ok()); + assert!(!builtin_storage_schemes(AccessMode::Read).contains(&"memory")); + let err = factory_result("memory:///path", AccessMode::Read).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: memory"), + "unexpected error: {err}" + ); + } + + #[test] + fn mixed_case_scheme_is_rejected_for_both_modes() { + // OpenDAL strips the scheme prefix from every path case-sensitively at open time + // (`S3://bucket/key` fails its `s3://bucket/` prefix check), so the factory must not + // admit what the open cannot serve. The JVM gates match the built-in set verbatim too. + for mode in [AccessMode::Read, AccessMode::Write] { + for path in ["S3://bucket/key", "File:///tmp/x", "GS://bucket/key"] { + let err = factory_result(path, mode).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {path} in {mode:?}: {err}" + ); + } + } + } + + #[test] + fn alias_scheme_is_matched_verbatim() { + // An alias location is opened as written, and OpenDAL's S3 backend checks it against a + // lowercase `scheme://bucket/` prefix, so a mixed-case alias that passed a + // case-insensitive gate here would only fail at open time. The list entry may be written + // in any case; the location's scheme must match its lowercase form exactly. + for listed in ["blob", " BLOB ", "minio,Blob"] { + let props = HashMap::from([( + "fs.comet.s3Compliant.schemes".to_string(), + listed.to_string(), + )]); + for mode in [AccessMode::Read, AccessMode::Write] { + assert!( + storage_factory_for("blob://bucket/key", &props, "test_cat", mode).is_ok(), + "blob://bucket/key must be admitted for {mode:?} with list {listed:?}" + ); + for path in [ + "BLOB://bucket/key", + "Blob://bucket/key", + "BLOB:///bucket/key", + ] { + let err = storage_factory_for(path, &props, "test_cat", mode) + .map(|_| ()) + .expect_err(&format!( + "{path} must be rejected for {mode:?} with list {listed:?}" + )) + .to_string(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {path} in {mode:?}: {err}" + ); + } + } + } + } + #[test] fn unknown_scheme_is_rejected() { let err = factory_result("hdfs://nn/db/table", AccessMode::Read).unwrap_err(); @@ -869,6 +977,43 @@ mod tests { assert!(has_explicit_s3_credentials(&blank)); } + #[test] + fn listed_schemes_are_accepted_by_storage_factory() { + for mode in [AccessMode::Read, AccessMode::Write] { + for scheme in builtin_storage_schemes(mode) { + let url = format!("{scheme}://bucket/path"); + assert!( + factory_result(&url, mode).is_ok(), + "{scheme} is listed for {mode:?} but the factory rejects it" + ); + } + } + assert!(builtin_storage_schemes(AccessMode::Read).contains(&"oss")); + assert!(!builtin_storage_schemes(AccessMode::Write).contains(&"oss")); + assert!(!builtin_storage_schemes(AccessMode::Read).contains(&"memory")); + assert!(builtin_storage_schemes(AccessMode::Write).contains(&"memory")); + } + + #[test] + fn unlisted_schemes_are_rejected_for_both_modes() { + let unlisted = [ + "hdfs", "abfs", "abfss", "wasb", "wasbs", "gcs", "http", "https", "azure", + ]; + for mode in [AccessMode::Read, AccessMode::Write] { + for scheme in unlisted { + let err = factory_result(&format!("{scheme}://bucket/path"), mode).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {scheme} in {mode:?}: {err}" + ); + assert!( + !builtin_storage_schemes(mode).contains(&scheme), + "{scheme} must not be listed for {mode:?}" + ); + } + } + } + #[test] fn scheme_of_extracts_scheme_from_all_uri_forms() { // Host-bearing and hostless/opaque vendor forms must resolve to the same scheme, so an diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index b712a65c115..5eb56aaec01 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -31,6 +31,7 @@ extern crate datafusion_comet_jni_bridge; use jni::{ objects::{JClass, JString}, + sys::{jboolean, jstring, JNI_FALSE}, EnvUnowned, }; use log::info; @@ -52,6 +53,9 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::cloud::s3::credential_bridge::AccessMode; +use crate::execution::operators::iceberg_common::builtin_storage_schemes; + pub mod alloc_accounting; pub mod cloud; pub mod comet_native_udf_bridge; @@ -270,6 +274,27 @@ pub fn tzdata_version() -> &'static str { chrono_tz::IANA_TZDB_VERSION } +/// JNI: the comma-joined built-in storage schemes the native Iceberg scan (`for_write` false) or +/// write (`for_write` true) path admits. This is the source of truth the JVM read and write gates +/// load. Opt-in S3-compliant alias schemes are not answered here; the JVM adds them from catalog +/// properties. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeBase_icebergStorageSchemes( + env: EnvUnowned, + _: JClass, + for_write: jboolean, +) -> jstring { + try_unwrap_or_throw(&env, |env| { + let access_mode = if for_write != JNI_FALSE { + AccessMode::Write + } else { + AccessMode::Read + }; + let joined = builtin_storage_schemes(access_mode).join(","); + Ok(env.new_string(joined)?.into_raw()) + }) +} + // Creates a default log4rs config, which logs to console with log level. fn default_logger_config(log_level: &str) -> CometResult { let console_append = ConsoleAppender::builder() diff --git a/native/core/src/parquet/objectstore/s3_blob_fs_support.rs b/native/core/src/parquet/objectstore/s3_blob_fs_support.rs index f455f8c3849..ab09e7d24c3 100644 --- a/native/core/src/parquet/objectstore/s3_blob_fs_support.rs +++ b/native/core/src/parquet/objectstore/s3_blob_fs_support.rs @@ -98,8 +98,10 @@ pub(crate) fn normalize_object_store_url( /// True if `scheme` is a configured s3-compliant alias (`fs.comet.s3Compliant.schemes`; empty or /// unset means none). These route to `AmazonS3` but are not recognized by /// `ObjectStoreScheme::parse`. `s3a` is excluded -- object_store knows it and callers special-case -/// it. Shared by the Parquet normalizer above and the Iceberg storage-factory gate -/// (`iceberg_common::storage_factory_for`) so both admit the same schemes. +/// it. Case-insensitive, which is safe for the Parquet normalizer above because it rewrites the +/// alias to `s3://` before anything opens the path. The Iceberg storage-factory gate opens the +/// recorded location as written, so it wraps this in `iceberg_common::is_iceberg_alias_scheme`, +/// which additionally requires the scheme to be lowercase. pub(crate) fn is_s3_compliant_alias_scheme( scheme: &str, object_store_configs: &HashMap, diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index ab4eee856ba..7952e736116 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -384,4 +384,16 @@ public static void releaseNative() throws Throwable { * @return the version compiled into libcomet */ public static native String getTzdataVersion(); + + /** + * The comma-joined URL schemes the native Iceberg storage factory can open, for reads ({@code + * forWrite} false) or writes ({@code forWrite} true). This is the authoritative list the JVM + * Iceberg scan and write gates load, so the planner never hardcodes (and drifts from) the set the + * native factory actually builds. Opt-in S3-compliant alias schemes are not included; the JVM + * adds them from {@code fs.comet.s3Compliant.schemes}. + * + * @param forWrite true for the write path's schemes, false for the scan path's + * @return the supported schemes joined by commas, e.g. "file,memory,gs,s3,s3a" + */ + public static native String icebergStorageSchemes(boolean forWrite); } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala new file mode 100644 index 00000000000..51f4b5f7019 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala @@ -0,0 +1,77 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import java.util.Locale + +import org.apache.spark.internal.Logging + +import org.apache.comet.NativeBase + +/** + * The storage schemes the native Iceberg storage factory publishes over JNI, so the JVM scan and + * write gates decline what native cannot open instead of failing at execution. Native is the only + * source of these lists. Every caller sits behind `isCometLoaded`, so a JVM where the library is + * not loaded never plans a native scan or write, and both sets are simply empty there. + */ +private[comet] object IcebergStorageSchemes extends Logging { + + lazy val read: Set[String] = load(forWrite = false) + lazy val write: Set[String] = load(forWrite = true) + + /** + * Splits the comma-joined native list, trimming and lowercasing each entry. A null or blank + * list throws: `builtin_storage_schemes` is a fixed non-empty constant natively, so a loaded + * library that publishes nothing can only be a build or marshalling bug, and a warning here + * would silently disable every native Iceberg scan and write. + */ + private[comet] def parse(joined: String, forWrite: Boolean): Set[String] = { + val schemes = Option(joined).toSeq + .flatMap(_.split(",")) + .map(_.trim.toLowerCase(Locale.ROOT)) + .filter(_.nonEmpty) + .toSet + if (schemes.isEmpty) { + val path = if (forWrite) "write" else "read" + throw new IllegalStateException( + s"Comet native library published no Iceberg $path scheme list") + } + schemes + } + + /** + * Loads one list over JNI, or answers the empty set with a warning when the library is not + * loaded; that branch never runs a plan, because every caller sits behind `isCometLoaded`. + * `isLoaded` is a parameter so the unloaded answer can be tested without unloading the library. + * Once loaded, anything but a non-empty list is a build bug that must fail loudly: a native + * fault while answering the probe propagates, and a blank answer throws in `parse`. + */ + private[comet] def load( + forWrite: Boolean, + isLoaded: Boolean = NativeBase.isLoaded): Set[String] = { + if (isLoaded) { + parse(NativeBase.icebergStorageSchemes(forWrite), forWrite) + } else { + val path = if (forWrite) "write" else "read" + logWarning(s"Comet native library is not loaded; the Iceberg $path scheme list is empty") + Set.empty + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index b80136314b7..c7c9a31909d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -47,7 +47,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.{CometConf, DataTypeSupport, NativeBase} import org.apache.comet.CometConf._ import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, isSpark35Plus, withFallbackReason, withFallbackReasons} -import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflection} +import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflection, IcebergStorageSchemes} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.parquet.CometParquetUtils.{encryptionEnabled, isEncryptionConfigSupported, readFieldId} import org.apache.comet.serde.operator.{CometIcebergNativeScan, CometNativeScan} @@ -1286,20 +1286,25 @@ object CometScanRule extends Logging { org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometScan") /** - * Schemes Comet's native Iceberg scan can actually open, mirroring the match arms in - * `native/core/src/execution/operators/iceberg_common.rs::storage_factory_for`. Deliberately - * NOT delegated to `isNativelyReadableScheme`: object_store recognizes schemes (http/https, - * azure, memory) that iceberg-rust's OpenDAL storage factory cannot build, and admitting them - * here turns a clean JVM fallback into a native runtime "Unsupported storage scheme" error. Add - * here what you add to `storage_factory_for` (currently Aliyun `oss` and GCS `gs`). - * S3-compliant aliases like `blob` are opt-in via `fs.comet.s3Compliant.schemes` (see - * `isIcebergReadableScheme`), not hardcoded, since the native planner opens them via S3. The - * write path keeps its own list (`CometIcebergNativeWrite.SupportedStorageSchemes`), which - * differs deliberately: it excludes `oss` (fails closed, see `storage_factory_for`) and - * includes `memory`. + * Schemes Comet's native Iceberg scan can open, loaded from the native storage factory over JNI + * so this gate cannot drift from `storage_factory_for`. Lazy so that constructing the rule does + * not touch the native library before `isCometLoaded` has been consulted. Opt-in aliases from + * `fs.comet.s3Compliant.schemes` are additive (see `isIcebergReadableScheme`); the write path + * loads its own set (`CometIcebergNativeWrite.SupportedStorageSchemes`). */ - private val icebergReadableSchemes: Set[String] = - Set("file", "s3", "s3a", "gs", "oss") + private lazy val icebergReadableSchemes: Set[String] = IcebergStorageSchemes.read + + /** + * True when the Iceberg scan gate admits `scheme`, written exactly as recorded. Native opens a + * location by its raw scheme, and OpenDAL's S3 backend checks it against a `scheme://bucket/` + * prefix whose scheme comes from `Url::parse` and so is lowercase: `S3://` is not `s3://`, and + * `BLOB://` is not an opted-in `blob`. The alias set is lowercase + * (`NativeConfig.parseSchemeSet`), so an alias is admitted only when the location writes it in + * lowercase. The Parquet gate stays case-insensitive because it rewrites alias URLs to `s3://` + * before anything opens them. + */ + private def isAdmittedIcebergScheme(scheme: String, s3CompliantSchemes: Set[String]): Boolean = + icebergReadableSchemes.contains(scheme) || s3CompliantSchemes.contains(scheme) /** * "Supported schemes: ..." suffix shared by the Iceberg scheme-fallback messages. Lists the @@ -1320,8 +1325,7 @@ object CometScanRule extends Logging { s3CompliantSchemes: Set[String]): Boolean = { val scheme = uri.getScheme if (scheme == null) return true - val lower = scheme.toLowerCase(Locale.ROOT) - icebergReadableSchemes.contains(lower) || s3CompliantSchemes.contains(lower) + isAdmittedIcebergScheme(scheme, s3CompliantSchemes) } /** @@ -1383,9 +1387,6 @@ object CometScanRule extends Logging { // hasOpenableAuthority); non-empty => decline. One example suffices for the message. var hostlessLocation: Option[String] = None - // Union of the two admitted scheme sets, built once so the per-file loop does one lookup. - val openableSchemes = icebergReadableSchemes ++ s3CompliantSchemes - // Classify one data/delete file location; see `icebergReadableSchemes` for why that allowlist // is narrower than the Parquet native gate. Runs per data and delete file, so the scheme and // the bucket are each derived once and threaded down. @@ -1396,9 +1397,8 @@ object CometScanRule extends Logging { // A schemeless local path routes to iceberg-rust's LocalFs and needs no host. val scheme = uri.getScheme if (scheme == null) return - val lower = scheme.toLowerCase(Locale.ROOT) - if (!openableSchemes.contains(lower)) { - unsupportedSchemes += lower + if (!isAdmittedIcebergScheme(scheme, s3CompliantSchemes)) { + unsupportedSchemes += scheme } else if (!hasOpenableAuthority(uri, s3CompliantSchemes)) { if (hostlessLocation.isEmpty) hostlessLocation = Some(rawPath) } else { diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index e55014812be..34c07ff8dd5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, Icebe import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason -import org.apache.comet.iceberg.{IcebergReflection, PositionDeltaWrite, ReplaceDataWrite} +import org.apache.comet.iceberg.{IcebergReflection, IcebergStorageSchemes, PositionDeltaWrite, ReplaceDataWrite} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.Operator @@ -86,12 +86,11 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // `timestamp_ns`, `geometry` or `geography` today, so those are declined in case it learns to. private val UnsupportedWriteTypeIds: Set[String] = Set("UUID", "VARIANT", "UNKNOWN", "TIMESTAMP_NANO", "GEOMETRY", "GEOGRAPHY") - // `oss` is deliberately absent: iceberg-rust has an OSS backend, but Comet does not forward - // `oss.*` catalog properties to it and no functional test covers the path, so an OSS write - // could silently drop endpoint/credential configuration. Fail closed until it is covered. - // `gs` is additionally gated on the resolved FileIO (`requireGcsFileIOForGcsDataLocation`). - private val SupportedStorageSchemes: Set[String] = - Set("file", "memory", "s3", "s3a", "gs") + // Loaded from the native storage factory: `builtin_storage_schemes` in + // `native/core/src/execution/operators/iceberg_common.rs` is the single point of change and + // explains why `oss` is read-only and `memory` write-only. Lazy so constructing the serde does + // not touch the native library. `gs` is additionally gated on the resolved FileIO below. + private lazy val SupportedStorageSchemes: Set[String] = IcebergStorageSchemes.write // Supported schemes whose native backend is local and needs no host. Every other supported // scheme reads its bucket from the URL host (`requireSupportedStorageScheme`). private val LocalStorageSchemes: Set[String] = Set("file", "memory") diff --git a/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java new file mode 100644 index 00000000000..95ff00b818a --- /dev/null +++ b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java @@ -0,0 +1,51 @@ +/* + * 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. + */ + +package org.apache.comet.hadoop.fs; + +import java.net.URI; + +import org.apache.hadoop.fs.RawLocalFileSystem; + +/** + * A local-disk-backed FileSystem that reports the {@code wasb} scheme, so a test can create and + * read an Iceberg table under a {@code wasb://} warehouse without Azure. Used to assert that the + * Iceberg scan and write gates decline a table whose storage scheme the native Iceberg storage + * factory has no arm for, instead of claiming it and failing at execution. + */ +public class FakeWasbSchemeFileSystem extends RawLocalFileSystem { + + public static final String PREFIX = "wasb://fake-container"; + + public FakeWasbSchemeFileSystem() { + // Avoid `URI scheme is not "file"` error on + // RawLocalFileSystem$DeprecatedRawLocalFileStatus.getOwner + RawLocalFileSystem.useStatIfAvailable(); + } + + @Override + public String getScheme() { + return "wasb"; + } + + @Override + public URI getUri() { + return URI.create(PREFIX); + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 49238eee228..922b5d6073f 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -545,6 +545,20 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("fall-back: mixed-case data location scheme (native opens the location verbatim)") { + // OpenDAL strips the scheme prefix from a path case-sensitively, so `S3://` cannot be opened + // natively even though `s3://` can. The gate must match the scheme verbatim and decline it. + withDetectionCatalog { dir => + createTable( + dir, + "mixed_case_scheme", + partitionSpec = "", + properties = + Some("'write.data.path'='S3://nonexistent-bucket/iceberg/db/mixed_case_scheme'")) + assertUnsupportedContainsAllowingWriteFailure("mixed_case_scheme", "storage scheme", "S3") + } + } + test("fall-back: hostless hdfs:/ data location is read as hdfs, not file") { // Hadoop normalises `hdfs:///p` to `hdfs:/p`; with no `://` the gate used to call it `file`. withDetectionCatalog { dir => diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala index 5a14f9f592e..21011fae65d 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala @@ -27,11 +27,12 @@ import java.util.UUID import org.apache.commons.io.FileUtils import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, SaveMode} -import org.apache.spark.sql.comet.CometScanExec +import org.apache.spark.sql.comet.{CometIcebergNativeScanExec, CometIcebergWriteExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} -import org.apache.comet.CometConf -import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem} +import org.apache.comet.{CometConf, CometIcebergTestBase, ExtendedExplainInfo, NativeBase} +import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem, FakeWasbSchemeFileSystem} +import org.apache.comet.iceberg.IcebergStorageSchemes /** * Comet's native readers go through object_store, which understands a fixed set of URL schemes. A @@ -43,7 +44,7 @@ import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem} * native probe is cached per probe URL so an authorityless URL can't poison the authority-bearing * form of the same scheme. */ -class CometScanSchemeFallbackSuite extends CometTestBase { +class CometScanSchemeFallbackSuite extends CometTestBase with CometIcebergTestBase { private var fakeRootDir: File = _ @@ -54,6 +55,9 @@ class CometScanSchemeFallbackSuite extends CometTestBase { // cluster. `hdfs` is natively readable by default, so this scan must be CLAIMED, not declined. conf.set("spark.hadoop.fs.hdfs.impl", "org.apache.comet.hadoop.fs.FakeHdfsSchemeFileSystem") conf.set("spark.hadoop.fs.defaultFS", FakeHDFSFileSystem.PREFIX) + // Back `wasb` with a local FS so an Iceberg table can live under a `wasb://` warehouse. The + // native Iceberg storage factory has no arm for it, so that scan must be DECLINED, not claimed. + conf.set("spark.hadoop.fs.wasb.impl", classOf[FakeWasbSchemeFileSystem].getName) // Intentionally NOT setting CometConf.COMET_LIBHDFS_SCHEMES -- `fake` is not natively readable, // and `hdfs` must still be claimed by default (mirrors the native `is_hdfs_scheme` default). conf @@ -120,10 +124,65 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "gs://bucket/... must be admitted even after an authorityless gs:// URI was probed first") } + test("iceberg gate: the JNI scheme list parser trims, lowercases and rejects a blank list") { + assert( + IcebergStorageSchemes.parse("file, S3,,s3a,", forWrite = false) == Set("file", "s3", "s3a")) + // A loaded library publishes a fixed non-empty list, so a blank answer is a build bug and + // must fail loudly rather than quietly disable every native Iceberg scan and write. + Seq(("", false), (" , ", true), (null, false)).foreach { case (joined, forWrite) => + val path = if (forWrite) "write" else "read" + val e = intercept[IllegalStateException](IcebergStorageSchemes.parse(joined, forWrite)) + assert( + e.getMessage == s"Comet native library published no Iceberg $path scheme list", + s"unexpected message for ${Option(joined)}: ${e.getMessage}") + } + } + + test("iceberg gate: an unloaded native library yields no schemes") { + // Every caller sits behind `isCometLoaded`, so this answer never reaches a plan; it only + // guarantees that nothing hand-maintained stands in for the native list. + assert(IcebergStorageSchemes.load(forWrite = false, isLoaded = false) == Set.empty) + assert(IcebergStorageSchemes.load(forWrite = true, isLoaded = false) == Set.empty) + } + + test("iceberg gate: the scheme sets come from native and carry its mode-specific entries") { + // The lazy sets must be what the JNI probe answers, not an empty set from a load that ran + // before the library was ready. `oss` is read-only (no `oss.*` property forwarding for + // writes) and `memory` is write-only (a fresh in-process store per FileIO, so a read finds + // nothing); the Azure schemes and `gcs` have no storage factory arm at all. + assume(NativeBase.isLoaded, "Comet native library not loaded") + val read = IcebergStorageSchemes.read + val write = IcebergStorageSchemes.write + assert(read == IcebergStorageSchemes.parse(NativeBase.icebergStorageSchemes(false), false)) + assert(write == IcebergStorageSchemes.parse(NativeBase.icebergStorageSchemes(true), true)) + assert(read.nonEmpty, "native published no read schemes") + assert(write.nonEmpty, "native published no write schemes") + assert(read.contains("oss"), s"oss must be in the native read schemes ${read.toSeq.sorted}") + assert( + !write.contains("oss"), + s"oss must not be in the native write schemes ${write.toSeq.sorted}") + assert( + write.contains("memory"), + s"memory must be in the native write schemes ${write.toSeq.sorted}") + assert( + !read.contains("memory"), + s"memory must not be in the native read schemes ${read.toSeq.sorted}") + Seq("abfs", "abfss", "wasb", "wasbs", "gcs").foreach { scheme => + assert( + !read.contains(scheme), + s"$scheme must not be in the native read schemes ${read.toSeq.sorted}") + assert( + !write.contains(scheme), + s"$scheme must not be in the native write schemes ${write.toSeq.sorted}") + } + } + test("iceberg gate: builtin allowlist admitted, unbuildable schemes rejected") { - // The Iceberg gate is an explicit allowlist mirroring `storage_factory_for`'s arms (keep in - // lockstep). Narrower than the Parquet gate: object_store recognizes http/abfs/wasb but - // iceberg-rust can't build them, so reject up-front rather than fail during native setup. + // The Iceberg gate is the allowlist the native storage factory publishes over JNI. Narrower + // than the Parquet gate: object_store recognizes http and the Azure schemes, but the native + // Iceberg storage factory has no arm for them, so reject up-front rather than fail at + // execution. `memory` is write-only (a fresh in-process store per FileIO holds no table), and + // the built-in set matches verbatim because native opens a location by its raw scheme. Seq( "file:///tmp/key.parquet", "s3://bucket/key.parquet", @@ -135,6 +194,9 @@ class CometScanSchemeFallbackSuite extends CometTestBase { s"$u must be iceberg-readable; icebergReadableSchemes has regressed") } Seq( + "memory:///key.parquet", + "S3://bucket/key.parquet", + "File:///tmp/key.parquet", "http://bucket.example.com/key.parquet", "https://bucket.example.com/key.parquet", "abfs://container@acct/key.parquet", @@ -143,10 +205,31 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "wasbs://container@acct/key.parquet").foreach { u => assert( !CometScanRule.isIcebergReadableScheme(new URI(u), Set.empty), - s"$u must not be iceberg-readable; storage_factory_for has no matching arm") + s"$u must not be iceberg-readable; storage_factory_for rejects it for reads") } } + test("iceberg gate: an opt-in alias is matched verbatim, like the built-in schemes") { + // Native opens an alias location as written and OpenDAL's S3 backend checks it against a + // lowercase `scheme://bucket/` prefix, so a `BLOB://` location admitted here would only fail + // at execution. The alias set is lowercase, so a scheme written any other way is declined. + val schemes = Set("blob") + assert( + CometScanRule.isIcebergReadableScheme(new URI("blob://bucket/key.parquet"), schemes), + "blob://bucket/... must be admitted once blob is opted in") + Seq("BLOB://bucket/key.parquet", "Blob://bucket/key.parquet", "BLOB:///bucket/key.parquet") + .foreach { u => + assert( + !CometScanRule.isIcebergReadableScheme(new URI(u), schemes), + s"$u must be declined: native opens the location as written and the S3 backend " + + "rejects a scheme prefix that is not lowercase") + } + // The fallback reason names the lowercase alias, which is the form the scan would admit. + assert( + CometScanRule.icebergSupportedSchemesMessage(schemes).contains("blob"), + "the supported-schemes message must list the opted-in alias in its admitted form") + } + test("parquet gate: mixed-bucket alias scan is declined (single object store per partition)") { // Native planning registers one object store per FilePartition and strips the authority from // every file's object key, so files in a second bucket would be read from the first. A scan @@ -218,6 +301,67 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "without opt-in, blob gets no bucket promotion, so a hostless blob location is unopenable") } + test("iceberg scan and write decline a wasb table instead of failing at execution") { + // Regression guard for both scheme gates: the native Iceberg storage factory has no arm for + // `wasb`, so the scan and the write must each be declined with a reason naming the scheme + // and run in Spark, rather than be claimed and fail at execution with "Unsupported storage + // scheme: wasb". + assume(icebergAvailable, "Iceberg not available in classpath") + withTempIcebergDir { dir => + val warehouse = s"${FakeWasbSchemeFileSystem.PREFIX}${dir.getAbsolutePath}/warehouse" + withSQLConf( + "spark.sql.catalog.wasb_catalog" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.wasb_catalog.type" -> "hadoop", + "spark.sql.catalog.wasb_catalog.warehouse" -> warehouse, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "true", + // The VALUES source must be native, or the write is declined for its input before the + // write serde (and its scheme gate) is ever consulted. + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + spark.sql("CREATE TABLE wasb_catalog.db.wasb_table (id INT, name STRING) USING iceberg") + try { + val insertPlans = capturePlans(spark) { + spark.sql( + "INSERT INTO wasb_catalog.db.wasb_table VALUES (1, 'Alice'), (2, 'Bob'), (3, 'Carol')") + } + assert(insertPlans.nonEmpty, "no INSERT plan was captured") + val nativeWrites = insertPlans.flatMap(_.collectWithSubqueries { + case w: CometIcebergWriteExec => w + }) + assert( + nativeWrites.isEmpty, + "`wasb://` has no native Iceberg storage factory arm; the write must fall back to " + + s"Spark, but Comet claimed it:\n${insertPlans.mkString("\n--\n")}") + val writeReasons = insertPlans.flatMap(new ExtendedExplainInfo().getFallbackReasons) + assert( + writeReasons.exists(_.contains("unsupported storage scheme: wasb")), + s"the write must be declined for its scheme, got reasons: $writeReasons") + + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + "SELECT * FROM wasb_catalog.db.wasb_table ORDER BY id", + "wasb") + val nativeScans = cometPlan.collect { case s: CometIcebergNativeScanExec => s } + assert( + nativeScans.isEmpty, + "`wasb://` has no native Iceberg storage factory arm; the scan must fall back to " + + s"Spark, but Comet claimed it:\n$cometPlan") + // Nothing but the scheme may have caused the fallback: every reason names `wasb`, + // except the one the operator above records about its now-Spark BatchScanExec child. + val scanReasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + val unrelated = scanReasons.filterNot(_.contains("wasb")) + assert( + unrelated.forall(_.contains("BatchScanExec")), + s"fallback reasons other than the scheme were recorded: $unrelated") + } finally { + spark.sql("DROP TABLE wasb_catalog.db.wasb_table") + } + } + } + } + test("native scan claims hdfs:// when libhdfs.schemes is unset (native-default lockstep)") { // Native `is_hdfs_scheme` treats `hdfs` as readable when `fs.comet.libhdfs.schemes` is unset, // so the JVM gate must agree and CLAIM `hdfs://`. Guards the `case None => Set("hdfs")` default From f80042f78715e5a009352f1582f6c246a5f96131 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 7 Oct 2026 23:19:26 +0000 Subject: [PATCH 26/89] fix: Route collated array element comparisons through the codegen dispatcher (#6471) * fix: Route collated array element comparisons through the codegen dispatcher array_contains, arrays_overlap, array_distinct and array_union compare array elements by raw bytes in their native kernels, so they returned wrong results for strings with a non-UTF8_BINARY collation (for example UTF8_LCASE or UTF8_BINARY_RTRIM) when a dispatched collated cast fed them. Report these cases as Incompatible and mix in CodegenDispatchFallback so they run Spark's collation-aware comparison through the JVM codegen dispatcher by default. ArrayDistinct gets its own serde for this. * docs: Mark collation-routed array functions as Hybrid in expressions.md --- docs/source/user-guide/latest/expressions.md | 8 +- .../scala/org/apache/comet/serde/arrays.scala | 68 ++++++++++++-- .../array_element_equality_collation.sql | 90 +++++++++++++++++++ ..._equality_collation_allow_incompatible.sql | 45 ++++++++++ ...ay_element_equality_collation_disabled.sql | 49 ++++++++++ .../comet/CometMapExpressionSuite.scala | 24 +++-- 6 files changed, 265 insertions(+), 19 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 74232aac652..dffd6b49884 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -155,8 +155,8 @@ The tables below list every Spark built-in expression with its current status. | `array` | ✅ | Native | | | `array_append` | ✅ | Native | | | `array_compact` | ✅ | — | | -| `array_contains` | ✅ | Native | Float/double element arrays route through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible | -| `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | +| `array_contains` | ✅ | Hybrid | Float/double element arrays and non-default string collations route through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible | +| `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)). Non-default string collations fall back | | `array_except` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | | `array_insert` | ✅ | Native | | | `array_intersect` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | @@ -167,8 +167,8 @@ The tables below list every Spark built-in expression with its current status. | `array_prepend` | ✅ | — | | | `array_remove` | ✅ | Native | | | `array_repeat` | ✅ | Native | | -| `array_union` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | -| `arrays_overlap` | ✅ | Native | | +| `array_union` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)). Non-default string collations fall back | +| `arrays_overlap` | ✅ | Hybrid | Non-default string collations use the JVM codegen dispatcher | | `arrays_zip` | ✅ | Native | | | `element_at` | ✅ | Native | | | `flatten` | ✅ | Native | Binary/struct/map elements fall back | diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index b2a75f84782..5aeb6824346 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -115,15 +115,24 @@ object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBas object CometArrayContains extends CometExpressionSerde[ArrayContains] + with CometTypeShim with CodegenDispatchFallback { + override def hasConditionalNativeDefault: Boolean = true + private val floatingPointReason: String = "Spark compares array elements with ordering.equiv, so -0.0 matches +0.0 and all NaNs match " + "each other; Comet's native array_contains compares the raw Arrow values bitwise" - override def getIncompatibleReasons(): Seq[String] = Seq(floatingPointReason) + private val collationReason: String = + ArrayElementEqualitySupport.collationReason("array_contains") + + override def getIncompatibleReasons(): Seq[String] = Seq(floatingPointReason, collationReason) override def getSupportLevel(expr: ArrayContains): SupportLevel = expr.left.dataType match { + // See ArrayElementEqualitySupport: bytewise string comparison ignores the collation. + case ArrayType(elementType, _) if hasNonDefaultStringCollation(elementType) => + Incompatible(Some(collationReason)) // Native array_contains compares floating-point elements bitwise, disagreeing with Spark for // -0.0/+0.0 and NaN. Report Incompatible (not Unsupported) for float/double element types (at // any nesting level) so the expression routes through the JVM codegen dispatcher (Spark's own @@ -309,7 +318,19 @@ object CometArrayMin extends CometExpressionSerde[ArrayMin] with CodegenDispatch } } -object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] { +object CometArraysOverlap + extends CometExpressionSerde[ArraysOverlap] + with CodegenDispatchFallback { + + override def hasConditionalNativeDefault: Boolean = true + + override def getIncompatibleReasons(): Seq[String] = + Seq(ArrayElementEqualitySupport.collationReason("arrays_overlap")) + + // Both inputs share one element type, so checking the left side is enough. + override def getSupportLevel(expr: ArraysOverlap): SupportLevel = + ArrayElementEqualitySupport.getSupportLevel("arrays_overlap", expr.left.dataType) + override def convert( expr: ArraysOverlap, inputs: Seq[Attribute], @@ -538,6 +559,30 @@ object CometSlice extends CometExpressionSerde[Slice] { } } +/** + * Support level for native array kernels that compare elements for equality (membership, overlap, + * dedup). They compare strings by raw bytes, so a non-UTF8_BINARY collation (for example + * UTF8_LCASE, where 'a' equals 'A') gives wrong answers, and these cases report Incompatible. + * Serdes with a boolean result mix in CodegenDispatchFallback, so the JVM codegen dispatcher runs + * Spark's collation-aware comparison. Array-valued serdes fall back instead (see + * ArraySetSupport). + */ +private object ArrayElementEqualitySupport extends CometTypeShim { + def collationReason(name: String): String = + "Spark compares non-UTF8_BINARY collated string elements under their collation, while " + + s"Comet's native $name compares raw bytes" + + def collationSupportLevel(name: String, arrayType: DataType): Option[SupportLevel] = + if (hasNonDefaultStringCollation(arrayType)) { + Some(Incompatible(Some(collationReason(name)))) + } else { + None + } + + def getSupportLevel(name: String, arrayType: DataType): SupportLevel = + collationSupportLevel(name, arrayType).getOrElse(Compatible()) +} + private[comet] object ArraySetSupport { val floatingPointReason: String = "Floating-point elements match Spark's signed-zero and NaN semantics natively only on " + @@ -560,22 +605,33 @@ private[comet] object ArraySetSupport { Compatible() } } + + /** Collated string elements are also Incompatible; see ArrayElementEqualitySupport. */ + def supportLevel(name: String, dataType: DataType): SupportLevel = + ArrayElementEqualitySupport + .collationSupportLevel(name, dataType) + .getOrElse(supportLevel(dataType)) + + def incompatibleReasons(name: String): Seq[String] = + Seq(floatingPointReason, ArrayElementEqualitySupport.collationReason(name)) } // Use projection fallback to avoid codegen dispatch overhead for array-valued results. // The native implementation remains available through opt-in. object CometArrayDistinct extends CometScalarFunction[ArrayDistinct]("array_distinct") { - override def getIncompatibleReasons(): Seq[String] = Seq(ArraySetSupport.floatingPointReason) + override def getIncompatibleReasons(): Seq[String] = + ArraySetSupport.incompatibleReasons("array_distinct") override def getSupportLevel(expr: ArrayDistinct): SupportLevel = - ArraySetSupport.supportLevel(expr.dataType) + ArraySetSupport.supportLevel("array_distinct", expr.dataType) } object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { - override def getIncompatibleReasons(): Seq[String] = Seq(ArraySetSupport.floatingPointReason) + override def getIncompatibleReasons(): Seq[String] = + ArraySetSupport.incompatibleReasons("array_union") override def getSupportLevel(expr: ArrayUnion): SupportLevel = - ArraySetSupport.supportLevel(expr.dataType) + ArraySetSupport.supportLevel("array_union", expr.dataType) override def convert( expr: ArrayUnion, diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql new file mode 100644 index 00000000000..e8073461829 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql @@ -0,0 +1,90 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=false +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=false +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=false +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=false + +-- array_contains, arrays_overlap, array_distinct and array_union compare array elements for +-- equality. The native kernels compare strings by raw bytes, so under a non-UTF8_BINARY collation +-- they would miss 'a' = 'A' (UTF8_LCASE) or 'x ' = 'x' (UTF8_BINARY_RTRIM). Collated +-- array_contains and arrays_overlap route through the JVM codegen dispatcher, which runs Spark's +-- collation-aware comparison. array_distinct and array_union return arrays, where dispatch costs +-- more than a projection fallback (see ArraySetSupport), so they fall back to Spark instead. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES + (1, 'a', 'A'), + (2, 'x ', 'x'), + (3, 'b', 'c'), + (4, NULL, 'a'), + (5, 'A', NULL) + +-- UTF8_BINARY strings stay on the native kernels. +query expect_native(array_contains,arrays_overlap,array_distinct,array_union) +SELECT id, array_contains(array(a), b), arrays_overlap(array(a), array(b)), + array_distinct(array(a, b)), array_union(array(a), array(b)) +FROM test_array_eq_collation + +-- UTF8_LCASE: 'a' and 'A' are equal. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)) +FROM test_array_eq_collation + +query expect_dispatch(arrays_overlap) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +-- UTF8_BINARY_RTRIM: 'x ' and 'x' are equal. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM)) +FROM test_array_eq_collation + +query expect_dispatch(arrays_overlap) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM), CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +-- Collation detection also recurses into struct fields of the array elements. +query expect_dispatch(array_contains,arrays_overlap) +SELECT id, + array_contains(array(named_struct('s', CAST(a AS STRING COLLATE UTF8_LCASE))), + named_struct('s', CAST(b AS STRING COLLATE UTF8_LCASE))), + arrays_overlap(array(named_struct('s', CAST(a AS STRING COLLATE UTF8_LCASE))), + array(named_struct('s', CAST(b AS STRING COLLATE UTF8_LCASE)))) +FROM test_array_eq_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql new file mode 100644 index 00000000000..7bedec72ce9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql @@ -0,0 +1,45 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=true +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=true +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=true +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=true + +-- allowIncompatible=true opts collated inputs into the native kernels, which compare raw bytes. +-- The data only holds values whose bytewise and collation-aware equality agree, so the answers +-- still match Spark; the point is that the native path is taken. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES (1, 'a', 'a'), (2, 'b', 'c'), (3, NULL, 'a') + +query expect_native(array_contains,arrays_overlap) +SELECT id, + array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)), + arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_native(array_distinct,array_union) +SELECT id, + array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))), + array_union(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql new file mode 100644 index 00000000000..e07ad2f8a2c --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql @@ -0,0 +1,49 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=false +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=false +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=false +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=false +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=false + +-- With the JVM codegen dispatcher disabled, collated array_contains and arrays_overlap have no +-- Spark-compatible Comet path and fall back to Spark. array_distinct and array_union fall back +-- whether or not the dispatcher is enabled. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES (1, 'a', 'A'), (2, 'x ', 'x'), (3, 'b', 'c'), (4, NULL, 'a') + +query expect_fallback(array_contains: spark.comet.exec.scalaUDF.codegen.enabled=false) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)) +FROM test_array_eq_collation + +query expect_fallback(arrays_overlap: spark.comet.exec.scalaUDF.codegen.enabled=false) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 98ef07b96fb..091f777dbe9 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -467,18 +467,24 @@ class CometMapExpressionSuite extends CometTestBase { } // The collated counterpart of the double-key `map_contains_key` case above: a nested - // `UTF8_LCASE`-keyed map would reach `array_contains` with bytewise comparison. Expansion declines - // the folded literal so the case-insensitive lookup stays on Spark. The outer lookup key is the - // dynamic `_1` so the inner map survives as a literal (a constant key would let Spark fold - // `map_keys` into a collated-string array literal instead, which never reaches this guard). + // `UTF8_LCASE`-keyed map would reach `array_contains` with bytewise comparison. Like the + // floating-point case, `CometArrayContains` reports collated element types Incompatible, so the + // default config codegen-dispatches it; `allowIncompatible=true` forces the native kernel so the + // expansion guard runs. Expansion declines the folded literal so the case-insensitive lookup + // stays on Spark. The outer lookup key is the dynamic `_1` so the inner map survives as a literal + // (a constant key would let Spark fold `map_keys` into a collated-string array literal instead, + // which never reaches this guard). test("map_contains_key over nested collated map keys falls back (multirow)") { assume(isSpark40Plus) + val query = "SELECT _1 AS id, map_contains_key(" + + "element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), _1), " + + "CAST('a1' AS STRING COLLATE UTF8_LCASE)) AS present FROM tbl" withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndFallbackReason( - "SELECT _1 AS id, map_contains_key(" + - "element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), _1), " + - "CAST('a1' AS STRING COLLATE UTF8_LCASE)) AS present FROM tbl", - "Unsupported data type MapType") + withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[ArrayContains]) -> "true") { + checkSparkAnswerAndFallbackReason(query, "Unsupported data type MapType") + } + // Under the default config the collated lookup is dispatched and matches Spark. + checkSparkAnswer(query) } } From 5db0efe69966b11d8af8e1f7c9c8c2ca9ce03a95 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 8 Oct 2026 12:49:06 +0000 Subject: [PATCH 27/89] perf: convert Spark's unsafe rows to Arrow straight from their memory (#6730) * perf: convert Spark's unsafe rows to Arrow straight from their memory ArrowWriter's row path read every row through Spark's generic getters. A string went through a UTF8String and Arrow's per-value setter, each array, map and struct allocated a view, and every element took a generic write with two or three virtual or interface calls. Unsafe rows now take a path of their own, and unsafe arrays are written in bulk: - Fixed-width values are read from their slot at the vector's width, and arrays of them are copied in one block, with the validity inverted from the array's null words. - Strings and binaries are copied from the offset and size in their slot, short ones a word at a time. A negative size or an offset overflow is rejected before anything is written or allocated, for a value and for every element of an array. - Decimals up to 18 digits are written from the unscaled long, and wider ones are range-checked in place. - Arrays, maps and structs point one reused view at each value. Closes #6721. * Simplify the unsafe-row writers - Use Arrow's fillEmpties and setValueLengthSafe in place of hand-written hole filling and buffer growth, and share one variable-width overflow check. - Read the fixed-width vector through its field in every per-value method, and have the base defaults delegate to write. - Move the columnar code StringWriter and BinaryWriter shared into their base. - Write unsafe arrays of decimals up to 18 digits from the unscaled long, range-checked as UnsafeArrayData.getDecimal checks it. - Mark map entries valid with one masked write rather than a write per bit, and call setOne only for long runs. - Share the tests' reject, root-comparison and fill helpers. * test: check that a decimal past its precision in an unsafe array still fails Reads an unsafe array of decimals at a narrower precision than it was written with, holding 10^p or -10^p, for both the unscaled long and the unscaled bytes. * test: write an unsafe array of each primitive type Each of the twelve vector types whose unsafe array elements are copied in one block now reaches that copy, at every element width. Only int and long arrays did before. * fix: give each nested value a view of its own again The array, map and struct writers pointed one reused unsafe view at every value they read. Each view then kept the last row it read reachable after that row's values were copied to Arrow, along with the row's memory, until the writer was dropped. They now read each value through Spark's getters, which allocate a view per value as the generic path does. That keeps about nine tenths of the gain on nested types. A test walks the writers' fields after each row and checks that the row's memory is not reachable from them. * refactor: share the variable-width growth step between the append paths writeVariableWidth, writeDictionaryVariableWidth and the unsafe array path each filled the offsets of values skipped since the last one set and grew the validity and offset buffers. They now call reserveValues, next to reserveData. A new test appends a dictionary-encoded string field as a column after its struct took the row path, which the dictionary path's growth step had no test for. The array of each primitive type no longer repeats the two that nestedTypes lists. --- .../comet/execution/arrow/ArrowWriters.scala | 660 ++++++++++++++---- .../arrow/CometArrowStreamSuite.scala | 2 +- .../arrow/CometArrowWriterSuite.scala | 354 +++++++++- .../arrow/CometStringWriterSuite.scala | 87 ++- 4 files changed, 927 insertions(+), 176 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 00c4de38f6d..53b4ae7888d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -29,7 +29,7 @@ import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ import org.apache.arrow.vector.util.OversizedAllocationException import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeArrayData, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeArrayData, UnsafeMapData, UnsafeRow} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} @@ -151,9 +151,17 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { def write(row: InternalRow): Unit = { var i = 0 - while (i < fields.length) { - fields(i).writeUnsafe(row, i) - i += 1 + row match { + case unsafe: UnsafeRow => + while (i < fields.length) { + fields(i).writeUnsafeRowField(unsafe, i) + i += 1 + } + case _ => + while (i < fields.length) { + fields(i).write(row, i) + i += 1 + } } count += 1 } @@ -212,26 +220,62 @@ private[arrow] object ArrowFieldWriter { // Below this many rows the per-value loop is faster. val MinOnHeapBulkCopyRows = 32 + // Below this many bytes, copying a word at a time beats Unsafe.copyMemory's checks and call. + // Most strings, and the elements of most arrays, are shorter. + private final val MaxWordCopyBytes = 64 + + /** Copies `length` bytes from `srcOffset` in `src` to the native address `dst`. */ + def copyMemory(src: AnyRef, srcOffset: Long, dst: Long, length: Long): Unit = { + if (length > MaxWordCopyBytes || !Platform.unaligned()) { + Platform.copyMemory(src, srcOffset, null, dst, length) + } else { + var i = 0L + while (i + 8 <= length) { + Platform.putLong(null, dst + i, Platform.getLong(src, srcOffset + i)) + i += 8 + } + if (i + 4 <= length) { + Platform.putInt(null, dst + i, Platform.getInt(src, srcOffset + i)) + i += 4 + } + if (i + 2 <= length) { + Platform.putShort(null, dst + i, Platform.getShort(src, srcOffset + i)) + i += 2 + } + if (i < length) { + Platform.putByte(null, dst + i, Platform.getByte(src, srcOffset + i)) + } + } + } + /** Spark's own writable vectors, whose storage layout the columnar fast paths rely on. */ def isSparkVector(input: ColumnVector): Boolean = input.isInstanceOf[OnHeapColumnVector] || input.isInstanceOf[OffHeapColumnVector] + // Arrow's setOne calls Unsafe.setMemory, which the JIT does not inline, so shorter runs of whole + // bytes are set one at a time. + private final val MinSetOneBytes = 64 + /** Marks bits `[start, start + numRows)` of `validity` as valid. */ def setValid(validity: ArrowBuf, start: Int, numRows: Int): Unit = { val end = start + numRows - var i = start - while (i < end && (i & 7) != 0) { - BitVectorHelper.setBit(validity, i.toLong) - i += 1 - } - val alignedEnd = end & ~7 - if (i < alignedEnd) { - validity.setOne((i >> 3).toLong, ((alignedEnd - i) >> 3).toLong) - i = alignedEnd - } - while (i < end) { - BitVectorHelper.setBit(validity, i.toLong) - i += 1 + val firstByte = (start + 7) >> 3 + val endByte = end >> 3 + if (firstByte >= endByte) { + // No whole byte, as for most of a map's entries. + writeBits(validity, start, -1L, numRows) + } else { + writeBits(validity, start, -1L, (firstByte << 3) - start) + if (endByte - firstByte >= MinSetOneBytes) { + validity.setOne(firstByte.toLong, (endByte - firstByte).toLong) + } else { + var b = firstByte + while (b < endByte) { + validity.setByte(b.toLong, 0xff) + b += 1 + } + } + writeBits(validity, endByte << 3, -1L, end - (endByte << 3)) } } @@ -282,6 +326,37 @@ private[arrow] object ArrowFieldWriter { } } + /** Writes the low `numBits` bits of `bits`, at most 64, to `[start, start + numBits)`. */ + def writeBits(buffer: ArrowBuf, start: Int, bits: Long, numBits: Int): Unit = { + val end = start + numBits + var remaining = bits + var i = start + while (i < end) { + val shift = i & 7 + val n = Math.min(8 - shift, end - i) + val mask = ((1 << n) - 1) << shift + val index = (i >> 3).toLong + buffer.setByte(index, (buffer.getByte(index) & ~mask) | ((remaining.toInt << shift) & mask)) + remaining >>>= n + i += n + } + } + + /** + * Sets the validity of the elements of `array` from bit `start` of `validity`. Spark keeps an + * unsafe array's null bits in 64-bit words after its element count, a set bit marking a null. + */ + def writeArrayValidity(validity: ArrowBuf, start: Int, array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + val nulls = array.getBaseOffset + 8 + var i = 0 + while (i < numElements) { + val nullBits = Platform.getLong(array.getBaseObject, nulls + (i >> 3)) + writeBits(validity, start + i, ~nullBits, Math.min(64, numElements - i)) + i += 64 + } + } + /** * Clears each bit of `validity` in `[start, start + numRows)` whose bit in `parent` is clear. * Whole bytes are combined: the bits before `start` in the first byte belong to rows already @@ -386,12 +461,7 @@ private[arrow] object ArrowFieldWriter { if (numRows == 0) { return } - if (vector.getLastSet < outStart - 1) { - vector.fillEmpties(outStart) - } - while (vector.getValueCapacity < outStart + numRows) { - vector.reallocValidityAndOffsetBuffers() - } + reserveValues(vector, outStart, numRows) val offsets = vector.getOffsetBuffer val dataStart = vector.getStartOffset(outStart) val hasNull = input.hasNull @@ -413,10 +483,7 @@ private[arrow] object ArrowFieldWriter { } next = offset + length end += length - if (end > Integer.MAX_VALUE) { - throw new OversizedAllocationException( - s"Arrow variable-width data would exceed ${Integer.MAX_VALUE} bytes") - } + checkDataEnd(end) } } offsets.setInt((outStart + i + 1).toLong * BaseVariableWidthVector.OFFSET_WIDTH, end.toInt) @@ -467,12 +534,7 @@ private[arrow] object ArrowFieldWriter { if (numRows == 0) { return } - if (vector.getLastSet < outStart - 1) { - vector.fillEmpties(outStart) - } - while (vector.getValueCapacity < outStart + numRows) { - vector.reallocValidityAndOffsetBuffers() - } + reserveValues(vector, outStart, numRows) val offsets = vector.getOffsetBuffer val ids = input.getDictionaryIds val hasNull = input.hasNull @@ -523,12 +585,30 @@ private[arrow] object ArrowFieldWriter { vector.setLastSet(outStart + numRows - 1) } - /** The data buffer of `vector`, grown to hold at least `end` bytes. */ - private def reserveData(vector: BaseVariableWidthVector, end: Long): ArrowBuf = { + /** Rejects variable-width data ending at `end`, past what Arrow's 32-bit offsets address. */ + def checkDataEnd(end: Long): Unit = { if (end > Integer.MAX_VALUE) { throw new OversizedAllocationException( s"Arrow variable-width data would exceed ${Integer.MAX_VALUE} bytes") } + } + + /** + * Grows the validity and offset buffers of `vector` to hold values `[outStart, outStart + + * numValues)`, after filling in the offsets of any values skipped since the last one set. + */ + def reserveValues(vector: BaseVariableWidthVector, outStart: Int, numValues: Int): Unit = { + if (vector.getLastSet < outStart - 1) { + vector.fillEmpties(outStart) + } + while (vector.getValueCapacity < outStart + numValues) { + vector.reallocValidityAndOffsetBuffers() + } + } + + /** The data buffer of `vector`, grown to hold at least `end` bytes. */ + def reserveData(vector: BaseVariableWidthVector, end: Long): ArrowBuf = { + checkDataEnd(end) if (vector.getDataBuffer.capacity < end) { vector.reallocDataBuffer(end) } @@ -610,8 +690,20 @@ private[arrow] abstract class ArrowFieldWriter { count += 1 } - def writeUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - write(input, ordinal) + /** + * Appends field `ordinal` of `row`. Most rows Spark produces are unsafe rows, and writers that + * can read one's memory directly override this. + */ + private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = write(row, ordinal) + + /** Appends every element of `array`, which writers that can copy them in bulk override. */ + private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + var i = 0 + while (i < numElements) { + write(array, i) + i += 1 + } } def writeCol(input: ColumnarArray): Unit = { @@ -691,16 +783,22 @@ private[arrow] abstract class ArrowFieldWriter { } } -private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWriter { +/** + * `vector` is `valueVector`. The methods called per value read it through this field rather than + * the subclass's accessor, a virtual call. + */ +private[arrow] abstract class FixedWidthArrowFieldWriter(vector: BaseFixedWidthVector) + extends ArrowFieldWriter { import ArrowFieldWriter._ override def valueVector: BaseFixedWidthVector protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit + // Arrow keeps a fixed-width vector's capacity in a field, so checking it per value is cheap. protected def ensureCapacity(inputNumElements: Int): Unit = { - while (valueVector.getValueCapacity < inputNumElements) { - valueVector.reAlloc() + while (vector.getValueCapacity < inputNumElements) { + vector.reAlloc() } } @@ -841,26 +939,75 @@ private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWrite } override def setNull(): Unit = { - valueVector.setNull(count) + vector.setNull(count) } protected def setNullUnsafe(): Unit = { - BitVectorHelper.unsetBit(valueVector.getValidityBuffer, count) + BitVectorHelper.unsetBit(vector.getValidityBuffer, count) } - override def writeUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - if (input.isNullAt(ordinal)) { + // The width of this type's values when Spark's unsafe formats hold the bits Arrow stores, at the + // vector's type width and little-endian: an unsafe row in an 8-byte slot, an unsafe array back + // to back. 0 for a type whose values need converting. + private val unsafeValueWidth: Int = vector match { + case _ if !LittleEndian => 0 + case _: TinyIntVector | _: SmallIntVector | _: IntVector | _: DateDayVector | + _: IntervalYearVector | _: Float4Vector | _: BigIntVector | _: TimeStampMicroTZVector | + _: TimeStampMicroVector | _: DurationVector | _: TimeNanoVector | _: Float8Vector => + vector.getTypeWidth + case _ => 0 + } + + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + ensureCapacity(count + 1) + if (row.isNullAt(ordinal)) { setNullUnsafe() + } else if (unsafeValueWidth == 0) { + setValueUnsafe(row, ordinal) } else { - setValueUnsafe(input, ordinal) + val target = vector.getDataBufferAddress + count.toLong * unsafeValueWidth + unsafeValueWidth match { + case 8 => Platform.putLong(null, target, row.getLong(ordinal)) + case 4 => Platform.putInt(null, target, row.getInt(ordinal)) + case 2 => Platform.putShort(null, target, row.getShort(ordinal)) + case _ => Platform.putByte(null, target, row.getByte(ordinal)) + } + BitVectorHelper.setBit(vector.getValidityBuffer, count.toLong) } count += 1 } + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + ensureCapacity(count + numElements) + if (unsafeValueWidth == 0) { + var i = 0 + while (i < numElements) { + if (array.isNullAt(i)) { + setNullUnsafe() + } else { + setValueUnsafe(array, i) + } + count += 1 + i += 1 + } + } else { + // The values are copied in one block, so what lands under a null element is whatever Spark + // left there, as with the columnar copies. + copyMemory( + array.getBaseObject, + array.getBaseOffset + UnsafeArrayData.calculateHeaderPortionInBytes(numElements), + vector.getDataBufferAddress + count.toLong * unsafeValueWidth, + numElements.toLong * unsafeValueWidth) + writeArrayValidity(vector.getValidityBuffer, count, array) + count += numElements + } + } + override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { ensureCapacity(count + numRows) if (copyValues(input, startRow, numRows)) { - writeValidity(valueVector.getValidityBuffer, count, input, startRow, numRows) + writeValidity(vector.getValidityBuffer, count, input, startRow, numRows) count += numRows } else { val slice = new ColumnarArray(input, startRow, numRows) @@ -901,7 +1048,7 @@ private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWrite } private[arrow] class BooleanWriter(val valueVector: BitVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, if (input.getBoolean(ordinal)) 1 else 0) @@ -949,10 +1096,25 @@ private[arrow] class BooleanWriter(val valueVector: BitVector) ArrowFieldWriter.writeValidity(valueVector.getValidityBuffer, count, input, startRow, numRows) count += numRows } + + // An unsafe array holds a byte per value, each of which becomes a bit. Null elements get a + // clear bit, as in writeColumnSlice. + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + ensureCapacity(count + numElements) + val values = valueVector.getDataBuffer + var i = 0 + while (i < numElements) { + ArrowFieldWriter.writeBit(values, count + i, !array.isNullAt(i) && array.getBoolean(i)) + i += 1 + } + ArrowFieldWriter.writeArrayValidity(valueVector.getValidityBuffer, count, array) + count += numElements + } } private[arrow] class ByteWriter(val valueVector: TinyIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getByte(ordinal)) @@ -964,7 +1126,7 @@ private[arrow] class ByteWriter(val valueVector: TinyIntVector) } private[arrow] class ShortWriter(val valueVector: SmallIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getShort(ordinal)) @@ -976,7 +1138,7 @@ private[arrow] class ShortWriter(val valueVector: SmallIntVector) } private[arrow] class IntegerWriter(val valueVector: IntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -988,7 +1150,7 @@ private[arrow] class IntegerWriter(val valueVector: IntVector) } private[arrow] class LongWriter(val valueVector: BigIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1000,7 +1162,7 @@ private[arrow] class LongWriter(val valueVector: BigIntVector) } private[arrow] class FloatWriter(val valueVector: Float4Vector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getFloat(ordinal)) @@ -1012,7 +1174,7 @@ private[arrow] class FloatWriter(val valueVector: Float4Vector) } private[arrow] class DoubleWriter(val valueVector: Float8Vector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getDouble(ordinal)) @@ -1031,7 +1193,7 @@ private[arrow] class DoubleWriter(val valueVector: Float8Vector) * that `getDecimal` would build is only built for a value that does not fit, to fail as before. */ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: Int, scale: Int) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { import ArrowFieldWriter.LittleEndian override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { @@ -1039,23 +1201,49 @@ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: In setValueUnsafe(input, ordinal) } + // Spark's unsafe rows and arrays hold up to 18 digits as the unscaled long, and more as the + // unscaled bytes, which are read in place. An unsafe row's getDecimal does not check the long + // against the precision, but an unsafe array's does, so a long that does not fit is left to it. override protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - if (precision > Decimal.MAX_LONG_DIGITS && LittleEndian && - (input.isInstanceOf[UnsafeRow] || input.isInstanceOf[UnsafeArrayData])) { - val bytes = input.getBinary(ordinal) - if (putBigEndianIfFits(count, bytes, Platform.BYTE_ARRAY_OFFSET.toLong, bytes.length)) { - BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) - return - } - } - val decimal = input.getDecimal(ordinal, precision, scale) - if (decimal.changePrecision(precision, scale)) { - setUnscaled(count, decimal) + val written = LittleEndian && (input match { + case row: UnsafeRow if precision <= Decimal.MAX_LONG_DIGITS => + putLong(valueAddress, row.getLong(ordinal)) + true + case array: UnsafeArrayData if precision <= Decimal.MAX_LONG_DIGITS => + val unscaled = array.getLong(ordinal) + val fits = DecimalWriter.fitsPrecision(unscaled >> 63, unscaled, precision) + if (fits) { + putLong(valueAddress, unscaled) + } + fits + case row: UnsafeRow => + putUnsafeBytesIfFits(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal)) + case array: UnsafeArrayData => + putUnsafeBytesIfFits(array.getBaseObject, array.getBaseOffset, array.getLong(ordinal)) + case _ => false + }) + if (written) { + BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) } else { - setNullUnsafe() + val decimal = input.getDecimal(ordinal, precision, scale) + if (decimal.changePrecision(precision, scale)) { + setUnscaled(count, decimal) + } else { + setNullUnsafe() + } } } + private def valueAddress: Long = + valueVector.getDataBufferAddress + count.toLong * DecimalVector.TYPE_WIDTH + + /** + * Writes the unscaled bytes that `offsetAndSize` locates in an unsafe row or array to value + * `count`, leaving validity alone, if they fit the precision. Returns whether it did. + */ + private def putUnsafeBytesIfFits(base: AnyRef, baseOffset: Long, offsetAndSize: Long): Boolean = + putBigEndianIfFits(count, base, baseOffset + (offsetAndSize >> 32), offsetAndSize.toInt) + /** Sets the value and validity of `index`, which must be within capacity. */ private def setUnscaled(index: Int, decimal: Decimal): Unit = { if (precision <= Decimal.MAX_LONG_DIGITS) { @@ -1082,7 +1270,7 @@ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: In * whether it did. */ private def putBigEndianIfFits(index: Int, base: AnyRef, offset: Long, length: Int): Boolean = { - if (length == 0 || length > DecimalVector.TYPE_WIDTH) { + if (length <= 0 || length > DecimalVector.TYPE_WIDTH) { return false } var low = if (Platform.getByte(base, offset) < 0) -1L else 0L @@ -1208,37 +1396,26 @@ private[arrow] object DecimalWriter { } } -private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowFieldWriter { +/** + * Writes strings or binaries. Spark's own vectors are copied in bulk, and its unsafe rows and + * arrays straight from their memory: both keep a value's offset and size in its fixed-length + * slot, so its bytes are copied without a `UTF8String` or Arrow's per-value setters. Other input + * goes through `setValue`. + */ +private[arrow] abstract class VariableWidthArrowFieldWriter extends ArrowFieldWriter { + import ArrowFieldWriter._ + + override def valueVector: BaseVariableWidthVector override def setNull(): Unit = { valueVector.setNull(count) } - override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val utf8 = input.getUTF8String(ordinal) - if (utf8.getBaseObject == null) { - val length = utf8.numBytes() - require(length >= 0, "String length must be non-negative") - valueVector.setValueLengthSafe(count, length) - - // Reservation can replace the buffer. Copy into its current address while Spark still owns - // the source bytes, without staging the off-heap payload in a JVM byte array. - val data = valueVector.getDataBuffer - val offset = valueVector.getStartOffset(count).toLong - require(offset >= 0 && offset + length <= data.capacity(), "Invalid Arrow string range") - utf8.writeToMemory(null, Math.addExact(data.memoryAddress(), offset)) - valueVector.setIndexDefined(count) - } else { - val utf8ByteBuffer = utf8.getByteBuffer - valueVector.setSafe(count, utf8ByteBuffer, utf8ByteBuffer.position(), utf8.numBytes()) - } - } - override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { input match { - case vector: WritableColumnVector if ArrowFieldWriter.isSparkVector(vector) => + case vector: WritableColumnVector if isSparkVector(vector) => if (vector.hasDictionary) { - ArrowFieldWriter.writeDictionaryVariableWidth( + writeDictionaryVariableWidth( valueVector, count, vector, @@ -1246,7 +1423,7 @@ private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowF numRows, dictionaryCache) } else { - ArrowFieldWriter.writeVariableWidth(valueVector, count, vector, startRow, numRows) + writeVariableWidth(valueVector, count, vector, startRow, numRows) } count += numRows case _ => @@ -1262,6 +1439,102 @@ private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowF super.reset() dictionaryCache.clear() } + + // The value is checked before Arrow grows anything, so one that cannot be written changes + // nothing. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + val offsetAndSize = row.getLong(ordinal) + val start = valueVector.getStartOffset(valueVector.getLastSet + 1) + val length = (checkedEnd(start.toLong, offsetAndSize) - start).toInt + // Grows the buffers, and gives null rows written since the last value their empty offsets. + valueVector.setValueLengthSafe(count, length) + copyMemory( + row.getBaseObject, + row.getBaseOffset + (offsetAndSize >> 32), + valueVector.getDataBuffer.memoryAddress + start, + length.toLong) + BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) + } + count += 1 + } + + // Every element is checked before anything changes, as one value is, so the buffers grow once + // and the bytes are copied without further checks. + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + if (numElements == 0) { + return + } + val start = valueVector.getStartOffset(valueVector.getLastSet + 1) + var end = start.toLong + var i = 0 + while (i < numElements) { + if (!array.isNullAt(i)) { + end = checkedEnd(end, array.getLong(i)) + } + i += 1 + } + reserveValues(valueVector, count, numElements) + reserveData(valueVector, end) + val base = array.getBaseObject + val baseOffset = array.getBaseOffset + val data = valueVector.getDataBuffer.memoryAddress + val offsets = valueVector.getOffsetBuffer + var offset = start.toLong + i = 0 + while (i < numElements) { + if (!array.isNullAt(i)) { + val offsetAndSize = array.getLong(i) + val length = offsetAndSize.toInt + copyMemory(base, baseOffset + (offsetAndSize >> 32), data + offset, length.toLong) + offset += length + } + offsets.setInt((count + i + 1).toLong * BaseVariableWidthVector.OFFSET_WIDTH, offset.toInt) + i += 1 + } + writeArrayValidity(valueVector.getValidityBuffer, count, array) + valueVector.setLastSet(count + numElements - 1) + count += numElements + } + + /** The end of a value written at `start` with the size packed into `offsetAndSize`. */ + private def checkedEnd(start: Long, offsetAndSize: Long): Long = { + val length = offsetAndSize.toInt + if (length < 0) { + val kind = if (valueVector.isInstanceOf[VarCharVector]) "String" else "Binary" + throw new IllegalArgumentException(s"$kind length must be non-negative") + } + val end = start + length + checkDataEnd(end) + end + } +} + +private[arrow] class StringWriter(val valueVector: VarCharVector) + extends VariableWidthArrowFieldWriter { + + override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { + val utf8 = input.getUTF8String(ordinal) + if (utf8.getBaseObject == null) { + val length = utf8.numBytes() + require(length >= 0, "String length must be non-negative") + valueVector.setValueLengthSafe(count, length) + + // Reservation can replace the buffer. Copy into its current address while Spark still owns + // the source bytes, without staging the off-heap payload in a JVM byte array. + val data = valueVector.getDataBuffer + val offset = valueVector.getStartOffset(count).toLong + require(offset >= 0 && offset + length <= data.capacity(), "Invalid Arrow string range") + utf8.writeToMemory(null, Math.addExact(data.memoryAddress(), offset)) + valueVector.setIndexDefined(count) + } else { + val utf8ByteBuffer = utf8.getByteBuffer + valueVector.setSafe(count, utf8ByteBuffer, utf8ByteBuffer.position(), utf8.numBytes()) + } + } } private[arrow] class LargeStringWriter(val valueVector: LargeVarCharVector) @@ -1279,45 +1552,13 @@ private[arrow] class LargeStringWriter(val valueVector: LargeVarCharVector) } } -private[arrow] class BinaryWriter(val valueVector: VarBinaryVector) extends ArrowFieldWriter { - - override def setNull(): Unit = { - valueVector.setNull(count) - } +private[arrow] class BinaryWriter(val valueVector: VarBinaryVector) + extends VariableWidthArrowFieldWriter { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { val bytes = input.getBinary(ordinal) valueVector.setSafe(count, bytes, 0, bytes.length) } - - override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { - input match { - case vector: WritableColumnVector if ArrowFieldWriter.isSparkVector(vector) => - if (vector.hasDictionary) { - ArrowFieldWriter.writeDictionaryVariableWidth( - valueVector, - count, - vector, - startRow, - numRows, - dictionaryCache) - } else { - ArrowFieldWriter.writeVariableWidth(valueVector, count, vector, startRow, numRows) - } - count += numRows - case _ => - super.writeColumnSlice(input, startRow, numRows) - } - } - - private val dictionaryCache = new DictionaryCache - - override private[arrow] def startInputBatch(): Unit = dictionaryCache.clear() - - override def reset(): Unit = { - super.reset() - dictionaryCache.clear() - } } private[arrow] class LargeBinaryWriter(val valueVector: LargeVarBinaryVector) @@ -1334,7 +1575,7 @@ private[arrow] class LargeBinaryWriter(val valueVector: LargeVarBinaryVector) } private[arrow] class DateWriter(val valueVector: DateDayVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -1346,7 +1587,7 @@ private[arrow] class DateWriter(val valueVector: DateDayVector) } private[arrow] class TimestampWriter(val valueVector: TimeStampMicroTZVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1358,7 +1599,7 @@ private[arrow] class TimestampWriter(val valueVector: TimeStampMicroTZVector) } private[arrow] class TimestampNTZWriter(val valueVector: TimeStampMicroVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1370,7 +1611,7 @@ private[arrow] class TimestampNTZWriter(val valueVector: TimeStampMicroVector) } private[arrow] class TimeNanoWriter(val valueVector: TimeNanoVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1387,13 +1628,49 @@ private[arrow] class ArrayWriter(val valueVector: ListVector, val elementWriter: override def setNull(): Unit = {} override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val array = input.getArray(ordinal) + input.getArray(ordinal) match { + case unsafe: UnsafeArrayData => + writeElements(unsafe) + case array => + val numElements = array.numElements() + valueVector.startNewValue(count) + var i = 0 + while (i < numElements) { + elementWriter.write(array, i) + i += 1 + } + valueVector.endValue(count, numElements) + } + } + + // Each unsafe array gets a view of its own from Spark's getter. A view reused across values + // would keep the last row it read reachable after that row's values were copied. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeElements(row.getArray(ordinal)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 - valueVector.startNewValue(count) - while (i < array.numElements()) { - elementWriter.write(array, i) + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeElements(array.getArray(i)) + } + count += 1 i += 1 } + } + + private def writeElements(array: UnsafeArrayData): Unit = { + valueVector.startNewValue(count) + elementWriter.writeArrayElements(array) valueVector.endValue(count, array.numElements()) } @@ -1449,11 +1726,48 @@ private[arrow] class StructWriter( } override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val struct = input.getStruct(ordinal, children.length) + input.getStruct(ordinal, children.length) match { + case unsafe: UnsafeRow => + writeFields(unsafe) + case struct => + var i = 0 + valueVector.setIndexDefined(count) + while (i < struct.numFields) { + children(i).write(struct, i) + i += 1 + } + } + } + + // Each unsafe struct gets a view of its own, as in ArrayWriter. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeFields(row.getStruct(ordinal, children.length)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeFields(array.getStruct(i, children.length)) + } + count += 1 + i += 1 + } + } + + private def writeFields(struct: UnsafeRow): Unit = { valueVector.setIndexDefined(count) - while (i < struct.numFields) { - children(i).write(struct, i) + var i = 0 + while (i < children.length) { + children(i).writeUnsafeRowField(struct, i) i += 1 } } @@ -1524,19 +1838,65 @@ private[arrow] class MapWriter( override def setNull(): Unit = {} override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val map = input.getMap(ordinal) - valueVector.startNewValue(count) - val keys = map.keyArray() - val values = map.valueArray() + input.getMap(ordinal) match { + case unsafe: UnsafeMapData => + writeEntries(unsafe) + case map => + val numElements = map.numElements() + valueVector.startNewValue(count) + val keys = map.keyArray() + val values = map.valueArray() + var i = 0 + while (i < numElements) { + structVector.setIndexDefined(keyWriter.count) + keyWriter.write(keys, i) + valueWriter.write(values, i) + i += 1 + } + valueVector.endValue(count, numElements) + } + } + + // Each unsafe map gets a view of its own, as in ArrayWriter. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeEntries(row.getMap(ordinal)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 - while (i < map.numElements()) { - structVector.setIndexDefined(keyWriter.count) - keyWriter.write(keys, i) - valueWriter.write(values, i) + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeEntries(array.getMap(i)) + } + count += 1 i += 1 } + } + + private def writeEntries(map: UnsafeMapData): Unit = { + val numElements = map.numElements() + valueVector.startNewValue(count) + if (numElements > 0) { + setEntriesValid(keyWriter.count, numElements) + keyWriter.writeArrayElements(map.keyArray()) + valueWriter.writeArrayElements(map.valueArray()) + } + valueVector.endValue(count, numElements) + } - valueVector.endValue(count, map.numElements()) + /** Marks entries `[start, start + numEntries)` valid. */ + private def setEntriesValid(start: Int, numEntries: Int): Unit = { + // Grows the validity buffer to the last entry. + structVector.setIndexDefined(start + numEntries - 1) + ArrowFieldWriter.setValid(structVector.getValidityBuffer, start, numEntries) } // Like ArrayWriter.writeColumnSlice, with the keys and values in Spark's two child vectors. @@ -1547,9 +1907,7 @@ private[arrow] class MapWriter( val numEntries = ArrowFieldWriter.writeListOffsets(valueVector, count, vector, startRow, numRows) if (numEntries > 0) { - // Grows the entries' validity buffer, then marks every entry valid. - structVector.setIndexDefined(entryStart + numEntries - 1) - ArrowFieldWriter.setValid(structVector.getValidityBuffer, entryStart, numEntries) + setEntriesValid(entryStart, numEntries) } val keys = vector.getChild(0) val values = vector.getChild(1) @@ -1600,7 +1958,7 @@ private[arrow] class NullWriter(val valueVector: NullVector) extends ArrowFieldW } private[arrow] class IntervalYearWriter(val valueVector: IntervalYearVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -1612,7 +1970,7 @@ private[arrow] class IntervalYearWriter(val valueVector: IntervalYearVector) } private[arrow] class DurationWriter(val valueVector: DurationVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1624,7 +1982,7 @@ private[arrow] class DurationWriter(val valueVector: DurationVector) } private[arrow] class IntervalMonthDayNanoWriter(val valueVector: IntervalMonthDayNanoVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { val ci = input.getInterval(ordinal) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index ef73522577a..34bb32d908f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -249,7 +249,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { val allocator = new RootAllocator(Long.MaxValue) val numRows = 32 class CountingFixedWidthWriter(override val valueVector: BaseFixedWidthVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { var scalarWrites: Int = 0 override def setValue(input: SpecializedGetters, ordinal: Int): Unit = scalarWrites += 1 override protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit = diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala index 294f9152985..400d9aad797 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala @@ -19,23 +19,28 @@ package org.apache.spark.sql.comet.execution.arrow +import java.lang.reflect.Modifier import java.math.{BigDecimal => JavaBigDecimal, BigInteger} +import java.util.{ArrayDeque, Collections, IdentityHashMap} +import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.util.Random import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers -import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{DecimalVector, FieldVector, IntVector, ValueVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.memory.{ArrowBuf, RootAllocator} +import org.apache.arrow.vector.{BaseVariableWidthVector, DecimalVector, FieldVector, IntVector, ValueVector, VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.{ListVector, StructVector} -import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection} -import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeArrayData, UnsafeMapData, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus @@ -124,6 +129,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { /** * Fills rows `[0, n)`. Array and map rows are laid out back to back, or from the last row * backwards when `reversed`, and null rows get no offsets at all, as in Spark's readers. + * Collections hold fewer than `maxLength` elements. */ private def fill( v: WritableColumnVector, @@ -131,18 +137,19 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { n: Int, rnd: Random, nullFraction: Double, - reversed: Boolean): Unit = dataType match { + reversed: Boolean, + maxLength: Int = 6): Unit = dataType match { case st: StructType => (0 until n).foreach(i => if (rnd.nextDouble() < nullFraction) v.putNull(i)) st.fields.zipWithIndex.foreach { case (field, ordinal) => val child = v.getChild(ordinal) child.reserve(n) - fill(child, field.dataType, n, rnd, nullFraction, reversed) + fill(child, field.dataType, n, rnd, nullFraction, reversed, maxLength) // Spark's Parquet reader nulls the fields of a null struct. Other producers need not. (0 until n).foreach(i => if (v.isNullAt(i) && rnd.nextBoolean()) child.putNull(i)) } case _: ArrayType | _: MapType => - val lengths = Array.fill(n)(rnd.nextInt(6)) + val lengths = Array.fill(n)(rnd.nextInt(maxLength)) val children = dataType match { case _: ArrayType => Seq(v.arrayData()) case _ => Seq(v.getChild(0), v.getChild(1)) @@ -159,10 +166,10 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } dataType match { case ArrayType(elementType, _) => - fill(v.arrayData(), elementType, offset, rnd, nullFraction, reversed) + fill(v.arrayData(), elementType, offset, rnd, nullFraction, reversed, maxLength) case MapType(keyType, valueType, _) => - fill(v.getChild(0), keyType, offset, rnd, 0.0, reversed) - fill(v.getChild(1), valueType, offset, rnd, nullFraction, reversed) + fill(v.getChild(0), keyType, offset, rnd, 0.0, reversed, maxLength) + fill(v.getChild(1), valueType, offset, rnd, nullFraction, reversed, maxLength) } case _ => (0 until n).foreach { i => @@ -206,6 +213,14 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + /** Both roots must hold the same rows, compared vector by vector. */ + private def assertSameRoots(expected: VectorSchemaRoot, actual: VectorSchemaRoot): Unit = { + actual.getRowCount shouldBe expected.getRowCount + expected.getFieldVectors.asScala.zip(actual.getFieldVectors.asScala).foreach { case (e, a) => + assertSameVectors(e, a, "") + } + } + /** Each vector in the tree must match, leaf values under null parents included. */ private def assertSameVectors( expected: ValueVector, @@ -222,12 +237,20 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { a.getElementStartIndex(i) shouldBe e.getElementStartIndex(i) a.getElementEndIndex(i) shouldBe e.getElementEndIndex(i) case (_: StructVector, _) => - case _ if !expected.isNull(i) => - (expected.getObject(i), actual.getObject(i)) match { - case (e: Array[Byte], a: Array[Byte]) => a.toSeq shouldBe e.toSeq - case (e, a) => a shouldBe e - } case _ => + // A reader takes each value's start from the end of the one before, null or not. + // The paths may differ in the bytes under a null struct, but never go backwards. + actual match { + case a: BaseVariableWidthVector => + a.getStartOffset(i + 1) should be >= a.getStartOffset(i) + case _ => + } + if (!expected.isNull(i)) { + (expected.getObject(i), actual.getObject(i)) match { + case (e: Array[Byte], a: Array[Byte]) => a.toSeq shouldBe e.toSeq + case (e, a) => a shouldBe e + } + } } } } @@ -269,9 +292,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { rowWriter.finish() columnar.getRowCount shouldBe length - rows.getFieldVectors.asScala.zip(columnar.getFieldVectors.asScala).foreach { case (e, a) => - assertSameVectors(e, a, "") - } + assertSameRoots(rows, columnar) } finally { columnar.close() rows.close() @@ -488,6 +509,35 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + test("a decimal past its precision in an unsafe array still fails the row path") { + // Written at a wider precision and read at precision p, so the array holds 10^p or -10^p, the + // smallest values past p. Up to 18 digits the array holds the unscaled long, and past that the + // unscaled bytes. + Seq( + (DecimalType(18, 2), DecimalType(5, 2), "1000.00"), + (DecimalType(38, 0), DecimalType(20, 0), "100000000000000000000")).foreach { + case (written, read, magnitude) => + Seq(magnitude, s"-$magnitude").foreach { value => + val decimal = Decimal(new JavaBigDecimal(value), written.precision, written.scale) + val row = UnsafeProjection.create(new StructType().add("a", ArrayType(written)))( + new GenericInternalRow(Array[Any](new GenericArrayData(Array[Any](decimal))))) + val allocator = new RootAllocator(Long.MaxValue) + val root = VectorSchemaRoot.create( + Utils.toArrowSchema(new StructType().add("a", ArrayType(read)), "UTC"), + allocator) + try { + val writer = ArrowWriter.create(root, 1) + withClue(s"$read $value: ") { + intercept[ArithmeticException](writer.write(row)) + } + } finally { + root.close() + allocator.close() + } + } + } + } + test("a narrow decimal with more digits than its precision passes through") { // Spark's getDecimal does not check an int- or long-backed value against the precision, so // neither path does. Arrow's BigDecimal setter, which the writer used before, threw instead. @@ -588,6 +638,52 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + test("a dictionary-encoded field appends as a column after its struct took the row path") { + // A struct with a collection field takes the row path in a batch that holds null structs, + // which leaves the offsets of its string field's trailing nulls for the next value to fill. A + // dictionary-encoded field has to fill them, and grow its buffers, before it appends a column. + val st = new StructType().add("a", ArrayType(IntegerType)).add("s", StringType) + val schema = new StructType().add("st", st) + val rnd = new Random(5) + val batches = Seq((0.5, 100), (0.0, 5000), (0.5, 100), (0.0, 5000)).map { + case (nullFraction, n) => + val v = newVector(n, st, offHeap = false) + (0 until n).foreach { i => + if (rnd.nextDouble() < nullFraction || (nullFraction > 0 && i == n - 1)) { + v.putNull(i) + } + } + fill(v.getChild(0), ArrayType(IntegerType), n, rnd, nullFraction = 0.0, reversed = false) + fillDictionary(v.getChild(1), StringType, n, rnd, nullFraction = 0.0) + // Spark's Parquet reader nulls the fields of a null struct. + (0 until n).foreach { i => + if (v.isNullAt(i)) { + v.getChild(0).putNull(i) + v.getChild(1).putNull(i) + } + } + new ColumnarBatch(Array[ColumnVector](v), n) + } + val allocator = new RootAllocator(Long.MaxValue) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val columnar = VectorSchemaRoot.create(arrowSchema, allocator) + val rows = VectorSchemaRoot.create(arrowSchema, allocator) + try { + val columnarWriter = ArrowWriter.create(columnar, 1) + batches.foreach(b => columnarWriter.writeColumns(b, 0, b.numRows())) + columnarWriter.finish() + val rowWriter = ArrowWriter.create(rows, 1) + batches.foreach(b => (0 until b.numRows()).foreach(i => rowWriter.write(b.getRow(i)))) + rowWriter.finish() + assertSameRoots(rows, columnar) + } finally { + columnar.close() + rows.close() + allocator.close() + batches.foreach(_.close()) + } + } + test("a struct that switches between the columnar and row paths across appended batches") { // A struct with an array or a map field takes the row path only in batches that hold null // structs, so its fields' writers keep appending where the other path left off. @@ -627,10 +723,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { val rowWriter = ArrowWriter.create(rows, batches.map(_.numRows()).sum) batches.foreach(b => (0 until b.numRows()).foreach(i => rowWriter.write(b.getRow(i)))) rowWriter.finish() - columnar.getRowCount shouldBe rows.getRowCount - rows.getFieldVectors.asScala.zip(columnar.getFieldVectors.asScala).foreach { - case (e, a) => assertSameVectors(e, a, "") - } + assertSameRoots(rows, columnar) } finally { columnar.close() rows.close() @@ -640,4 +733,221 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } } + + // Nested shapes whose unsafe forms take paths of their own: an array of each primitive type not + // in `nestedTypes`, whose elements are either copied in one block or converted one at a time, + // collections inside collections, and a struct wider than one word of null bits. + private val moreNestedTypes: Seq[DataType] = + primitiveTypes.map(ArrayType(_)).filterNot(nestedTypes.contains) ++ Seq( + ArrayType(MapType(StringType, IntegerType)), + MapType(StringType, ArrayType(StringType)), + MapType(LongType, DecimalType(9, 2)), + new StructType() + .add("m", MapType(IntegerType, StringType)) + .add("s", new StructType().add("x", BinaryType).add("y", DecimalType(18, 4))), + StructType( + (0 until 70).map(i => StructField(s"f$i", if (i % 7 == 3) StringType else LongType)))) + + /** + * Writes `numRows` rows, `row(i)` through the generic row path and its unsafe projection + * through the unsafe one, the latter copied to off-heap memory when `offHeap`, as Spark's + * off-heap pages hold rows. A third writer alternates between the two kinds of row, so each + * path appends where the other left off, and a fourth writes generic rows holding the unsafe + * row's values, as a copied row holds unsafe arrays. All of them must write the same vectors. + */ + private def assertUnsafeRowsMatchGeneric(numRows: Int, schema: StructType, offHeap: Boolean)( + row: Int => InternalRow): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val roots = Seq.fill(4)(VectorSchemaRoot.create(arrowSchema, allocator)) + val buffers = ArrayBuffer.empty[ArrowBuf] + try { + // Every writer starts undersized, so each buffer has to grow. + val writers = roots.map(ArrowWriter.create(_, 1)) + val project = UnsafeProjection.create(schema) + (0 until numRows).foreach { i => + val generic = row(i) + val projected = project(generic) + val unsafe = if (offHeap) { + val buffer = allocator.buffer(math.max(projected.getSizeInBytes, 8).toLong) + buffers += buffer + projected.writeToMemory(null, buffer.memoryAddress()) + val copy = new UnsafeRow(projected.numFields()) + copy.pointTo(null, buffer.memoryAddress(), projected.getSizeInBytes) + copy + } else { + projected + } + writers(0).write(generic) + writers(1).write(unsafe) + writers(2).write(if (i % 3 == 1) generic else unsafe) + writers(3).write(new GenericInternalRow(Array.tabulate[Any](schema.length) { ordinal => + unsafe.get(ordinal, schema(ordinal).dataType) + })) + } + writers.foreach(_.finish()) + roots.head.getRowCount shouldBe numRows + roots.tail.foreach(assertSameRoots(roots.head, _)) + } finally { + roots.foreach(_.close()) + buffers.foreach(_.close()) + allocator.close() + } + } + + /** Fills `n` rows of `schema` as Spark's readers lay them out, and checks them as above. */ + private def assertFilledRowsMatch( + schema: StructType, + n: Int, + nullFraction: Double, + offHeap: Boolean, + maxLength: Int = 6): Unit = { + val rnd = new Random(schema.hashCode) + val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false)) + try { + schema.fields.zip(vectors).foreach { case (field, v) => + fill(v, field.dataType, n, rnd, nullFraction, reversed = false, maxLength) + } + val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n) + assertUnsafeRowsMatchGeneric(n, schema, offHeap)(batch.getRow) + } finally { + vectors.foreach(_.close()) + } + } + + for (offHeap <- Seq(false, true); nullFraction <- Seq(0.0, 0.2, 1.0)) { + test(s"unsafe rows match the generic row path: offHeap=$offHeap, nulls=$nullFraction") { + val types = primitiveTypes ++ nestedTypes ++ moreNestedTypes + val schemas = types.map(t => new StructType().add("c", t)) :+ + // A row wider than one word of null bits, with every type in it. + StructType(types.zipWithIndex.flatMap { case (t, i) => + Seq(StructField(s"a$i", t), StructField(s"b$i", t), StructField(s"c$i", t)) + }) + schemas.foreach { schema => + withClue(s"${schema.simpleString}: ") { + assertFilledRowsMatch(schema, numRows, nullFraction, offHeap) + } + } + } + } + + test("unsafe rows match the generic row path with collections past 64 elements") { + // An unsafe array keeps its null bits in 64-bit words, so long collections span several, + // and their elements land at every bit offset of the Arrow validity buffer. + val types = (nestedTypes ++ moreNestedTypes).collect { case t @ (_: ArrayType | _: MapType) => + t + } + for (dataType <- types; nullFraction <- Seq(0.0, 0.3)) { + withClue(s"$dataType nulls=$nullFraction: ") { + val schema = new StructType().add("c", dataType) + assertFilledRowsMatch(schema, 24, nullFraction, offHeap = false, maxLength = 150) + } + } + } + + /** + * Whether `target` is reachable from `root` through the fields of the Arrow writers and of any + * Spark unsafe views they hold. Arrow's vectors hold only Arrow memory, so they are skipped. + */ + private def reachable(root: AnyRef, target: AnyRef): Boolean = { + val writerPackage = classOf[ArrowWriter].getPackage.getName + "." + val seen = Collections.newSetFromMap(new IdentityHashMap[AnyRef, java.lang.Boolean]) + val pending = new ArrayDeque[AnyRef] + def pushFields(o: AnyRef): Unit = { + var c: Class[_] = o.getClass + while (c != null) { + c.getDeclaredFields.foreach { field => + if (!field.getType.isPrimitive && !Modifier.isStatic(field.getModifiers)) { + field.setAccessible(true) + val value = field.get(o) + if (value != null) { + pending.push(value) + } + } + } + c = c.getSuperclass + } + } + pending.push(root) + while (!pending.isEmpty) { + val o = pending.pop() + if (o eq target) { + return true + } + if (seen.add(o)) { + o match { + case objects: Array[AnyRef] => objects.foreach(x => if (x != null) pending.push(x)) + case _: UnsafeArrayData | _: UnsafeMapData | _: UnsafeRow => pushFields(o) + case _ if o.getClass.getName.startsWith(writerPackage) => pushFields(o) + case _ => + } + } + } + false + } + + test("the row path keeps no reference to an unsafe row once it is written") { + // Arrays, maps and structs at the top level, inside one another and as array elements. A view + // reused across values would keep the last row it read, and the row's memory, reachable. + val schema = new StructType() + .add("a", ArrayType(ArrayType(IntegerType))) + .add("s", new StructType().add("x", StringType).add("y", ArrayType(LongType))) + .add("m", MapType(StringType, new StructType().add("z", BinaryType))) + .add("as", ArrayType(new StructType().add("w", IntegerType))) + .add("am", ArrayType(MapType(IntegerType, StringType))) + val n = 20 + val rnd = new Random(7) + val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false)) + val allocator = new RootAllocator(Long.MaxValue) + val root = VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), allocator) + try { + schema.fields.zip(vectors).foreach { case (field, v) => + fill(v, field.dataType, n, rnd, nullFraction = 0.0, reversed = false) + } + val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n) + val project = UnsafeProjection.create(schema) + val writer = ArrowWriter.create(root, n) + (0 until n).foreach { i => + val row = project(batch.getRow(i)).copy() + writer.write(row) + withClue(s"row $i: ") { + reachable(writer, row.getBaseObject) shouldBe false + } + } + } finally { + root.close() + allocator.close() + vectors.foreach(_.close()) + } + } + + test("unsafe rows copy strings and binaries of every length") { + // Short values are copied a word at a time and long ones in one call, so this covers both + // sides of the threshold and every remainder. + val schema = new StructType() + .add("s", StringType) + .add("b", BinaryType) + .add("a", ArrayType(StringType)) + .add("m", MapType(StringType, BinaryType)) + val rows = (0 to 140).map { length => + if (length % 17 == 5) { + new GenericInternalRow(4) + } else { + val bytes = Array.tabulate[Byte](length)(i => (i * 31 + length).toByte) + val text = UTF8String.fromBytes(bytes) + val half = UTF8String.fromBytes(bytes, length / 3, length / 2) + new GenericInternalRow( + Array[Any]( + text, + bytes, + new GenericArrayData(Array[Any](text, null, half)), + new ArrayBasedMapData( + new GenericArrayData(Array[Any](text, UTF8String.fromString(s"key-$length"))), + new GenericArrayData(Array[Any](null, bytes))))) + } + } + Seq(false, true).foreach { offHeap => + assertUnsafeRowsMatchGeneric(rows.size, schema, offHeap)(rows) + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala index 1e4230fb199..3f23851d76f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala @@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.arrow import java.nio.charset.StandardCharsets.UTF_8 import java.nio.file.{Files, Paths} +import scala.reflect.ClassTag import scala.util.Using import org.scalatest.funsuite.AnyFunSuite @@ -32,9 +33,10 @@ import org.apache.arrow.memory.{ArrowBuf, OutOfMemoryException, RootAllocator} import org.apache.arrow.vector.{VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.util.OversizedAllocationException import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeArrayData, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.types.{StringType, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, StringType, StructField, StructType} import org.apache.spark.unsafe.types.UTF8String class CometStringWriterSuite extends AnyFunSuite with Matchers { @@ -299,4 +301,85 @@ class CometStringWriterSuite extends AnyFunSuite with Matchers { allocator.getAllocatedMemory shouldBe 0L } } + + private def unsafeRow(value: UTF8String): UnsafeRow = + UnsafeProjection.create(schema)(row(value)).copy() + + /** An unsafe array of `values`, in a row of its own. */ + private def unsafeArray(values: UTF8String*): UnsafeArrayData = { + val arraySchema = StructType(Seq(StructField("texts", ArrayType(StringType)))) + UnsafeProjection + .create(arraySchema)( + new GenericInternalRow(Array[Any](new GenericArrayData(values.toArray[Any])))) + .copy() + .getArray(0) + } + + /** `offsetAndSize`, as an unsafe row or array packs it, with the size made negative. */ + private def negativeSize(offsetAndSize: Long): Long = + (offsetAndSize & ~0xffffffffL) | 0xffffffffL + + /** + * Runs `write` from row `index(vector)` of a vector holding one value at row 0, whose offsets + * then end at `end`, and checks that it throws `E` having changed and allocated nothing. + */ + private def assertRejected[E <: Exception: ClassTag](end: Int, index: VarCharVector => Int)( + write: StringWriter => Unit): E = { + Using.resource(new RootAllocator(1024 * 1024)) { allocator => + val rejected = Using.Manager { use => + val vector = use(new VarCharVector("text", allocator)) + vector.allocateNew(8, 4) + vector.setSafe(0, Array[Byte](42)) + vector.getOffsetBuffer.setInt(4L, end) + val writer = new StringWriter(vector) + val at = index(vector) + writer.count = at + val allocated = allocator.getAllocatedMemory + val e = intercept[E](write(writer)) + allocator.getAllocatedMemory shouldBe allocated + writer.count shouldBe at + vector.getLastSet shouldBe 0 + vector.getOffsetBuffer.getInt(4L) shouldBe end + vector.getOffsetBuffer.getInt(8L) shouldBe 0 + vector.getDataBuffer.getByte(0L) shouldBe 42.toByte + if (at < vector.getValueCapacity) vector.isNull(at) shouldBe true + e + }.get + allocator.getAllocatedMemory shouldBe 0L + rejected + } + } + + test("unsafe row string offset overflow throws before anything changes or is allocated") { + // Unlike Arrow's setters, the unsafe paths check before growing the offsets. + val nearlyFull = Int.MaxValue - 5 + val value = unsafeRow(UTF8String.fromBytes(Array[Byte](1, 2, 3, 4, 5, 6))) + Seq[VarCharVector => Int](_ => 1, _.getValueCapacity).foreach { index => + assertRejected[OversizedAllocationException](nearlyFull, index) { + _.writeUnsafeRowField(value, 0) + } + } + // An array's elements are all checked before any of them is written. + val fitsThenOverflows = unsafeArray( + UTF8String.fromBytes(Array[Byte](1, 2)), + null, + UTF8String.fromBytes(Array[Byte](3, 4, 5, 6))) + Seq[VarCharVector => Int](_ => 1, _.getValueCapacity - 1).foreach { index => + assertRejected[OversizedAllocationException](nearlyFull, index) { + _.writeArrayElements(fitsThenOverflows) + } + } + } + + test("unsafe row negative string length is rejected before reserving or copying") { + val value = unsafeRow(UTF8String.fromString("abc")) + value.setLong(0, negativeSize(value.getLong(0))) + val array = unsafeArray(UTF8String.fromString("ab"), UTF8String.fromString("cd")) + array.setLong(1, negativeSize(array.getLong(1))) + Seq[StringWriter => Unit](_.writeUnsafeRowField(value, 0), _.writeArrayElements(array)) + .foreach { write => + assertRejected[IllegalArgumentException](1, _ => 1)(write).getMessage should + include("String length must be non-negative") + } + } } From 7d294535e87367d819c4e2a23422cd1158ac31b8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 8 Oct 2026 13:17:03 +0000 Subject: [PATCH 28/89] fix: match Spark 4.2 date_trunc overflow for SECOND and MILLISECOND (#6740) * fix: match Spark 4.2 date_trunc overflow for SECOND and MILLISECOND Spark 4.2.0 (SPARK-56663) raises `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, where earlier versions wrap the Long subtraction. The serde now tells the native kernel which behavior to follow, and the wrapping fixture from #5956 is split into a 4.1-and-earlier file and a 4.2 file. * perf: keep values-only truncation for safe non-UTC SECOND/MILLISECOND dictionaries SECOND and MILLISECOND ignore the timezone, so when they do not wrap they need key masking only for a value near the lower bound. Previously a non-UTC dictionary fell through to the unconditional non-UTC masking branch and scanned every key. --- .../expression-audits/datetime_funcs.md | 1 + .../src/execution/expressions/temporal.rs | 7 +- native/proto/src/proto/expr.proto | 5 + native/spark-expr/benches/timestamp_trunc.rs | 10 +- .../src/datetime_funcs/timestamp_trunc.rs | 15 +- native/spark-expr/src/kernels/temporal.rs | 323 ++++++++++++------ .../org/apache/comet/serde/datetime.scala | 5 + .../expressions/datetime/trunc_timestamp.sql | 6 +- .../trunc_timestamp_fine_overflow.sql | 41 +++ .../trunc_timestamp_fine_overflow_spark42.sql | 49 +++ 10 files changed, 357 insertions(+), 105 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql diff --git a/docs/source/contributor-guide/expression-audits/datetime_funcs.md b/docs/source/contributor-guide/expression-audits/datetime_funcs.md index dc88f63915f..d75661a6502 100644 --- a/docs/source/contributor-guide/expression-audits/datetime_funcs.md +++ b/docs/source/contributor-guide/expression-audits/datetime_funcs.md @@ -35,6 +35,7 @@ ## date_trunc +- Spark 4.2.0 raises `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound. - Performance (tuned 2026-10-06, PR [#5956](https://github.com/apache/datafusion-comet/pull/5956)): scalar timestamp truncation skips dictionary key masking for infallible fine units and safe UTC/NTZ values with substantial repetition. A values-only range check retains masking near the lower timestamp bound; non-UTC calendar paths and high-cardinality coarse batches retain masking to avoid errors and unnecessary calendar work on unused entries. About 19x faster for MICROSECOND with 65,536 keys and 32 distinct values. Matched Criterion coverage includes sparse/dense/all-NULL keys, high-cardinality YEAR batches, and used/unused overflow values. Benchmark: `benches/timestamp_trunc.rs`. ## dayname diff --git a/native/core/src/execution/expressions/temporal.rs b/native/core/src/execution/expressions/temporal.rs index 3c27c3b8263..61a9c36edd9 100644 --- a/native/core/src/execution/expressions/temporal.rs +++ b/native/core/src/execution/expressions/temporal.rs @@ -158,7 +158,12 @@ impl ExpressionBuilder for TruncTimestampBuilder { let format = planner.create_expr(expr.format.as_ref().unwrap(), input_schema)?; let timezone = expr.timezone.clone(); - Ok(Arc::new(TimestampTruncExpr::new(child, format, timezone))) + Ok(Arc::new(TimestampTruncExpr::new( + child, + format, + timezone, + expr.wrap_second_millisecond_overflow, + ))) } } diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 096773af7e5..adcf2d64dfe 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -597,6 +597,11 @@ message TruncTimestamp { Expr format = 1; Expr child = 2; string timezone = 3; + // Whether SECOND and MILLISECOND truncation wraps when the result would fall below the + // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 (SPARK-56663) + // raises `long overflow` like every other unit. This tracks the Spark version Comet is + // running against. + bool wrap_second_millisecond_overflow = 4; } message Subquery { diff --git a/native/spark-expr/benches/timestamp_trunc.rs b/native/spark-expr/benches/timestamp_trunc.rs index 1ad974a01e1..654e30e6393 100644 --- a/native/spark-expr/benches/timestamp_trunc.rs +++ b/native/spark-expr/benches/timestamp_trunc.rs @@ -56,6 +56,7 @@ fn criterion_benchmark(c: &mut Criterion) { Arc::new(Column::new("a", 0)), lit(format), timezone.to_string(), + true, ); for rows in ROW_COUNTS { for (null_ratio, tag) in NULL_RATIOS.into_iter().chain([(0.875, "dense")]) { @@ -110,6 +111,7 @@ fn criterion_benchmark(c: &mut Criterion) { Arc::new(Column::new("a", 0)), lit(format), "UTC".into(), + true, ); let cardinality_tag = if cardinality == rows { "high" } else { "low" }; group.bench_with_input( @@ -136,8 +138,12 @@ fn criterion_benchmark(c: &mut Criterion) { true, )])); let batch = RecordBatch::try_new(schema, vec![Arc::new(input)]).unwrap(); - let expr = - TimestampTruncExpr::new(Arc::new(Column::new("a", 0)), lit("YEAR"), "UTC".into()); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit("YEAR"), + "UTC".into(), + true, + ); assert_eq!(expr.evaluate(&batch).is_err(), referenced); group.bench_with_input( BenchmarkId::from_parameter(format!("YEAR/{rows}/{tag}")), diff --git a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs index 91a83b77e50..2835fab6bd2 100644 --- a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs +++ b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs @@ -41,6 +41,9 @@ pub struct TimestampTruncExpr { /// `Arc` so it can be cheaply cloned onto Arrow `Timestamp` data types without /// reallocating, and parsed once into a `chrono::TimeZone` per batch. timezone: Arc, + /// Whether SECOND and MILLISECOND truncation wraps below the smallest timestamp, as Spark + /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 does (SPARK-56663). + wrap_second_millisecond_overflow: bool, } impl Hash for TimestampTruncExpr { @@ -48,6 +51,7 @@ impl Hash for TimestampTruncExpr { self.child.hash(state); self.format.hash(state); self.timezone.hash(state); + self.wrap_second_millisecond_overflow.hash(state); } } impl PartialEq for TimestampTruncExpr { @@ -55,6 +59,7 @@ impl PartialEq for TimestampTruncExpr { self.child.eq(&other.child) && self.format.eq(&other.format) && self.timezone.eq(&other.timezone) + && self.wrap_second_millisecond_overflow == other.wrap_second_millisecond_overflow } } @@ -63,11 +68,13 @@ impl TimestampTruncExpr { child: Arc, format: Arc, timezone: String, + wrap_second_millisecond_overflow: bool, ) -> Self { TimestampTruncExpr { child, format, timezone: Arc::from(timezone), + wrap_second_millisecond_overflow, } } } @@ -100,6 +107,7 @@ impl PhysicalExpr for TimestampTruncExpr { let format = self.format.evaluate(batch)?; let output_type = output_type(×tamp.data_type()); let tz = &self.timezone; + let wrap = self.wrap_second_millisecond_overflow; let resolve_tz = |ts: ArrayRef| -> datafusion::common::Result { // For TimestampNTZ (Timestamp(Microsecond, None)), skip timezone conversion. // NTZ values are timezone-independent and truncation should operate directly on the @@ -121,7 +129,7 @@ impl PhysicalExpr for TimestampTruncExpr { }; match (timestamp, format) { (ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(Some(format)))) => { - let result = timestamp_trunc_dyn(&resolve_tz(ts)?, format)?; + let result = timestamp_trunc_dyn(&resolve_tz(ts)?, format, wrap)?; Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Array(ts), ColumnarValue::Array(formats)) => { @@ -129,7 +137,8 @@ impl PhysicalExpr for TimestampTruncExpr { Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => { - let result = timestamp_trunc_dyn(&resolve_tz(ts_scalar.to_array()?)?, format)?; + let result = + timestamp_trunc_dyn(&resolve_tz(ts_scalar.to_array()?)?, format, wrap)?; let scalar = ScalarValue::try_from_array(&relabel(result)?, 0)?; Ok(ColumnarValue::Scalar(scalar)) } @@ -153,6 +162,7 @@ impl PhysicalExpr for TimestampTruncExpr { Arc::clone(&children[0]), Arc::clone(&self.format), self.timezone.to_string(), + self.wrap_second_millisecond_overflow, ))) } } @@ -195,6 +205,7 @@ mod tests { Arc::new(Column::new("ts", 0)), Arc::new(Literal::new(Utf8(Some("HOUR".to_string())))), "Asia/Kolkata".to_string(), + false, ); let declared = expr.data_type(&schema).unwrap(); let ColumnarValue::Array(result) = expr.evaluate(&batch).unwrap() else { diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index c3735d741e6..60bf0e735ce 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -666,9 +666,13 @@ where /// timezone or no timezone. The array may be a dictionary array. /// /// format is a scalar string specifying the format to apply to the timestamp value. +/// +/// wrap_second_millisecond_overflow selects how SECOND and MILLISECOND truncation treats a +/// result below the smallest timestamp: Spark before 4.2 wraps, 4.2 raises. pub(crate) fn timestamp_trunc_dyn( array: &dyn Array, format: String, + wrap_second_millisecond_overflow: bool, ) -> Result { match array.data_type().clone() { DataType::Dictionary(_, _) => { @@ -676,8 +680,17 @@ pub(crate) fn timestamp_trunc_dyn( array => { // Safe, low-cardinality truncation can operate only on distinct values, // rather than scanning every key just to discover unused values. - if !timestamp_trunc_dictionary_needs_mask(array.values(), array.len(), &format)? { - let values = timestamp_trunc_dyn(array.values(), format)?; + if !timestamp_trunc_dictionary_needs_mask( + array.values(), + array.len(), + &format, + wrap_second_millisecond_overflow, + )? { + let values = timestamp_trunc_dyn( + array.values(), + format, + wrap_second_millisecond_overflow, + )?; return Ok(Arc::new(array.with_values(values))); } // Dictionary values can outlive the rows that reference them (e.g. after @@ -693,7 +706,11 @@ pub(crate) fn timestamp_trunc_dyn( } else { Arc::clone(array.values()) }; - let truncated_values = timestamp_trunc_dyn(values.as_ref(), format)?; + let truncated_values = timestamp_trunc_dyn( + values.as_ref(), + format, + wrap_second_millisecond_overflow, + )?; Ok(Arc::new(array.with_values(truncated_values))) } dt => return_compute_error_with!("timestamp_trunc does not support", dt), @@ -702,7 +719,7 @@ pub(crate) fn timestamp_trunc_dyn( _ => { downcast_temporal_array!( array => { - timestamp_trunc(array, format) + timestamp_trunc(array, format, wrap_second_millisecond_overflow) .map(|a| Arc::new(a) as ArrayRef) } dt => return_compute_error_with!("timestamp_trunc does not support", dt), @@ -714,18 +731,36 @@ pub(crate) fn timestamp_trunc_dyn( /// Keep key masking for fallible values and for coarse units with many distinct values, where /// masking can avoid expensive calendar work on unused entries. Fine arithmetic processes the /// physical values regardless of validity, so infallible fine units never need a key scan. +/// MICROSECOND is always infallible. SECOND and MILLISECOND ignore the timezone and, when they do +/// not wrap, fail only near the lower bound, so only such a value needs the key scan. fn timestamp_trunc_dictionary_needs_mask( values: &dyn Array, keys_len: usize, format: &str, + wrap_second_millisecond_overflow: bool, ) -> Result { let Some(values) = values.as_any().downcast_ref::() else { return Ok(true); }; + // Truncation moves backwards by at most a year. Above this conservative lower-bound margin, + // UTC/NTZ truncation cannot underflow i64 and all values can be evaluated without a key scan. + const LOWER_SAFE_MICROS: i64 = i64::MIN + 370 * MICROS_PER_DAY; + let has_value_near_lower_bound = || { + values + .iter() + .flatten() + .any(|micros| micros < LOWER_SAFE_MICROS) + }; let granularity = normalize_timestamp_trunc_format(format)?; - if matches!(granularity, "microsecond" | "millisecond" | "second") { + if granularity == "microsecond" + || (wrap_second_millisecond_overflow && matches!(granularity, "millisecond" | "second")) + { return Ok(false); } + if matches!(granularity, "millisecond" | "second") { + // These units ignore the timezone, so a non-UTC value needs no masking of its own. + return Ok(has_value_near_lower_bound()); + } if values.timezone().is_some_and(|tz| !is_utc_timezone(tz)) { // Unused non-UTC values can also hit chrono boundary panics, not just Result errors. return Ok(true); @@ -736,13 +771,7 @@ fn timestamp_trunc_dictionary_needs_mask( // Dense NULLs/high cardinality otherwise benefit from masking before truncation. return Ok(true); } - // Truncation moves backwards by at most a year. Above this conservative lower-bound margin, - // UTC/NTZ truncation cannot underflow i64 and all values can be evaluated without a key scan. - const LOWER_SAFE_MICROS: i64 = i64::MIN + 370 * MICROS_PER_DAY; - Ok(values - .iter() - .flatten() - .any(|micros| micros < LOWER_SAFE_MICROS)) + Ok(has_value_near_lower_bound()) } /// Convert microseconds since epoch to NaiveDateTime @@ -1069,6 +1098,7 @@ fn timestamp_trunc_fine_boundary( array: &TimestampMicrosecondArray, granularity: &str, unit: i64, + wrap_second_millisecond_overflow: bool, ) -> Result { let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); for value in array.iter() { @@ -1076,11 +1106,15 @@ fn timestamp_trunc_fine_boundary( None => builder.append_null(), Some(micros) => { let remainder = micros.rem_euclid(unit); - let truncated = if matches!(granularity, "second" | "millisecond") { - // Spark uses unchecked Long subtraction for these timezone-independent units. + let truncated = if wrap_second_millisecond_overflow + && matches!(granularity, "second" | "millisecond") + { + // Before 4.2, Spark uses unchecked Long subtraction for these + // timezone-independent units. micros.wrapping_sub(remainder) } else { - // MINUTE/HOUR/DAY use Spark's exact instant-to-microseconds conversion. + // MINUTE/HOUR/DAY use Spark's exact instant-to-microseconds conversion, and + // Spark 4.2 (SPARK-56663) checks the SECOND/MILLISECOND subtraction too. micros.checked_sub(remainder).ok_or_else(|| { SparkError::Internal(format!( "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" @@ -1097,6 +1131,7 @@ fn timestamp_trunc_fine_boundary( fn timestamp_trunc_upstream( array: &TimestampMicrosecondArray, format: &str, + wrap_second_millisecond_overflow: bool, ) -> Result { let granularity = normalize_timestamp_trunc_format(format)?; @@ -1109,7 +1144,7 @@ fn timestamp_trunc_upstream( // for HOUR/DAY (without a nanosecond limit) and integer calendar code for coarse units. // The unlabelled call still applies the coarse-unit range guard and calendar fallback. let input = array.clone().with_timezone_opt(None::>); - let result = timestamp_trunc_upstream(&input, format)?; + let result = timestamp_trunc_upstream(&input, format, wrap_second_millisecond_overflow)?; return Ok(result.with_timezone_opt(array.timezone())); } @@ -1138,7 +1173,12 @@ fn timestamp_trunc_upstream( let lower = i64::MIN + (unit - i64::MIN.rem_euclid(unit)).rem_euclid(unit); // Check physical values too: DataFusion processes NULL slots before restoring validity. if array.values().iter().any(|micros| *micros < lower) { - return timestamp_trunc_fine_boundary(array, granularity, unit); + return timestamp_trunc_fine_boundary( + array, + granularity, + unit, + wrap_second_millisecond_overflow, + ); } } @@ -1227,6 +1267,7 @@ where pub(crate) fn timestamp_trunc( array: &PrimitiveArray, format: String, + wrap_second_millisecond_overflow: bool, ) -> Result where T: ArrowTemporalType + ArrowNumericType, @@ -1239,6 +1280,7 @@ where .downcast_ref::() .expect("TimestampMicrosecond type mismatch"), &format, + wrap_second_millisecond_overflow, ), dt => return_compute_error_with!( "Unsupported input type '{:?}' for function 'timestamp_trunc'", @@ -1414,8 +1456,8 @@ where mod tests { use super::{ naive_to_micros, normalize_timestamp_trunc_format, ntz_trunc_fn_for_format, - timestamp_trunc_coarse_micros, timestamp_trunc_ntz, MICROS_PER_DAY, - TIMESTAMP_TRUNC_ALIASES, + timestamp_trunc_coarse_micros, timestamp_trunc_dictionary_needs_mask, timestamp_trunc_ntz, + MICROS_PER_DAY, TIMESTAMP_TRUNC_ALIASES, }; use crate::kernels::temporal::{ date_trunc, date_trunc_array_fmt_dyn, date_trunc_dyn, timestamp_trunc, @@ -1483,7 +1525,7 @@ mod tests { ) .with_timezone_opt(timezone); assert_eq!( - timestamp_trunc(&input, format.to_string()).unwrap(), + timestamp_trunc(&input, format.to_string(), true).unwrap(), expected ); } @@ -1810,7 +1852,7 @@ mod tests { .with_timezone_utc(); for format in ["MILLISECONDS", "invalid", " DAY ", ""] { let SparkError::Internal(message) = - timestamp_trunc(&invalid_input, format.to_string()).unwrap_err() + timestamp_trunc(&invalid_input, format.to_string(), true).unwrap_err() else { panic!("expected an internal unsupported-format error"); }; @@ -1974,13 +2016,13 @@ mod tests { ]) .with_timezone_opt(timezone); assert_eq!( - timestamp_trunc(&input, format.to_string()).unwrap(), + timestamp_trunc(&input, format.to_string(), true).unwrap(), expected ); for (alias, canonical) in TIMESTAMP_TRUNC_ALIASES { if canonical == normalize_timestamp_trunc_format(format).unwrap() { assert_eq!( - timestamp_trunc(&input, alias.to_lowercase()).unwrap(), + timestamp_trunc(&input, alias.to_lowercase(), true).unwrap(), expected ); } @@ -1995,7 +2037,7 @@ mod tests { let input = TimestampMicrosecondArray::from(vec![Some(i64::MIN)]).with_timezone_opt(timezone); for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { - let error = timestamp_trunc(&input, format.to_string()).unwrap_err(); + let error = timestamp_trunc(&input, format.to_string(), true).unwrap_err(); assert!(error.to_string().contains("out of range"), "{error}"); } } @@ -2045,7 +2087,7 @@ mod tests { for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { let expected = timestamp_trunc_ntz(&input, format.to_string()).unwrap(); assert_eq!( - timestamp_trunc(&input, format.to_string()).unwrap(), + timestamp_trunc(&input, format.to_string(), true).unwrap(), expected ); } @@ -2065,7 +2107,7 @@ mod tests { Arc::new(expected_values), ) .unwrap(); - let result = timestamp_trunc_dyn(&input, "YEAR".to_string()).unwrap(); + let result = timestamp_trunc_dyn(&input, "YEAR".to_string(), true).unwrap(); assert_eq!(result.as_ref(), &expected as &dyn Array); } @@ -2075,7 +2117,10 @@ mod tests { for values in [vec![], vec![None, None]] { let input = TimestampMicrosecondArray::from(values).with_timezone_opt(timezone); for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { - assert_eq!(timestamp_trunc(&input, format.to_string()).unwrap(), input); + assert_eq!( + timestamp_trunc(&input, format.to_string(), true).unwrap(), + input + ); } } } @@ -2117,7 +2162,7 @@ mod tests { usize::from(!expected.is_empty()), expected.len().saturating_sub(1), ); - let result = timestamp_trunc_dyn(&input, "YEAR".into()).unwrap(); + let result = timestamp_trunc_dyn(&input, "YEAR".into(), true).unwrap(); // Unused entries can be masked, so compare the decoded logical rows. let result = arrow::compute::cast( result.as_ref(), @@ -2131,7 +2176,7 @@ mod tests { Arc::new(values.clone()), ) .unwrap(); - assert!(timestamp_trunc_dyn(&used, "YEAR".into()) + assert!(timestamp_trunc_dyn(&used, "YEAR".into(), true) .unwrap_err() .to_string() .contains("long overflow")); @@ -2155,9 +2200,9 @@ mod tests { ) .unwrap(); for format in ["YEAR", "WEEK"] { - let result = timestamp_trunc_dyn(&input, format.into()).unwrap(); + let result = timestamp_trunc_dyn(&input, format.into(), true).unwrap(); let decoded = arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); - let expected = timestamp_trunc(&values.slice(0, 1), format.into()).unwrap(); + let expected = timestamp_trunc(&values.slice(0, 1), format.into(), true).unwrap(); let decoded = decoded .as_any() .downcast_ref::() @@ -2168,32 +2213,79 @@ mod tests { } #[test] - fn test_timestamp_trunc_dictionary_infallible_unused_extremes() { + fn test_timestamp_trunc_dictionary_fine_mask_decision() { + // SECOND/MILLISECOND ignore the timezone, so a dictionary of safe values is truncated + // through its values alone in either mode. Only a value that can fail needs the key scan. for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { - let values = TimestampMicrosecondArray::from(vec![ + let safe = TimestampMicrosecondArray::from(vec![ Some(instant_micros("2024-05-17T12:34:56Z")), - Some(i64::MIN), None, ]) .with_timezone_opt(timezone); - for keys in [ - vec![Some(0), None, Some(2), Some(0)], - vec![None, None], - vec![], - ] { - let input = DictionaryArray::::try_new( - Int32Array::from(keys), - Arc::new(values.clone()), - ) - .unwrap(); - let decoded_input = arrow::compute::cast(&input, values.data_type()).unwrap(); - for format in ["MICROSECOND", "MILLISECOND", "SECOND"] { - let result = timestamp_trunc_dyn(&input, format.into()).unwrap(); - let decoded = - arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); - let expected = - timestamp_trunc_dyn(decoded_input.as_ref(), format.into()).unwrap(); - assert_eq!(decoded.as_ref(), expected.as_ref()); + let near_min = + TimestampMicrosecondArray::from(vec![Some(i64::MIN)]).with_timezone_opt(timezone); + for format in ["SECOND", "MILLISECOND"] { + for wrap in [true, false] { + assert!( + !timestamp_trunc_dictionary_needs_mask(&safe, 8192, format, wrap).unwrap() + ); + } + assert!( + !timestamp_trunc_dictionary_needs_mask(&near_min, 8192, format, true).unwrap() + ); + assert!( + timestamp_trunc_dictionary_needs_mask(&near_min, 8192, format, false).unwrap() + ); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_fine_unused_extremes() { + // Unused MIN entries must not raise, whether SECOND/MILLISECOND wrap (Spark before 4.2) + // or fail (Spark 4.2). + for wrap in [true, false] { + for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { + let values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MIN), + None, + ]) + .with_timezone_opt(timezone); + for keys in [ + vec![Some(0), None, Some(2), Some(0)], + vec![None, None], + vec![], + ] { + let input = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(values.clone()), + ) + .unwrap(); + let decoded_input = arrow::compute::cast(&input, values.data_type()).unwrap(); + for format in ["MICROSECOND", "MILLISECOND", "SECOND"] { + let result = timestamp_trunc_dyn(&input, format.into(), wrap).unwrap(); + let decoded = + arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); + let expected = + timestamp_trunc_dyn(decoded_input.as_ref(), format.into(), wrap) + .unwrap(); + assert_eq!(decoded.as_ref(), expected.as_ref()); + } + } + if !wrap { + // A used MIN entry still raises. + let used = DictionaryArray::::try_new( + Int32Array::from(vec![0, 1]), + Arc::new(values.clone()), + ) + .unwrap(); + for format in ["MILLISECOND", "SECOND"] { + assert!(timestamp_trunc_dyn(&used, format.into(), wrap) + .unwrap_err() + .to_string() + .contains("long overflow")); + } } } } @@ -2201,44 +2293,46 @@ mod tests { #[test] fn test_timestamp_trunc_fine_poisoned_nulls() { - for timezone in [None, Some("UTC")] { - let input = TimestampMicrosecondArray::new( - vec![ - i64::MIN, - instant_micros("2024-05-17T12:34:56Z"), - i64::MIN + 1, - i64::MAX, - ] - .into(), - Some(arrow::buffer::NullBuffer::from(vec![ - false, true, false, true, - ])), - ) - .with_timezone_opt(timezone); - let clean = TimestampMicrosecondArray::from(vec![ - None, - Some(instant_micros("2024-05-17T12:34:56Z")), - None, - Some(i64::MAX), - ]) - .with_timezone_opt(timezone); - for format in [ - "MINUTE", - "HOUR", - "DAY", - "SECOND", - "MILLISECOND", - "MICROSECOND", - ] { - assert_eq!( - timestamp_trunc(&input, format.into()).unwrap(), - timestamp_trunc(&clean, format.into()).unwrap() - ); - let slice = input.slice(2, 2); - assert_eq!( - timestamp_trunc(&slice, format.into()).unwrap(), - timestamp_trunc(&clean.slice(2, 2), format.into()).unwrap() - ); + for wrap in [true, false] { + for timezone in [None, Some("UTC")] { + let input = TimestampMicrosecondArray::new( + vec![ + i64::MIN, + instant_micros("2024-05-17T12:34:56Z"), + i64::MIN + 1, + i64::MAX, + ] + .into(), + Some(arrow::buffer::NullBuffer::from(vec![ + false, true, false, true, + ])), + ) + .with_timezone_opt(timezone); + let clean = TimestampMicrosecondArray::from(vec![ + None, + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + Some(i64::MAX), + ]) + .with_timezone_opt(timezone); + for format in [ + "MINUTE", + "HOUR", + "DAY", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ] { + assert_eq!( + timestamp_trunc(&input, format.into(), wrap).unwrap(), + timestamp_trunc(&clean, format.into(), wrap).unwrap() + ); + let slice = input.slice(2, 2); + assert_eq!( + timestamp_trunc(&slice, format.into(), wrap).unwrap(), + timestamp_trunc(&clean.slice(2, 2), format.into(), wrap).unwrap() + ); + } } } } @@ -2255,7 +2349,7 @@ mod tests { for micros in [i64::MIN, i64::MIN + 1, boundary - 1] { let input = TimestampMicrosecondArray::from(vec![micros]).with_timezone_opt(timezone); - assert!(timestamp_trunc(&input, format.into()) + assert!(timestamp_trunc(&input, format.into(), true) .unwrap_err() .to_string() .contains("long overflow")); @@ -2264,14 +2358,17 @@ mod tests { .with_timezone_opt(timezone); let expected = TimestampMicrosecondArray::from(vec![boundary, boundary]) .with_timezone_opt(timezone); - assert_eq!(timestamp_trunc(&input, format.into()).unwrap(), expected); + assert_eq!( + timestamp_trunc(&input, format.into(), true).unwrap(), + expected + ); } } } #[test] fn test_timestamp_trunc_fine_spark_wrapping() { - // Spark uses unchecked Long subtraction for SECOND/MILLISECOND, even at MIN. + // Before 4.2, Spark uses unchecked Long subtraction for SECOND/MILLISECOND, even at MIN. for (format, expected) in [ ("SECOND", 9_223_372_036_854_551_616), ("MILLISECOND", 9_223_372_036_854_775_616), @@ -2281,7 +2378,41 @@ mod tests { .with_timezone_opt(timezone); let expected = TimestampMicrosecondArray::from(vec![Some(expected), None]) .with_timezone_opt(timezone); - assert_eq!(timestamp_trunc(&input, format.into()).unwrap(), expected); + assert_eq!( + timestamp_trunc(&input, format.into(), true).unwrap(), + expected + ); + } + } + } + + #[test] + fn test_timestamp_trunc_fine_spark42_overflow() { + // Spark 4.2 (SPARK-56663) checks the SECOND/MILLISECOND subtraction like every other + // unit. First representable boundary for each unit. + for (format, boundary) in [ + ("SECOND", -9_223_372_036_854_000_000), + ("MILLISECOND", -9_223_372_036_854_775_000), + ] { + for timezone in [None, Some("UTC"), Some("Asia/Kolkata")] { + for micros in [i64::MIN, i64::MIN + 1, boundary - 1] { + let input = TimestampMicrosecondArray::from(vec![Some(micros), None]) + .with_timezone_opt(timezone); + assert!(timestamp_trunc(&input, format.into(), false) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + let input = + TimestampMicrosecondArray::from(vec![Some(boundary), Some(boundary + 1), None]) + .with_timezone_opt(timezone); + let expected = + TimestampMicrosecondArray::from(vec![Some(boundary), Some(boundary), None]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.into(), false).unwrap(), + expected + ); } } } @@ -2428,7 +2559,7 @@ mod tests { )); let input_keys = input.keys().clone(); - let result = timestamp_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = timestamp_trunc_dyn(&input, "MONTH".to_string(), true).unwrap(); let result = result .as_any() .downcast_ref::>() @@ -2608,7 +2739,7 @@ mod tests { let array = TimestampMicrosecondArray::from(vec![ts_utc_micros]).with_timezone("America/Denver"); - let result = timestamp_trunc(&array, "QUARTER".to_string()).unwrap(); + let result = timestamp_trunc(&array, "QUARTER".to_string(), true).unwrap(); // 2023-10-01 00:00:00 MDT = 2023-10-01 06:00:00 UTC let expected_utc_micros: i64 = 1696140000 * 1_000_000; diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index ed5faa019ec..eaecffac875 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, Dou import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.CometGetDateField.CometGetDateField import org.apache.comet.serde.ExprOuterClass.Expr @@ -643,6 +644,10 @@ object CometTruncTimestamp builder.setChild(childExpr.get) builder.setFormat(formatExpr.get) builder.setTimezone(timeZone.get) + // Spark 4.2.0 (SPARK-56663) raises `long overflow` when SECOND or MILLISECOND truncation + // falls below the smallest timestamp; earlier versions wrap the Long subtraction. The + // native side has to match whichever Spark we are running against. + builder.setWrapSecondMillisecondOverflow(!isSpark42Plus) Some( ExprOuterClass.Expr diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql index 642777cd6ad..662774307ea 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql @@ -163,7 +163,5 @@ SELECT date_trunc('HOUR', ts) FROM test_trunc_ts_overflow query expect_error(long overflow) SELECT date_trunc('DAY', ts) FROM test_trunc_ts_overflow --- Spark intentionally wraps Long subtraction for SECOND/MILLISECOND at the lower bound. -query expect_native(date_trunc) -SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) -FROM test_trunc_ts_overflow +-- SECOND/MILLISECOND at the lower bound depend on the Spark version. See +-- trunc_timestamp_fine_overflow.sql and trunc_timestamp_fine_overflow_spark42.sql. diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql new file mode 100644 index 00000000000..dc8a5478d18 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql @@ -0,0 +1,41 @@ +-- 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. + +-- Before 4.2, Spark wraps the Long subtraction when SECOND or MILLISECOND truncation falls below +-- Long.MinValue. Spark 4.2 (SPARK-56663) raises instead, which +-- trunc_timestamp_fine_overflow_spark42.sql covers. +-- MaxSparkVersion: 4.1 + +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Config: spark.sql.session.timeZone=UTC +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- ConfigMatrix: parquet.enable.dictionary=false,true + +statement +CREATE TABLE test_trunc_ts_fine_overflow(ts timestamp) USING parquet + +-- Long.MinValue, and an ordinary value. +statement +INSERT INTO test_trunc_ts_fine_overflow VALUES + (timestamp_micros(-9223372036854775808)), + (timestamp('2024-05-17 12:34:56.123456')) + +query expect_native(date_trunc) +SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) +FROM test_trunc_ts_fine_overflow diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql new file mode 100644 index 00000000000..60f917b8ef7 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql @@ -0,0 +1,49 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Spark 4.2 (SPARK-56663) raises when SECOND or MILLISECOND truncation falls below +-- Long.MinValue, like every other unit. trunc_timestamp_fine_overflow.sql covers 4.1 and +-- earlier, where Spark wraps the Long subtraction instead. +-- MinSparkVersion: 4.2 + +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Config: spark.sql.session.timeZone=UTC +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- ConfigMatrix: parquet.enable.dictionary=false,true + +statement +CREATE TABLE test_trunc_ts_fine_overflow(ts timestamp) USING parquet + +-- Long.MinValue, and an ordinary value. +statement +INSERT INTO test_trunc_ts_fine_overflow VALUES + (timestamp_micros(-9223372036854775808)), + (timestamp('2024-05-17 12:34:56.123456')) + +-- Shows date_trunc runs natively here, so the errors below come from Comet, not a fallback. +query expect_native(date_trunc) +SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) +FROM test_trunc_ts_fine_overflow +WHERE ts > timestamp_micros(-9223372036854775808) + +query expect_error(long overflow) +SELECT date_trunc('SECOND', ts) FROM test_trunc_ts_fine_overflow + +query expect_error(long overflow) +SELECT date_trunc('MILLISECOND', ts) FROM test_trunc_ts_fine_overflow From 7f9b0aefc8cfb0c5c2412bc9823a5eb1d8d14576 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:37:57 +0000 Subject: [PATCH 29/89] test: preserve regex escapes in SQL fixtures (#6745) --- .../expressions/string/regexp_extract.sql | 29 ++++++++++--------- .../expressions/string/regexp_extract_all.sql | 25 ++++++++-------- .../expressions/string/regexp_instr.sql | 21 +++++++------- .../string/regexp_replace_java.sql | 27 ++++++++--------- .../string/regexp_replace_rust.sql | 13 +++++---- .../expressions/string/rlike_auto_native.sql | 18 ++++++------ .../expressions/string/rlike_java.sql | 25 ++++++++-------- 7 files changed, 82 insertions(+), 76 deletions(-) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql index adea5b30ec5..81f01248444 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_extract via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_extract(s string) USING parquet @@ -24,33 +25,33 @@ statement INSERT INTO test_regexp_extract VALUES ('abc123def'), ('no match'), (NULL), ('xyz789'), ('hello world'), ('aa') -- group 0: entire match -query -SELECT regexp_extract(s, '\d+', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'\d+', 0) FROM test_regexp_extract -- group 1: first capturing group -query -SELECT regexp_extract(s, '([a-z]+)(\d+)', 1) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract -- group 2: second capturing group -query -SELECT regexp_extract(s, '([a-z]+)(\d+)', 2) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract -- no match returns empty string -query +query expect_dispatch(regexp_extract) SELECT regexp_extract(s, 'NOMATCH', 0) FROM test_regexp_extract -- backreference pattern (Java-only) -query -SELECT regexp_extract(s, '(\w)\1', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'(\w)\1', 0) FROM test_regexp_extract -- lookahead (Java-only) -query -SELECT regexp_extract(s, 'abc(?=\d)', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'abc(?=\d)', 0) FROM test_regexp_extract -- embedded flags (Java-only) -query +query expect_dispatch(regexp_extract) SELECT regexp_extract(s, '(?i)HELLO', 0) FROM test_regexp_extract -- literal arguments -query -SELECT regexp_extract('abc123', '(\d+)', 1), regexp_extract('no digits', '(\d+)', 1), regexp_extract(NULL, '(\d+)', 1) +query expect_dispatch(regexp_extract) +SELECT regexp_extract('abc123', r'(\d+)', 1), regexp_extract('no digits', r'(\d+)', 1), regexp_extract(NULL, r'(\d+)', 1) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql index 880c7515a28..af3ada87c9c 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_extract_all via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_extract_all(s string) USING parquet @@ -24,29 +25,29 @@ statement INSERT INTO test_regexp_extract_all VALUES ('abc123def456'), ('no match'), (NULL), ('100-200-300'), ('hello world') -- group 0: all entire matches -query -SELECT regexp_extract_all(s, '\d+', 0) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'\d+', 0) FROM test_regexp_extract_all -- group 1: first capturing group from each match -query -SELECT regexp_extract_all(s, '([a-z]+)(\d+)', 1) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract_all -- group 2: second capturing group from each match -query -SELECT regexp_extract_all(s, '([a-z]+)(\d+)', 2) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract_all -- no match returns empty array -query +query expect_dispatch(regexp_extract_all) SELECT regexp_extract_all(s, 'NOMATCH', 0) FROM test_regexp_extract_all -- backreference pattern (Java-only) -query -SELECT regexp_extract_all(s, '(\d)\1', 0) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'(\d)\1', 0) FROM test_regexp_extract_all -- embedded flags (Java-only) -query +query expect_dispatch(regexp_extract_all) SELECT regexp_extract_all(s, '(?i)[A-Z]+', 0) FROM test_regexp_extract_all -- literal arguments -query -SELECT regexp_extract_all('abc123def456', '(\d+)', 1), regexp_extract_all('no digits', '(\d+)', 1), regexp_extract_all(NULL, '(\d+)', 1) +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all('abc123def456', r'(\d+)', 1), regexp_extract_all('no digits', r'(\d+)', 1), regexp_extract_all(NULL, r'(\d+)', 1) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql index d4754ef5c22..71cf1c9dab1 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_instr via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_instr(s string) USING parquet @@ -24,25 +25,25 @@ statement INSERT INTO test_regexp_instr VALUES ('abc123def'), ('no match'), (NULL), ('123xyz'), ('hello world'), ('aa') -- basic: position of first digit sequence -query -SELECT regexp_instr(s, '\d+', 0) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'\d+', 0) FROM test_regexp_instr -- group 1 (still returns position of entire match per Spark semantics) -query -SELECT regexp_instr(s, '([a-z]+)(\d+)', 1) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_instr -- no match returns 0 -query +query expect_dispatch(regexp_instr) SELECT regexp_instr(s, 'NOMATCH', 0) FROM test_regexp_instr -- backreference pattern (Java-only) -query -SELECT regexp_instr(s, '(\w)\1', 0) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'(\w)\1', 0) FROM test_regexp_instr -- embedded flags (Java-only) -query +query expect_dispatch(regexp_instr) SELECT regexp_instr(s, '(?i)HELLO', 0) FROM test_regexp_instr -- literal arguments -query -SELECT regexp_instr('abc123', '\d+', 0), regexp_instr('no digits', '\d+', 0), regexp_instr(NULL, '\d+', 0) +query expect_dispatch(regexp_instr) +SELECT regexp_instr('abc123', r'\d+', 0), regexp_instr('no digits', r'\d+', 0), regexp_instr(NULL, r'\d+', 0) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql index c0e4f955606..211740e10a2 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_replace via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_replace_java(s string) USING parquet @@ -23,28 +24,28 @@ CREATE TABLE test_regexp_replace_java(s string) USING parquet statement INSERT INTO test_regexp_replace_java VALUES ('100-200'), ('abc'), (''), (NULL), ('phone 123-456-7890'), ('aabbcc') -query -SELECT regexp_replace(s, '\d+', 'X') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+', 'X') FROM test_regexp_replace_java -query -SELECT regexp_replace(s, '\d+', 'X', 1) FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+', 'X', 1) FROM test_regexp_replace_java -- backreference in replacement -query -SELECT regexp_replace(s, '(\d+)-(\d+)', '$2-$1') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'(\d+)-(\d+)', '$2-$1') FROM test_regexp_replace_java -- backreference in pattern (Java-only) -query -SELECT regexp_replace(s, '(\w)\1', 'Z') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'(\w)\1', 'Z') FROM test_regexp_replace_java -- lookahead (Java-only) -query -SELECT regexp_replace(s, '\d+(?=-)', 'X') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+(?=-)', 'X') FROM test_regexp_replace_java -- embedded flags (Java-only) -query +query expect_dispatch(regexp_replace) SELECT regexp_replace(s, '(?i)ABC', 'X') FROM test_regexp_replace_java -- literal arguments -query -SELECT regexp_replace('100-200', '(\d+)', 'X'), regexp_replace('abc', '(\d+)', 'X'), regexp_replace(NULL, '(\d+)', 'X') +query expect_dispatch(regexp_replace) +SELECT regexp_replace('100-200', r'(\d+)', 'X'), regexp_replace('abc', r'(\d+)', 'X'), regexp_replace(NULL, r'(\d+)', 'X') diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql index 6aabf1de5ef..2f5c4c94294 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql @@ -17,6 +17,7 @@ -- Test regexp_replace() with Rust regexp engine -- Opt in to the native (rust) regexp engine. Without this, it runs through the codegen dispatcher. +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. -- Config: spark.comet.expression.RegExpReplace.allowIncompatible=true statement @@ -25,12 +26,12 @@ CREATE TABLE test_regexp_replace_enabled(s string) USING parquet statement INSERT INTO test_regexp_replace_enabled VALUES ('100-200'), ('abc'), (''), (NULL), ('phone 123-456-7890') -query -SELECT regexp_replace(s, '(\d+)', 'X') FROM test_regexp_replace_enabled +query expect_native(regexp_replace) +SELECT regexp_replace(s, r'(\d+)', 'X') FROM test_regexp_replace_enabled -query -SELECT regexp_replace(s, '(\d+)', 'X', 1) FROM test_regexp_replace_enabled +query expect_native(regexp_replace) +SELECT regexp_replace(s, r'(\d+)', 'X', 1) FROM test_regexp_replace_enabled -- literal + literal + literal -query -SELECT regexp_replace('100-200', '(\d+)', 'X'), regexp_replace('abc', '(\d+)', 'X'), regexp_replace(NULL, '(\d+)', 'X') +query expect_native(regexp_replace) +SELECT regexp_replace('100-200', r'(\d+)', 'X'), regexp_replace('abc', r'(\d+)', 'X'), regexp_replace(NULL, r'(\d+)', 'X') diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql index e93b18491e9..4fe2792b8eb 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql @@ -16,8 +16,8 @@ -- under the License. -- Default settings: in-subset literal patterns run on the native Rust path. --- Routing is asserted in CometRegExpJvmSuite / CometRegexParitySuite; this file --- only checks result equality with Spark. +-- Each query checks both native routing and result equality with Spark. +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_rlike_auto(s string) USING parquet @@ -25,20 +25,20 @@ CREATE TABLE test_rlike_auto(s string) USING parquet statement INSERT INTO test_rlike_auto VALUES ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab'), ('foo'), ('bar'), ('a+b') -query +query expect_native(rlike) SELECT s RLIKE 'abc[0-9]+' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '[a-zA-Z_][a-zA-Z0-9_]*' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '(foo|bar){1,3}' FROM test_rlike_auto -query -SELECT s RLIKE 'a\+b' FROM test_rlike_auto +query expect_native(rlike) +SELECT s RLIKE r'a\+b' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '' FROM test_rlike_auto -query +query expect_native(rlike) SELECT 'hello' RLIKE '[a-z]+', '12345' RLIKE '[0-9]+', '' RLIKE '', NULL RLIKE 'a' diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql index bf86973b232..7ea90ba90dc 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql @@ -15,7 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- Test RLIKE via JVM regex engine +-- Test RLIKE dispatcher patterns and the native empty-pattern case +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_rlike_java(s string) USING parquet @@ -23,27 +24,27 @@ CREATE TABLE test_rlike_java(s string) USING parquet statement INSERT INTO test_rlike_java VALUES ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab') -query -SELECT s RLIKE '^\d+$' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'^\d+$' FROM test_rlike_java -query +query expect_dispatch(rlike) SELECT s RLIKE '^[a-z]+$' FROM test_rlike_java -query +query expect_native(rlike) SELECT s RLIKE '' FROM test_rlike_java -- backreference (Java-only) -query -SELECT s RLIKE '^(\w)\1$' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'^(\w)\1$' FROM test_rlike_java -- lookahead (Java-only) -query -SELECT s RLIKE 'abc(?=\d)' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'abc(?=\d)' FROM test_rlike_java -- embedded flags (Java-only) -query +query expect_dispatch(rlike) SELECT s RLIKE '(?i)hello' FROM test_rlike_java --- literal arguments +-- Literal arguments mix native (empty pattern) and dispatched RLIKE, so only check parity. query -SELECT 'hello' RLIKE '^[a-z]+$', '12345' RLIKE '^\d+$', '' RLIKE '', NULL RLIKE 'a' +SELECT 'hello' RLIKE '^[a-z]+$', '12345' RLIKE r'^\d+$', '' RLIKE '', NULL RLIKE 'a' From f9dd86d482d136426990186c2f1d42eef6a1ca59 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 8 Oct 2026 15:49:12 +0000 Subject: [PATCH 30/89] docs: add an FAQ page with an answer on vectorized Java/Scala UDFs (#6778) * docs: add an FAQ page with an answer on vectorized Java/Scala UDFs The About section gets a Frequently Asked Questions page. Its first question is whether Comet plans an API for vectorized Java/Scala UDFs like the Rust UDF API. The answer gives what the prototype in #6697 found: for simple functions the vectorized form saved about 3 ns per row once the code generator's overheads were removed, which did not justify the API's cost to users. It also says where a vectorized function could still pay off, and points to #6694 for those use cases. * docs: shorten the vectorized UDF FAQ answer --- docs/source/about/faq.md | 37 +++++++++++++++++++++++++++++++++++++ docs/source/about/index.md | 5 ++++- 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 docs/source/about/faq.md diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md new file mode 100644 index 00000000000..f7313fac28b --- /dev/null +++ b/docs/source/about/faq.md @@ -0,0 +1,37 @@ + + +# Frequently Asked Questions + +## Does Comet plan to add an API for vectorized Java/Scala UDFs similar to the Rust UDF API? + +Not at the moment. Comet already runs ordinary Scala and Java UDFs in its pipeline without code +changes, compiling each call into a loop over a whole batch that the JVM optimizes well. See +[Scala UDF and Java UDF Support](../user-guide/latest/scala_java_udfs.md). + +A prototype ([#6697](https://github.com/apache/datafusion-comet/pull/6697)) found that, once the code +generator's overheads were reduced, a vectorized rewrite of a simple function such as `x + 1` saved +only about 3 ns per row. That does not justify asking users to rewrite their functions against +Comet's relocated Arrow classes and rebuild them for every Comet release, so the effort is going into +the code generator instead. + +If you have a use case that a function of one row handles poorly, such as working on the UTF-8 bytes +of strings or calling a library that processes whole batches, please describe it in +[#6694](https://github.com/apache/datafusion-comet/issues/6694). For native speed, see the +experimental [Rust UDF API](../user-guide/latest/rust_udfs.md). diff --git a/docs/source/about/index.md b/docs/source/about/index.md index c5bddb38d19..fc30f765c7f 100644 --- a/docs/source/about/index.md +++ b/docs/source/about/index.md @@ -20,10 +20,12 @@ under the License. # About Background on the project: how Comet compares to similar accelerators, how releases -are versioned, and links to the Apache Software Foundation resources behind it. +are versioned, answers to common questions, and links to the Apache Software Foundation resources +behind it. - [Comparison with Gluten](gluten_comparison.md) — how Comet differs from other Spark accelerators - [Versioning Policy](versioning_policy.md) — what our version numbers mean and our release cadence +- [Frequently Asked Questions](faq.md) — answers to common questions about Comet's features and plans - [Blog Posts and Talks](blogs-and-talks.md) — external writing and conference talks about Comet - [ASF Links](../asf/index.md) — licensing, donations, security reporting, and the code of conduct @@ -34,6 +36,7 @@ are versioned, and links to the Apache Software Foundation resources behind it. Comparison with Gluten Versioning Policy +Frequently Asked Questions Blog Posts and Talks ASF Links <../asf/index> ``` From 27b388f6202fbc5303b283843396e55c956330da Mon Sep 17 00:00:00 2001 From: Unik Dahal <61407386+unikdahal@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:50:35 +0000 Subject: [PATCH 31/89] fix: make Iceberg delete-file reflection failures fatal (#5515) * fix: harden Iceberg delete-file reflection handling * test: cover Iceberg delete-file serde failures --------- Co-authored-by: Andy Grove --- .../operator/CometIcebergNativeScan.scala | 197 +++++++++--------- .../CometIcebergNativeScanSuite.scala | 167 ++++++++++++++- 2 files changed, 263 insertions(+), 101 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index d7865833776..335f02ff034 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -282,6 +282,19 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit deleteFile: Any): Option[AnyRef] = IcebergReflection.findMethod(clazz, methodName).flatMap(m => Option(m.invoke(deleteFile))) + /** + * Reads equality-delete field IDs for the native serde path. The accessor is required on every + * supported Iceberg DeleteFile API, so a missing method or invocation failure must propagate. A + * null return is distinct: Iceberg uses it for files without equality keys. + */ + private def requiredEqualityFieldIds( + deleteFileClass: Class[_], + deleteFile: Any): java.util.List[Integer] = { + val method = IcebergReflection.getMethod(deleteFileClass, "equalityFieldIds") + val ids = method.invoke(deleteFile).asInstanceOf[java.util.List[Integer]] + if (ids == null) new java.util.ArrayList[Integer]() else ids + } + /** * Extracts delete files from an Iceberg FileScanTask as a list (for deduplication). * @@ -301,97 +314,15 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) - deletes.asScala.map { deleteFile => - // The path is the one essential field. A delete file we cannot locate cannot be applied, - // and silently skipping it would leak deleted rows, so treat a missing path as fatal. - val deletePath = IcebergReflection - .extractFileLocation(contentFileClass, deleteFile) - .getOrElse( - throw new RuntimeException( - "Neither location() nor path() is declared on this Iceberg version's " + - "ContentFile -- cannot extract delete file path from FileScanTask")) - - val deleteBuilder = OperatorOuterClass.IcebergDeleteFile.newBuilder() - deleteBuilder.setFilePathIdx(internPath(deletePath)) - - val contentType = - try { - val contentMethod = IcebergReflection.getMethod(deleteFileClass, "content") - val content = contentMethod.invoke(deleteFile) - content.toString match { - case IcebergReflection.ContentTypes.POSITION_DELETES => - IcebergReflection.ContentTypes.POSITION_DELETES - case IcebergReflection.ContentTypes.EQUALITY_DELETES => - IcebergReflection.ContentTypes.EQUALITY_DELETES - case other => other - } - } catch { - case _: Exception => - IcebergReflection.ContentTypes.POSITION_DELETES - } - deleteBuilder.setContentType(contentType) - - // "PARQUET", or "PUFFIN" for a V3 deletion vector. iceberg-rust selects its - // deletion-vector reader on this, and a wrong value reads a Puffin blob as Parquet, so an - // undeterminable format is fatal: by serde time there is no fallback left. - val fileFormat = IcebergReflection - .getFileFormat(contentFileClass, deleteFile) - .getOrElse( - throw new RuntimeException( - "ContentFile.format() is not declared on this Iceberg version -- cannot tell a " + - "deletion vector from a Parquet delete file")) - deleteBuilder.setFileFormat(fileFormat) - - val specId = - try { - val specIdMethod = IcebergReflection.getMethod(deleteFileClass, "specId") - specIdMethod.invoke(deleteFile).asInstanceOf[Int] - } catch { - case _: Exception => 0 - } - deleteBuilder.setPartitionSpecId(specId) - - try { - val equalityIdsMethod = - IcebergReflection.getMethod(deleteFileClass, "equalityFieldIds") - val equalityIds = equalityIdsMethod - .invoke(deleteFile) - .asInstanceOf[java.util.List[Integer]] - deleteBuilder.addAllEqualityIds(equalityIds) - } catch { - case _: Exception => - } - - // Gated on the format, not on the accessors returning a value: Iceberg also sets - // referencedDataFile on file-scoped Parquet position deletes, where iceberg-rust ignores - // it. Forwarding it there would serialize a data-file path per delete file that nothing - // reads, and would invite keying deletion-vector detection on the field instead of on - // fileFormat, which is the only discriminator. - if (fileFormat.equalsIgnoreCase(IcebergReflection.FileFormats.PUFFIN)) { - deletionVectorField(deleteFileClass, "referencedDataFile", deleteFile) - .foreach(p => deleteBuilder.setReferencedDataFile(p.asInstanceOf[String])) - deletionVectorField(deleteFileClass, "contentOffset", deleteFile) - .foreach(o => deleteBuilder.setContentOffset(o.asInstanceOf[java.lang.Long])) - deletionVectorField(deleteFileClass, "contentSizeInBytes", deleteFile) - .foreach(s => deleteBuilder.setContentSizeInBytes(s.asInstanceOf[java.lang.Long])) - } - - // recordCount is declared on ContentFile, so it is present on every supported Iceberg - // version and a lookup failure is a real defect rather than an old-version absence. - // iceberg-rust rejects a deletion vector without it, since it checks the count against - // the cardinality it decodes from the blob. - deleteBuilder.setRecordCount( - IcebergReflection - .getMethod(contentFileClass, "recordCount") - .invoke(deleteFile) - .asInstanceOf[java.lang.Long]) - - // Encrypted delete files carry a plaintext StandardKeyMetadata blob; forward it verbatim. - // Unencrypted delete files leave the field unset. - keyMetadataBytes(keyMetadataMethod, deleteFile).foreach(deleteBuilder.setKeyMetadata) - - deleteBuilder.build() - }.toSeq + deletes.asScala + .map( + serializeDeleteFile( + _, + contentFileClass, + deleteFileClass, + keyMetadataMethod, + internPath)) + .toSeq } catch { case e: Exception => val msg = @@ -402,6 +333,85 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } + /** + * Serializes a single Iceberg DeleteFile to protobuf. + * + * Required delete metadata must never fall back to guessed values after native execution has + * been selected. Missing accessors and invocation failures therefore propagate. + */ + private[operator] def serializeDeleteFile( + deleteFile: Any, + contentFileClass: Class[_], + deleteFileClass: Class[_], + keyMetadataMethod: java.lang.reflect.Method, + internPath: String => Int): OperatorOuterClass.IcebergDeleteFile = { + val deletePath = IcebergReflection + .extractFileLocation(contentFileClass, deleteFile) + .getOrElse( + throw new RuntimeException( + "Neither location() nor path() is declared on this Iceberg version's " + + "ContentFile -- cannot extract delete file path from FileScanTask")) + + val deleteBuilder = OperatorOuterClass.IcebergDeleteFile.newBuilder() + deleteBuilder.setFilePathIdx(internPath(deletePath)) + + val contentMethod = IcebergReflection.getMethod(deleteFileClass, "content") + val contentType = contentMethod.invoke(deleteFile).toString + deleteBuilder.setContentType(contentType) + + // "PARQUET", or "PUFFIN" for a V3 deletion vector. iceberg-rust selects its + // deletion-vector reader on this, and a wrong value reads a Puffin blob as Parquet, so an + // undeterminable format is fatal: by serde time there is no fallback left. + val fileFormat = IcebergReflection + .getFileFormat(contentFileClass, deleteFile) + .getOrElse( + throw new RuntimeException( + "ContentFile.format() is not declared on this Iceberg version -- cannot tell a " + + "deletion vector from a Parquet delete file")) + deleteBuilder.setFileFormat(fileFormat) + + val specIdMethod = IcebergReflection.getMethod(deleteFileClass, "specId") + deleteBuilder.setPartitionSpecId(specIdMethod.invoke(deleteFile).asInstanceOf[Int]) + + val equalityFieldIds = requiredEqualityFieldIds(deleteFileClass, deleteFile) + val isEqualityDelete = contentType == IcebergReflection.ContentTypes.EQUALITY_DELETES + if (isEqualityDelete && equalityFieldIds.isEmpty) { + throw new IllegalStateException( + s"Iceberg equality delete file '$deletePath' has no equality field IDs") + } + deleteBuilder.addAllEqualityIds(equalityFieldIds) + + // Gated on the format, not on the accessors returning a value: Iceberg also sets + // referencedDataFile on file-scoped Parquet position deletes, where iceberg-rust ignores + // it. Forwarding it there would serialize a data-file path per delete file that nothing + // reads, and would invite keying deletion-vector detection on the field instead of on + // fileFormat, which is the only discriminator. + if (fileFormat.equalsIgnoreCase(IcebergReflection.FileFormats.PUFFIN)) { + deletionVectorField(deleteFileClass, "referencedDataFile", deleteFile) + .foreach(p => deleteBuilder.setReferencedDataFile(p.asInstanceOf[String])) + deletionVectorField(deleteFileClass, "contentOffset", deleteFile) + .foreach(o => deleteBuilder.setContentOffset(o.asInstanceOf[java.lang.Long])) + deletionVectorField(deleteFileClass, "contentSizeInBytes", deleteFile) + .foreach(s => deleteBuilder.setContentSizeInBytes(s.asInstanceOf[java.lang.Long])) + } + + // recordCount is declared on ContentFile, so it is present on every supported Iceberg + // version and a lookup failure is a real defect rather than an old-version absence. + // iceberg-rust rejects a deletion vector without it, since it checks the count against + // the cardinality it decodes from the blob. + deleteBuilder.setRecordCount( + IcebergReflection + .getMethod(contentFileClass, "recordCount") + .invoke(deleteFile) + .asInstanceOf[java.lang.Long]) + + // Encrypted delete files carry a plaintext StandardKeyMetadata blob; forward it verbatim. + // Unencrypted delete files leave the field unset. + keyMetadataBytes(keyMetadataMethod, deleteFile).foreach(deleteBuilder.setKeyMetadata) + + deleteBuilder.build() + } + /** * Serializes partition spec and data from an Iceberg FileScanTask. * @@ -1134,10 +1144,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // dropped ones from the table's schema history (mirrors Iceberg-Java's // DeleteFilter.fileProjection). val equalityFieldIds = deletes.asScala.flatMap { df => - IcebergReflection - .getEqualityFieldIds(deleteFileClass, df) - .asScala - .map(_.asInstanceOf[java.lang.Integer].intValue()) + requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) }.toSeq if (equalityFieldIds.nonEmpty) { IcebergReflection diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala index a11fc34d9a2..6ff5daf27fe 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala @@ -19,19 +19,19 @@ package org.apache.comet.serde.operator +import java.lang.reflect.InvocationTargetException + import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers +import org.apache.iceberg.{DeleteFile, FileContent} import org.apache.iceberg.expressions.Expressions import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType, StructType} -/** - * Unit tests for [[CometIcebergNativeScan.hadoopToIcebergS3Properties]]. The pinned iceberg-rust - * S3 parser reads ONLY global `s3.*` keys (never `s3.bucket.*`), so the function drops per-bucket - * keys and promotes just the TARGET bucket's keys to global `s3.*`. Pure-function assertions, so - * a lightweight `AnyFunSuite` (no Spark session) suffices. - */ +import org.apache.comet.iceberg.IcebergReflection + +/** Unit tests for Iceberg native-scan serde helpers that do not require a Spark session. */ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { test("complex type null residuals are not serialized") { @@ -239,4 +239,159 @@ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { out("s3.endpoint") shouldBe "https://global.example.com" out.values.toSet should not contain "https://some.example.com" } + + private def keyMetadataMethod(clazz: Class[_]) = clazz.getMethod("keyMetadata") + + private def serializeDeleteFile(file: AnyRef) = + CometIcebergNativeScan.serializeDeleteFile( + file, + file.getClass, + file.getClass, + keyMetadataMethod(file.getClass), + _ => 0) + + test("required Iceberg delete-file accessors are present") { + Seq("content", "specId", "equalityFieldIds").foreach { accessor => + IcebergReflection.findMethod(classOf[DeleteFile], accessor).isDefined shouldBe true + } + } + + test("Iceberg delete content names match native serde literals") { + FileContent.POSITION_DELETES.toString shouldBe IcebergReflection.ContentTypes.POSITION_DELETES + FileContent.EQUALITY_DELETES.toString shouldBe IcebergReflection.ContentTypes.EQUALITY_DELETES + } + + test("position-delete file with null equality ids serializes without equality ids") { + val proto = serializeDeleteFile(new PositionDeleteFile) + proto.getContentType shouldBe "POSITION_DELETES" + proto.getPartitionSpecId shouldBe 7 + proto.getEqualityIdsCount shouldBe 0 + proto.getFilePathIdx shouldBe 0 + } + + test("equality-delete file serializes declared equality ids") { + val proto = serializeDeleteFile(new EqualityDeleteFile) + proto.getContentType shouldBe "EQUALITY_DELETES" + proto.getEqualityIdsCount shouldBe 2 + proto.getEqualityIds(0) shouldBe 3 + proto.getEqualityIds(1) shouldBe 5 + } + + test("equality-delete file with null equality ids is fatal") { + val ex = + intercept[IllegalStateException](serializeDeleteFile(new EqualityDeleteFileWithNullIds)) + ex.getMessage shouldBe + "Iceberg equality delete file 's3://bucket/eq-null-ids.parquet' has no equality field IDs" + } + + test("equality-delete file with empty equality ids is fatal") { + val ex = + intercept[IllegalStateException](serializeDeleteFile(new EqualityDeleteFileWithEmptyIds)) + ex.getMessage shouldBe + "Iceberg equality delete file 's3://bucket/eq-empty-ids.parquet' has no equality field IDs" + } + + test("content invocation failure propagates instead of defaulting to position deletes") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingContentDeleteFile)) + ex.getCause.getMessage shouldBe "content boom" + } + + test("spec id invocation failure propagates instead of defaulting to zero") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingSpecIdDeleteFile)) + ex.getCause.getMessage shouldBe "spec boom" + } + + test("equality-id invocation failure propagates instead of dropping ids") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingEqualityIdsDeleteFile)) + ex.getCause.getMessage shouldBe "ids boom" + } + + test("missing content accessor is fatal") { + assertThrows[NoSuchMethodException](serializeDeleteFile(new NoContentAccessorDeleteFile)) + } + + test("missing equality-id accessor is fatal") { + assertThrows[NoSuchMethodException](serializeDeleteFile(new NoEqualityIdsAccessorDeleteFile)) + } + + test("missing delete-file path accessor is fatal") { + val ex = intercept[RuntimeException](serializeDeleteFile(new NoPathAccessorDeleteFile)) + ex.getMessage should include("Neither location() nor path() is declared") + } + + private abstract class BaseDeleteFile { + def location(): String + def format(): String = "PARQUET" + def recordCount(): java.lang.Long = java.lang.Long.valueOf(1L) + def keyMetadata(): java.nio.ByteBuffer = null + } + + private class PositionDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/pos-delete.parquet" + def content(): String = "POSITION_DELETES" + def specId(): Int = 7 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class EqualityDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-delete.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = + java.util.List.of(Integer.valueOf(3), Integer.valueOf(5)) + } + + private class EqualityDeleteFileWithNullIds extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-null-ids.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class EqualityDeleteFileWithEmptyIds extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-empty-ids.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = java.util.List.of[Integer]() + } + + private class ThrowingContentDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = throw new RuntimeException("content boom") + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class ThrowingSpecIdDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "POSITION_DELETES" + def specId(): Int = throw new RuntimeException("spec boom") + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class ThrowingEqualityIdsDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = throw new RuntimeException("ids boom") + } + + private class NoContentAccessorDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class NoEqualityIdsAccessorDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + } + + private class NoPathAccessorDeleteFile { + def keyMetadata(): java.nio.ByteBuffer = null + } } From 7a1c042f37c0a4e0cad97ba16aaf558ad30dcb69 Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Thu, 8 Oct 2026 16:03:17 +0000 Subject: [PATCH 32/89] test: enable SPARK-57298 collect_set tests (#5228) --- dev/diffs/4.2.0.diff | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 3335dfb3713..e53ec96c3f5 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -464,7 +464,7 @@ index 085dbcd8046..3090d321b6c 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 5b8154d2900..eb3de47f387 100644 +index 5b8154d2900..d42374b127f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -478,18 +478,7 @@ index 5b8154d2900..eb3de47f387 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -667,7 +668,9 @@ class DataFrameAggregateSuite extends SharedSparkSession - df.selectExpr("sort_array(collect_set(b) RESPECT NULLS)"), Seq(Row(Seq(null, 2)))) - } - -- test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types") { -+ test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types", -+ // Comet's native collect_set does not yet normalize NaN / -0.0. -+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4966")) { - checkAnswer( - sql("SELECT collect_set(v) FROM VALUES (double('NaN')), (double('NaN')) AS t(v)"), - Row(Seq(Double.NaN))) -@@ -932,7 +935,9 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -932,7 +933,9 @@ class DataFrameAggregateSuite extends SharedSparkSession case _ => false }.isDefined) } else { @@ -500,7 +489,7 @@ index 5b8154d2900..eb3de47f387 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -941,7 +946,8 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -941,7 +944,8 @@ class DataFrameAggregateSuite extends SharedSparkSession val objHashAggOrSortAggPlan = stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) if (useObjectHashAgg) { @@ -510,7 +499,7 @@ index 5b8154d2900..eb3de47f387 100644 } else { assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) } -@@ -991,12 +997,12 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -991,12 +995,12 @@ class DataFrameAggregateSuite extends SharedSparkSession assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { From 46d6d7c403c4f5339d1034e1981414964784d603 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:10:37 +0000 Subject: [PATCH 33/89] test: add optimizer controls and config isolation to SQL fixtures (#6731) --- .../contributor-guide/sql-file-tests.md | 76 ++++- .../array/folded_array_map_literals.sql | 33 +++ .../expressions/math/abs_null_literal.sql | 29 ++ .../comet/CometArrayExpressionSuite.scala | 21 +- .../apache/comet/CometSqlFileTestSuite.scala | 278 +++++++++++++++--- .../org/apache/comet/SqlFileTestParser.scala | 52 +++- .../apache/comet/SqlFileTestParserSuite.scala | 79 +++++ 7 files changed, 508 insertions(+), 60 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql diff --git a/docs/source/contributor-guide/sql-file-tests.md b/docs/source/contributor-guide/sql-file-tests.md index 57cc717547f..62bff931bb1 100644 --- a/docs/source/contributor-guide/sql-file-tests.md +++ b/docs/source/contributor-guide/sql-file-tests.md @@ -98,6 +98,9 @@ Sets a Spark SQL config for all queries in the file. -- Config: spark.sql.ansi.enabled=true ``` +The first `=` separates the key from the value. Values containing `=`, such as base64 +encryption keys, are preserved. An empty value is allowed; the key must be non-empty. + #### `ConfigMatrix` Runs the entire file once per combination of values. Multiple `ConfigMatrix` lines produce a @@ -117,6 +120,47 @@ sql-file: expressions/conditional/in_set.sql [spark.sql.optimizer.inSetConversio Only add a `ConfigMatrix` directive when there is a real reason to run the test under multiple configurations. Do not add `ConfigMatrix` directives speculatively. +Values are comma-separated and must be non-empty. As with `Config`, any `=` within a +value is preserved. A matrix value overrides a `Config` value for the same key. + +#### `ExcludeRules` + +Adds fully qualified optimizer rule names to the exclusion list. Constant folding remains +disabled by default. + +```sql +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation +``` + +Repeated lines are combined, and duplicate rules are removed. Rules declared through +`Config: spark.sql.optimizer.excludedRules=...` are also retained; a matrix selects the +value of that config before `ExcludeRules` and the default exclusion are added. Unknown +rules and rules Spark does not permit excluding fail the fixture rather than silently +leaving the optimizer unchanged. + +Exclude `NullPropagation` when the test needs a NULL argument to reach an expression. +Disabling ConstantFolding alone does not prevent Spark from replacing `abs(NULL)` with a +NULL literal. + +#### `ConstantFolding` + +Lets Spark fold constants for a fixture that tests folded complex literals. + +```sql +-- ConstantFolding: enabled +``` + +`disabled` is the default. With `enabled`, an all-literal `array(...)`, `map(...)` or +`named_struct(...)` can become a complex `Literal` before Comet sees it. This exercises +Comet's literal handling, rather than its expression constructors. Ordinary all-literal +function queries may instead fold away entirely; use column inputs when those functions +must execute in Comet. Combining `enabled` with an explicit ConstantFolding exclusion +is rejected as contradictory. Repeated `ConstantFolding` directives must agree. + +Malformed configuration directives fail that fixture with its file path and source line +instead of being treated as ordinary comments. Other fixtures still run. + #### `MinSparkVersion` Skips the file when running on a Spark version older than the specified version. @@ -148,6 +192,34 @@ statement INSERT INTO my_table VALUES (1, 2.0), (3, 4.0), (NULL, NULL) ``` +Use `SET` statements to change SQL configs between queries: + +```sql +-- Starts with the fixture default, ANSI disabled. +query +SELECT CAST(v AS INT) FROM my_strings + +statement +SET spark.sql.ansi.enabled=true + +query expect_error(CAST_INVALID_INPUT) +SELECT CAST(v AS INT) FROM my_strings +``` + +All SQLConf changes are restored when the fixture finishes, including on failure and +between matrix combinations. Keys changed by `SET` need not also appear in `Config`. +`RESET` is restored too, but it removes the selected setting for the remaining queries; +`RESET` without a key also resets the harness's ANSI and optimizer settings. A mid-file +`SET spark.sql.optimizer.excludedRules=...` replaces the whole exclusion list, including +the default ConstantFolding exclusion. + +Enabling the codegen dispatcher through `SET` still requires a non-error sentinel when +the file contains `expect_error`, just as enabling it through `Config` does. Run that +sentinel under the same settings as the error query. + +This isolation covers SQL configs, not `USE`, temporary views or functions, or SQL +session variables. Manage those separately. + ### Queries A `query` block executes a SELECT and compares results between Spark and Comet. The query @@ -345,7 +417,9 @@ planning, so Comet would never see them. However, `CometSqlFileTestSuite` automa disables constant folding (by excluding `ConstantFolding` from the optimizer rules), so all-literal queries are evaluated by Comet's native engine. This means you can use the default `query` mode for all-literal cases and they will be tested natively just like -column-based queries. +column-based queries, subject to other optimizer rewrites such as `NullPropagation` and +Comet's own plan-time literal handling. `ConstantFolding: enabled` opts out of the default +exclusion for fixtures that deliberately test folded literals. #### Cover edge cases diff --git a/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql b/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql new file mode 100644 index 00000000000..f63cee958d1 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql @@ -0,0 +1,33 @@ +-- 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. + +-- ConstantFolding: enabled + +statement +CREATE TABLE folded_array_map_literals(id INT, value BIGINT) USING parquet + +statement +INSERT INTO folded_array_map_literals SELECT CAST(id AS INT), id FROM range(0, 3, 1, 1) + +-- The whole array folds to one ArrayType(MapType(...), containsNull=true) literal. +-- The NULL element and the rebuilt populated map must retain identical map types. +-- Three rows in one file exercise broadcasting the folded literal across a batch. +query +SELECT array(map(1, 2), NULL) AS arr FROM folded_array_map_literals + +query +SELECT array(NULL, map(1, 2)) AS arr FROM folded_array_map_literals diff --git a/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql b/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql new file mode 100644 index 00000000000..1f8a03b0929 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql @@ -0,0 +1,29 @@ +-- 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. + +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation + +statement +CREATE TABLE abs_null_literal(v INT) USING parquet + +statement +INSERT INTO abs_null_literal VALUES (1), (-2), (NULL) + +-- Excluding ConstantFolding alone is insufficient: NullPropagation would replace +-- abs(NULL) with a NULL literal. Pin abs itself so that such a rewrite fails this test. +query expect_native(abs) +SELECT abs(CAST(NULL AS INT)) FROM abs_null_literal diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 5cbaf538316..c464746c8c7 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1379,12 +1379,10 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - // The tests below deliberately live in this suite, not a `sql-tests` fixture: constant folding is - // enabled by default here, so each `map(...)` collapses to a MapType Literal and the outer - // `array(...)` reaches `CometCreateArray` with folded-Literal children, which `CometLiteral` - // rebuilds as an equivalent `CreateMap` of primitive literals -- the folded-literal expansion path - // under review. `CometSqlFileTestSuite` force-disables `ConstantFolding`, so an equivalent SQL - // fixture would only exercise the constructor path (which `create_array.sql` already covers). + // These tests exercise folded complex literals, including CometLiteral's expansion to native + // constructors. SQL fixtures need `-- ConstantFolding: enabled` to reach that same path. + // folded_array_map_literals.sql covers the mixed map/NULL array; the remaining conversions are + // tracked in #6628. test("array of folded map literals with array values (multirow)") { withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { checkSparkAnswerAndOperator( @@ -1432,17 +1430,6 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - // A NULL element sits next to a populated one inside a single folded - // `ArrayType(MapType(IntegerType, IntegerType, false), containsNull = true)` literal. The null - // slot serializes as a typed null literal carrying the declared map type, so the rebuilt sibling - // has to report that same type for `make_array` to accept the pair. - test("folded array mixing a map literal and a NULL element (multirow)") { - withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndOperator("SELECT array(map(1, 2), NULL) AS arr FROM tbl") - checkSparkAnswerAndOperator("SELECT array(NULL, map(1, 2)) AS arr FROM tbl") - } - } - // A map value that is itself a map or an array of maps is built by Spark's own generated code // inside `CometCreateMap`, so it keeps its declared nullability all the way to the consumer. test("folded map literals with nested map values (multirow)") { diff --git a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala index b1770b72bfa..72ffb23ebaf 100644 --- a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala @@ -21,9 +21,13 @@ package org.apache.comet import java.io.File +import scala.util.{Failure, Success, Try} + import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{Abs, CreateArray, Literal} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.ArrayType class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -87,15 +91,56 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Disable constant folding so that literal expressions are evaluated by Comet's // native engine rather than being folded away by Spark's optimizer at plan time. - private val constantFoldingExcluded = Seq( - "spark.sql.optimizer.excludedRules" -> - "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") + private val constantFoldingRule = "org.apache.spark.sql.catalyst.optimizer.ConstantFolding" // Most SQL fixtures here predate Spark 4 ANSI default and expect non-ANSI semantics // (silent overflow/null on bad input). Individual files can opt in via their own // --CONFIG line, which appears later in the pair list and wins. private val ansiDisabled = Seq(SQLConf.ANSI_ENABLED.key -> "false") + /** Apply file settings and restore every SQLConf change, including mid-file SET and RESET. */ + private def withFixtureConfigs(file: SqlTestFile)(body: => Unit): Unit = { + val conf = spark.sessionState.conf + val original = conf.getAllConfs + val optimizerKey = SQLConf.OPTIMIZER_EXCLUDED_RULES.key + // Matrix settings are appended to configs, so the last value for a key wins as before. + val configuredRules = file.configs.reverse + .collectFirst { case (`optimizerKey`, value) => + value + } + .toSeq + .flatMap(_.split(",")) + val rules = (configuredRules ++ file.excludedRules ++ + (if (file.constantFoldingEnabled) Seq.empty else Seq(constantFoldingRule))) + .map(_.trim) + .filter(_.nonEmpty) + .distinct + require( + !file.constantFoldingEnabled || !rules.contains(constantFoldingRule), + "ConstantFolding: enabled conflicts with an explicit ConstantFolding exclusion") + // Spark silently ignores unknown and non-excludable rules. A fixture must fail instead of + // claiming to test a path that its optimizer settings never reached. + val optimizer = spark.sessionState.optimizer + val knownRules = optimizer.defaultBatches.flatMap(_.rules).map(_.ruleName).toSet + rules.foreach { rule => + require(knownRules.contains(rule), s"Unknown optimizer rule: $rule") + require( + !optimizer.nonExcludableRules.contains(rule), + s"Cannot exclude optimizer rule: $rule") + } + val configs = ansiDisabled ++ file.configs.filterNot(_._1 == optimizerKey) :+ + (optimizerKey -> rules.mkString(",")) + try { + withSQLConf(configs: _*)(body) + } finally { + val current = conf.getAllConfs + (current.keySet -- original.keySet).foreach(conf.unsetConf) + original.foreach { case (key, value) => + if (current.get(key) != Some(value)) conf.setConfString(key, value) + } + } + } + /** * Pin the sentinel-query convention for fixtures that route an `expect_error` through the * codegen dispatcher. See [[ExpectError]] for the failure mode this guards against. @@ -103,9 +148,9 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { private def requireSentinelForCodegenExpectError( relativePath: String, file: SqlTestFile): Unit = { - val codegenFlagOn = file.configs.exists { case (k, v) => - k == CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key && v.equalsIgnoreCase("true") - } + val codegenFlagOn = spark.sessionState.conf.getAllConfs + .get(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key) + .exists(_.equalsIgnoreCase("true")) if (!codegenFlagOn) return val hasExpectError = file.records.exists { case SqlQuery(_, _: ExpectError, _) => true @@ -132,9 +177,7 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { } private def runTestFile(relativePath: String, file: SqlTestFile): Unit = { - requireSentinelForCodegenExpectError(relativePath, file) - val allConfigs = ansiDisabled ++ file.configs ++ constantFoldingExcluded - withSQLConf(allConfigs: _*) { + withFixtureConfigs(file) { withTable(file.tables: _*) { file.records.foreach { case SqlStatement(sql, line) => @@ -167,6 +210,8 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { case Ignore(reason) => logInfo(s"IGNORED query ($reason): $sql") case ExpectError(pattern) => + // Check the effective setting after any mid-file SET or RESET. + requireSentinelForCodegenExpectError(relativePath, file) val (sparkError, cometError) = checkSparkAnswerMaybeThrows(spark.sql(sql)) assert( sparkError.isDefined, @@ -192,39 +237,198 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("SQL fixture optimizer directives control the optimized expression") { + val default = SqlFileTestParser.parse(Seq.empty) + withFixtureConfigs(default) { + val plan = spark.sql("SELECT array(1, 2)").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists(_.isInstanceOf[CreateArray]))) + val nullPlan = spark.sql("SELECT abs(CAST(NULL AS INT))").queryExecution.optimizedPlan + assert(!nullPlan.expressions.exists(_.exists(_.isInstanceOf[Abs]))) + } + val folded = SqlFileTestParser.parse(Seq("-- ConstantFolding: enabled")) + withFixtureConfigs(folded) { + val plan = spark.sql("SELECT array(1, 2)").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists { + case Literal(_, _: ArrayType) => true + case _ => false + })) + assert(!plan.expressions.exists(_.exists(_.isInstanceOf[CreateArray]))) + } + val excluded = SqlFileTestParser.parse( + Seq("-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation")) + withFixtureConfigs(excluded) { + val plan = spark.sql("SELECT abs(CAST(NULL AS INT))").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists(_.isInstanceOf[Abs]))) + assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES.key).contains(constantFoldingRule)) + } + } + + test("SQL fixture optimizer Config and matrix settings retain precedence") { + val optimizerKey = SQLConf.OPTIMIZER_EXCLUDED_RULES.key + val nullPropagation = "org.apache.spark.sql.catalyst.optimizer.NullPropagation" + val eliminateSorts = "org.apache.spark.sql.catalyst.optimizer.EliminateSorts" + val file = SqlFileTestParser.parse( + Seq( + s"-- Config: $optimizerKey=$nullPropagation", + s"-- ConfigMatrix: $optimizerKey=$eliminateSorts,$nullPropagation", + s"-- ExcludeRules: $nullPropagation, $nullPropagation")) + val original = spark.sessionState.conf.getAllConfs + configMatrix(file.configMatrix).foreach { matrix => + withFixtureConfigs(file.copy(configs = file.configs ++ matrix)) { + val selected = matrix.head._2 + assert( + spark.conf.get(optimizerKey).split(",").toSeq == + Seq(selected, nullPropagation, constantFoldingRule).distinct) + } + assert(spark.sessionState.conf.getAllConfs == original) + } + } + + test("SQL fixture optimizer directives reject ineffective exclusions") { + Seq( + "org.apache.spark.sql.catalyst.optimizer.NoSuchRule" -> "Unknown optimizer rule", + spark.sessionState.optimizer.nonExcludableRules.head -> "Cannot exclude optimizer rule") + .foreach { case (rule, message) => + val file = SqlFileTestParser.parse(Seq(s"-- ExcludeRules: $rule")) + assert( + intercept[IllegalArgumentException](withFixtureConfigs(file) {}).getMessage + .contains(message)) + } + val conflict = SqlFileTestParser.parse( + Seq("-- ConstantFolding: enabled", s"-- ExcludeRules: $constantFoldingRule")) + assert( + intercept[IllegalArgumentException](withFixtureConfigs(conflict) {}).getMessage + .contains("conflicts")) + } + + test("SQL fixture SET and RESET restore existing and previously unset keys") { + val key = "spark.comet.sqlFixture.isolation" + val original = spark.sessionState.conf.getAllConfs + assert(!original.contains(key)) + val file = SqlFileTestParser.parse(Seq("statement", s"SET $key=first")) + runTestFile("set-isolation", file) + assert(spark.sessionState.conf.getAllConfs == original) + withSQLConf(key -> "original") { + val before = spark.sessionState.conf.getAllConfs + runTestFile("set-existing", file) + assert(spark.sessionState.conf.getAllConfs == before) + runTestFile("reset-existing", SqlFileTestParser.parse(Seq("statement", s"RESET $key"))) + assert(spark.sessionState.conf.getAllConfs == before) + runTestFile("reset-all", SqlFileTestParser.parse(Seq("statement", "RESET"))) + assert(spark.sessionState.conf.getAllConfs == before) + } + } + + test("SQL fixture settings are restored after a failing statement or query") { + val key = "spark.comet.sqlFixture.isolation" + Seq("statement", "query").foreach { directive => + val original = spark.sessionState.conf.getAllConfs + val file = SqlFileTestParser.parse( + Seq( + "statement", + s"SET $key=changed", + "", + directive, + "SELECT * FROM missing_sql_fixture_isolation_table")) + intercept[RuntimeException](runTestFile("failure-isolation", file)) + assert(spark.sessionState.conf.getAllConfs == original) + } + } + + test("SQL fixture SET changes take effect between queries") { + val file = SqlFileTestParser.parse( + Seq( + "statement", + "CREATE TABLE sql_fixture_set(v STRING) USING parquet", + "", + "statement", + "INSERT INTO sql_fixture_set VALUES ('invalid'), ('1')", + "", + "query", + "SELECT CAST(v AS INT) FROM sql_fixture_set", + "", + "statement", + "SET spark.sql.ansi.enabled=true", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST(v AS INT) FROM sql_fixture_set")) + val original = spark.sessionState.conf.getAllConfs + runTestFile("set-between-queries", file) + assert(spark.sessionState.conf.getAllConfs == original) + } + + test("SQL fixture SET cannot bypass the dispatcher error sentinel") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false", + "statement", + s"SET ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST('invalid' AS INT)")) + val original = spark.sessionState.conf.getAllConfs + val error = intercept[Exception](runTestFile("set-dispatcher-sentinel", file)) + assert(error.getMessage.contains("missing a non-error")) + assert(spark.sessionState.conf.getAllConfs == original) + } + + test("SQL fixture dispatcher sentinel follows SET before an error query") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true", + "statement", + s"SET ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST('invalid' AS INT)")) + val original = spark.sessionState.conf.getAllConfs + runTestFile("set-disable-dispatcher", file) + assert(spark.sessionState.conf.getAllConfs == original) + } + // Discover and register all .sql test files discoverTestFiles(testResourceDir).foreach { file => val relativePath = testResourceDir.toURI.relativize(file.toURI).getPath - val parsed = SqlFileTestParser.parse(file) - val combinations = configMatrix(parsed.configMatrix) - - // Skip tests that fall outside the file's declared Spark version range. - val skip = skipReason(parsed) - - if (combinations.size <= 1) { - // No matrix or single combination - test(s"sql-file: $relativePath") { - skip match { - case Some(reason) => - logInfo(s"SKIPPED ($reason): $relativePath") - case None => - val effectiveConfigs = parsed.configs ++ combinations.headOption.getOrElse(Seq.empty) - runTestFile(relativePath, parsed.copy(configs = effectiveConfigs)) - } - } - } else { - // Multiple combinations: generate one test per combination - combinations.foreach { matrixConfigs => - val label = matrixConfigs.map { case (k, v) => s"$k=$v" }.mkString(", ") - test(s"sql-file: $relativePath [$label]") { - skip match { - case Some(reason) => - logInfo(s"SKIPPED ($reason): $relativePath") - case None => - runTestFile(relativePath, parsed.copy(configs = parsed.configs ++ matrixConfigs)) + Try(SqlFileTestParser.parse(file)) match { + case Failure(error) => + // A malformed file must not prevent every other fixture from being registered. + test(s"sql-file: $relativePath") { throw error } + case Success(parsed) => + val combinations = configMatrix(parsed.configMatrix) + + // Skip tests that fall outside the file's declared Spark version range. + val skip = skipReason(parsed) + + if (combinations.size <= 1) { + // No matrix or single combination + test(s"sql-file: $relativePath") { + skip match { + case Some(reason) => + logInfo(s"SKIPPED ($reason): $relativePath") + case None => + val effectiveConfigs = + parsed.configs ++ combinations.headOption.getOrElse(Seq.empty) + runTestFile(relativePath, parsed.copy(configs = effectiveConfigs)) + } + } + } else { + // Multiple combinations: generate one test per combination + combinations.foreach { matrixConfigs => + val label = matrixConfigs.map { case (k, v) => s"$k=$v" }.mkString(", ") + test(s"sql-file: $relativePath [$label]") { + skip match { + case Some(reason) => + logInfo(s"SKIPPED ($reason): $relativePath") + case None => + runTestFile( + relativePath, + parsed.copy(configs = parsed.configs ++ matrixConfigs)) + } + } } } - } } } } diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala index 2b7c573bb19..f7e653c2e13 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala @@ -104,6 +104,10 @@ case class ExpectError(pattern: String) extends QueryAssertionMode * Optional maximum Spark version this test applies to (e.g. "3.4"). The test is skipped on * newer versions. Useful for paired fixtures where each version range has its own expected * error class or output format. + * @param excludedRules + * Additional optimizer rule names to exclude, alongside the default ConstantFolding exclusion. + * @param constantFoldingEnabled + * Whether to let Spark fold constants before Comet sees the query. */ case class SqlTestFile( configs: Seq[(String, String)], @@ -111,12 +115,18 @@ case class SqlTestFile( records: Seq[SqlTestRecord], tables: Seq[String], minSparkVersion: Option[String] = None, - maxSparkVersion: Option[String] = None) + maxSparkVersion: Option[String] = None, + excludedRules: Seq[String] = Seq.empty, + constantFoldingEnabled: Boolean = false) object SqlFileTestParser { - private val ConfigPattern = """--\s*Config:\s*(.+)=(.+)""".r - private val ConfigMatrixPattern = """--\s*ConfigMatrix:\s*(.+)=(.+)""".r + private val ConfigPattern = """--\s*Config:\s*([^=]+)=(.*)""".r + private val ConfigMatrixPattern = """--\s*ConfigMatrix:\s*([^=]+)=(.*)""".r + private val ExcludeRulesPattern = """--\s*ExcludeRules:\s*(.*)""".r + private val ConstantFoldingPattern = """--\s*ConstantFolding:\s*(.*)""".r + private val ConfigDirectivePattern = + """--\s*(Config|ConfigMatrix|ExcludeRules|ConstantFolding):.*""".r private val MinSparkVersionPattern = """--\s*MinSparkVersion:\s*(.+)""".r private val MaxSparkVersionPattern = """--\s*MaxSparkVersion:\s*(.+)""".r private val CreateTablePattern = """(?i)CREATE\s+TABLE\s+(\w+)""".r.unanchored @@ -125,6 +135,9 @@ object SqlFileTestParser { val source = Source.fromFile(file, "UTF-8") try { parse(source.getLines().toSeq) + } catch { + case e: IllegalArgumentException => + throw new IllegalArgumentException(s"${file.getPath}: ${e.getMessage}", e) } finally { source.close() } @@ -135,6 +148,8 @@ object SqlFileTestParser { var configMatrix = Seq.empty[(String, Seq[String])] var minSparkVersion: Option[String] = None var maxSparkVersion: Option[String] = None + var excludedRules = Seq.empty[String] + var constantFoldingEnabled: Option[Boolean] = None val records = Seq.newBuilder[SqlTestRecord] val tables = Seq.newBuilder[String] @@ -144,13 +159,38 @@ object SqlFileTestParser { line match { case ConfigPattern(key, value) => + require(key.trim.nonEmpty, s"Empty Config key at line ${lineIdx + 1}") configs :+= (key.trim -> value.trim) lineIdx += 1 case ConfigMatrixPattern(key, values) => - configMatrix :+= (key.trim -> values.split(",").map(_.trim).toSeq) + require(key.trim.nonEmpty, s"Empty ConfigMatrix key at line ${lineIdx + 1}") + val choices = values.split(",", -1).map(_.trim).toSeq + require(choices.forall(_.nonEmpty), s"Empty ConfigMatrix value at line ${lineIdx + 1}") + configMatrix :+= (key.trim -> choices) lineIdx += 1 + case ExcludeRulesPattern(rules) => + val names = splitNames(rules) + require(names.nonEmpty, s"Empty ExcludeRules at line ${lineIdx + 1}") + excludedRules ++= names + lineIdx += 1 + + case ConstantFoldingPattern(mode) => + require( + mode == "enabled" || mode == "disabled", + s"Invalid ConstantFolding mode '$mode' at line ${lineIdx + 1}: " + + "expected enabled or disabled") + val enabled = mode == "enabled" + require( + constantFoldingEnabled.forall(_ == enabled), + s"Conflicting ConstantFolding directives at line ${lineIdx + 1}") + constantFoldingEnabled = Some(enabled) + lineIdx += 1 + + case ConfigDirectivePattern(name) => + throw new IllegalArgumentException(s"Malformed $name directive at line ${lineIdx + 1}") + case MinSparkVersionPattern(version) => minSparkVersion = Some(version.trim) lineIdx += 1 @@ -188,7 +228,9 @@ object SqlFileTestParser { records.result(), tables.result(), minSparkVersion, - maxSparkVersion) + maxSparkVersion, + excludedRules.distinct, + constantFoldingEnabled.getOrElse(false)) } private val FallbackPattern = """query\s+expect_fallback\((.+)\)""".r diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala index c356c3c37a6..0d107395ba7 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala @@ -20,6 +20,8 @@ package org.apache.comet import java.io.File +import java.nio.charset.StandardCharsets +import java.nio.file.Files import org.scalatest.funsuite.AnyFunSuite @@ -78,6 +80,83 @@ class SqlFileTestParserSuite extends AnyFunSuite { } } + test("configuration values preserve equals signs and empty Config values") { + val parsed = SqlFileTestParser.parse( + Seq( + "-- Config: fixture.key = MDEyMzQ1Njc4OTAxMjM0NQ==", + "-- Config: fixture.other = a=b=c", + "-- Config: fixture.empty=", + "-- ConfigMatrix: fixture.matrix = first==, second=part=")) + assert( + parsed.configs == Seq( + "fixture.key" -> "MDEyMzQ1Njc4OTAxMjM0NQ==", + "fixture.other" -> "a=b=c", + "fixture.empty" -> "")) + assert(parsed.configMatrix == Seq("fixture.matrix" -> Seq("first==", "second=part="))) + } + + test("optimizer directives preserve defaults and combine repeated exclusions") { + val defaults = SqlFileTestParser.parse(Seq.empty) + assert(!defaults.constantFoldingEnabled) + assert(defaults.excludedRules.isEmpty) + val parsed = SqlFileTestParser.parse( + Seq( + "-- ConstantFolding: enabled", + "-- ExcludeRules: first.Rule, second.Rule,", + "-- ExcludeRules: second.Rule, third.Rule")) + assert(parsed.constantFoldingEnabled) + assert(parsed.excludedRules == Seq("first.Rule", "second.Rule", "third.Rule")) + assert(!SqlFileTestParser.parse(Seq("-- ConstantFolding: disabled")).constantFoldingEnabled) + } + + test("malformed configuration directives fail with their source line") { + Seq( + "-- Config: no_assignment", + "-- Config: =value", + "-- ConfigMatrix: no_assignment", + "-- ConfigMatrix: key=", + "-- ConfigMatrix: key=first,,second", + "-- ConfigMatrix: key=first,", + "-- ExcludeRules:", + "-- ExcludeRules: , ,", + "-- ConstantFolding: true").foreach { directive => + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(Seq("-- comment", directive)) + } + assert(error.getMessage.contains("line 2"), directive) + } + // Prose mentioning a directive is still an ordinary comment unless it includes the colon. + assert( + SqlFileTestParser.parse(Seq("-- ConstantFolding, so literals remain visible")) == + SqlTestFile(Seq.empty, Seq.empty, Seq.empty, Seq.empty)) + } + + test("repeated ConstantFolding directives must agree") { + Seq("enabled", "disabled").foreach { mode => + val directive = s"-- ConstantFolding: $mode" + val parsed = SqlFileTestParser.parse(Seq(directive, directive)) + assert(parsed.constantFoldingEnabled == (mode == "enabled")) + val other = if (mode == "enabled") "disabled" else "enabled" + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(Seq(directive, s"-- ConstantFolding: $other")) + } + assert(error.getMessage.contains("Conflicting ConstantFolding directives at line 2")) + } + } + + test("malformed fixture files report their path and source line") { + val tempDir = Files.createDirectories(new File(System.getProperty("java.io.tmpdir")).toPath) + val path = Files.createTempFile(tempDir, "comet-malformed-fixture-", ".sql") + try { + Files.write(path, "-- comment\n-- Config: no_assignment".getBytes(StandardCharsets.UTF_8)) + val error = intercept[IllegalArgumentException](SqlFileTestParser.parse(path.toFile)) + assert(error.getMessage.contains(path.toString)) + assert(error.getMessage.contains("line 2")) + } finally { + Files.delete(path) + } + } + test("statements and queries preserve SQL, source lines and tables for cleanup") { val lines = Seq( "-- fixture comment", From 9876862df98ee10fbded5365a4abeaa4a2b022b9 Mon Sep 17 00:00:00 2001 From: dustin Date: Thu, 8 Oct 2026 16:22:54 +0000 Subject: [PATCH 34/89] perf: keep the sort-merge join when the forced hash join's build side is too large (#6384) * perf: keep the sort-merge join when the forced hash join's build side is too large With spark.comet.exec.forceShuffledHashJoin on, RewriteJoin turned every eligible SortMergeJoinExec into a shuffled hash join and read the build side's size estimate only to pick the build side. DataFusion's hash join has to hold that side in memory, so one large join failed the task while the rest of the query would have gained from the rewrite. The rule now rewrites only when the build side's estimate is under a limit and otherwise keeps the SortMergeJoinExec, which still runs natively, recording the reason as plan info. The limit is Spark's own rule for choosing a shuffled hash join, spark.sql.autoBroadcastJoinThreshold times the initial shuffle partition count, with Spark's 10 MB default threshold standing in when broadcasts are disabled. spark.comet.exec.forceShuffledHashJoin.maxBuildSize sets a fixed limit, and a non-positive value removes it, which is what the hash join benchmark configurations and the macOS benchmarking guide now pin so they keep measuring every join. * fix: restore the sort a kept sort-merge join needs after a hash join rewrite When RewriteJoin keeps a sort-merge join above a join it has already turned into a hash join on the same key, the kept join reads unsorted input. EnsureRequirements placed no sort between the two because the inner sort-merge join's ordering satisfied the outer one, and the rewrite then removed the inner join's sorts. The native sort-merge join does not check its input ordering, so matches were silently dropped. After the rewrite, add a local SortExec wherever a child no longer satisfies its parent's required ordering, the same check EnsureRequirements makes. This covers joins kept by the build-size limit and the LeftSemi and ExistenceJoin joins the rewrite already keeps, which lost rows the same way. * test: cover a window over a join the hash join rewrite reordered A window partitioned by the join key gets no sort of its own, because the sort-merge join's output ordering already satisfies it. Once RewriteJoin turns that join into a hash join, the window needs the ordering restored, or the native window fails its partition-ordering assertion. The test runs with AQE on and off and checks the join became a hash join. --- benchmarks/pyspark/run_all_benchmarks.sh | 4 + benchmarks/tpc/engines/comet-hashjoin.toml | 2 + .../tpc/engines/comet-iceberg-hashjoin.toml | 2 + .../contributor-guide/benchmarking_macos.md | 1 + .../contributor-guide/memory_management.md | 2 +- .../user-guide/latest/tuning/operators.md | 8 + .../scala/org/apache/comet/CometConf.scala | 16 + .../apache/comet/rules/CometExecRule.scala | 9 +- .../org/apache/comet/rules/RewriteJoin.scala | 95 ++++- .../apache/comet/exec/CometJoinSuite.scala | 339 +++++++++++++++++- 10 files changed, 458 insertions(+), 20 deletions(-) diff --git a/benchmarks/pyspark/run_all_benchmarks.sh b/benchmarks/pyspark/run_all_benchmarks.sh index 452ea8d1ff8..4e43641874a 100755 --- a/benchmarks/pyspark/run_all_benchmarks.sh +++ b/benchmarks/pyspark/run_all_benchmarks.sh @@ -67,6 +67,7 @@ $SPARK_HOME/bin/spark-submit \ # Run Comet JVM shuffle echo "" echo ">>> Running COMET JVM shuffle benchmark..." +# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite. $SPARK_HOME/bin/spark-submit \ --master "$SPARK_MASTER" \ --executor-memory "$EXECUTOR_MEMORY" \ @@ -86,6 +87,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.shuffle.mode=jvm \ --conf spark.comet.shuffle.mode=jvm \ --conf spark.comet.exec.replaceSortMergeJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \ --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \ --conf spark.comet.cast.allowIncompatible=true \ @@ -96,6 +98,7 @@ $SPARK_HOME/bin/spark-submit \ # Run Comet Native shuffle echo "" echo ">>> Running COMET NATIVE shuffle benchmark..." +# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite. $SPARK_HOME/bin/spark-submit \ --master "$SPARK_MASTER" \ --executor-memory "$EXECUTOR_MEMORY" \ @@ -114,6 +117,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.explain.fallback.enabled=true \ --conf spark.comet.shuffle.mode=native \ --conf spark.comet.exec.replaceSortMergeJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \ --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \ --conf spark.comet.cast.allowIncompatible=true \ diff --git a/benchmarks/tpc/engines/comet-hashjoin.toml b/benchmarks/tpc/engines/comet-hashjoin.toml index 202dcad914b..50e184b049b 100644 --- a/benchmarks/tpc/engines/comet-hashjoin.toml +++ b/benchmarks/tpc/engines/comet-hashjoin.toml @@ -31,4 +31,6 @@ driver_class_path = ["$COMET_JAR"] "spark.plugins" = "org.apache.spark.CometPlugin" "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" "spark.comet.exec.replaceSortMergeJoin" = "true" +# Measure the hash join alone: no build-side size limit on the rewrite. +"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1" "spark.comet.expression.Cast.allowIncompatible" = "true" diff --git a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml b/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml index 421d3201315..4a563613108 100644 --- a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml +++ b/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml @@ -34,6 +34,8 @@ driver_class_path = ["$COMET_JAR", "$ICEBERG_JAR"] "spark.plugins" = "org.apache.spark.CometPlugin" "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" "spark.comet.exec.replaceSortMergeJoin" = "true" +# Measure the hash join alone: no build-side size limit on the rewrite. +"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1" "spark.comet.expression.Cast.allowIncompatible" = "true" "spark.comet.enabled" = "true" "spark.comet.exec.enabled" = "true" diff --git a/docs/source/contributor-guide/benchmarking_macos.md b/docs/source/contributor-guide/benchmarking_macos.md index 13ed2c26a27..4f3541811ce 100644 --- a/docs/source/contributor-guide/benchmarking_macos.md +++ b/docs/source/contributor-guide/benchmarking_macos.md @@ -154,6 +154,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.enabled=true \ --conf spark.comet.shuffle.enabled=true \ --conf spark.comet.exec.forceShuffledHashJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ $DF_BENCH/runners/datafusion-comet/tpcbench.py \ --benchmark tpch \ --data $BENCH_DATA/ \ diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 7bd16d880f6..e3283fe26cd 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -631,4 +631,4 @@ A checklist for triaging an executor OOM kill: `batch_size * columns`, and wide or deeply nested schemas amplify it. 4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so `spark.comet.exec.forceShuffledHashJoin=true` converts a spillable sort-merge join into one that - is not. + is not, for build sides under `spark.comet.exec.forceShuffledHashJoin.maxBuildSize`. diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 644f35483ba..97baabfccad 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -30,6 +30,14 @@ to configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`. Comet does to test with both for your specific workloads. To configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`, set `spark.comet.exec.forceShuffledHashJoin=true`. +The conversion only happens when the build side is under a size limit, and a join whose build side has no statistics +is left as a `SortMergeJoin`. The size is Spark's planning estimate of the build side, or under AQE the materialized +shuffle size of the build side. By default the limit is Spark's own rule for choosing a `ShuffledHashJoin`: +`spark.sql.autoBroadcastJoinThreshold` times the initial shuffle partition count +(`spark.sql.adaptive.coalescePartitions.initialPartitionNum` when AQE and partition coalescing are both on and it is +set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive threshold, Spark's default +threshold of 10 MB is used instead. Set `spark.comet.exec.forceShuffledHashJoin.maxBuildSize` to a size in bytes to use a fixed +limit, or to a non-positive value to convert every eligible join regardless of size. ### Join Runtime Filters diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 248772fb6c4..a187f8fa5f8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -559,6 +559,22 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_FORCE_SHJ_MAX_BUILD_SIZE: OptionalConfigEntry[Long] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.forceShuffledHashJoin.maxBuildSize") + .category(CATEGORY_EXEC) + .doc(s"The build side size below which `${COMET_FORCE_SHJ.key}` converts a " + + "SortMergeJoin to ShuffledHashJoin. The size is Spark's planning estimate of the build " + + "child, or under AQE the materialized shuffle size of the build side. A build side at " + + "or over this size, or one with no statistics, keeps the SortMergeJoin. When unset, " + + "Spark's own rule applies: `spark.sql.autoBroadcastJoinThreshold` times the initial " + + "shuffle partition count (`spark.sql.adaptive.coalescePartitions.initialPartitionNum` " + + "when AQE and partition coalescing are both on and it is set, else " + + "`spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive " + + "threshold, Spark's default threshold of 10 MB is used instead. A " + + s"non-positive value removes the limit. $TUNING_GUIDE.") + .bytesConf(ByteUnit.BYTE) + .createOptional + val COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregate.skipPartial.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index fdf37b4e2bf..a47f1c2a6b4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -994,10 +994,13 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } else { val normalizedPlan = normalizeDivisors(plan) + // The rewrite removes the sorts of each join it converts, so put back any sort a kept + // operator above still needs. This runs before transform(), which converts the sorts it + // adds like any other SortExec. val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - normalizedPlan.transformUp { case p => - RewriteJoin.rewrite(p) - } + RewriteJoin.restoreRequiredOrdering(normalizedPlan.transformUp { case p => + RewriteJoin.rewrite(p, conf) + }) } else { normalizedPlan } diff --git a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala index 4ffd34451b4..5f543cc84aa 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala @@ -19,13 +19,16 @@ package org.apache.comet.rules +import org.apache.spark.sql.catalyst.expressions.SortOrder import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.execution.{SortExec, SparkPlan} import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} /** * Adapted from equivalent rule in Apache Gluten. @@ -64,7 +67,69 @@ object RewriteJoin extends JoinSelectionHelper { case _ => plan } - def rewrite(plan: SparkPlan): SparkPlan = plan match { + /** + * Adds a local sort under any operator whose child no longer has the ordering the operator + * requires, using the same check as EnsureRequirements. EnsureRequirements puts no sort above a + * sort-merge join whose output ordering already satisfies its parent, so when rewrite turns + * that join into a hash join and removes its sorts, a kept parent such as another sort-merge + * join on the same key would read unsorted input. Plans that already satisfy their required + * orderings are returned unchanged. + */ + def restoreRequiredOrdering(plan: SparkPlan): SparkPlan = plan.transformUp { case p => + p.withNewChildren(p.children.zip(p.requiredChildOrdering).map { + case (child, required) if !SortOrder.orderingSatisfies(child.outputOrdering, required) => + SortExec(required, global = false, child = child) + case (child, _) => child + }) + } + + /** + * The largest build side the rewrite converts, or None when the limit is switched off. Unset + * means Spark's own rule for choosing a shuffled hash join by size: the broadcast threshold + * times the initial shuffle partition count. A non-positive threshold only disables broadcast + * joins and says nothing about the hash table an executor can hold, so Spark's default + * threshold stands in for it. + */ + private def maxBuildSize(conf: SQLConf): Option[BigInt] = + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(conf) match { + case Some(limit) if limit <= 0 => None + case Some(limit) => Some(BigInt(limit)) + case None => + val threshold = conf.autoBroadcastJoinThreshold + val perPartition = + if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + Some(BigInt(perPartition) * conf.numShufflePartitions) + } + + /** + * Why the build side is too large to rewrite, or None when it fits. Like Spark's + * canBuildLocalHashMapBySize, the size must be strictly under the limit. The hash join cannot + * spill, so a build side without statistics is refused as well. Spark reports an unknown size + * as Long.MaxValue, which fails the comparison on its own. + */ + private def buildSideOverLimit( + smj: SortMergeJoinExec, + buildSide: BuildSide, + conf: SQLConf): Option[String] = + maxBuildSize(conf).flatMap { limit => + smj.logicalLink match { + case Some(join: Join) => + val buildSize = buildSide match { + case BuildLeft => join.left.stats.sizeInBytes + case BuildRight => join.right.stats.sizeInBytes + } + if (buildSize >= limit) { + Some( + s"build side size estimate of $buildSize bytes is not under the limit of " + + s"$limit bytes") + } else { + None + } + case _ => Some("no statistics are available for the build side") + } + } + + def rewrite(plan: SparkPlan, conf: SQLConf): SparkPlan = plan match { case smj: SortMergeJoinExec => getSmjBuildSide(smj) match { case Some(BuildRight) @@ -77,15 +142,23 @@ object RewriteJoin extends JoinSelectionHelper { s"BuildRight with ${smj.joinType} is not supported") plan case Some(buildSide) => - ShuffledHashJoinExec( - smj.leftKeys, - smj.rightKeys, - smj.joinType, - buildSide, - smj.condition, - removeSort(smj.left), - removeSort(smj.right), - smj.isSkewJoin) + buildSideOverLimit(smj, buildSide, conf) match { + case Some(reason) => + // The sort-merge join still runs natively, so this is information for extended + // explain rather than a fallback reason. + withInfo(smj, s"Cannot rewrite SortMergeJoin to HashJoin: $reason") + plan + case None => + ShuffledHashJoinExec( + smj.leftKeys, + smj.rightKeys, + smj.joinType, + buildSide, + smj.condition, + removeSort(smj.left), + removeSort(smj.right), + smj.isSkewJoin) + } case _ => plan } case _ => plan diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index eec16205a11..369c2517636 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,16 +31,20 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} +import org.apache.comet.rules.RewriteJoin class CometJoinSuite extends CometTestBase { @@ -890,6 +894,331 @@ class CometJoinSuite extends CometTestBase { } } + // The MERGE hint makes Spark plan a SortMergeJoin, which forceShuffledHashJoin then rewrites. + private val forcedSmjQuery = + "SELECT /*+ MERGE(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + + // tbl_b is the smaller side, so the rewrite builds on the right. + private def withForcedSmjTables(f: => Unit): Unit = { + withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") { + withParquetTable((0 until 10).map(i => (i % 10, i + 2)), "tbl_b") { + f + } + } + } + + private def rewriteReason(reason: String): String = + s"Cannot rewrite SortMergeJoin to HashJoin: $reason" + + // Checks that the SortMergeJoin and its sorts survived and that the kept join, which still + // runs natively, is not reported as a fallback. Returns Spark's size estimate of the build + // side, read through the same logical link the rewrite used. + private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = { + val joins = collect(plan) { case j: CometSortMergeJoinExec => j } + assert(joins.size == 1, plan) + assert(collect(plan) { case s: CometSortExec => s }.size == 2, plan) + assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, plan) + val kept = joins.head + assert(!hasFallbackReason(kept) && !hasFallbackReason(kept.originalPlan), plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(!reasons.exists(_.contains("SortMergeJoin")), reasons) + kept.originalPlan.logicalLink match { + case Some(join: Join) => join.right.stats.sizeInBytes + case other => fail(s"Expected a Join logical link, got $other") + } + } + + // The refusal is recorded as information on the kept join and rendered as COMET-INFO. + private def assertKeptWithInfo(plan: SparkPlan, expected: String): Unit = { + val join = collect(plan) { case j: CometSortMergeJoinExec => j }.head + assert(join.getTagValue(CometExplainInfo.EXTENSION_INFO).exists(_.contains(expected)), plan) + val verbose = new ExtendedExplainInfo().generateVerboseInfo(plan) + assert(verbose.contains("[COMET-INFO:") && verbose.contains(expected), verbose) + } + + private def assertRewrittenToHashJoin(plan: SparkPlan): Unit = { + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.map(_.buildSide) == Seq(BuildRight), plan) + assert(collect(plan) { case s: CometSortExec => s }.isEmpty, plan) + assert(collect(plan) { case j: CometSortMergeJoinExec => j }.isEmpty, plan) + } + + private def assertKeptWithReason(plan: SparkPlan, limit: Long): Unit = { + val estimate = assertSortMergeJoinKept(plan) + assert(estimate >= limit, estimate) + assertKeptWithInfo( + plan, + rewriteReason( + s"build side size estimate of $estimate bytes is not under the limit of $limit bytes")) + } + + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin keeps SortMergeJoin over maxBuildSize, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, 1) + } + } + } + + test(s"forceShuffledHashJoin rewrites SortMergeJoin under maxBuildSize, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1g") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + } + + test("forceShuffledHashJoin matches Spark's strict size rule at the limit") { + // AQE off, so the estimate read from the first plan is the one the later runs compare. + withForcedSmjTables { + var estimate = BigInt(0) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + estimate = assertSortMergeJoinKept(cometPlan) + } + // A build side exactly at the limit is not under it, as in canBuildLocalHashMapBySize. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> estimate.toString) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, estimate.toLong) + } + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> (estimate + 1).toString) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin with a non-positive maxBuildSize rewrites regardless of size") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin without maxBuildSize uses Spark's shuffled hash join size rule") { + // The limit is autoBroadcastJoinThreshold times the shuffle partition count. + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, 2) + } + } + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1g") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin uses Spark's default threshold when broadcasts are disabled") { + // Disabling broadcasts says nothing about the hash table an executor can hold, so the limit + // falls back to Spark's default threshold times the partition count. AQE is off so the + // decision uses the planning estimate, which fileCompressionFactor scales past 10 MB without + // writing a 10 MB table. + val defaultThreshold = SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + assert(defaultThreshold == 10L * 1024 * 1024) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + withSQLConf(SQLConf.FILE_COMPRESSION_FACTOR.key -> "1000000") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, defaultThreshold) + } + } + } + } + + test("forceShuffledHashJoin keeps a SortMergeJoin that has no statistics") { + val left = spark.range(1).toDF("a").queryExecution.sparkPlan + val right = spark.range(1).toDF("b").queryExecution.sparkPlan + def smj: SortMergeJoinExec = + SortMergeJoinExec(left.output, right.output, Inner, None, left, right) + assert(smj.logicalLink.isEmpty) + + val kept = smj + assert(RewriteJoin.rewrite(kept, spark.sessionState.conf) eq kept) + assert(!hasFallbackReason(kept)) + val info = kept.getTagValue(CometExplainInfo.EXTENSION_INFO) + assert( + info.exists(_.contains(rewriteReason("no statistics are available for the build side"))), + info) + + withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + val rewritten = RewriteJoin.rewrite(smj, spark.sessionState.conf) + assert(rewritten.isInstanceOf[ShuffledHashJoinExec], rewritten) + } + } + + // Both joins are on the same key, so EnsureRequirements puts no sort between them: the inner + // sort-merge join's output ordering satisfies the outer one. When the inner join is rewritten + // to a hash join and the outer one is kept, the outer join needs that ordering restored. + private def withChainedJoinTables(midRows: Int, midKeys: Int)(f: => Unit): Unit = { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withParquetTable((0 until midRows).map(i => (i % midKeys, i)), "mid") { + f + } + } + } + } + + // Whether a hash join is reachable without leaving the stage. + private def reachesHashJoinInStage(plan: SparkPlan): Boolean = plan match { + case _: CometHashJoinExec => true + case _: ShuffleExchangeLike | _: QueryStageExec => false + case p => p.children.exists(reachesHashJoinInStage) + } + + // Whether a sort-merge join input is a sort over a hash join in the same stage. A Spark + // sort-merge join reads the sort through codegen wrappers and a columnar to row transition. + private def isSortOverHashJoin(input: SparkPlan): Boolean = input match { + case s: CometSortExec => reachesHashJoinInStage(s.child) + case s: SortExec => reachesHashJoinInStage(s.child) + case t: ColumnarToRowTransition => isSortOverHashJoin(t.child) + case w: WholeStageCodegenExec => isSortOverHashJoin(w.child) + case a: InputAdapter => isSortOverHashJoin(a.child) + case _ => false + } + + // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join + // through a sort. + private def checkSortRestoredOverHashJoin(query: String): Unit = { + val (_, cometPlan) = checkSparkAnswer(query) + val inputs = collect(cometPlan) { + case j: CometSortMergeJoinExec => j.children + case j: SortMergeJoinExec => j.children + }.flatten + assert(inputs.exists(isSortOverHashJoin), cometPlan) + } + + private val chainedJoinQuery = "SELECT big._1, big._2, small._2, mid._2 FROM big " + + "JOIN small ON big._1 = small._1 JOIN mid ON big._1 = mid._1" + + test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=false") { + // The default limit, 3000 bytes times two partitions, falls between the two build sides. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") { + withChainedJoinTables(midRows = 3000, midKeys = 100) { + checkSortRestoredOverHashJoin(chainedJoinQuery) + } + } + } + + test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=true") { + // Under AQE the decision uses the materialized shuffle sizes, so mid is made large enough + // to stay over the limit while small stays under it. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1000000", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") { + withChainedJoinTables(midRows = 300000, midKeys = 30000) { + checkSortRestoredOverHashJoin(chainedJoinQuery) + } + } + } + + // The rewrite never converts a LeftSemi or ExistenceJoin that builds on the right, whatever + // the build size, so these joins are kept above the rewritten inner join with no size limit. + // Each predicate also refers to small, so neither the filter nor the semi join is pushed below + // the inner join. + for (adaptive <- Seq(false, true)) { + def withKeptSemiJoinConf(f: => Unit): Unit = withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100)(f) + } + + test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") { + withKeptSemiJoinConf { + checkSortRestoredOverHashJoin( + "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + + "WHERE small._2 < 0 OR EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1)") + } + } + + test(s"forceShuffledHashJoin restores ordering for a kept LeftSemi join, AQE=$adaptive") { + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.PushLeftSemiLeftAntiThroughJoin") { + withKeptSemiJoinConf { + checkSortRestoredOverHashJoin( + "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + + "WHERE EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1 AND mid._2 > small._2)") + } + } + } + } + + // A window partitioned by the join key gets no sort of its own, since the sort-merge join's + // output ordering satisfies it, so the window needs that ordering restored too. + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100) { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + + "FROM big JOIN small ON big._1 = small._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + } + } + } + } + test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") { withSQLConf( SQLConf.PREFER_SORTMERGEJOIN.key -> "false", From e8c7ee739d3530b309d38388ba1c41decdb8abea Mon Sep 17 00:00:00 2001 From: Peter Lee Date: Thu, 8 Oct 2026 16:45:26 +0000 Subject: [PATCH 35/89] =?UTF-8?q?fix:=20honor=20Spark=E2=80=99s=20commit?= =?UTF-8?q?=20protocol=20in=20Spark=203.x=20native=20writes=20(#4746)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: honor Spark’s commit protocol in Spark 3.x native writes * ci: register legacy native write commit protocol tests * refactor: drop the native Parquet writer's work_dir file naming Both native writers now take the exact file name from Spark's commit protocol, so nothing sets work_dir, job_id or task_attempt_id any more. Remove the Some(work_dir) arm from ParquetWriterExec::execute, reserve the three proto fields, and correct the comments that still described the Spark 3.x writer naming its own files. * fix: handle empty inputs in the Spark 3.x native writer Mirror CometWriteFilesExec and FileFormatWriter: swap a zero-partition child for a single empty partition so the output still carries a schema (SPARK-23271), and only write a file from partition 0 or a partition that has rows. Move createTaskContext to the companion object so the task closure no longer captures the exec node and its child subtree. The SPARK-23271 and empty-partition tests now run on Spark 3.x too. * fix: match Spark 3.x write admission and cover empty scans * docs: say native writes over an empty relation stay native With the empty-relation guard gone from CometDataWritingCommand, the user guide and the CometWriteFilesExec comment still described Parquet writes over an empty relation falling back to Spark's writer. Both native write paths now swap a zero-partition input for one empty partition, so such a write stays native and still leaves a readable schema-only file. * fix: take the Spark 3.x native write codec from the job configuration The Spark 3.x task named its file after the codec in the job configuration, because getFileExtension reads it from there. The native writer used the codec that convert resolved at planning time. The two agree today only because both copy the precedence rule of ParquetOptions. Read the codec with CodecConfig.from(taskContext) next to setOutputPath, as CometWriteFilesExec.executeTask does, so the file name and the footer agree by construction. A new CometNativeWriteSuite test sets a job codec that differs from the planned one and checks the name and the footer. The mixed-case codec test now checks the file name on Spark 3.x too. * test: make the injected cleanup failure throw from SQLMetric.add too Since #6416, CometMetricNode reports a sum metric with add instead of set. The test metric threw only from set, so the two cleanup-failure tests in CometNativeWriteSuite stopped injecting their failure and failed on Spark 3.4 and 3.5. * ci: pre-compile Spark SQL test classes in the manual writer workflow The Parquet writer test step caps the sbt heap at 1 GB for the forked test JVM, but this workflow compiles Spark first, and sbt ran out of heap while it compiled sql/core. Compile the SQL test classes in an earlier step with a 3 GB heap. This branch had the same step in July, before a history rewrite dropped it. * test: build CometNativeWriteExec with the merged constructor The merge with main brings in two callers written against the old constructor. #5957's RevertNativeForTransitionHeavyStagesSuite relied on the removed committer and jobTrackerID defaults, and #6247's write_reserving_from helper still passed the removed work_dir, job_id and task_attempt_id arguments to ParquetWriterExec::try_new. CleanupFailingNativeWriteExec also needs #5957's originalPlan argument. --------- Co-authored-by: Andy Grove --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .github/workflows/spark_sql_writer_tests.yml | 7 + .../latest/compatibility/operators.md | 6 +- docs/source/user-guide/latest/operators.md | 2 +- .../src/execution/operators/parquet_writer.rs | 72 +--- native/core/src/execution/planner/write.rs | 3 - native/proto/src/proto/operator.proto | 30 +- .../operator/CometDataWritingCommand.scala | 94 ++--- .../serde/operator/NativeWriteUtils.scala | 13 +- .../sql/comet/CometNativeWriteExec.scala | 396 +++++++----------- .../spark/sql/comet/CometWriteFilesExec.scala | 8 +- .../parquet/CometParquetWriterSuite.scala | 81 ++-- ...tNativeForTransitionHeavyStagesSuite.scala | 5 +- .../sql/comet/CometNativeWriteSuite.scala | 390 +++++++++++++++++ ...CometEmptyRelationParquetWriterSuite.scala | 8 +- 16 files changed, 681 insertions(+), 436 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 69f0fa89612..debfccf311e 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -460,6 +460,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 0b525419943..742fd100d09 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -164,6 +164,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase diff --git a/.github/workflows/spark_sql_writer_tests.yml b/.github/workflows/spark_sql_writer_tests.yml index b26ff88e864..6b5fc7f6d65 100644 --- a/.github/workflows/spark_sql_writer_tests.yml +++ b/.github/workflows/spark_sql_writer_tests.yml @@ -115,6 +115,13 @@ jobs: spark-short-version: ${{ inputs.spark-version }} skip-native-build: true + - name: Pre-compile Spark SQL test classes + run: | + cd apache-spark + rm -rf /root/.m2/repository/org/apache/parquet + NOLINT_ON_COMPILE=true build/sbt -Dsbt.log.noformat=true -mem 3072 \ + 'sql/Test/compile' + - name: Run Parquet writer tests run: | cd apache-spark diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 4c65b2ee6c5..40732ced79f 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -29,8 +29,10 @@ executed. Supported parent joins and aggregates remain eligible for native execution. Global aggregates still return one row (`COUNT = 0`, `SUM = NULL`), and grouped aggregates return no rows. Independent operator restrictions and aggregate buffer compatibility checks still apply. -Parquet writes whose input plans contain an empty relation use Spark's writer to preserve -readable empty output files and their schema metadata. + +A native Parquet write over a native empty relation stays native. Like Spark's writer, it runs one +task for the empty input, so the output still gets a schema-only Parquet file that readers can +infer the schema from. ## In-Memory Cache diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 0155ebb8dae..acefd4f0890 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -64,7 +64,7 @@ omitted from the tables below and may be reconsidered based on demand: | `FileSourceScanExec` | ✅ | Parquet only. Some types and configurations fall back. See [Parquet Scan Compatibility](compatibility/scans.md). | | `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet Scan Compatibility](compatibility/scans.md) and the [Iceberg Guide](iceberg.md). | | `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). | -| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support and writer fallback. | +| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support, including native Parquet writes. | | `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. | | `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index c96cb17020e..cc2da616e54 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -234,18 +234,9 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Where this task writes. When `work_dir` is set (the Spark 3.x `CometNativeWriteExec` - /// path) this is the write's output directory and is unused; the file name is derived from - /// `work_dir`. Otherwise (Spark 4.0+, `CometWriteFilesExec`) it is the exact path of the file - /// to write, chosen by the JVM commit protocol and used verbatim - this operator then never - /// derives file names of its own. + /// The exact path of the file this task writes, chosen by Spark's commit protocol on the JVM + /// side and used verbatim - this operator never derives file names of its own. output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only. - work_dir: Option, - /// Job ID for tracking this write operation - job_id: Option, - /// Task attempt ID for this specific task - task_attempt_id: Option, /// Compression codec compression: ParquetCompression, /// Partition ID (from Spark TaskContext) @@ -268,9 +259,6 @@ impl ParquetWriterExec { pub fn try_new( input: Arc, output_path: String, - work_dir: Option, - job_id: Option, - task_attempt_id: Option, compression: ParquetCompression, partition_id: i32, column_names: Vec, @@ -290,9 +278,6 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, output_path, - work_dir, - job_id, - task_attempt_id, compression, partition_id, column_names, @@ -478,9 +463,6 @@ impl ExecutionPlan for ParquetWriterExec { 1 => Ok(Arc::new(ParquetWriterExec::try_new( Arc::clone(&children[0]), self.output_path.clone(), - self.work_dir.clone(), - self.job_id.clone(), - self.task_attempt_id, self.compression.clone(), self.partition_id, self.column_names.clone(), @@ -510,8 +492,6 @@ impl ExecutionPlan for ParquetWriterExec { .register(context.memory_pool()); let input = self.input.execute(partition, context)?; let input_schema = self.input.schema(); - let work_dir = self.work_dir.clone(); - let task_attempt_id = self.task_attempt_id; let compression = self.compression.to_parquet()?; let column_names = self.column_names.clone(); @@ -530,19 +510,8 @@ impl ExecutionPlan for ParquetWriterExec { Arc::new(Schema::new(fields)) }); - let part_file = match &work_dir { - // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. - None => self.output_path.clone(), - // Spark 3.x hands over a working directory instead and expects the writer to name the - // file; that branch goes away with Spark 3.x support. - Some(work_dir) => match task_attempt_id { - Some(attempt_id) => format!( - "{}/part-{:05}-{:05}.parquet", - work_dir, self.partition_id, attempt_id - ), - None => format!("{}/part-{:05}.parquet", work_dir, self.partition_id), - }, - }; + // The JVM commit protocol has already chosen the exact file to write. + let part_file = self.output_path.clone(); // Configure writer properties let props = WriterProperties::builder() @@ -675,11 +644,10 @@ mod tests { ); } - /// Spark 4.0+ hands over the exact file to write rather than a working directory. The writer - /// must use that path verbatim - Spark's commit protocol owns naming and staging, and + /// The JVM hands over the exact file to write. The writer must use that path verbatim - Spark's commit protocol owns naming and staging, and /// committers that track individual files depend on the name it chose. #[tokio::test] - async fn test_parquet_writer_uses_output_path_verbatim_without_work_dir() -> Result<()> { + async fn test_parquet_writer_uses_output_path_verbatim() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); let batch = RecordBatch::try_new( Arc::clone(&schema), @@ -696,9 +664,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, output_path, - None, // work_dir: Spark 4.0+ path - None, - None, ParquetCompression::None, // A non-zero partition id must not leak into the file name. 3, @@ -763,13 +728,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["required_id".to_string(), "values".to_string()], @@ -832,13 +794,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["values".to_string()], @@ -903,9 +862,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, format!("file://{}", file.display()), - None, - None, - None, ParquetCompression::None, 0, column_names, @@ -1160,16 +1116,14 @@ mod tests { let memory_exec = Arc::new(DataSourceExec::new(Arc::new(memory_source_config))); // Create ParquetWriterExec with DataSourceExec as input - let output_path = "unused".to_string(); - let work_dir = "hdfs://namenode:9000/user/test_parquet_writer_exec".to_string(); + let output_path = + "hdfs://namenode:9000/user/test_parquet_writer_exec/part-00000-00123.parquet" + .to_string(); let column_names = vec!["id".to_string(), "name".to_string()]; let parquet_writer = ParquetWriterExec::try_new( memory_exec, output_path, - Some(work_dir), - None, // job_id - Some(123), // task_attempt_id ParquetCompression::None, 0, // partition_id column_names, diff --git a/native/core/src/execution/planner/write.rs b/native/core/src/execution/planner/write.rs index 1a3f9674cd1..5665939afb2 100644 --- a/native/core/src/execution/planner/write.rs +++ b/native/core/src/execution/planner/write.rs @@ -107,9 +107,6 @@ impl OperatorBuilder for ParquetWriterBuilder { let parquet_writer = Arc::new(ParquetWriterExec::try_new( Arc::clone(&child.native_plan), writer.output_path.clone(), - writer.work_dir.clone(), - writer.job_id.clone(), - writer.task_attempt_id, codec, planner.partition(), writer.column_names.clone(), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 527461f47a5..97c65f0da36 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -785,7 +785,7 @@ message IcebergWriteCommon { } // Single Iceberg write operator. Per-task fields are populated by the JVM exec wrapper at task -// launch (matching how `ParquetWriter.task_attempt_id` is filled in `CometNativeWriteExec`). +// launch (matching how `ParquetWriter.output_path` is filled in by the Parquet write execs). message IcebergWrite { IcebergWriteCommon common = 1; @@ -920,27 +920,19 @@ message ShuffleWriter { } message ParquetWriter { - // Where this task writes. Two shapes, selected by whether `work_dir` is set: - // - // - Spark 4.0+ (CometWriteFilesExec): `work_dir` is unset and `output_path` is the - // fully-qualified path of the Parquet file to write, set per task from - // FileCommitProtocol.newTaskTempFile. Naming and staging are owned by Spark's commit - // protocol so that task-attempt isolation, speculative execution, and committers that track - // individual files (S3A magic, streaming manifest) all behave as they do for Spark's own - // writer. The native writer uses this path verbatim. - // - Spark 3.x (CometNativeWriteExec): `work_dir` is set and the native writer derives the file - // name from it, the partition id and the task attempt id. `output_path` is the write's - // output directory and is unused natively. Goes away with Spark 3.x support. + // The fully-qualified path of the Parquet file this task writes, set per task from + // FileCommitProtocol.newTaskTempFile by CometWriteFilesExec (Spark 4.0+) or + // CometNativeWriteExec (Spark 3.x). Naming and staging are owned by Spark's commit protocol so + // that task-attempt isolation, speculative execution, and committers that track individual + // files (S3A magic, streaming manifest) all behave as they do for Spark's own writer. The native + // writer uses this path verbatim. string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; - // Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only; see - // output_path above. - optional string work_dir = 5; - // Job ID for tracking this write operation - optional string job_id = 6; - // Task attempt ID for this specific task - optional int32 task_attempt_id = 7; + // Formerly work_dir, job_id and task_attempt_id, which let the native writer name its own file + // under a working directory. + reserved 5, 6, 7; + reserved "work_dir", "job_id", "task_attempt_id"; // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object // stores. diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index a56dea862c4..545cbc774ae 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -19,17 +19,21 @@ package org.apache.comet.serde.operator +import java.util.UUID + import scala.jdk.CollectionConverters._ -import org.apache.spark.SparkException -import org.apache.spark.sql.comet.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.hadoop.mapreduce.Job +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.spark.internal.io.FileCommitProtocol +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec} import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.util.SerializableConfiguration import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -55,25 +59,18 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec case cmd: InsertIntoHadoopFsRelationCommand => cmd.fileFormat match { case _: ParquetFileFormat => - // AQE can replace the write input with a zero-partition empty relation. Keep - // Spark's writer, which creates an empty task to preserve the output file schema. - // The native writer only maps existing partitions; see #5303. This guard is - // conservative: an empty relation below an exchange can have nonzero partitions - // at the write input. Revisit the guard when native empty-file handling is fixed. - if (hasEmptyRelationInput(op.child)) { - return Unsupported(Some( - "Parquet writes with empty-relation inputs require Spark's empty-file handling")) - } - if (!cmd.outputPath.toString.startsWith("file:") && !cmd.outputPath.toString .startsWith("hdfs:")) { return Unsupported(Some("Supported output filesystems: local, HDFS")) } + // Spark 3.x HadoopMapReduceCommitProtocol.getFilename hardcodes "part"; + // mapreduce.output.basename is only honored on Spark 4.0+. Custom committer paths + // are still checked by checkNativeWriteDestination before native execution. NativeWriteUtils - // This writer names its own files `part--.parquet`, so the - // prefix is fixed rather than read from `mapreduce.output.basename`. - .escapedHdfsDestination(cmd.outputPath.toString, "part") + .escapedHdfsDestination( + cmd.outputPath.toString, + NativeWriteUtils.DEFAULT_BASE_OUTPUT_NAME) .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -98,14 +95,6 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec } } - private def hasEmptyRelationInput(plan: SparkPlan): Boolean = plan match { - case _: CometEmptyRelationExec => true - case wrapper: CometScanWrapper => hasEmptyRelationInput(wrapper.originalPlan) - case stage: QueryStageExec => hasEmptyRelationInput(stage.plan) - case reused: ReusedExchangeExec => hasEmptyRelationInput(reused.child) - case _ => plan.children.exists(hasEmptyRelationInput) - } - override def convert( op: DataWritingCommandExec, builder: Operator.Builder, @@ -123,6 +112,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec val outputPath = cmd.outputPath.toString + // Planning-time value only. CometNativeWriteExec replaces it per task with the codec that + // names the file. val plannedCodec = NativeWriteUtils.parseCompressionCodec(cmd.options) val codec = NativeWriteUtils.protoCompressionCodec(plannedCodec) match { case Some(codec) => codec @@ -140,8 +131,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec cmd.query.schema.fields.toIndexedSeq, Some( op.session.sessionState.conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))).asJava) - // Note: work_dir, job_id, and task_attempt_id will be set at execution time - // in CometNativeWriteExec, as they depend on the Spark task context + // CometNativeWriteExec replaces output_path with the committer's exact task filename + // at execution time. // Collect S3/cloud storage configurations val session = op.session @@ -189,29 +180,30 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec other } - // Create FileCommitProtocol for atomic writes - val jobId = java.util.UUID.randomUUID().toString - val committer = - try { - // Use Spark's SQLHadoopMapReduceCommitProtocol - val committerClass = - classOf[org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol] - val constructor = - committerClass.getConstructor(classOf[String], classOf[String], classOf[Boolean]) - Some( - constructor - .newInstance( - jobId, - outputPath, - java.lang.Boolean.FALSE // dynamicPartitionOverwrite = false for now - ) - .asInstanceOf[org.apache.spark.internal.io.FileCommitProtocol]) - } catch { - case e: Exception => - throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}") - } - - CometNativeWriteExec(nativeOp, op, childPlan, outputPath, cmd.mode, committer, jobId) + val session = op.session + val job = Job.getInstance(session.sessionState.newHadoopConfWithOptions(cmd.options)) + job.setOutputKeyClass(classOf[Void]) + job.setOutputValueClass(classOf[InternalRow]) + FileOutputFormat.setOutputPath(job, cmd.outputPath) + val outputWriterFactory = + cmd.fileFormat.prepareWrite(session, job, CaseInsensitiveMap(cmd.options), cmd.query.schema) + + val committer = FileCommitProtocol.instantiate( + session.sessionState.conf.fileCommitProtocolClass, + UUID.randomUUID().toString, + outputPath, + dynamicPartitionOverwrite = false) + job.getConfiguration.set("spark.sql.sources.writeJobUUID", UUID.randomUUID().toString) + + CometNativeWriteExec( + nativeOp, + op, + childPlan, + outputPath, + cmd.mode, + committer, + new SerializableConfiguration(job.getConfiguration), + outputWriterFactory) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 5246099fdcb..a0ebdea520f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -141,12 +141,13 @@ object NativeWriteUtils { * byte for byte (see [[hdfsPathDivergence]] for why it may not): * * - the destination directory, and - * - `fileNamePrefix`, the basename every file name is built from. On Spark 4.0+ that is - * `mapreduce.output.basename`, which `HadoopMapReduceCommitProtocol.getFilename` - * interpolates into `--`; on 3.x Comet names the files itself and - * the basename is always the literal `part`. A basename holding `?` or `#` is the dangerous - * one: the native URL parser truncates there, so *every* task writes a file with the same - * truncated name and they overwrite each other during commit. + * - `fileNamePrefix`, the basename every file name is built from. Spark 4.0+ + * `HadoopMapReduceCommitProtocol.getFilename` reads `mapreduce.output.basename` and + * interpolates it into `--`; Spark 3.x hardcodes `part` instead. + * Both native writers take the exact file name from the commit protocol. A basename holding + * `?` or `#` is dangerous on 4.0+: the native URL parser truncates there, so every task + * writes a file with the same truncated name and they overwrite each other during commit. + * Custom committer paths are checked by [[checkNativeWriteDestination]]. * * The basename is checked by running [[hdfsPathDivergence]] over the path it produces rather * than over the name alone. Everything else `getFilename` interpolates -- the split number, the diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index bb2c560ba64..200c7831321 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -19,13 +19,16 @@ package org.apache.spark.sql.comet -import scala.jdk.CollectionConverters._ +import java.util.Date +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext, TaskAttemptID, TaskID, TaskType} import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.parquet.hadoop.codec.CodecConfig import org.apache.spark.TaskContext -import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec} +import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage import org.apache.spark.rdd.RDD import org.apache.spark.sql.SaveMode import org.apache.spark.sql.catalyst.InternalRow @@ -35,42 +38,23 @@ import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.command.DataWritingCommandExec -import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.{OutputWriterFactory, WriteFilesExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils +import org.apache.spark.util.{SerializableConfiguration, Utils} -import com.google.protobuf.CodedOutputStream - -import org.apache.comet.CometExecIterator import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.NativeWriteUtils /** - * Comet physical operator for native Parquet write operations with FileCommitProtocol support. - * - * This operator writes data to Parquet files using the native Comet engine. It integrates with - * Spark's FileCommitProtocol to provide atomic writes with proper staging and commit semantics. + * Native Parquet writes on Spark 3.x, where Comet replaces the entire write command. * - * The implementation includes support for Spark's file commit protocol through work_dir, job_id, - * and task_attempt_id parameters that can be set in the operator. When work_dir is set, files are - * written to a temporary location that can be atomically committed later. + * Both execution entry points run Spark's job and task commit lifecycle. Native code writes the + * exact filename supplied by the configured commit protocol, using the prepared Hadoop job. Spark + * 4.0+ uses CometWriteFilesExec and leaves the surrounding command with Spark instead. * - * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) * @param originalPlan * The JVM data-writing command restored when Comet reverts a transition-heavy stage - * @param child - * The child operator providing the data to write - * @param outputPath - * The path where the Parquet file will be written - * @param mode - * The Spark SaveMode governing target-exists behavior (Append / Overwrite / ErrorIfExists / - * Ignore). Comet takes over Spark's DataWritingCommandExec so we must apply these semantics - * here - a direct port of Spark's InsertIntoHadoopFsRelationCommand.run() doInsertion logic. - * @param committer - * FileCommitProtocol for atomic writes. If None, files are written directly. - * @param jobTrackerID - * Unique identifier for this write job */ case class CometNativeWriteExec( nativeOp: Operator, @@ -78,8 +62,10 @@ case class CometNativeWriteExec( child: SparkPlan, outputPath: String, mode: SaveMode, - committer: Option[FileCommitProtocol] = None, - jobTrackerID: String = Utils.createTempDir().getName) + committer: FileCommitProtocol, + serializableHadoopConf: SerializableConfiguration, + outputWriterFactory: OutputWriterFactory, + jobTrackerID: String = SparkHadoopWriterUtils.createJobTrackerID(new Date())) extends CometNativeExec with UnaryExecNode { @@ -97,19 +83,8 @@ case class CometNativeWriteExec( s"${getClass.getSimpleName} expected one reverted input but received ${newChildren.size}") } - // Accumulator to collect TaskCommitMessages from all tasks - // Must be eagerly initialized on driver, not lazy - @transient private val taskCommitMessagesAccum = - sparkContext.collectionAccumulator[FileCommitProtocol.TaskCommitMessage]("taskCommitMessages") - - override def serializedPlanOpt: SerializedPlan = { - val size = nativeOp.getSerializedSize - val bytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(bytes) - nativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - SerializedPlan(Some(bytes)) - } + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) @@ -122,219 +97,132 @@ case class CometNativeWriteExec( "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) override def doExecute(): RDD[InternalRow] = { - // Setup job if committer is present - committer.foreach { c => - val jobContext = createJobContext() - c.setupJob(jobContext) - } - - // Execute the native write with commit protocol - val resultRDD = doExecuteColumnar() - - // Force execution by consuming all batches - resultRDD - .mapPartitions { iter => - iter.foreach(_.close()) - Iterator.empty - } - .count() - - // Extract write statistics from metrics - val filesWritten = metrics("files_written").value - val bytesWritten = metrics("bytes_written").value - val rowsWritten = metrics("rows_written").value - - // Collect TaskCommitMessages from accumulator - val commitMessages = taskCommitMessagesAccum.value.asScala.toSeq - - // Commit job with collected TaskCommitMessages - committer.foreach { c => - val jobContext = createJobContext() - try { - c.commitJob(jobContext, commitMessages) - logInfo( - s"Successfully committed write job to $outputPath: " + - s"$filesWritten files, $bytesWritten bytes, $rowsWritten rows") - } catch { - case e: Exception => - logError("Failed to commit job, aborting", e) - c.abortJob(jobContext) - throw e - } - } - - // Return empty RDD as write operations don't return data + executeWriteAndCommit() sparkContext.emptyRDD[InternalRow] } override def doExecuteColumnar(): RDD[ColumnarBatch] = { - // Comet replaces DataWritingCommandExec entirely, so Spark's - // InsertIntoHadoopFsRelationCommand.run() never runs. That method is where Spark handles - // SaveMode semantics (path-exists check, delete-before-Overwrite, Ignore short-circuit) - - // port the non-partitioned, non-catalog branch of that logic here. See Spark 3.5's - // InsertIntoHadoopFsRelationCommand.run doInsertion match. This runs on the driver before - // any executor tasks fire, mirroring where Spark does the delete. + executeWriteAndCommit() + sparkContext.emptyRDD[ColumnarBatch] + } + + private def executeWriteAndCommit(): Unit = { if (!prepareOutputPathForMode()) { logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") - return sparkContext.emptyRDD[ColumnarBatch] + return } - // Get the input data from the child operator - val childRDD = if (child.supportsColumnar) { + val job = Job.getInstance(new Configuration(serializableHadoopConf.value)) + // Like FileFormatWriter, only abort after setupJob has succeeded. + committer.setupJob(job) + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Include configuration changes made by setupJob in the task contexts. + val commitMessages = runNativeWriteJob(new SerializableConfiguration(job.getConfiguration)) + committer.commitJob(job, commitMessages.toSeq) + logInfo( + s"Successfully committed native write job to $outputPath: " + + s"${metrics("files_written").value} files, " + + s"${metrics("bytes_written").value} bytes, ${metrics("rows_written").value} rows") + })(catchBlock = committer.abortJob(job)) + } + + private def runNativeWriteJob( + hadoopConf: SerializableConfiguration): Array[TaskCommitMessage] = { + val childRDD: RDD[ColumnarBatch] = if (child.supportsColumnar) { child.executeColumnar() } else { - // If child doesn't support columnar, convert to columnar child.execute().mapPartitionsInternal { _ => - // TODO this could delegate to CometRowToColumnar, but maybe Comet - // does not need to support this case? throw new UnsupportedOperationException( "Row-based child operators not yet supported for native write") } } - // Capture metadata before the transformation - val numPartitions = childRDD.getNumPartitions - val numOutputCols = child.output.length + // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, + // so the output directory would carry no schema for readers. Swap in a single empty + // partition, as Spark's FileFormatWriter and CometWriteFilesExec do. + val writeRDD = if (childRDD.getNumPartitions == 0) { + sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) + } else { + childRDD + } + + // The task closure below must not touch `this`, which holds `nativeOp` plus the whole + // converted child subtree; capture what it needs by value instead, as CometWriteFilesExec + // does. + val numPartitions = writeRDD.getNumPartitions + val childSchema = CometUtils.fromAttributes(child.output) + val capturedNativeOp = nativeOp val capturedCommitter = committer val capturedJobTrackerID = jobTrackerID - val capturedNativeOp = nativeOp - val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks - - // Execute native write operation with task-level commit protocol - childRDD.mapPartitionsInternal { iter => - val partitionId = org.apache.spark.TaskContext.getPartitionId() - val taskAttemptId = org.apache.spark.TaskContext.get().taskAttemptId() - - // Setup task-level commit protocol if provided - val (workDir, taskContext, _) = capturedCommitter - .map { committer => - val taskContext = - createTaskContext(capturedJobTrackerID, partitionId, taskAttemptId.toInt) - - // Setup task - this creates the temporary working directory - committer.setupTask(taskContext) - - // Get the work directory for temp files - // Spark 4.1 made the (taskContext, dir, ext: String) overload throw by default; - // the FileNameSpec overload is the supported one and exists in 3.4+. - val workPath = committer.newTaskTempFile(taskContext, None, FileNameSpec("", "")) - val workDir = new Path(workPath).getParent.toString - - (Some(workDir), Some((committer, taskContext)), null) - } - .getOrElse((None, None, null)) - - // Modify the native operator to include task-specific parameters - val modifiedNativeOp = if (workDir.isDefined) { - val parquetWriter = capturedNativeOp.getParquetWriter.toBuilder - .setWorkDir(workDir.get) - .setJobId(capturedJobTrackerID) - .setTaskAttemptId(taskAttemptId.toInt) - .build() - - capturedNativeOp.toBuilder.setParquetWriter(parquetWriter).build() - } else { - capturedNativeOp - } - - val nativeMetrics = CometMetricNode.fromCometPlan(this) - // Register before CometExecIterator so completion listeners run after iterator close - // (Spark runs task completion callbacks in reverse registration order). - Option(TaskContext.get()).foreach(nativeMetrics.reportNativeWriteOutputMetrics) - - val size = modifiedNativeOp.getSerializedSize - val planBytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(planBytes) - modifiedNativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - - val execIterator = new CometExecIterator( - CometExec.newIterId, - CometArrowStream.inputObjects( - iter, - CometUtils.fromAttributes(child.output), - "CometNativeWriteExec"), - numOutputCols, - planBytes, - nativeMetrics, - numPartitions, - partitionId, - None, - Seq.empty) - - // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage - new Iterator[ColumnarBatch] { - private var completed = false - private var thrownException: Option[Throwable] = None - - override def hasNext: Boolean = { - val result = - try { - execIterator.hasNext - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - - if (!result && !completed) { - handleTaskEnd() - } - - result - } - - override def next(): ColumnarBatch = { - try { - execIterator.next() - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - } - - private def handleTaskEnd(): Unit = { - if (!completed) { - completed = true - - // Handle commit or abort based on whether an exception was thrown - taskContext.foreach { case (committer, ctx) => - try { - if (thrownException.isEmpty) { - // Commit the task and add message to accumulator - val message = committer.commitTask(ctx) - capturedAccumulator.add(message) - logDebug(s"Task ${ctx.getTaskAttemptID} committed successfully") - } else { - // Abort the task - committer.abortTask(ctx) - val exMsg = thrownException.get.getMessage - logWarning(s"Task ${ctx.getTaskAttemptID} aborted due to exception: $exMsg") - } - } catch { - case e: Exception => - // Log the commit/abort exception but don't mask the original exception - logError(s"Error during task commit/abort: ${e.getMessage}", e) - if (thrownException.isEmpty) { - // If no original exception, propagate the commit/abort exception - throw e - } + val writerFactory = outputWriterFactory + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val commitMessages = new Array[TaskCommitMessage](numPartitions) + + sparkContext.runJob( + writeRDD, + (context: TaskContext, batches: Iterator[ColumnarBatch]) => { + val taskContext = + CometNativeWriteExec.createTaskContext(hadoopConf.value, capturedJobTrackerID, context) + capturedCommitter.setupTask(taskContext) + // Guard filename allocation, native iterator construction, execution, cleanup and commit. + // Spark's helper preserves the original error if abortTask also fails. + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Mirrors FileFormatWriter's EmptyDirectoryDataWriter case: an empty input still writes + // one file from partition 0 so that the output carries the schema, but every other + // empty partition produces no file at all. + if (context.partitionId() == 0 || batches.hasNext) { + val extension = writerFactory.getFileExtension(taskContext) + val filePath = capturedCommitter.newTaskTempFile( + taskContext, + None, + FileNameSpec("", "-c000" + extension)) + NativeWriteUtils.checkNativeWriteDestination(filePath) + // `extension` names the codec in the job configuration. The writer uses that codec + // too, so the file name and the footer agree. + val codec = CodecConfig.from(taskContext).getCodec + val protoCodec = NativeWriteUtils + .protoCompressionCodec(codec.name()) + .getOrElse(throw new UnsupportedOperationException( + s"Comet's native Parquet writer cannot write $codec")) + val writer = capturedNativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .setCompression(protoCodec) + .build() + val taskOp = capturedNativeOp.toBuilder.setParquetWriter(writer).build() + + // Register before the iterator so this listener runs after its cleanup. + nativeMetrics.reportNativeWriteOutputMetrics(context) + val execIterator = CometExec.getCometIterator( + CometArrowStream.inputObjects(batches, childSchema, "CometNativeWriteExec"), + childSchema.length, + taskOp, + nativeMetrics, + numPartitions, + context.partitionId(), + None, + Seq.empty) + + // Close before committing. A failed write must remain the primary error even if + // native teardown (including the final metrics update) also throws. + Utils.tryWithSafeFinally { + while (execIterator.hasNext) { + execIterator.next().close() } + } { + execIterator.close() } } - } - } - } - } - - /** Create a JobContext for the write job */ - private def createJobContext(): Job = { - val job = Job.getInstance() - job.setJobID(new org.apache.hadoop.mapreduce.JobID(jobTrackerID, 0)) - job + // `hasNext` already ran an empty partition's child to completion, so there is nothing + // to drain or release when no file was written. + capturedCommitter.commitTask(taskContext) + })(catchBlock = capturedCommitter.abortTask(taskContext)) + }, + writeRDD.partitions.indices, + (index, message: TaskCommitMessage) => { + committer.onTaskCommit(message) + commitMessages(index) = message + }) + commitMessages } /** @@ -347,7 +235,7 @@ case class CometNativeWriteExec( */ private def prepareOutputPathForMode(): Boolean = { val path = new Path(outputPath) - val hadoopConf = sparkContext.hadoopConfiguration + val hadoopConf = serializableHadoopConf.value val fs = path.getFileSystem(hadoopConf) val qualifiedOutputPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) @@ -361,11 +249,7 @@ case class CometNativeWriteExec( true case SaveMode.Overwrite => if (fs.exists(qualifiedOutputPath)) { - val deleted = committer match { - case Some(c) => c.deleteWithJob(fs, qualifiedOutputPath, true) - case None => fs.delete(qualifiedOutputPath, true) - } - if (!deleted) { + if (!committer.deleteWithJob(fs, qualifiedOutputPath, true)) { throw QueryExecutionErrors.cannotClearOutputDirectoryError(qualifiedOutputPath) } } @@ -375,15 +259,27 @@ case class CometNativeWriteExec( } } - /** Create a TaskAttemptContext for a specific task */ +} + +object CometNativeWriteExec { + + /** + * Match FileFormatWriter's Hadoop task identifiers and configuration. Kept off the exec node so + * that the task closure does not capture it. + */ private def createTaskContext( - jobId: String, - partitionId: Int, - attemptNumber: Int): TaskAttemptContext = { - val job = Job.getInstance() - val taskAttemptID = new TaskAttemptID( - new TaskID(new org.apache.hadoop.mapreduce.JobID(jobId, 0), TaskType.REDUCE, partitionId), - attemptNumber) - new TaskAttemptContextImpl(job.getConfiguration, taskAttemptID) + conf: Configuration, + jobTrackerID: String, + context: TaskContext): TaskAttemptContext = { + val hadoopConf = new Configuration(conf) + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, context.stageId()) + val taskId = new TaskID(jobId, TaskType.MAP, context.partitionId()) + val attemptId = new TaskAttemptID(taskId, context.taskAttemptId().toInt & Integer.MAX_VALUE) + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskId.toString) + hadoopConf.set("mapreduce.task.attempt.id", attemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + new TaskAttemptContextImpl(hadoopConf, attemptId) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala index 305593e609d..7721cb7c2a5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -129,10 +129,10 @@ case class CometWriteFilesExec( // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, // so the output directory would carry no schema for readers. Spark's own WriteFilesExec swaps - // in a dummy single-partition RDD for exactly this case. AQE reaches it by collapsing a - // completed empty shuffle into a CometEmptyRelationExec, which is why - // CometDataWritingCommand declines empty-relation inputs on the Spark 3.x path (#5303): its - // writer has nowhere to put this swap. See CometEmptyRelationParquetWriterSuite. + // in a dummy single-partition RDD for exactly this case, and CometNativeWriteExec does the + // same on the Spark 3.x path (#5303). AQE reaches it by collapsing a completed empty shuffle + // into a CometEmptyRelationExec, so a write over a native empty relation stays native and + // still leaves a schema-only file. See CometEmptyRelationParquetWriterSuite. val writeRDD = if (childRDD.getNumPartitions == 0) { sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) } else { diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 5b01a603534..466f6cad9e6 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -613,13 +613,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { checkAnswer(spark.read.parquet(outputPath), df.collect()) assertParquetCodec(outputPath, CompressionCodecName.GZIP) - if (isSpark40Plus) { - // Spark names the file; Comet fills it. The extension is the only externally visible - // statement of the codec, so it has to agree with the footer. (On 3.x the native writer - // invents a name with no codec suffix, so there is nothing to compare.) - listPartFileNames(outputPath).foreach { name => - assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") - } + // Spark names the file; Comet fills it. The extension is the only externally visible + // statement of the codec, so it has to agree with the footer. + listPartFileNames(outputPath).foreach { name => + assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") } } } @@ -1237,16 +1234,13 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } // --------------------------------------------------------------------------------------------- - // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in - // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the - // `WriteFilesExecBase` trait. See CometWriteFilesExec. + // Commit-protocol checks run on both writers. Tests requiring the surrounding Spark write + // command are gated to Spark 4.0+, where Comet replaces only WriteFilesExec. // --------------------------------------------------------------------------------------------- test("write creates a _SUCCESS marker") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/2985 - the marker comes from - // HadoopMapReduceCommitProtocol.commitJob, which only runs because Comet leaves - // InsertIntoHadoopFsRelationCommand in the plan. + // HadoopMapReduceCommitProtocol.commitJob on both native writer paths. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath withTempPath { srcDir => @@ -1266,7 +1260,6 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("written file names follow Spark's naming convention") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // The file name comes from FileCommitProtocol.newTaskTempFile and must be used verbatim: // part---c..parquet. Committers that track individual files // and tools that parse these names depend on it. @@ -1295,12 +1288,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } - test("a custom output basename is honored on local storage") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // `escapedHdfsDestination` gates the basename on HDFS only, where the native URL parser would - // rename the file out from under the committer. Local writes hand the path to the native - // writer verbatim, so this is the control that keeps that guard from being widened: `part%foo` - // is declined on HDFS and has to keep working here. + test("local output basenames match Spark's version-specific commit protocol") { + // Spark 4.0+ honors the basename, including literal '%' on local storage where the path + // reaches the native writer verbatim. Spark 3.x ignores the option and always uses "part". Seq("out", "part%foo").foreach { basename => withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath @@ -1316,9 +1306,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } val written = new File(outputPath).listFiles().map(_.getName).filter(_.endsWith(".parquet")) + val expectedBasename = if (isSpark40Plus) basename else "part" assert( - written.nonEmpty && written.forall(_.startsWith(s"$basename-")), - s"expected every data file to be named '$basename-...', found: " + + written.nonEmpty && written.forall(_.startsWith(s"$expectedBasename-")), + s"expected every data file to be named '$expectedBasename-...', found: " + written.mkString(", ")) checkAnswer(spark.read.parquet(outputPath), df) } @@ -1350,8 +1341,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("dynamic partition overwrite falls back to Spark") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // A dynamic overwrite is a partitioned write, which CometWriteFiles declines - but the + // A dynamic overwrite is a partitioned write, which both native writers decline - but the // consequence of getting it wrong is silent data loss across untouched partitions, so assert // the fallback and the semantics explicitly rather than relying on the partitioning check. withTempPath { dir => @@ -1411,14 +1401,14 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("empty input still writes a schema-only file (SPARK-23271)") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // An empty input must still leave a schema behind for downstream readers: `spark.read.parquet` // of the output must see the write's schema, not fail. Comet reaches this in two ways - if the - // native child has one partition producing no batches, the partition-0 branch of executeTask - // writes a metadata-only file; if it produces zero partitions, doExecuteWrite swaps in a dummy - // single-partition RDD to get to the same branch. This test exercises the first; the - // zero-partition swap is reached by an AQE-collapsed empty relation and is covered by - // CometEmptyRelationParquetWriterSuite. + // native child has one partition producing no batches, the partition-0 branch of the write + // task writes a metadata-only file; if it produces zero partitions, the writer swaps in a + // dummy single-partition RDD to get to the same branch. Both CometWriteFilesExec (Spark 4.0+) + // and CometNativeWriteExec (Spark 3.x) do this. This test exercises the first; the + // zero-partition swap is covered by the empty-directory test below and, on Spark 4.0+, + // by CometEmptyRelationParquetWriterSuite. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val sourcePath = new File(dir, "source.parquet").getAbsolutePath @@ -1439,6 +1429,31 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a zero-partition native scan writes a schema-only file (#5303)") { + withTempPath { dir => + val source = new File(dir, "empty-source") + assert(source.mkdirs()) + val output = new File(dir, "output") + withNativeWriter { + val empty = spark.read.schema("id INT, name STRING").parquet(source.getAbsolutePath) + val scans = collect(empty.queryExecution.executedPlan) { case scan: CometNativeScanExec => + scan + } + assert(scans.size == 1, empty.queryExecution.executedPlan.toString) + assert(scans.head.executeColumnar().getNumPartitions == 0) + + val plan = captureWritePlan(p => empty.write.parquet(p), output.getAbsolutePath) + assertHasCometNativeWriteExec(plan) + assert(new File(output, "_SUCCESS").isFile) + assert(listPartFileNames(output.getAbsolutePath).size == 1) + assert(!new File(output, "_temporary").exists()) + val readBack = spark.read.parquet(output.getAbsolutePath) + assert(readBack.schema == empty.schema) + assert(readBack.collect().isEmpty) + } + } + } + test("INSERT INTO ... SELECT writes the target table's column names") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3426, which is Spark's own @@ -1487,8 +1502,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("an empty partition writes no file and still commits") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // executeTask's `sparkPartitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile + // The write task's `partitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile // altogether and still commit the task, matching FileFormatWriter's EmptyDirectoryDataWriter. // Hash-partitioning into eight and keeping a single id leaves at most one partition with // rows, so at most two files can appear: that one, plus partition 0's schema-only file when @@ -1547,8 +1561,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("a failing task aborts and cleans up its staging file") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // CometWriteFilesExec.executeTask must call committer.abortTask and rethrow. Injecting the + // Both native writers must call committer.abortTask and rethrow. Injecting the // failure through the commit protocol rather than the data lets the write get as far as // creating a staging file, so the cleanup is actually observable. // diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index d0ca67c60d5..40281ce4ea1 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala @@ -117,7 +117,10 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { command, child, outputPath = "/tmp/unused-native-write", - mode = SaveMode.Overwrite) + mode = SaveMode.Overwrite, + committer = null, + serializableHadoopConf = null, + outputWriterFactory = null) } private def assertRestoredParquetWrite(reverted: SparkPlan): WriteFilesExec = { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala new file mode 100644 index 00000000000..098b1c16bc5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -0,0 +1,390 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.io.{File, IOException} +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.jdk.CollectionConverters._ +import scala.util.Using + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.hadoop.mapreduce.{JobContext, TaskAttemptContext} +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.parquet.hadoop.{ParquetFileReader, ParquetOutputFormat} +import org.apache.parquet.hadoop.metadata.CompressionCodecName +import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.spark.TaskContext +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage +import org.apache.spark.internal.io.FileNameSpec +import org.apache.spark.sql.{CometTestBase, DataFrame, SaveMode} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, SQLHadoopMapReduceCommitProtocol} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{StringType, StructField} +import org.apache.spark.util.SerializableConfiguration + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.{Incompatible, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{schema2Proto, CometDataWritingCommand, NativeWriteUtils} + +/** Exercises the retained Spark 3.x writer, including both terminal execution entry points. */ +class CometNativeWriteSuite extends CometTestBase { + import NativeWriteCommitProtocol._ + + private def withWriter(partitions: Int = 1)( + f: (CometNativeWriteExec, DataFrame, File) => Unit): Unit = { + assume(!isSpark40Plus, "Spark 4.0+ uses CometWriteFilesExec") + withTempPath { dir => + val source = new File(dir, "source").getAbsolutePath + val output = new File(dir, "output with % space") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 20, 1, partitions).write.parquet(source) + } + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key -> classOf[NativeWriteCommitProtocol].getName, + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> "0") { + val data = spark.read.parquet(source) + val sourcePlan = data.queryExecution.executedPlan + val child = sourcePlan + .collectFirst { case scan: CometNativeScanExec => + scan + } + .getOrElse(fail(s"Expected a native source scan: $sourcePlan")) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path(output.toURI), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map("compression" -> "gzip", "comet.test.option" -> "present"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, child) + val nativeOp = CometDataWritingCommand.convert(command, Operator.newBuilder()).get + val writer = CometDataWritingCommand + .createExec(nativeOp, command) + .asInstanceOf[CometNativeWriteExec] + assert(writer.committer.isInstanceOf[NativeWriteCommitProtocol]) + reset() + try f(writer, data, output) + finally reset() + } + } + } + + test("Spark 3.x HDFS admission ignores the configurable output basename") { + assume(!isSpark40Plus, "Spark 4.0+ honors mapreduce.output.basename") + val data = spark.range(1) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path("hdfs://localhost:8020/output"), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map(NativeWriteUtils.BASE_OUTPUT_NAME -> "out%?#"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, data.queryExecution.executedPlan) + assert(CometDataWritingCommand.getSupportLevel(command).isInstanceOf[Incompatible]) + val escapedPath = cmd.copy(outputPath = new Path("hdfs://localhost:8020/out%?#")) + assert( + CometDataWritingCommand + .getSupportLevel(command.copy(cmd = escapedPath)) + .isInstanceOf[Unsupported]) + } + + Seq(false, true).foreach { columnar => + test( + s"configured commit protocol completes the ${if (columnar) "columnar" else "row"} write") { + withWriter(2) { (writer, data, output) => + val resultCount = + if (columnar) writer.executeColumnar().count() else writer.execute().count() + assert(resultCount == 0) + assert(events.asScala.head == "setupJob") + assert(events.asScala.last == "commitJob") + assert(!events.asScala.exists(_.startsWith("abort"))) + assert(writer.metrics("rows_written").value == 20, "the writer must execute natively") + val files = allocated.asScala.map(p => new Path(p).getName).toSet + assert(files.nonEmpty) + assert(files.forall(n => n.startsWith("chosen % ") && n.endsWith(".gz.parquet"))) + assert(output.list().filter(_.endsWith(".parquet")).toSet == files) + assert(new File(output, "_SUCCESS").isFile) + assert(!new File(output, "_temporary").exists()) + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + } + + test("the footer codec matches the codec in the file name") { + withWriter() { (writer, data, output) => + // The planner chose gzip, but the job configuration says zstd. The job configuration also + // names the file. + val conf = new Configuration(writer.serializableHadoopConf.value) + conf.set(ParquetOutputFormat.COMPRESSION, CompressionCodecName.ZSTD.name()) + writer.copy(serializableHadoopConf = new SerializableConfiguration(conf)).execute().count() + + val files = output.listFiles().filter(_.getName.endsWith(".parquet")) + assert(files.nonEmpty) + files.foreach { file => + assert(file.getName.endsWith(".zstd.parquet"), s"Expected a zstd file: ${file.getName}") + val input = HadoopInputFile.fromPath(new Path(file.toURI), conf) + Using.resource(ParquetFileReader.open(input)) { reader => + val codecs = reader.getFooter.getBlocks.asScala + .flatMap(_.getColumns.asScala) + .map(_.getCodec) + .toSet + assert(codecs == Set(CompressionCodecName.ZSTD), s"${file.getName} uses $codecs") + } + } + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + + test("native iterator creation failure aborts the task and job") { + withWriter() { (writer, _, output) => + // Inject an unknown native config after the child scan has initialized, so the writer's + // CometExecIterator constructor fails in createPlan, before the drain/close block is entered. + failAt = "createIterator" + failAbort = true + val error = intercept[Exception](writer.execute().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "createPlan"))) + assert(causes.exists(e => Option(e.getMessage).exists(_.contains("no_such_namespace")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("native execution failure survives throwing cleanup and abort callbacks") { + withWriter() { (writer, _, output) => + // A mismatched output type fails when writing the first native batch, after the native + // stream is initialized. This also lets releasePlan exercise its final metrics callback. + val badWriter = writer.nativeOp.getParquetWriter.toBuilder + .clearOutputSchema() + .addAllOutputSchema(schema2Proto(Seq(StructField("id", StringType))).asJava) + .build() + val badPlan = writer.nativeOp.toBuilder.setParquetWriter(badWriter).build() + failAbort = true + val failingWriter = new CleanupFailingNativeWriteExec(writer.copy(nativeOp = badPlan)) + val error = intercept[Exception](failingWriter.executeColumnar().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "executePlan"))) + assert(causes.exists(e => + Option(e.getMessage).exists(_.contains("Failed to rename batch columns")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected cleanup failure"))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("a cleanup failure aborts instead of committing") { + withWriter() { (writer, _, output) => + val error = + intercept[Exception](new CleanupFailingNativeWriteExec(writer).execute().count()) + assert(allErrors(error).exists(_.getMessage == "injected cleanup failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + Seq("commitTask", "onTaskCommit", "commitJob").foreach { phase => + test(s"$phase failure aborts and preserves the commit failure") { + withWriter() { (writer, _, output) => + failAt = phase + failAbort = true + val error = intercept[Exception](writer.execute().count()) + assert(allErrors(error).exists(_.getMessage == s"injected $phase failure")) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assert(events.contains("abortJob")) + if (phase == "commitTask") { + assert(events.contains("abortTask")) + assert( + allErrors(error).exists( + _.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + } + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert( + written.asScala.nonEmpty, + "the failure must occur after the native file was written") + } + } + } + + private def allErrors(error: Throwable): Seq[Throwable] = + Seq(error) ++ Option(error.getCause).toSeq.flatMap(allErrors) ++ + error.getSuppressed.toSeq.flatMap(allErrors) + + private def assertAborted(output: File): Unit = { + assert(events.contains("abortTask")) + assert(events.contains("abortJob")) + assert(!events.contains("commitJob")) + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert(allocated.asScala.nonEmpty, "the test must reach task file allocation") + } +} + +/** + * Local-mode probes are reset for each write; executor and driver callbacks share this object. + */ +object NativeWriteCommitProtocol { + val events = new ConcurrentLinkedQueue[String]() + val allocated = new ConcurrentLinkedQueue[String]() + val written = new ConcurrentLinkedQueue[String]() + val received = new ConcurrentLinkedQueue[TaskCommitMessage]() + @volatile var failAt: String = "" + @volatile var failAbort = false + + def reset(): Unit = { + events.clear() + allocated.clear() + written.clear() + received.clear() + failAt = "" + failAbort = false + } +} + +class NativeWriteCommitProtocol(jobId: String, path: String, dynamicPartitionOverwrite: Boolean) + extends SQLHadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite) { + import NativeWriteCommitProtocol._ + private var taskFile: String = _ + + override def setupJob(context: JobContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getOutputFormatClass == classOf[ParquetOutputFormat[_]]) + assert(FileOutputFormat.getOutputPath(context) == new Path(path)) + assert(context.getConfiguration.get("spark.sql.sources.writeJobUUID") != null) + super.setupJob(context) + context.getConfiguration.set("comet.test.setupJob", "present") + events.add("setupJob") + } + + override def setupTask(context: TaskAttemptContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getConfiguration.get("comet.test.setupJob") == "present") + super.setupTask(context) + events.add("setupTask") + } + + override def newTaskTempFile( + context: TaskAttemptContext, + dir: Option[String], + spec: FileNameSpec): String = { + val sparkPath = new Path(super.newTaskTempFile(context, dir, spec)) + taskFile = new Path(sparkPath.getParent, "chosen % " + sparkPath.getName).toString + allocated.add(taskFile) + if (failAt == "createIterator") { + val properties = TaskContext.get().getLocalProperties + properties.setProperty(CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key, "true") + properties.setProperty("spark.comet.datafusion.no_such_namespace.option", "1") + } + taskFile + } + + override def commitTask(context: TaskAttemptContext): TaskCommitMessage = { + val file = new Path(taskFile) + assert(file.getFileSystem(context.getConfiguration).getFileStatus(file).getLen > 0) + written.add(taskFile) + events.add("commitTask") + if (failAt == "commitTask") throw new IOException("injected commitTask failure") + new TaskCommitMessage((context.getTaskAttemptID.toString, super.commitTask(context))) + } + + override def onTaskCommit(message: TaskCommitMessage): Unit = { + events.add("onTaskCommit") + received.add(message) + if (failAt == "onTaskCommit") throw new IOException("injected onTaskCommit failure") + super.onTaskCommit(message.obj.asInstanceOf[(String, TaskCommitMessage)]._2) + } + + override def commitJob(context: JobContext, messages: Seq[TaskCommitMessage]): Unit = { + assert(messages.nonEmpty) + assert(messages.toSet == received.asScala.toSet) + assert( + messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._1).distinct.size == + messages.size) + events.add("commitJob") + if (failAt == "commitJob") throw new IOException("injected commitJob failure") + super.commitJob(context, messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._2)) + } + + override def abortTask(context: TaskAttemptContext): Unit = { + events.add("abortTask") + super.abortTask(context) + if (failAbort) throw new IOException("injected abortTask failure") + } + + override def abortJob(context: JobContext): Unit = { + events.add("abortJob") + super.abortJob(context) + if (failAbort) throw new IOException("injected abortJob failure") + } +} + +/** + * With periodic updates disabled, this metric throws only during native iterator teardown. + * CometMetricNode reports a sum metric with `add`, so both update methods throw. + */ +class CleanupFailingWriteMetric extends SQLMetric("sum") { + override def copy(): SQLMetric = new CleanupFailingWriteMetric + override def set(value: Long): Unit = throw new IOException("injected cleanup failure") + override def add(value: Long): Unit = throw new IOException("injected cleanup failure") +} + +class CleanupFailingNativeWriteExec(writer: CometNativeWriteExec) + extends CometNativeWriteExec( + writer.nativeOp, + writer.originalPlan, + writer.child, + writer.outputPath, + writer.mode, + writer.committer, + writer.serializableHadoopConf, + writer.outputWriterFactory, + writer.jobTrackerID) { + override lazy val metrics: Map[String, SQLMetric] = { + val failingMetric = new CleanupFailingWriteMetric + sparkContext.register(failingMetric) + writer.metrics.updated("rows_written", failingMetric) + } +} diff --git a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala index 04585612ab0..c739e420315 100644 --- a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala @@ -64,12 +64,8 @@ class CometEmptyRelationParquetWriterSuite extends CometParquetWriterTestBase { } if (nativeEmpty) { assert(collect(plan) { case e: CometEmptyRelationExec => e }.nonEmpty) - // A native empty relation is a zero-partition RDD, which is what - // CometDataWritingCommand declines on the Spark 3.x path: that writer only maps - // existing partitions, so no task runs and no file is written. Spark 4.0+ goes - // through the WriteFilesExec seam instead, where CometWriteFilesExec swaps in a - // dummy single-partition RDD exactly as Spark's own WriteFilesExec does, so - // partition 0 still writes the schema-only file the readback above needs. + // CometWriteFilesExec replaces the zero-partition native empty relation + // with one empty partition, so a task writes the schema-only file above. assertHasCometNativeWriteExec(plan) } else { // Spark's own EmptyRelationExec is not a Comet operator, so From 7e3cbc09f27cba1a767156bca1a4e80888237c10 Mon Sep 17 00:00:00 2001 From: dustin Date: Thu, 8 Oct 2026 17:46:13 +0000 Subject: [PATCH 36/89] fix: retry a shuffle page allocation that Spark failed after dropping the task's entry (#6403) * fix: retry a shuffle page allocation that Spark failed after dropping the task's entry Spark's ExecutionMemoryPool registers a task's memoryForTask entry once, before its wait loop, and reads it on every pass, while releaseMemory removes the entry when the task's balance reaches zero. When the JVM shuffle allocator's page or pointer array request was parked below the task's minimum share and Comet's native memory consumer released the task's last bytes, the request woke to NoSuchElementException "key not found: " and failed the task, because the shuffle writers only expect SparkOutOfMemoryError. CometUnifiedShuffleMemoryAllocator now retries allocate and allocateArray when Spark throws that exception from ExecutionMemoryPool. The retry registers the task again and waits as the first attempt would have. Nothing is counted or allocated before the exception, so retrying cannot leak. After three attempts it refuses the request with the same SparkOutOfMemoryError it uses for a refused page, with the last exception as its cause. The Spark issue is SPARK-59827. * docs: say which callers the page allocation retry covers Only the shuffle allocator's callers are guarded; Spark's own operators in the task can still hit the removed entry until Spark re-registers a waiting task (apache/spark#59103). * docs: point the shuffle allocator retry at SPARK-59444 apache/spark#59103 (SPARK-59827) was closed as a duplicate of apache/spark#58747 (SPARK-59444), which fixes the same race in ExecutionMemoryPool, so link that issue and PR instead. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + docs/source/contributor-guide/jvm_shuffle.md | 4 +- .../contributor-guide/memory_management.md | 11 +- .../CometUnifiedShuffleMemoryAllocator.java | 87 +++++- ...etUnifiedShuffleMemoryAllocatorSuite.scala | 285 ++++++++++++++++++ 6 files changed, 386 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index debfccf311e..2bb8fd363d1 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -505,6 +505,7 @@ jobs: org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnifiedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 742fd100d09..148dde4784e 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -209,6 +209,7 @@ jobs: org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnifiedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index dca7144a01a..474764b90e2 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -197,7 +197,9 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies depends on Spark's memory mode. - Off-heap mode gets `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`. When it cannot acquire a page the writer spills to - disk. + disk. An allocation that loses the task's entry in Spark's execution pool while it waits + ([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)) is retried a few times before + it counts as a page that could not be acquired. - On-heap mode gets `CometUnboundedShuffleMemoryAllocator`, which keeps no budget and so never refuses a page. Nothing bounds these allocations, and memory pressure never triggers a spill. That mode exists only so the Spark SQL tests can run against Comet. See diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index e3283fe26cd..196e0d6608f 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -136,7 +136,16 @@ allocator before reading anything into the split. **The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance` returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the -same task like any other allocation. +same task like any other allocation. A page or pointer array allocation that waits in Spark for +memory can wake up to find the task's entry gone from Spark's execution pool, because another +consumer of the task released its last bytes in the meantime, and Spark then throws a +`NoSuchElementException` ("key not found") instead of a grant +([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). The allocator retries that +case, since the failed call was granted nothing, and after three attempts throws +`SparkOutOfMemoryError`, which the shuffle writers handle like any other refused page. Only the +shuffle allocator's callers are guarded this way. Spark's operators in the same task, such as its +sorters and aggregates, can still hit the exception until Spark re-registers a waiting task in +`ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). Which allocator each call site uses, and who ends up charged for the bytes: diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index b7c7c58848e..e3a2269cacc 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -20,10 +20,18 @@ package org.apache.spark.shuffle.comet; import java.io.IOException; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.function.Supplier; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.apache.spark.memory.MemoryConsumer; import org.apache.spark.memory.MemoryMode; +import org.apache.spark.memory.SparkOutOfMemoryError; import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.unsafe.array.LongArray; import org.apache.spark.unsafe.memory.MemoryBlock; /** @@ -38,6 +46,12 @@ */ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { + private static final Logger logger = + LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class); + + /** Attempts at an allocation whose task entry Spark keeps losing, see retryMissingTaskEntry. */ + private static final int MAX_ALLOCATE_ATTEMPTS = 3; + CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) { @@ -54,7 +68,78 @@ public long spill(long l, MemoryConsumer memoryConsumer) throws IOException { } public synchronized MemoryBlock allocate(long required) { - return this.allocatePage(required); + return retryMissingTaskEntry(required, () -> this.allocatePage(required)); + } + + /** + * Allocates a pointer array, retrying like {@link #allocate}. The inherited method asks Spark for + * the page directly rather than through allocate, and the sorters' pointer arrays come through + * here. + */ + @Override + public LongArray allocateArray(long size) { + return retryMissingTaskEntry(size * 8L, () -> super.allocateArray(size)); + } + + /** + * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's + * entry. Spark removes a task's entry from its execution pool when the task's balance reaches + * zero, and an allocation that was waiting for memory then fails when it wakes up with a + * NoSuchElementException ("key not found: " and the task id). Releases take no task monitor, so a + * native consumer or another JVM consumer of the same task can empty the balance while an + * allocation waits. The failed call was granted nothing and took no page, because Spark only + * drops the entry at a zero balance, so trying again is safe: it registers the task again and + * waits for its share as Spark would have. After a few failed attempts the allocation fails with + * a SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused + * by the last attempt's exception. The Spark issue is SPARK-59444. + */ + private T retryMissingTaskEntry(long required, Supplier allocation) { + for (int attempt = 1; ; attempt++) { + try { + return allocation.get(); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e)) { + throw e; + } + if (attempt >= MAX_ALLOCATE_ATTEMPTS) { + logger.warn( + "Lost the task's execution memory entry in Spark on {} attempts to allocate {} " + + "bytes, refusing the allocation", + attempt, + required, + e); + SparkOutOfMemoryError error = + new SparkOutOfMemoryError( + "UNABLE_TO_ACQUIRE_MEMORY", + Map.of( + "requestedBytes", String.valueOf(required), + "receivedBytes", String.valueOf(0))); + error.initCause(e); + throw error; + } + logger.info( + "Lost the task's execution memory entry in Spark while waiting to allocate {} bytes, " + + "trying again", + required); + } + } + } + + /** + * Whether {@code e} is Spark's execution pool failing to find the task's entry. The task id is + * not available here, so this matches the message prefix and the pool's stack frame. + */ + private static boolean isMissingTaskEntry(NoSuchElementException e) { + String message = e.getMessage(); + if (message == null || !message.startsWith("key not found: ")) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; } public synchronized long free(MemoryBlock block) { diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala new file mode 100644 index 00000000000..b978c272802 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala @@ -0,0 +1,285 @@ +/* + * 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. + */ + +package org.apache.spark.shuffle.comet + +import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable.ArrayBuffer + +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.{LogEvent, LoggerContext} +import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkFunSuite, TaskContext, TaskContextImpl} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, UnifiedMemoryManager} + +class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { + + private val PageSize = 20L + private val TimeoutSeconds = 10L + private val MaxAllocateAttempts = 3 + + test("a page waiting in Spark survives a release that empties the task's balance") { + checkWaitingAllocationSurvives(20L)(allocatePage) + } + + test("a pointer array waiting in Spark survives a release that empties the task's balance") { + checkWaitingAllocationSurvives(24L)(allocatePointerArray) + } + + test("a NoSuchElementException other than Spark's missing task entry is rethrown") { + val otherMessage = noSuchElement("no entry for task 0", fromExecutionMemoryPool = true) + val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false) + for (error <- Seq(otherMessage, notFromPool); allocation <- allocations) { + val taskMemoryManager = new FailingTaskMemoryManager(() => error, failures = Int.MaxValue) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + val thrown = intercept[NoSuchElementException](allocation(allocator)) + assert(thrown eq error) + } + assert(events.isEmpty, messages(events)) + assert(taskMemoryManager.calls.get == 1) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + test("an allocation that keeps losing its task entry fails with Spark's out of memory error") { + for ((allocation, requested) <- allocations.zip(Seq("20", "24"))) { + // A new exception on every call, so the test can tell which one the allocator keeps. + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + val thrown = intercept[SparkOutOfMemoryError](allocation(allocator)) + assert(thrown.getErrorClass == "UNABLE_TO_ACQUIRE_MEMORY") + assert(thrown.getMessageParameters.get("requestedBytes") == requested) + assert(thrown.getMessageParameters.get("receivedBytes") == "0") + // The last attempt's exception carries the task id and where Spark was waiting. + assert(thrown.getCause eq taskMemoryManager.errors.last) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) + assert(events.last.getThrown eq taskMemoryManager.errors.last) + assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(taskMemoryManager.errors.distinct.size == MaxAllocateAttempts) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + test("an allocation that loses its task entry twice is granted on the last attempt") { + val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true) + for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) { + val taskMemoryManager = + new FailingTaskMemoryManager(() => error, failures = MaxAllocateAttempts - 1) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + assert(allocation(allocator) == requested) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO), messages(events)) + assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + /** Allocates a 20 byte page, frees it and returns its size. */ + private def allocatePage(allocator: CometUnifiedShuffleMemoryAllocator): Long = { + val page = allocator.allocate(20L) + val size = page.size() + allocator.free(page) + size + } + + /** Allocates a 3 entry (24 byte) pointer array, frees it and returns its size in bytes. */ + private def allocatePointerArray(allocator: CometUnifiedShuffleMemoryAllocator): Long = { + val array = allocator.allocateArray(3L) + val size = array.memoryBlock().size() + allocator.freeArray(array) + size + } + + private def allocations: Seq[CometUnifiedShuffleMemoryAllocator => Long] = + Seq(allocatePage _, allocatePointerArray _) + + /** + * The task holds 10 bytes of a 100 byte pool through a native consumer and another task holds + * 90. `allocate` asks for `requested` bytes on another thread and waits in Spark below the + * task's minimum share of 25 while the native consumer releases the task's last 10 bytes, which + * removes the task's entry from Spark's pool. Once the other task frees its memory the + * allocation must be granted in full. + */ + private def checkWaitingAllocationSurvives(requested: Long)( + allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val native = new CometTaskMemoryManager(1L, 0L) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. + assert(otherTask.acquireMemory(90L) == 90L) + assert(native.acquireMemory(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val thread = new Thread(() => + try granted.set(allocate(allocator)) + catch { case t: Throwable => failure.set(t) }) + thread.setDaemon(true) + + val events = logEvents(Level.INFO) { + try { + thread.start() + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(thread), s"the allocation is ${thread.getState}") + native.releaseMemory(10L) + } finally { + // Free the other task's memory so that the allocation does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + } + + assert(!thread.isAlive, s"the allocation is ${thread.getState}") + assert(failure.get == null, s"the allocation failed: ${failure.get}") + assert(granted.get == requested) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + // The first attempt lost the task's entry and the allocator tried again. + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + } + } + + private def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** A 100 byte off-heap execution pool. */ + private def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** + * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the + * result of `newError`, recording each one in `errors`. + */ + private class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) + extends TaskMemoryManager(offHeapMemoryManager(), 0L) { + val calls = new AtomicInteger() + val errors = new ArrayBuffer[Throwable]() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + val call = calls.incrementAndGet() + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (call > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + if (call <= failures) { + val error = newError() + errors += error + throw error + } + super.acquireExecutionMemory(required, consumer) + } + } + + private def noSuchElement( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + + /** The events that the allocator logs at `level` or above. */ + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { + val appender = new LogAppender("unified shuffle memory allocator") + appender.setThreshold(level) + val loggers = Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName) + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) + try withLogAppender(appender, loggers, Some(level))(f) + finally { + // For a logger with no config of its own, withLogAppender adds one and never removes it. + // It copies the root config's additivity, which is off, so left in place it would keep the + // logger's events out of the test log for the rest of the run. + unconfigured.foreach(context.getConfiguration.removeLogger) + context.updateLoggers() + } + appender.loggingEvents.toSeq + } + + private def messages(events: Seq[LogEvent]): String = + events.map(_.getMessage.getFormattedMessage).mkString("\n") + + private def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { + val taskContext = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = 0L, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + + TaskContext.setTaskContext(taskContext) + try { + f(taskMemoryManager) + } finally { + try { + taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + TaskContext.unset() + } + } + } +} From 9a6241f140e9fd576ba9467ee26c2e21cf46ec02 Mon Sep 17 00:00:00 2001 From: ChenChen Lai <72776271+0lai0@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:03:07 +0000 Subject: [PATCH 37/89] comet-collation-suite-4.1-plus-shim (#4864) Co-authored-by: Andy Grove --- .../org/apache/spark/sql/CometCollationSuite.scala | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename spark/src/test/{spark-4.1 => spark-4.1+}/org/apache/spark/sql/CometCollationSuite.scala (100%) diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala b/spark/src/test/spark-4.1+/org/apache/spark/sql/CometCollationSuite.scala similarity index 100% rename from spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala rename to spark/src/test/spark-4.1+/org/apache/spark/sql/CometCollationSuite.scala From 412468e2817207eff0c1907d1c4aac8c3f06680d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 8 Oct 2026 20:01:23 +0000 Subject: [PATCH 38/89] docs: expand the FAQ and refresh the Gluten comparison (#6783) * docs: expand the FAQ and refresh the Gluten comparison * docs: correct the shuffle manager and PyArrow UDF answers in the FAQ When spark.shuffle.manager does not name Comet's shuffle manager and spark.comet.shuffle.enabled is true (the default), Comet disables itself instead of only running shuffles in Spark. Say so, and name the setting that keeps Comet enabled. State that the mapInArrow and mapInPandas acceleration needs Spark 4.0 or later, and name spark.comet.exec.pyarrowUDF.enabled. --- docs/source/about/faq.md | 174 ++++++++++++++++++++++++- docs/source/about/gluten_comparison.md | 97 ++++++++------ docs/source/about/index.md | 6 +- 3 files changed, 236 insertions(+), 41 deletions(-) diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md index f7313fac28b..25edfe2a816 100644 --- a/docs/source/about/faq.md +++ b/docs/source/about/faq.md @@ -19,7 +19,147 @@ # Frequently Asked Questions -## Does Comet plan to add an API for vectorized Java/Scala UDFs similar to the Rust UDF API? +Answers to common questions about choosing, installing, and running Comet. If your question is not +answered here, see [Where can I ask questions or report bugs?](#where-can-i-ask-questions-or-report-bugs) + +## Choosing Comet + +### How does Comet compare to other open-source Spark accelerators? + +Several open-source projects speed up Spark by running query plans in native code. They differ in +the engine they build on and the hardware they need. + +- [Apache Gluten](https://gluten.apache.org/) is the closest to Comet in design: a Spark plugin that + passes the parts of a plan it supports to a native engine, such as the C++ engine Velox. See + [Comparison of Comet and Gluten](gluten_comparison.md). +- [Apache Auron (incubating)](https://auron.apache.org/), formerly Blaze, is also a Spark plugin + built on DataFusion. Comet is developed within the DataFusion project itself. +- The [NVIDIA cuDF plugin for Apache Spark](https://github.com/NVIDIA/cudf-spark), formerly the + RAPIDS Accelerator, runs queries on NVIDIA GPUs. Comet runs on the CPUs your cluster already has. + +Gluten and Auron are also adding Flink support, while Comet focuses on Spark alone. Which +accelerator is fastest depends on the workload, so we recommend benchmarking your own jobs. + +### Will Comet return the same results as Spark? + +That is the goal, and Comet runs Apache Spark's own SQL test suite in CI to check it. By default, +Comet uses implementations that match Spark, including running Spark's own generated code for +expressions whose native implementations behave differently. Those native implementations are +opt-in, one expression at a time. The known differences that remain, such as the handling of strings +that contain invalid UTF-8, are listed in the +[Compatibility Guide](../user-guide/latest/compatibility/index.md). If you find a difference that is +not documented, please [file an issue](https://github.com/apache/datafusion-comet/issues). + +### Which Spark and Java versions does Comet support? + +Spark 3.4, 3.5, 4.0, 4.1, and 4.2, on Java 17 or later. Spark 3.4 is deprecated. The +[installation guide](../user-guide/latest/installation.md#supported-spark-versions) lists the exact +Spark, Java, and Scala versions, and the +[Spark Version Compatibility](../user-guide/latest/compatibility/spark-versions.md) page lists known +issues for each Spark version. + +Comet drops a Spark version some time after the Spark project stops maintaining it. Spark 3.3 and +earlier are no longer supported, and Comet 1.1.0 dropped Java 11. Java 8 cannot be supported because +Apache Arrow's Java library no longer supports it. If you cannot upgrade Spark yet, stay on an earlier +Comet release that supports your version, as the +[versioning policy](versioning_policy.md#apache-spark-version-support) describes. + +### Does Comet work with vendor distributions of Spark? + +Comet is built and tested against open-source Apache Spark only. Spark distributions from cloud +providers and other vendors can change the internal Spark APIs that Comet calls and add operators +that Comet does not recognize, so Comet may fail with errors such as `NoSuchMethodError`, or leave +much of the work to Spark. To use Comet in the cloud, run open-source Apache Spark there, for example +on Kubernetes as described in the [Kubernetes guide](../user-guide/latest/kubernetes.md). + +### Which platforms do the published jars support? + +The jars in Maven Central include native libraries for Linux on amd64 and arm64, and need glibc 2.31 +or newer. They do not load on distributions with an older glibc, such as RHEL 7 and CentOS 7. On +those systems and on macOS, [build Comet from source](../user-guide/latest/source.md). + +For performance, the libraries target CPUs that are common in data centers: `x86-64-v3`, which +includes AVX2, on amd64, and Arm Neoverse N1 or newer cores on arm64. If Comet crashes with `SIGILL` +(illegal instruction), the CPU is older than that. Build Comet from source for it, and please +[open an issue](https://github.com/apache/datafusion-comet/issues) describing your environment. + +A jar you build yourself is smaller than the published one because it contains the native library +for one platform only. + +## Running Comet + +### Why do I get a `ClassNotFoundException` for `CometShuffleManager`? + +Spark could not find the Comet jar when it created the shuffle manager. On Spark 3.4 and 3.5, +executors create the shuffle manager before they load jars passed with `--jars` or `--packages`, so +the jar must be on the startup classpath. Either copy it into `$SPARK_HOME/jars` on every machine, or +set `spark.driver.extraClassPath` and `spark.executor.extraClassPath` to its path, as the +[installation guide](../user-guide/latest/installation.md#run-spark-shell-with-comet-enabled) shows. + +These two settings take a local file path, not an `hdfs://` or other URL. The JVM silently ignores +classpath entries that do not exist, so the jar must be at that path on every machine that runs a +driver or an executor. + +### How much memory does Comet need? + +Comet needs Spark's off-heap memory: without `spark.memory.offHeap.enabled=true`, Comet disables +itself and logs a warning. Comet's native operators share the off-heap pool, sized by +`spark.memory.offHeap.size`, with Spark. When the pool runs out, operators that can spill, such as +sorts, aggregations, and shuffle writes, spill to disk, and operators that cannot spill fail the +task. + +Comet also uses memory that the pool does not track, which has to fit in +`spark.executor.memoryOverhead`. If the cluster manager kills executors for exceeding their memory +limit, raise the overhead rather than the pool. Each executor logs its native memory use while Comet +runs, which you can use to size these settings. See +[Memory Tuning](../user-guide/latest/tuning/memory.md). + +### Why isn't my query faster with Comet? + +First check how much of the query Comet runs. Set `spark.comet.explain.fallback.enabled=true` to log +the reasons that parts of each query run in Spark, and see +[Understanding Comet Plans](../user-guide/latest/understanding-comet-plans.md) to read the plan. To +estimate how much of a workload Comet would accelerate without changing how it runs, set +`spark.comet.explain.planOnly.enabled=true`. + +Common causes of a small speedup are: + +- **Fallbacks inside a stage.** Each switch between Comet and Spark converts data between columnar + and row formats, which can cost more than Comet saves. See + [Reducing Row/Columnar Conversion Overhead](../user-guide/latest/tuning/transitions.md). +- **Comet's shuffle manager is not set.** If `spark.shuffle.manager` does not name Comet's shuffle + manager, Comet disables itself and logs a warning, unless `spark.comet.shuffle.enabled=false`, in + which case Comet runs but every shuffle runs in Spark. See + [Shuffle Tuning](../user-guide/latest/tuning/shuffle.md). +- **Too little memory.** Native operators spill to disk when the off-heap pool is too small. See + [Memory Tuning](../user-guide/latest/tuning/memory.md). +- **Little computation to accelerate.** Queries that spend most of their time reading from storage, + listing files, or planning, or that process little data, have less to gain. + +## Feature Support + +### Does Comet support Delta Lake, Apache Hudi, or Apache Paimon tables? + +Not yet. Comet accelerates [Apache Iceberg](../user-guide/latest/iceberg.md) tables, with a native +reader that is enabled by default and +[experimental native writes](../user-guide/latest/iceberg-writes.md). It does not accelerate scans of +Delta Lake, Hudi, or Paimon tables, so Spark reads those. Native Delta Lake reads are in development; +see the [roadmap](../contributor-guide/roadmap.md#delta-lake-support). Hudi and Paimon are not on the +roadmap. + +### Does Comet accelerate PySpark jobs and Python UDFs? + +PySpark DataFrame and SQL queries produce the same physical plans as Scala, so Comet accelerates them +the same way. Set Comet up as the [installation guide](../user-guide/latest/installation.md) +describes. There is no pip package. + +For Python UDFs, Comet has experimental support for `mapInArrow` and `mapInPandas` on Spark 4.0 and +later. It keeps their data in Arrow format instead of converting it to rows and back, and is disabled +by default. To enable it, set `spark.comet.exec.pyarrowUDF.enabled=true`; see +[PyArrow UDF Acceleration](../user-guide/latest/pyarrow-udfs.md). Scalar `@pandas_udf` functions are +not accelerated yet, and Python UDFs that do not use Arrow are not planned. + +### Does Comet plan to add an API for vectorized Java/Scala UDFs similar to the Rust UDF API? Not at the moment. Comet already runs ordinary Scala and Java UDFs in its pipeline without code changes, compiling each call into a loop over a whole batch that the JVM optimizes well. See @@ -35,3 +175,35 @@ If you have a use case that a function of one row handles poorly, such as workin of strings or calling a library that processes whole batches, please describe it in [#6694](https://github.com/apache/datafusion-comet/issues/6694). For native speed, see the experimental [Rust UDF API](../user-guide/latest/rust_udfs.md). + +### Does Comet support Structured Streaming? + +No. Comet targets batch queries and leaves streaming queries to Spark, and streaming support is not +on the roadmap. See [Spark Operator Support](../user-guide/latest/operators.md#not-currently-planned). + +### Can I use Comet with a remote shuffle service? + +Yes, with [Apache Celeborn](https://celeborn.apache.org/). Comet accelerates scans, filters, and +other operators, and Celeborn's own Spark integration handles the shuffle. With the released Celeborn +0.6.x and 0.7.x clients, Comet's native shuffle cannot run through Celeborn. See +[Using Comet with Apache Celeborn](../user-guide/latest/celeborn.md). Native shuffle with Apache +Uniffle is not supported yet ([#4913](https://github.com/apache/datafusion-comet/issues/4913)). + +## Project and Community + +### Is there a roadmap? + +Yes. The [Comet Roadmap](../contributor-guide/roadmap.md) describes the major items that need +coordination between contributors. + +### Where can I ask questions or report bugs? + +- Ask questions in [GitHub Discussions](https://github.com/apache/datafusion-comet/discussions), or + in the Comet channels of the DataFusion + [Slack and Discord](https://datafusion.apache.org/contributor-guide/communication.html). +- Report bugs and request features in + [GitHub issues](https://github.com/apache/datafusion-comet/issues). +- Report security vulnerabilities privately, following the + [ASF security process](https://www.apache.org/security/), not in a public issue. +- Contributors also meet weekly on a video call; see + [Regular public meetings](../contributor-guide/contributing.md#regular-public-meetings). diff --git a/docs/source/about/gluten_comparison.md b/docs/source/about/gluten_comparison.md index 735662c52c4..56fb5340cc9 100644 --- a/docs/source/about/gluten_comparison.md +++ b/docs/source/about/gluten_comparison.md @@ -30,7 +30,9 @@ Comet and Gluten have very similar architectures. Both are Spark plugins that tr a serialized representation and pass the serialized plan to native code for execution. Gluten serializes the plans using the Substrait format and has an extensible architecture that supports execution -against multiple engines. Velox and Clickhouse are currently supported, but Velox is more widely used. +against multiple engines. Gluten 1.7.0 supports Velox and ClickHouse. A third backend, Bolt, a C++ engine from +ByteDance that is derived from Velox, has been added to Gluten's main branch but is not yet in a release. The rest of +this page compares Comet with Gluten's Velox backend. Comet serializes the plans in a proprietary Protocol Buffer format. Execution is delegated to Apache DataFusion. Comet does not plan to support multiple engines, but rather focus on a tight integration between Spark and DataFusion. @@ -87,7 +89,7 @@ Comet supports Spark 3.4, 3.5, 4.0, 4.1, and 4.2 in production builds. See the [Spark version compatibility guide]: /user-guide/latest/compatibility/spark-versions.md -Gluten supports Spark 3.3, 3.4, 3.5, 4.0, and 4.1. +Gluten 1.7.0 supports Spark 3.3, 3.4, 3.5, 4.0, and 4.1. ## ANSI Mode @@ -98,8 +100,10 @@ Comet implements ANSI semantics for the expressions it supports natively, includ ANSI cast behavior, and `try_*` variants. Queries running with `spark.sql.ansi.enabled=true` continue to be accelerated. See the [Comet Compatibility Guide] for details on which expressions have full ANSI coverage. -The Gluten Velox backend documents that ANSI mode is not supported and that any query executed with ANSI enabled -will fall back to vanilla Spark. See the [Gluten Velox limitations] page for the current status. +The Gluten Velox backend documents that ANSI mode is not supported: by default, any query executed with ANSI enabled +falls back to vanilla Spark. Setting `spark.gluten.sql.ansiFallback.enabled=false` makes Gluten attempt to run such +queries natively, but Gluten's developer documentation notes that the results do not yet match Spark. See the +[Gluten Velox limitations] page for the current status. [Gluten Velox limitations]: https://apache.github.io/gluten/velox-backend-limitations.html @@ -111,16 +115,22 @@ fraction of a workload that runs natively. Both projects can accelerate queries against Apache Iceberg tables, but they take different approaches and Gluten covers a broader set of table formats overall. -Comet provides a native Iceberg scan built on iceberg-rust. It has been tested with Iceberg 1.5 through 1.10 and -supports Iceberg spec v1 and v2, schema evolution, time travel and branch reads, positional and equality deletes -on merge-on-read tables, REST catalogs, and S3-compatible object storage. Iceberg writes still go through Spark. -Comet does not currently provide native integrations for Delta Lake, Hudi, or Paimon. See the -[Comet Iceberg guide] for the full list of supported features and known limitations. +Comet provides a native Iceberg scan built on iceberg-rust. It has been tested with Iceberg 1.5 and 1.8 through 1.11 +and supports Iceberg spec v1, v2, and v3, schema evolution, time travel and branch reads, positional and equality +deletes and deletion vectors, encrypted v3 tables, REST catalogs, and S3-compatible object storage. Comet can also +write Iceberg data files natively through iceberg-rust, as an experimental feature that is disabled by default (see +[Comet Iceberg writes]). Comet does not currently provide native integrations for Delta Lake, Hudi, or Paimon. See +the [Comet Iceberg guide] for the full list of supported features and known limitations. [Comet Iceberg guide]: /user-guide/latest/iceberg.md +[Comet Iceberg writes]: /user-guide/latest/iceberg-writes.md -Gluten ships dedicated modules for Iceberg, Delta Lake (2.0 through 4.0), Hudi, and Paimon. Users who need native -acceleration for Delta, Hudi, or Paimon will find broader coverage in Gluten today. +Gluten ships dedicated modules for Iceberg, Delta Lake (2.3 through 4.1, depending on the Spark version), Hudi, and +Paimon. Its [Iceberg documentation] says that reads of unpartitioned tables are offloaded, including tables with +position deletes, while reads of partitioned tables mostly fall back to Spark, as do equality deletes and all writes. +Users who need native acceleration for Delta, Hudi, or Paimon will find broader coverage in Gluten today. + +[Iceberg documentation]: https://github.com/apache/gluten/blob/v1.7.0/docs/get-started/VeloxIceberg.md ## Compatibility @@ -131,17 +141,17 @@ Spark are disabled by default, but users can opt in. See the [Comet Compatibilit [Comet Compatibility Guide]: https://datafusion.apache.org/comet/user-guide/latest/compatibility/index.html Gluten also aims to provide compatibility with Spark, and includes a subset of the Spark SQL tests in its own test -suite. See the Gluten [Velox backend limitations] page for known gaps, including notes on case sensitivity, regular -expression dialect (RE2 vs `java.util.regex`), NaN handling, and timestamp encodings. +suite. See the Gluten [Velox backend limitations] page for known gaps, such as differences in the regular expression +dialect (RE2 vs `java.util.regex`). [Velox backend limitations]: https://apache.github.io/gluten/velox-backend-limitations.html ## Codegen Dispatch -Comet has a feature called the codegen dispatcher that has no direct equivalent in Gluten. When an expression has a -native implementation with known semantic differences from Spark, or has no native implementation at all, Comet can -run Spark's own generated code for that expression inside the native pipeline. Data is passed to the JVM in Arrow -format, evaluated using Spark's byte-exact logic, and returned to the native pipeline for the rest of the query. +When an expression has a native implementation with known semantic differences from Spark, or has no native +implementation at all, Comet can run Spark's own generated code for that expression inside the native pipeline. Data +is passed to the JVM in Arrow format, evaluated using Spark's byte-exact logic, and returned to the native pipeline for +the rest of the query. Comet calls this the codegen dispatcher. This matters for two reasons: @@ -149,14 +159,15 @@ This matters for two reasons: expressions, JSON functions, and some datetime and array functions) route through the codegen dispatcher by default, so they produce results that match Spark exactly. The faster native path is opt-in per expression for users who accept its differences. -- **Fewer full fallbacks.** Because the dispatcher keeps evaluation inside the native pipeline instead of falling - back to whole-stage Spark execution, a single unsupported or incompatible expression does not force an entire - query stage back onto vanilla Spark. This keeps more of the plan running natively than a strict native-or-fallback - model would. +- **Fewer fallbacks.** Because the dispatcher keeps evaluation inside the native pipeline, a single unsupported or + incompatible scalar expression does not force its operator back onto vanilla Spark. -By contrast, Gluten falls back to vanilla Spark for expressions and operators that its Velox backend does not -support. See the [Comet Compatibility Guide] for more detail on how the codegen dispatcher works and which -expressions use it. +Gluten's closest equivalent is partial projection, which has been enabled by default since Gluten 1.3.0. When a +projection contains expressions that the backend cannot run, including Scala and Hive UDFs, Gluten splits the +projection so that vanilla Spark evaluates only those expressions and the rest of the projection runs natively. +Partial projection applies only to projections, so, for example, a filter that contains an unsupported expression +still falls back to Spark. See the [Comet Compatibility Guide] for more detail on how the codegen dispatcher works and +which expressions use it. ## Regular Expression Compatibility @@ -167,11 +178,13 @@ evaluate `rlike`, `regexp_replace`, `regexp_extract`, `regexp_extract_all`, `reg compatibility with Spark by default, including every pattern feature the JVM engine supports. Comet also offers a faster native Rust regex engine as an opt-in per expression for users who can accept its semantic differences. -Gluten's Velox backend evaluates regular expressions using the C++ RE2 engine. RE2 uses a different dialect from -`java.util.regex` and deliberately omits features such as backreferences and lookaround, so it cannot fully match -Spark's regex semantics regardless of configuration. This is a fundamental consequence of using a native C++ regex -engine rather than the JVM engine, and it is documented in the [Gluten Velox limitations] page. Comet's codegen -dispatcher sidesteps this class of incompatibility entirely by keeping Spark's own engine in the loop. +Gluten's Velox backend evaluates regular expressions using the C++ RE2 engine, which uses a different dialect from +`java.util.regex` and deliberately omits features such as backreferences and lookaround. Gluten falls back to Spark +for patterns that RE2 cannot compile, so those return correct results without acceleration. Patterns that RE2 accepts +run natively and can differ from Spark in edge cases: for example, `\s` does not match the vertical tab character. +Setting `spark.gluten.sql.fallbackRegexpExpressions=true` sends every regular expression function to Spark. See the +[Gluten Velox limitations] page for details. Comet's codegen dispatcher avoids this trade-off by running Spark's own +regex engine inside the native pipeline. ## Scala and Java UDF Support @@ -183,9 +196,10 @@ keep running natively while only the UDF itself is evaluated on the JVM. This co complex nested types and UDFs composed with other Catalyst expressions and higher-order functions. This means users can adopt Comet without giving up their existing Scala and Java UDFs, and without rewriting them in -the engine's native language. Gluten's Velox backend supports native C++ UDFs but does not run arbitrary existing -JVM UDFs in the native pipeline, so plans containing them typically fall back to Spark. See the -[Comet Scala and Java UDF guide] for the full list of supported and unsupported cases. +the engine's native language. Hive UDFs are not covered and still fall back to Spark. Gluten's Velox backend supports +native C++ UDFs, and its partial projection (see [Codegen Dispatch](#codegen-dispatch)) evaluates Scala and Hive UDFs +in a projection with vanilla Spark while the rest of the projection runs natively. A UDF in a filter makes the filter +fall back to Spark. See the [Comet Scala and Java UDF guide] for the full list of supported and unsupported cases. [Comet Scala and Java UDF guide]: /user-guide/latest/scala_java_udfs.md @@ -199,8 +213,15 @@ overall performance, with Gluten finishing roughly 9% faster than Comet across t The headline number masks wide per-query variation: some queries were significantly faster on Gluten (for example, large fact-table joins where Gluten uses a shuffled hash join strategy), while others were significantly faster on -Comet (for example, CPU-bound scan and aggregate pipelines). We expect Comet performance to continue improving over -time and for this gap to close. +Comet (for example, CPU-bound scan and aggregate pipelines). + +AWS Labs has also published separate runs of each accelerator against the same Spark baseline on the same TPC-DS 3TB +Parquet workload. [Comet 1.0.0] ran it 1.57× faster than Spark (and 1.75× faster on Iceberg tables), and +[Gluten 1.6.0] ran it 1.63× faster. We expect Comet performance to continue improving over time and for this gap to +close. + +[Comet 1.0.0]: https://awslabs.github.io/data-on-eks/docs/benchmarks/spark-datafusion-comet-benchmark +[Gluten 1.6.0]: https://awslabs.github.io/data-on-eks/docs/benchmarks/spark-gluten-velox-benchmark Although TPC-DS and TPC-H are good benchmarks for operators such as joins and aggregates, they don't necessarily represent real-world queries, especially for ETL use cases. For example, there are limited complex types involved @@ -216,7 +237,7 @@ management capabilities vs the complexities around installing C++ dependencies. Comet and Gluten are both good solutions for accelerating Spark jobs, and independent benchmarking shows they deliver similar overall performance. Comet currently has an edge for users on Spark 4.0 with ANSI mode enabled, for -users who want a fully native Iceberg scan path, and through its codegen dispatcher, which keeps partially compatible -or unsupported expressions running inside the native pipeline rather than falling back to Spark. Gluten holds a -small performance lead in the AWS Labs TPC-DS benchmark and offers broader native integration with Delta Lake, Hudi, -and Paimon, plus a second backend in ClickHouse. We recommend trying both to see which is the best fit for your needs. +Iceberg workloads, and through its codegen dispatcher, which runs partially compatible or unsupported expressions with +Spark's exact semantics inside the native pipeline. Gluten holds a small performance lead in the AWS Labs TPC-DS +benchmarks and offers broader native integration with Delta Lake, Hudi, and Paimon, plus a second backend in +ClickHouse. We recommend trying both to see which is the best fit for your needs. diff --git a/docs/source/about/index.md b/docs/source/about/index.md index fc30f765c7f..5573a9103c8 100644 --- a/docs/source/about/index.md +++ b/docs/source/about/index.md @@ -23,9 +23,11 @@ Background on the project: how Comet compares to similar accelerators, how relea are versioned, answers to common questions, and links to the Apache Software Foundation resources behind it. -- [Comparison with Gluten](gluten_comparison.md) — how Comet differs from other Spark accelerators +- [Comparison with Gluten](gluten_comparison.md) — how Comet differs from Apache Gluten, the most + similar accelerator - [Versioning Policy](versioning_policy.md) — what our version numbers mean and our release cadence -- [Frequently Asked Questions](faq.md) — answers to common questions about Comet's features and plans +- [Frequently Asked Questions](faq.md) — answers to common questions about choosing, installing, and + running Comet - [Blog Posts and Talks](blogs-and-talks.md) — external writing and conference talks about Comet - [ASF Links](../asf/index.md) — licensing, donations, security reporting, and the code of conduct From 14e49a0a3637408fc4b8496f91f768814e865e21 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 04:28:19 +0000 Subject: [PATCH 39/89] perf: compare floats in Spark's order without normalizing copies of the operands (#6691) * perf: compare floats in Spark's order without normalizing copies of the operands Since #6447 every native comparison of FLOAT or DOUBLE operands follows Spark's SQL ordering, but it got there by normalizing each operand that is not a literal into a copy and running Arrow's kernels on the copies. For arrays and structs, <=>, <, <=, > and >= deep-copied both operands in every batch. A new SparkComparison expression compares in place instead. Flat floats go through branch-free kernels over the value buffers, column against column or against a scalar, for =, <>, <, <=, >, >=, <=> and IS DISTINCT FROM. Lists and structs use spark_comparator for the ordering operators and spark_equality for the null-safe ones. The nested comparisons, nested = and <> included, now compare only the rows where neither side is null, iterating the validity bitmap, rather than testing both null buffers on every row. Planning is unchanged: a literal is still normalized while the plan is built, a float column compared with a literal in a scan's data filters is still a plain BinaryExpr so that Parquet pruning recognizes it, and operands of other types, such as dictionary-encoded floats, keep the normalize-then-compare path. * fix: split float comparison bitmaps with as_chunks for clippy 1.99 CI's rust-test job lints with the current stable toolchain, Rust 1.99.0, whose clippy adds chunks_exact_to_as_chunks. It rejected the three chunks_exact(64) calls in pack and pack_pairs under -D warnings, so the job failed before any test ran. as_chunks::<64>() yields the same 64-value chunks as arrays, which also drops the try_into().unwrap() conversions. * perf: keep CASE WHEN eager when a branch compares floats `is_infallible` in `case_when.rs` accepted the `BinaryExpr` of normalized operands that a float comparison used to be, but not the `SparkComparison` that replaced it. A CASE WHEN or IF with a float comparison in a later WHEN or in a branch value therefore fell back from eager to lazy evaluation. `SparkComparison::is_infallible` now reports a comparison of two operands of the same type as unable to fail, which covers every flat Float32 or Float64 comparison. The nested comparators reject a dictionary-encoded leaf against a plain one, which the logical type check in `spark_comparison` lets through, so a nested comparison counts only when the types are equal. The `is_infallible` test builds comparisons through `spark_comparison`, so a change in what the planner builds shows up there. The contributor guide describes comparing in place and lists the float comparison kernels. --- .../contributor-guide/floating_point.md | 23 +- .../latest/compatibility/floating-point.md | 17 +- native/core/src/execution/planner.rs | 50 +- native/spark-expr/benches/float_comparison.rs | 90 ++- .../spark-expr/benches/nested_comparison.rs | 29 +- native/spark-expr/src/array_funcs/mod.rs | 2 +- .../src/array_funcs/nested_comparison.rs | 669 ++++++++++++++++-- .../src/conditional_funcs/case_when.rs | 30 +- .../spark-expr/src/float_semantics/compare.rs | 147 +++- .../spark-expr/src/float_semantics/kernels.rs | 413 +++++++++++ native/spark-expr/src/float_semantics/mod.rs | 16 +- .../src/float_semantics/normalize.rs | 2 +- 12 files changed, 1289 insertions(+), 199 deletions(-) create mode 100644 native/spark-expr/src/float_semantics/kernels.rs diff --git a/docs/source/contributor-guide/floating_point.md b/docs/source/contributor-guide/floating_point.md index c113b86b8d5..165b280c634 100644 --- a/docs/source/contributor-guide/floating_point.md +++ b/docs/source/contributor-guide/floating_point.md @@ -92,6 +92,7 @@ Its module documentation is the reference. In short: | Helper | Rule | Use it to | | ---------------------------------------------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------- | | `compare_floats`, `float_lt`, `float_gt` | SQL ordering | Compare values in a kernel that returns the original bits, as `array_min` and `greatest` do | +| `compare_float_arrays`, `compare_float_array_scalar` | SQL ordering | Compare a `FLOAT` or `DOUBLE` array with another or with a scalar in place, as `SparkComparison` does | | `spark_comparator`, `spark_equality` | SQL ordering, at any depth | Compare arrays and structs in place | | `normalize_float`, `normalize_floats`, `normalize_nested_floats` | `NormalizeNaNAndZero` | Normalize values before an Arrow kernel sorts, row-encodes, hashes or compares them | | `NormalizeNaNAndZero`, `NormalizeNestedFloats` | `NormalizeNaNAndZero` | Wrap a key or operand expression. `wrap_if_needed` skips other types and keys that Spark already normalized | @@ -112,15 +113,19 @@ normalized and the output keeps the original values. sort orders them the Spark way. - `spark_comparison` in `native/spark-expr/src/array_funcs/nested_comparison.rs` builds every native `=`, `<>`, `<=>`, `<`, `<=`, `>`, `>=` and `IS DISTINCT FROM`, whichever operator evaluates it. It - normalizes float operands, and folds a literal while the plan is built so that it stays a - literal. Nested `=` and `<>` compare in place with `spark_equality`. -- A scan's pushed-down data filters leave a float column compared with a literal other than NaN - unwrapped (`FloatOperands::Raw`) so that Parquet pruning still recognizes it, but only while the - reader prunes with them without filtering rows. A bloom filter probe hashes the literal's bits, so - `=` against either zero becomes `= -0.0 OR = 0.0`. With - `spark.comet.parquet.rowFilterPushdown.enabled=true` they normalize both sides, because the - reader drops the rows a filter rejects, and a raw column would reject a stored NaN that Spark - matches ([#6702](https://github.com/apache/datafusion-comet/issues/6702)). Spark's own reader + compares the operands as they are, without normalized copies: `FLOAT` and `DOUBLE` with + `compare_float_arrays` and `compare_float_array_scalar`, and arrays and structs with + `spark_comparator`, or with `spark_equality` for `=`, `<>`, `<=>` and `IS DISTINCT FROM`. It + normalizes a literal while the plan is built so that it stays a literal. Operands of other types, + such as dictionary-encoded floats, are normalized and compared with Arrow's kernels. +- A scan's pushed-down data filters keep a float column compared with a literal other than NaN as a + plain `BinaryExpr` (`FloatOperands::Raw`) so that Parquet pruning still recognizes it, but only + while the reader prunes with them without filtering rows. A bloom filter probe hashes the + literal's bits, so `=` against either zero becomes `= -0.0 OR = 0.0`. With + `spark.comet.parquet.rowFilterPushdown.enabled=true` they compare in Spark's ordering like every + other comparison, because the reader drops the rows a filter rejects, and a raw column would + reject a stored NaN that Spark matches + ([#6702](https://github.com/apache/datafusion-comet/issues/6702)). Spark's own reader keeps the two zeros apart in its dictionary and bloom filters, so it can skip a row group that Comet reads. A test of this pruning writes out the expected rows instead of comparing them with Spark's. diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index f9e257dba3f..adf8518d42d 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -24,17 +24,18 @@ However, one exception is comparison. Spark does not normalize NaN and zero when because they are handled well in Spark (e.g., `SQLOrderingUtil.compareFloats`). But the comparison functions of arrow-rs used by DataFusion do not normalize NaN and zero (e.g., [arrow::compute::kernels::cmp::eq](https://docs.rs/arrow/latest/arrow/compute/kernels/cmp/fn.eq.html#)). For `FLOAT` and `DOUBLE` comparisons (`=`, `<>`, `<=>`, `<`, `<=`, `>` and `>=`), Comet -normalizes both operands before native execution, including noncanonical NaN literals. This -applies wherever a comparison appears: projections, filters, aggregate arguments and `FILTER` -clauses, join conditions, sort keys, and generator arguments. +compares the operands in Spark's order natively, including noncanonical NaNs: `-0.0` equals +`0.0`, all NaNs are equal, and NaN sorts above every other value. This applies wherever a +comparison appears: projections, filters, aggregate arguments and `FILTER` clauses, join +conditions, sort keys, and generator arguments. A native Parquet scan prunes row groups and pages with its data filters. So that this pruning still applies, a `FLOAT` or `DOUBLE` column compared with a constant other than NaN in a data filter is compared without normalizing the column, and the filter above the scan evaluates the -comparison again with Spark's semantics. Every other comparison in a data filter is normalized. +comparison again with Spark's semantics. Every other comparison in a data filter follows Spark's semantics. Bloom filters store `-0.0` and `0.0` separately, so `=` against either zero checks them for both. With `spark.comet.parquet.rowFilterPushdown.enabled=true` the scan also drops the rows its data -filters reject, so every comparison in a data filter is normalized, and a `FLOAT` or `DOUBLE` +filters reject, so every comparison in a data filter follows Spark's semantics, and a `FLOAT` or `DOUBLE` comparison in a data filter does not prune row groups or pages ([#6702](https://github.com/apache/datafusion-comet/issues/6702)). @@ -58,10 +59,10 @@ or structs; see [#6019](https://github.com/apache/datafusion-comet/issues/6019). For arrays and structs containing `FLOAT` or `DOUBLE`, native `=`, `<>`, `IN`, and `NOT IN` compare signed zeros as equal and all NaN representations as equal, matching Spark. This also covers single-candidate membership that Spark rewrites into equality. `<=>`, `<`, `<=`, `>`, and -`>=` normalize the floating-point values inside both operands first, so they follow Spark's order -too, in which NaN sorts above every other value. +`>=` follow Spark's order too, in which NaN sorts above every other value. -Equality and dynamic membership compare nested elements directly and stop at the first mismatch. +Comparisons and dynamic membership compare nested elements directly and stop at the first +mismatch. Constant membership sets use normalized comparison values for static lookup. These operations preserve SQL null semantics and do not change the values returned by projections. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 221d6b05360..b88b477931d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1159,10 +1159,10 @@ impl PhysicalPlanner { /// they reject, which Spark's Filter above the scan then never sees. /// /// Without row-level pushdown the filters only prune, and pruning only recognizes a column - /// compared with a literal, so that shape keeps the raw column ([`FloatOperands::Raw`]). With - /// it, every operand is normalized and float comparisons give up pruning: a raw column would - /// drop a stored NaN that Spark matches, such as one with the sign bit set, which Arrow orders - /// below every other value. + /// compared with a literal, so that shape stays a plain comparison of the raw column + /// ([`FloatOperands::Raw`]). With it, every float comparison follows Spark's ordering and + /// gives up pruning: a raw column would drop a stored NaN that Spark matches, such as one with + /// the sign bit set, which Arrow orders below every other value. fn data_filter_float_operands(&self) -> FloatOperands { if self .session_ctx @@ -5882,12 +5882,13 @@ mod tests { assert_eq!(0, filter_exec.additional_native_plans.len()); } - /// Comparisons normalize float operands, except in the data filters that a scan pushes into - /// the Parquet reader, where pruning has to see the column itself. + /// Comparisons follow Spark's float ordering without normalizing their operands, except in the + /// data filters that a scan pushes into the Parquet reader, where pruning has to see a plain + /// comparison of the column. #[test] fn scan_data_filters_compare_float_columns_directly() { use datafusion::physical_expr::expressions::BinaryExpr; - use datafusion_comet_spark_expr::NormalizeNaNAndZero; + use datafusion_comet_spark_expr::SparkComparison; let double = spark_expression::DataType { type_id: 6, type_info: None, @@ -5909,30 +5910,29 @@ mod tests { })))), }))); let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, true)])); - let left_operand = |expr: Arc| { - let comparison = expr.downcast_ref::().expect("a comparison"); - Arc::clone(comparison.left()) - }; let planner = PhysicalPlanner::default(); let comparison = planner.create_expr(&expr, Arc::clone(&schema)).unwrap(); - assert!(left_operand(comparison) - .downcast_ref::() - .is_some()); + let comparison = comparison + .downcast_ref::() + .expect("a Spark comparison"); + assert!(comparison.left().downcast_ref::().is_some()); let data_filter = planner .create_data_filter(&expr, schema, planner.data_filter_float_operands()) .unwrap(); - assert!(left_operand(data_filter).downcast_ref::().is_some()); + let data_filter = data_filter + .downcast_ref::() + .expect("a plain comparison"); + assert!(data_filter.left().downcast_ref::().is_some()); } /// With row-level pushdown the Parquet reader drops the rows a data filter rejects, so the - /// filter normalizes the column as well as the literal. A raw column would drop a stored NaN - /// whose bits differ from the normalized literal, and a stored NaN with the sign bit set - /// under any ordering comparison, both of which Spark matches. + /// filter compares the column in Spark's ordering rather than as it is. A raw column would drop + /// a stored NaN whose bits differ from the normalized literal, and a stored NaN with the sign + /// bit set under any ordering comparison, both of which Spark matches. #[test] - fn scan_data_filters_normalize_float_columns_with_row_level_pushdown() { + fn scan_data_filters_follow_spark_ordering_with_row_level_pushdown() { use arrow::array::{AsArray, BooleanArray}; - use datafusion::physical_expr::expressions::BinaryExpr; - use datafusion_comet_spark_expr::{FloatOperands, NormalizeNaNAndZero}; + use datafusion_comet_spark_expr::{FloatOperands, SparkComparison}; let double = spark_expression::DataType { type_id: 6, type_info: None, @@ -5985,13 +5985,7 @@ mod tests { planner.data_filter_float_operands(), ) .unwrap(); - let comparison = data_filter - .downcast_ref::() - .expect("a comparison"); - assert!(comparison - .left() - .downcast_ref::() - .is_some()); + assert!(data_filter.downcast_ref::().is_some()); let matched = data_filter.evaluate(&batch).unwrap().into_array(3).unwrap(); assert_eq!(matched.as_boolean(), &BooleanArray::from(expected.to_vec())); } diff --git a/native/spark-expr/benches/float_comparison.rs b/native/spark-expr/benches/float_comparison.rs index 7a411e05b59..aacadec6550 100644 --- a/native/spark-expr/benches/float_comparison.rs +++ b/native/spark-expr/benches/float_comparison.rs @@ -15,10 +15,11 @@ // specific language governing permissions and limitations // under the License. -//! Compare Spark's float comparisons, which normalize each float operand first, with DataFusion's -//! on the same data, and time `normalize_floats` on its own. Both comparisons return the same -//! answer except on a NaN with the sign bit set, which is checked before timing; special-value -//! semantics belong in tests. +//! Compare Spark's float comparisons, which follow Spark's SQL ordering without normalizing their +//! operands, with comparing normalized copies of the operands and with DataFusion's comparison, on +//! the same data, and time `normalize_floats` on its own. All three return the same answer except +//! that DataFusion's differs on a NaN with the sign bit set, which is checked before timing; +//! special-value semantics belong in tests. use arrow::array::{Array, ArrayRef, Float64Array}; use arrow::datatypes::{DataType, Field, Schema}; @@ -28,7 +29,9 @@ use datafusion::common::ScalarValue; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion::physical_expr::PhysicalExpr; -use datafusion_comet_spark_expr::{normalize_floats, spark_comparison, FloatOperands}; +use datafusion_comet_spark_expr::{ + normalize_floats, spark_comparison, FloatOperands, NormalizeNaNAndZero, +}; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -72,34 +75,55 @@ fn criterion_benchmark(c: &mut Criterion) { let b: Arc = Arc::new(Column::new("b", 1)); let literal: Arc = Arc::new(Literal::new(ScalarValue::Float64(Some(500.0)))); - for (shape, right) in [("column_column", b), ("column_literal", literal)] { - let datafusion: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&a), - Operator::Lt, - Arc::clone(&right), - )); - let spark = spark_comparison( - Arc::clone(&a), - Operator::Lt, - right, - &schema, - FloatOperands::Normalize, - ) - .unwrap(); - let evaluate = |expr: &Arc| { - expr.evaluate(&batch) - .unwrap() - .into_array(ROWS) - .unwrap() - .to_data() - }; - if special.map(f64::to_bits) != Some(NEGATIVE_NAN.to_bits()) { - assert_eq!(evaluate(&spark), evaluate(&datafusion)); - } - for (engine, expr) in [("spark", &spark), ("datafusion", &datafusion)] { - group.bench_function(BenchmarkId::new(format!("{shape}_{engine}"), data), |b| { - b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) - }); + for (op_name, op) in [ + ("lt", Operator::Lt), + ("eq", Operator::Eq), + ("not_distinct", Operator::IsNotDistinctFrom), + ] { + for (shape, right) in [("column_column", &b), ("column_literal", &literal)] { + let datafusion: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&a), op, Arc::clone(right))); + // Comparing normalized copies of the operands, as `spark_comparison` did before it + // compared them in place. It normalized a literal while planning, and 500.0 is + // already normal. + let normalize = |operand: &Arc| -> Arc { + if operand.downcast_ref::().is_some() { + Arc::clone(operand) + } else { + NormalizeNaNAndZero::wrap_if_needed(Arc::clone(operand), &schema).unwrap() + } + }; + let normalized: Arc = + Arc::new(BinaryExpr::new(normalize(&a), op, normalize(right))); + let spark = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(right), + &schema, + FloatOperands::Normalize, + ) + .unwrap(); + let evaluate = |expr: &Arc| { + expr.evaluate(&batch) + .unwrap() + .into_array(ROWS) + .unwrap() + .to_data() + }; + assert_eq!(evaluate(&spark), evaluate(&normalized)); + if special.map(f64::to_bits) != Some(NEGATIVE_NAN.to_bits()) { + assert_eq!(evaluate(&spark), evaluate(&datafusion)); + } + for (engine, expr) in [ + ("spark", &spark), + ("normalized", &normalized), + ("datafusion", &datafusion), + ] { + group.bench_function( + BenchmarkId::new(format!("{op_name}_{shape}_{engine}"), data), + |b| b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())), + ); + } } } let column = Arc::clone(batch.column(0)); diff --git a/native/spark-expr/benches/nested_comparison.rs b/native/spark-expr/benches/nested_comparison.rs index 85e1905c19d..3bfd0a34d38 100644 --- a/native/spark-expr/benches/nested_comparison.rs +++ b/native/spark-expr/benches/nested_comparison.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! Compare the base DataFusion path, PR #6073's eager normalization, and Spark equality. -//! Ordinary finite inputs keep the answers identical across all three implementations. +//! Compare the base DataFusion path, eager normalization of the operands (PR #6073's for `IN`, and +//! the one `spark_comparison` used for ordering comparisons before it compared them in place), and +//! Spark's comparisons. Ordinary finite inputs keep the answers identical across all three. use arrow::array::{ArrayRef, Float64Array, ListArray}; use arrow::buffer::{NullBuffer, OffsetBuffer}; @@ -65,11 +66,23 @@ fn expression(version: &str, mode: &str, batch: &RecordBatch) -> Arc = Arc::new(Column::new("a", 0)); let b: Arc = Arc::new(Column::new("b", 1)); - if mode == "eq" { - return if version == "new" { - spark_comparison(a, Operator::Eq, b, &schema, FloatOperands::Normalize).unwrap() - } else { - Arc::new(BinaryExpr::new(a, Operator::Eq, b)) + let normalize = |e: Arc| NormalizeNestedFloats::wrap_if_needed(e, &schema); + let comparison = match mode { + "eq" => Some(Operator::Eq), + "lt" => Some(Operator::Lt), + "not_distinct" => Some(Operator::IsNotDistinctFrom), + _ => None, + }; + if let Some(op) = comparison { + return match version { + "new" => spark_comparison(a, op, b, &schema, FloatOperands::Normalize).unwrap(), + // `=` already compared in place before ordering comparisons did. + "head" if op != Operator::Eq => Arc::new(BinaryExpr::new( + normalize(a).unwrap(), + op, + normalize(b).unwrap(), + )), + _ => Arc::new(BinaryExpr::new(a, op, b)), }; } let literal: Arc = Arc::new(Literal::new( @@ -118,7 +131,7 @@ fn benchmark(c: &mut Criterion) { Field::new("b", b.data_type().clone(), true), ])); let batch = RecordBatch::try_new(schema, vec![a, b]).unwrap(); - for mode in ["dynamic", "constant", "mixed", "eq"] { + for mode in ["dynamic", "constant", "mixed", "eq", "lt", "not_distinct"] { let expected = expression("base", mode, &batch) .evaluate(&batch) .unwrap() diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index e79d8f18a2e..18d35d7ac0a 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -42,7 +42,7 @@ pub use flatten::SparkFlatten; pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; pub use list_positions::ListPositionsExpr; -pub use nested_comparison::{spark_comparison, spark_in_list, FloatOperands}; +pub use nested_comparison::{spark_comparison, spark_in_list, FloatOperands, SparkComparison}; pub use sequence::spark_sequence; pub use size::{spark_size, SparkSizeFunc}; pub use sort_array::SparkSortArray; diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index e945eb8f269..585c244c0eb 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -15,15 +15,17 @@ // specific language governing permissions and limitations // under the License. -//! Spark comparisons for floating-point values, and Spark equality for nested floating-point -//! values without materializing normalized columns. +//! Spark comparisons for floating-point values, flat or nested in lists and structs, evaluated +//! without materializing normalized columns. use crate::float_semantics::{ - is_nested_with_float_leaf, normalize_comparison_operand, normalize_nested_floats, - spark_equality, NormalizeNestedFloats, + compare_float_array_scalar, compare_float_arrays, comparison_with_nulls, + is_nested_with_float_leaf, normalize_comparison_operand, normalize_float_scalar, + normalize_floats, normalize_nested_floats, spark_comparator_ignoring_nulls, + spark_equality_ignoring_nulls, NormalizeNestedFloats, }; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer}; use arrow::compute::{not, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; @@ -31,6 +33,7 @@ use datafusion::common::{internal_err, DFSchema, Result, ScalarValue}; use datafusion::logical_expr::{ColumnarValue, Operator}; use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, InListExpr, Literal}; use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr_common::datum::apply_cmp; use datafusion::physical_expr_common::physical_expr::is_volatile; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -55,27 +58,45 @@ impl Operand { }) } + /// The SQL nulls of a comparison of `self` with `other`, neither of them a null scalar: the + /// rows where either operand is null. + fn nulls(&self, other: &Self) -> Option { + match (self.scalar, other.scalar) { + (true, true) => None, + (true, false) => other.array.nulls().cloned(), + (false, true) => self.array.nulls().cloned(), + (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), + } + } + + /// Whether either operand is a null scalar, which makes every row of a comparison null. + fn has_null_scalar(&self, other: &Self) -> bool { + (self.scalar && self.array.is_null(0)) || (other.scalar && other.array.is_null(0)) + } + + /// The nulls of this operand over `len` rows: a null scalar is null on every row. + fn broadcast_nulls(&self, len: usize) -> Option { + if self.scalar { + self.array.is_null(0).then(|| NullBuffer::new_null(len)) + } else { + self.array.nulls().cloned() + } + } + fn equal(&self, other: &Self, rows: usize) -> Result { let scalar = self.scalar && other.scalar; let len = if scalar { 1 } else { rows }; // Scalars use a single value; only the Boolean result is broadcast across rows. - let scalar_null = - (self.scalar && self.array.is_null(0)) || (other.scalar && other.array.is_null(0)); - let result = if scalar_null + let result = if self.has_null_scalar(other) || self.array.data_type() == &DataType::Null || other.array.data_type() == &DataType::Null { BooleanArray::new_null(len) } else { // Inner nulls take part in structural equality. Outer SQL nulls are handled here. - let equal = spark_equality(self.array.as_ref(), other.array.as_ref())?; - let nulls = match (self.scalar, other.scalar) { - (true, true) => None, - (true, false) => other.array.nulls().cloned(), - (false, true) => self.array.nulls().cloned(), - (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), - }; - let values = BooleanBuffer::collect_bool(len, |row| { + let equal = spark_equality_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let nulls = self.nulls(other); + let values = collect_valid(len, nulls.as_ref(), |row| { equal( if self.scalar { 0 } else { row }, if other.scalar { 0 } else { row }, @@ -83,17 +104,90 @@ impl Operand { }); BooleanArray::new(values, nulls) }; - if scalar { - Ok(ColumnarValue::Scalar(ScalarValue::Boolean( - if result.is_null(0) { - None + Ok(into_value(result, scalar)) + } + + /// `self op other` for an ordering operator, or for `IS [NOT] DISTINCT FROM`, in the ordering + /// of `spark_comparator`. The null-safe operators test equality with `spark_equality` instead, + /// which tells lists of different lengths apart without comparing their elements. Both compare + /// only the rows where neither side is null, and then apply the nulls of the two sides. + fn compare(&self, other: &Self, op: Operator, rows: usize) -> Result { + let scalar = self.scalar && other.scalar; + let len = if scalar { 1 } else { rows }; + let (left_scalar, right_scalar) = (self.scalar, other.scalar); + let index = move |row: usize, scalar: bool| if scalar { 0 } else { row }; + let result = match op { + Operator::IsDistinctFrom | Operator::IsNotDistinctFrom => { + let equal = + spark_equality_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let distinct = op == Operator::IsDistinctFrom; + let (left_nulls, right_nulls) = + (self.broadcast_nulls(len), other.broadcast_nulls(len)); + let valid = NullBuffer::union(left_nulls.as_ref(), right_nulls.as_ref()); + let values = collect_valid(len, valid.as_ref(), |row| { + equal(index(row, left_scalar), index(row, right_scalar)) != distinct + }); + comparison_with_nulls(op, values, left_nulls.as_ref(), right_nulls.as_ref()) + } + _ => { + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return internal_err!("Unsupported operator for a nested comparison: {op}"); + } + if self.has_null_scalar(other) { + BooleanArray::new_null(len) } else { - Some(result.value(0)) - }, - ))) + let compare = + spark_comparator_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let compare_row = + |row: usize| compare(index(row, left_scalar), index(row, right_scalar)); + let nulls = self.nulls(other); + let valid = nulls.as_ref(); + // One loop per operator, so that the test on the ordering is inlined. + let values = match op { + Operator::Lt => collect_valid(len, valid, |row| compare_row(row).is_lt()), + Operator::LtEq => collect_valid(len, valid, |row| compare_row(row).is_le()), + Operator::Gt => collect_valid(len, valid, |row| compare_row(row).is_gt()), + _ => collect_valid(len, valid, |row| compare_row(row).is_ge()), + }; + BooleanArray::new(values, nulls) + } + } + }; + Ok(into_value(result, scalar)) + } +} + +/// `test(row)` for each of `len` rows, as a bitmap. A row that `valid` marks null is not tested +/// and comes out unset, which skips comparing the values under a null list or struct without +/// testing the null buffers on every row. +fn collect_valid( + len: usize, + valid: Option<&NullBuffer>, + test: impl Fn(usize) -> bool, +) -> BooleanBuffer { + let Some(valid) = valid.filter(|valid| valid.null_count() > 0) else { + return BooleanBuffer::collect_bool(len, test); + }; + let mut words = vec![0u64; len.div_ceil(64)]; + for row in valid.valid_indices() { + words[row / 64] |= (test(row) as u64) << (row % 64); + } + BooleanBuffer::new(Buffer::from_vec(words), 0, len) +} + +/// `result` as an array, or as a scalar when both operands were scalars. +fn into_value(result: BooleanArray, scalar: bool) -> ColumnarValue { + if scalar { + ColumnarValue::Scalar(ScalarValue::Boolean(if result.is_null(0) { + None } else { - Ok(ColumnarValue::Array(Arc::new(result))) - } + Some(result.value(0)) + })) + } else { + ColumnarValue::Array(Arc::new(result)) } } @@ -217,31 +311,190 @@ impl PhysicalExpr for NestedPredicate { } } +/// A comparison of two Float32 or Float64 operands, or of two lists or structs with a float leaf, +/// in Spark's SQL ordering: `-0.0` equals `0.0`, all NaNs are equal, and NaN sorts above every +/// other value. It reads the operands as they are, without normalized copies of them. +#[derive(Debug, Eq)] +pub struct SparkComparison { + left: Arc, + op: Operator, + right: Arc, + /// Whether the operands are lists or structs rather than Float32 or Float64 values. + nested: bool, +} + +impl SparkComparison { + pub fn left(&self) -> &Arc { + &self.left + } + + pub fn op(&self) -> Operator { + self.op + } + + pub fn right(&self) -> &Arc { + &self.right + } + + /// Whether the comparison cannot fail for any value: when both operands have the same type. + /// The float kernels and the nested comparators fail only on operands of different types, + /// such as a dictionary-encoded leaf against a plain one, which the logical type check in + /// [`spark_comparison`] lets through for lists and structs. + pub(crate) fn is_infallible(&self, input_schema: &Schema) -> bool { + match ( + self.left.data_type(input_schema), + self.right.data_type(input_schema), + ) { + (Ok(left), Ok(right)) => left == right, + _ => false, + } + } +} + +impl PartialEq for SparkComparison { + fn eq(&self, other: &Self) -> bool { + self.left.eq(&other.left) + && self.op == other.op + && self.right.eq(&other.right) + && self.nested == other.nested + } +} + +impl Hash for SparkComparison { + fn hash(&self, state: &mut H) { + self.left.hash(state); + self.op.hash(state); + self.right.hash(state); + self.nested.hash(state); + } +} + +impl Display for SparkComparison { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "SparkComparison({} {} {})", + self.left, self.op, self.right + ) + } +} + +impl PhysicalExpr for SparkComparison { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn data_type(&self, _: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, schema: &Schema) -> Result { + if matches!( + self.op, + Operator::IsDistinctFrom | Operator::IsNotDistinctFrom + ) { + return Ok(false); + } + Ok(self.left.nullable(schema)? || self.right.nullable(schema)?) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let left = self.left.evaluate(batch)?; + let right = self.right.evaluate(batch)?; + if self.nested { + Operand::new(left)?.compare(&Operand::new(right)?, self.op, batch.num_rows()) + } else { + compare_flat(self.op, left, right) + } + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let Ok([left, right]) = <[_; 2]>::try_from(children) else { + return internal_err!("SparkComparison expects two children"); + }; + Ok(Arc::new(Self { + left, + op: self.op, + right, + nested: self.nested, + })) + } +} + +/// `left op right` for Float32 or Float64 operands. An operand in any other layout, which these +/// types do not produce, is normalized and compared with Arrow's kernels instead. +fn compare_flat(op: Operator, left: ColumnarValue, right: ColumnarValue) -> Result { + let result = match (&left, &right) { + (ColumnarValue::Array(l), ColumnarValue::Array(r)) => { + compare_float_arrays(op, l.as_ref(), r.as_ref())? + } + (ColumnarValue::Array(l), ColumnarValue::Scalar(r)) => { + compare_float_array_scalar(op, l.as_ref(), r)? + } + (ColumnarValue::Scalar(l), ColumnarValue::Array(r)) => match op.swap() { + Some(swapped) => compare_float_array_scalar(swapped, r.as_ref(), l)?, + None => None, + }, + (ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) => { + if let Some(result) = + compare_float_arrays(op, l.to_array()?.as_ref(), r.to_array()?.as_ref())? + { + return Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)); + } + None + } + }; + match result { + Some(result) => Ok(ColumnarValue::Array(Arc::new(result))), + None => apply_cmp(op, &normalize_value(left), &normalize_value(right)), + } +} + +fn normalize_value(value: ColumnarValue) -> ColumnarValue { + match value { + ColumnarValue::Array(array) => ColumnarValue::Array(normalize_floats(&array)), + ColumnarValue::Scalar(value) => ColumnarValue::Scalar(normalize_float_scalar(value)), + } +} + /// How [`spark_comparison`] treats floating-point operands. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FloatOperands { - /// Normalize them, so that the comparison follows Spark's SQL ordering. + /// Compare them in Spark's SQL ordering. A literal is normalized while planning, and other + /// operands are compared as they are. Normalize, - /// Leave a Float32 or Float64 column compared with a literal other than NaN as it is, and - /// normalize every other operand. Only a scan's pushed-down data filters use this, and only - /// when the Parquet reader prunes with them but does not filter rows: Parquet pruning - /// recognizes a column compared with a literal but not a normalized column, and Spark's - /// Filter above the scan applies Spark's semantics to every row. A NaN literal is normalized - /// with the column, giving up pruning, because a stored NaN can have other bits than the - /// literal and a writer can leave NaNs out of the column statistics. With row-level pushdown - /// the reader would drop the rows such a comparison rejects, including a stored NaN with the - /// sign bit set, which Arrow orders below every other value, so the data filters use - /// [`FloatOperands::Normalize`] there instead. + /// Leave a Float32 or Float64 column compared with a literal other than NaN to a plain + /// [`BinaryExpr`] of the two as they are, and compare every other shape in Spark's SQL + /// ordering. Only a scan's pushed-down data filters use this, and only when the Parquet reader + /// prunes with them but does not filter rows: Parquet pruning recognizes a column compared with + /// a literal but not a Spark comparison, and Spark's Filter above the scan applies Spark's + /// semantics to every row. A NaN literal gets a Spark comparison too, giving up pruning, + /// because a stored NaN can have other bits than the literal and a writer can leave NaNs out + /// of the column statistics. With row-level pushdown the reader would drop the rows a plain + /// comparison rejects, including a stored NaN with the sign bit set, which Arrow orders below + /// every other value, so the data filters use [`FloatOperands::Normalize`] there instead. Raw, } /// Builds a comparison with Spark's SQL ordering for floats, in which `-0.0` equals `0.0`, all /// NaNs are equal and NaN sorts above every other value, at any depth of a list or struct. /// -/// Arrow compares floats by IEEE 754 total order instead, so each float operand is normalized -/// first with [`normalize_comparison_operand`], after which the two orders agree. Nested `=` and -/// `<>` compare with `spark_equality` instead, without building normalized copies of the nested -/// values. Any other operator, such as `AND`, builds a plain [`BinaryExpr`]. +/// Arrow compares floats by IEEE 754 total order instead, so a [`SparkComparison`] compares +/// Float32 and Float64 operands, and lists and structs with a float leaf, itself, without +/// normalized copies of them: flat floats with branch-free kernels over the value buffers, and +/// nested values with `spark_comparator`. Nested `=` and `<>` compare with `spark_equality`, which +/// also serves `IN`. A literal is normalized while planning. Operands of any other type, such as +/// dictionary-encoded floats, go through [`normalize_comparison_operand`] and a plain +/// [`BinaryExpr`], and so does any other operator, such as `AND`. /// /// The planner reconciles the nullability of nested operands before calling this. pub fn spark_comparison( @@ -260,23 +513,41 @@ pub fn spark_comparison( } // An operand whose type does not resolve against this schema falls back to the plain // comparison, the way `reconcile_nested_comparison_types` already leaves such operands alone. - let (Ok(left_type), Ok(_)) = (left.data_type(schema), right.data_type(schema)) else { + let (Ok(left_type), Ok(right_type)) = (left.data_type(schema), right.data_type(schema)) else { return Ok(Arc::new(BinaryExpr::new(left, op, right))); }; - if matches!(op, Eq | NotEq) && is_nested_with_float_leaf(&left_type) { - validate_types(&left, std::slice::from_ref(&right), schema)?; - return Ok(Arc::new(NestedPredicate { - value: left, - candidates: vec![right], - negated: op == NotEq, - membership: false, - })); + if is_nested_with_float_leaf(&left_type) { + if matches!(op, Eq | NotEq) { + validate_types(&left, std::slice::from_ref(&right), schema)?; + return Ok(Arc::new(NestedPredicate { + value: left, + candidates: vec![right], + negated: op == NotEq, + membership: false, + })); + } + if DFSchema::datatype_is_logically_equal(&left_type, &right_type) { + return Ok(Arc::new(SparkComparison { + left: normalize_literal(left, schema)?, + op, + right: normalize_literal(right, schema)?, + nested: true, + })); + } } if float_operands == FloatOperands::Raw { if let Some(comparison) = raw_float_comparison(&left, op, &right, schema) { return Ok(comparison); } } + if matches!(left_type, DataType::Float32 | DataType::Float64) && left_type == right_type { + return Ok(Arc::new(SparkComparison { + left: normalize_literal(left, schema)?, + op, + right: normalize_literal(right, schema)?, + nested: false, + })); + } Ok(Arc::new(BinaryExpr::new( normalize_comparison_operand(left, schema)?, op, @@ -284,6 +555,19 @@ pub fn spark_comparison( ))) } +/// Normalizes `expr` while planning if it is a literal, which keeps it a literal. Any other +/// operand is returned as is. +fn normalize_literal( + expr: Arc, + schema: &Schema, +) -> Result> { + if expr.downcast_ref::().is_some() { + normalize_comparison_operand(expr, schema) + } else { + Ok(expr) + } +} + /// The comparison that [`FloatOperands::Raw`] builds for a Float32 or Float64 column compared with /// a literal other than NaN: the column and the literal as they are. `None` for any other operands. /// @@ -1015,11 +1299,11 @@ mod tests { } /// A literal operand is normalized while planning, so the comparison stays `column op - /// literal`. With `FloatOperands::Raw`, as for a scan's data filters, the column and the - /// literal of that shape are left as they are. + /// literal`, and a column is compared as it is. With `FloatOperands::Raw`, as for a scan's + /// data filters, a float column compared with a literal stays a plain `BinaryExpr` of the two + /// as they are. #[test] fn float_operand_shapes() -> Result<()> { - use crate::float_semantics::NormalizeNaNAndZero; let batch = batch( Arc::new(Float64Array::from(vec![1.0])), Arc::new(Float64Array::from(vec![1.0])), @@ -1034,12 +1318,9 @@ mod tests { &schema, FloatOperands::Normalize, )?; - let binary = expr.downcast_ref::().unwrap(); - assert!(binary - .left() - .downcast_ref::() - .is_some()); - let folded = binary.right().downcast_ref::().unwrap(); + let comparison = expr.downcast_ref::().unwrap(); + assert!(comparison.left().downcast_ref::().is_some()); + let folded = comparison.right().downcast_ref::().unwrap(); assert!(matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == 0)); // With `FloatOperands::Raw`, a float column compared with a literal keeps both as they @@ -1065,8 +1346,8 @@ mod tests { assert!(Arc::ptr_eq(binary.left(), &negative_zero)); assert!(binary.right().downcast_ref::().is_some()); - // Any other shape is normalized even with `FloatOperands::Raw`, because a reader with - // row-level pushdown drops the rows it rejects. + // Any other shape follows Spark's ordering even with `FloatOperands::Raw`, because a + // reader with row-level pushdown drops the rows it rejects. let (a, b) = columns(); let expr = spark_comparison( Arc::clone(&a), @@ -1075,15 +1356,10 @@ mod tests { &schema, FloatOperands::Raw, )?; - let binary = expr.downcast_ref::().unwrap(); - assert!(binary - .left() - .downcast_ref::() - .is_some()); - assert!(binary - .right() - .downcast_ref::() - .is_some()); + let comparison = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(comparison.left(), &a)); + assert!(Arc::ptr_eq(comparison.right(), &b)); + assert_eq!(comparison.op(), Operator::Lt); // Other operators are not comparisons and keep their operands. let expr = spark_comparison( @@ -1096,16 +1372,31 @@ mod tests { let binary = expr.downcast_ref::().unwrap(); assert!(Arc::ptr_eq(binary.left(), &a)); assert!(Arc::ptr_eq(binary.right(), &b)); + + // Dictionary-encoded floats are normalized and compared by Arrow, as before. + let dictionary: ArrayRef = Arc::new(arrow::array::DictionaryArray::new( + Int32Array::from(vec![0, 1]), + Arc::new(Float64Array::from(vec![-0.0, 1.0])), + )); + let batch = self::batch(Arc::clone(&dictionary), dictionary); + let expr = spark_comparison( + Arc::clone(&a), + Operator::Lt, + Arc::clone(&b), + batch.schema().as_ref(), + FloatOperands::Normalize, + )?; + assert!(expr.downcast_ref::().is_some()); + assert_eq!(results(&expr, &batch), vec![Some(false), Some(false)]); Ok(()) } /// A bloom filter holds the bits of each value, so with `FloatOperands::Raw` an `=` against /// either zero probes the filter for both zeros, while other operators keep the literal. A - /// NaN literal is normalized together with the column instead, because no single literal - /// stands for every NaN a file can hold. + /// NaN literal gets a Spark comparison instead, because no single literal stands for every + /// NaN a file can hold. #[test] fn raw_float_comparison_literals() -> Result<()> { - use crate::float_semantics::NormalizeNaNAndZero; use datafusion::physical_expr::utils::{Guarantee, LiteralGuarantee}; use std::collections::HashSet; let schema = Schema::new(vec![ @@ -1161,12 +1452,9 @@ mod tests { &schema, FloatOperands::Raw, )?; - let binary = expr.downcast_ref::().unwrap(); - assert!(binary - .left() - .downcast_ref::() - .is_some()); - let folded = binary.right().downcast_ref::().unwrap(); + let comparison = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(comparison.left(), &column)); + let folded = comparison.right().downcast_ref::().unwrap(); assert!( matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == f64::NAN.to_bits()) ); @@ -1174,6 +1462,227 @@ mod tests { Ok(()) } + /// Columns that start at an offset, and scalar operands that only appear when the batch is + /// evaluated, so that planning cannot normalize them. + #[test] + fn float_operands_at_an_offset_and_runtime_scalars() -> Result<()> { + use crate::float_semantics::EDGE_VALUES; + let pairs: Vec<_> = EDGE_VALUES + .iter() + .flat_map(|&l| EDGE_VALUES.iter().map(move |&r| (l, r))) + .collect(); + // Pad both sides with a row in front, then slice it off. + let left: Vec<_> = std::iter::once(Some(5.0)) + .chain(pairs.iter().map(|p| p.0)) + .collect(); + let right: Vec<_> = std::iter::once(None) + .chain(pairs.iter().map(|p| p.1)) + .collect(); + for float32 in [false, true] { + let array = |values: &[Option]| -> ArrayRef { + let array: ArrayRef = if float32 { + Arc::new(arrow::array::Float32Array::from( + values + .iter() + .map(|v| v.map(|v| v as f32)) + .collect::>(), + )) + } else { + Arc::new(Float64Array::from(values.to_vec())) + }; + array.slice(1, values.len() - 1) + }; + let batch = batch(array(&left), array(&right)); + let schema = batch.schema(); + let (a, b) = columns(); + // A scalar that evaluation produces as it is, unlike a literal. + let runtime_scalar = |value: Option| -> Arc { + let value = if float32 { + ScalarValue::Float32(value.map(|v| v as f32)) + } else { + ScalarValue::Float64(value) + }; + Arc::new(Probe { + child: Arc::new(Literal::new(value)), + volatile: false, + calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + }) + }; + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, r)| spark_answer(op, float_ordering(l, r))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} b, f32={float32}"); + for &value in &EDGE_VALUES { + let expr = spark_comparison( + Arc::clone(&a), + op, + runtime_scalar(value), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, _)| spark_answer(op, float_ordering(l, value))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} {value:?}"); + let expr = spark_comparison( + runtime_scalar(value), + op, + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, _)| spark_answer(op, float_ordering(value, l))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "{value:?} {op} a"); + // Two scalars give a scalar. + for &other in &EDGE_VALUES { + let expr = spark_comparison( + runtime_scalar(value), + op, + runtime_scalar(other), + &schema, + FloatOperands::Normalize, + )?; + let ColumnarValue::Scalar(ScalarValue::Boolean(actual)) = + expr.evaluate(&batch)? + else { + panic!("{value:?} {op} {other:?} must be a Boolean scalar"); + }; + assert_eq!( + actual, + spark_answer(op, float_ordering(value, other)), + "{value:?} {op} {other:?}" + ); + } + } + } + } + Ok(()) + } + + /// An operand that evaluates to another layout than its Float32 or Float64 type, which the + /// planner does not produce, is normalized and compared by Arrow instead of failing. + #[test] + fn unexpected_layouts_fall_back_to_normalizing() -> Result<()> { + let dictionary: ArrayRef = Arc::new(arrow::array::DictionaryArray::new( + Int32Array::from(vec![0, 1, 0]), + Arc::new(Float64Array::from(vec![1.0, 2.0])), + )); + let result = compare_flat( + Operator::Lt, + ColumnarValue::Array(dictionary), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.5))), + )? + .into_array(3)?; + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + Ok(()) + } + + /// A list or struct literal on either side, a null one, and two literals, against the + /// orderings `spark_comparator` gives. + #[test] + fn nested_float_operands_with_literals() -> Result<()> { + use crate::float_semantics::spark_comparator; + let lists: ArrayRef = Arc::new(ListArray::from_iter_primitive::([ + Some(vec![Some(-0.0)]), + Some(vec![Some(0.0), Some(1.0)]), + Some(vec![Some(f64::from_bits(0xfff8_0000_0000_0000))]), + Some(vec![Some(f64::INFINITY)]), + Some(vec![None]), + Some(vec![]), + None, + ])); + let batch = batch(Arc::clone(&lists), Arc::clone(&lists)); + let schema = batch.schema(); + let (a, _) = columns(); + let null: Arc = + Arc::new(Literal::new(ScalarValue::try_from(lists.data_type())?)); + for op in COMPARISONS { + for row in 0..lists.len() { + let lit = literal(&lists, row); + let value = lit.evaluate(&batch)?.into_array(1)?; + let compare = spark_comparator(lists.as_ref(), value.as_ref())?; + let ordering = |i: usize| match (lists.is_null(i), value.is_null(0)) { + (false, false) => Some(Some(compare(i, 0))), + (true, true) => Some(None), + _ => None, + }; + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + assert!(expr.downcast_ref::().is_some() || !is_ordering(op)); + let expected: Vec<_> = (0..lists.len()) + .map(|i| spark_answer(op, ordering(i))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} row {row}"); + let reversed = spark_comparison( + Arc::clone(&lit), + op.swap().unwrap(), + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + assert_eq!( + results(&reversed, &batch), + expected, + "row {row} swapped {op} a" + ); + let both = spark_comparison( + Arc::clone(&lit), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + let ColumnarValue::Scalar(ScalarValue::Boolean(actual)) = both.evaluate(&batch)? + else { + panic!("two literals must give a Boolean scalar"); + }; + let ordering = if value.is_null(0) { + Some(None) + } else { + Some(Some(std::cmp::Ordering::Equal)) + }; + assert_eq!(actual, spark_answer(op, ordering), "row {row} {op} itself"); + } + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&null), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = (0..lists.len()) + .map(|i| spark_answer(op, if lists.is_null(i) { Some(None) } else { None })) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} NULL"); + } + Ok(()) + } + + fn is_ordering(op: Operator) -> bool { + !matches!(op, Operator::Eq | Operator::NotEq) + } + /// Ordering and null-safe comparisons of lists and structs match `spark_comparator`, which /// compares float leaves in Spark's order without normalizing them. #[test] @@ -1216,9 +1725,13 @@ mod tests { None, )) }; + // Also from an offset, which shifts the outer null buffers that the comparison skips. + let sliced = |array: &ArrayRef| array.slice(3, array.len() - 3); for (l, r) in [ (Arc::clone(&left), Arc::clone(&right)), (as_struct(&left), as_struct(&right)), + (sliced(&left), sliced(&right)), + (sliced(&as_struct(&left)), sliced(&as_struct(&right))), ] { let compare = spark_comparator(l.as_ref(), r.as_ref())?; let batch = batch(Arc::clone(&l), Arc::clone(&r)); diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index b128b165093..7dec89c61c0 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions}; +use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions, SparkComparison}; use arrow::array::{ downcast_primitive, Array, ArrayRef, AsArray, BooleanArray, GenericByteArray, PrimitiveArray, }; @@ -396,6 +396,8 @@ fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { } let node_is_infallible = if let Some(binary) = expr.downcast_ref::() { binary_is_infallible(binary, input_schema) + } else if let Some(comparison) = expr.downcast_ref::() { + comparison.is_infallible(input_schema) } else if let Some(cast) = expr.downcast_ref::() { cast.is_infallible(input_schema) } else if let Some(normalize) = expr.downcast_ref::() { @@ -836,6 +838,7 @@ fn merge_nulls(num_rows: usize, branches: &[Branch]) -> Option { #[cfg(test)] mod tests { use super::*; + use crate::{spark_comparison, FloatOperands}; use arrow::array::{ BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, StructArray, TimestampMicrosecondArray, @@ -1137,12 +1140,24 @@ mod tests { #[test] fn infallible_expressions() { + let list_of = |item: DataType| DataType::List(Arc::new(Field::new("item", item, true))); let schema = Schema::new(vec![ Field::new("i", DataType::Int64, true), Field::new("j", DataType::Int32, true), Field::new("d", DataType::Float64, true), + Field::new("e", DataType::Float64, true), Field::new("s", DataType::Utf8, true), Field::new("dec", DataType::Decimal128(10, 2), true), + Field::new("l", list_of(DataType::Float64), true), + Field::new("m", list_of(DataType::Float64), true), + Field::new( + "dictionary", + list_of(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Float64), + )), + true, + ), ]); let c = |name: &str| col(name, &schema).unwrap(); let cast = |e: Arc, to: DataType| -> Arc { @@ -1155,6 +1170,10 @@ mod tests { )) }; let infallible = |e: Arc| is_infallible(&e, &schema); + // A comparison as the planner builds it, which follows Spark's ordering for floats + let compare = |left: &str, op: Operator, right: Arc| { + spark_comparison(c(left), op, right, &schema, FloatOperands::Normalize).unwrap() + }; assert!(infallible(c("i"))); assert!(infallible(lit("x"))); @@ -1181,6 +1200,13 @@ mod tests { DataType::Float64, c("d") )))); + for op in [Operator::Lt, Operator::Eq, Operator::IsNotDistinctFrom] { + assert!(infallible(compare("d", op, lit(1.5))), "d {op} 1.5"); + assert!(infallible(compare("d", op, c("e"))), "d {op} e"); + } + for op in [Operator::Lt, Operator::IsNotDistinctFrom] { + assert!(infallible(compare("l", op, c("m"))), "l {op} m"); + } // Only a boolean can be negated, and only a float normalized assert!(!infallible(Arc::new(NotExpr::new(c("i"))))); @@ -1188,6 +1214,8 @@ mod tests { DataType::Int64, c("i") )))); + // A nested comparison fails on a dictionary-encoded leaf against a plain one + assert!(!infallible(compare("l", Operator::Lt, c("dictionary")))); // Integer division and remainder fail on a zero divisor assert!(!infallible(binary(c("i"), Operator::Divide, c("i")))); assert!(!infallible(binary(c("i"), Operator::Modulo, c("i")))); diff --git a/native/spark-expr/src/float_semantics/compare.rs b/native/spark-expr/src/float_semantics/compare.rs index 24c01e1e9df..a9d314cb730 100644 --- a/native/spark-expr/src/float_semantics/compare.rs +++ b/native/spark-expr/src/float_semantics/compare.rs @@ -17,6 +17,7 @@ use super::{compare_floats, has_float_leaf}; use arrow::array::{make_comparator, Array, AsArray, DynComparator, OffsetSizeTrait}; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; use datafusion::common::{internal_err, DFSchema, Result}; @@ -32,7 +33,7 @@ use std::ops::Range; /// arrays must have the same type, ignoring field names and nullability. pub fn spark_comparator(left: &dyn Array, right: &dyn Array) -> Result { check_types(left, right)?; - comparator(left, right, false) + comparator(left, right, false, true) } /// Builds a test of whether `left[i]` equals `right[j]` in the ordering of [`spark_comparator`]. @@ -42,7 +43,30 @@ pub fn spark_equality( right: &dyn Array, ) -> Result bool + Send + Sync>> { check_types(left, right)?; - let compare = comparator(left, right, true)?; + let compare = comparator(left, right, true, true)?; + Ok(Box::new(move |i, j| compare(i, j).is_eq())) +} + +/// [`spark_comparator`] for a caller that applies the top-level nulls of `left` and `right` +/// itself, as SQL comparisons do. A null list or struct is compared by the values under it instead +/// of sorting first, which saves testing both null buffers on every call. Nulls inside a list or +/// struct still sort first. +pub(crate) fn spark_comparator_ignoring_nulls( + left: &dyn Array, + right: &dyn Array, +) -> Result { + check_types(left, right)?; + comparator(left, right, false, false) +} + +/// [`spark_equality`] for a caller that applies the top-level nulls itself. See +/// [`spark_comparator_ignoring_nulls`]. +pub(crate) fn spark_equality_ignoring_nulls( + left: &dyn Array, + right: &dyn Array, +) -> Result bool + Send + Sync>> { + check_types(left, right)?; + let compare = comparator(left, right, true, false)?; Ok(Box::new(move |i, j| compare(i, j).is_eq())) } @@ -58,8 +82,15 @@ fn check_types(left: &dyn Array, right: &dyn Array) -> Result<()> { } } -/// With `equality` set, the comparator only has to tell equal from unequal values. -fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { +/// With `equality` set, the comparator only has to tell equal from unequal values. With `nulls` +/// unset, it ignores the nulls of `left` and `right` themselves, though not those of their +/// children. +fn comparator( + left: &dyn Array, + right: &dyn Array, + equality: bool, + nulls: bool, +) -> Result { if !has_float_leaf(left.data_type()) { let options = SortOptions { descending: false, @@ -68,38 +99,51 @@ fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result Ok(float_comparator::(left, right)), - (DataType::Float64, DataType::Float64) => Ok(float_comparator::(left, right)), - (DataType::List(_), DataType::List(_)) => list_comparator::(left, right, equality), + (DataType::Float32, DataType::Float32) => { + Ok(float_comparator::(left, right, nulls)) + } + (DataType::Float64, DataType::Float64) => { + Ok(float_comparator::(left, right, nulls)) + } + (DataType::List(_), DataType::List(_)) => { + list_comparator::(left, right, equality, nulls) + } (DataType::LargeList(_), DataType::LargeList(_)) => { - list_comparator::(left, right, equality) + list_comparator::(left, right, equality, nulls) } (DataType::FixedSizeList(_, _), DataType::FixedSizeList(_, _)) => { - fixed_size_list_comparator(left, right, equality) + fixed_size_list_comparator(left, right, equality, nulls) + } + (DataType::Struct(_), DataType::Struct(_)) => { + struct_comparator(left, right, equality, nulls) } - (DataType::Struct(_), DataType::Struct(_)) => struct_comparator(left, right, equality), (l, r) => internal_err!("Unsupported types for Spark comparison: {l} and {r}"), } } -fn float_comparator(left: &dyn Array, right: &dyn Array) -> DynComparator +fn float_comparator( + left: &dyn Array, + right: &dyn Array, + nulls: bool, +) -> DynComparator where T::Native: Float, { let l = left.as_primitive::().values().clone(); let r = right.as_primitive::().values().clone(); - nulls_first(left, right, move |i, j| compare_floats(l[i], r[j])) + nulls_first(left, right, nulls, move |i, j| compare_floats(l[i], r[j])) } fn list_comparator( left: &dyn Array, right: &dyn Array, equality: bool, + nulls: bool, ) -> Result { let (l, r) = (left.as_list::(), right.as_list::()); - let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality, true)?; let (l, r) = (l.offsets().clone(), r.offsets().clone()); - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { let left = l[i].as_usize()..l[i + 1].as_usize(); let right = r[j].as_usize()..r[j + 1].as_usize(); lexicographic(&compare, left, right, equality) @@ -110,11 +154,12 @@ fn fixed_size_list_comparator( left: &dyn Array, right: &dyn Array, equality: bool, + nulls: bool, ) -> Result { let (l, r) = (left.as_fixed_size_list(), right.as_fixed_size_list()); - let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality, true)?; let (l, r) = (l.value_length() as usize, r.value_length() as usize); - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { lexicographic(&compare, i * l..(i + 1) * l, j * r..(j + 1) * r, equality) })) } @@ -137,15 +182,20 @@ fn lexicographic( .unwrap_or(lengths) } -fn struct_comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { +fn struct_comparator( + left: &dyn Array, + right: &dyn Array, + equality: bool, + nulls: bool, +) -> Result { let fields = left .as_struct() .columns() .iter() .zip(right.as_struct().columns()) - .map(|(l, r)| comparator(l.as_ref(), r.as_ref(), equality)) + .map(|(l, r)| comparator(l.as_ref(), r.as_ref(), equality, true)) .collect::>>()?; - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { fields .iter() .map(|compare| compare(i, j)) @@ -154,18 +204,19 @@ fn struct_comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Res })) } -/// Orders nulls before values. A null slot never reaches `compare`, because the child values -/// under a null list or struct can be anything. +/// Orders nulls before values, unless `apply` is unset. A null slot then never reaches `compare`, +/// because the child values under a null list or struct can be anything. Without `apply`, the +/// caller masks the result for null slots: `compare` reads the values under them, which are in +/// bounds in any valid array. fn nulls_first( left: &dyn Array, right: &dyn Array, + apply: bool, compare: impl Fn(usize, usize) -> Ordering + Send + Sync + 'static, ) -> DynComparator { - let left = left.nulls().filter(|nulls| nulls.null_count() > 0).cloned(); - let right = right - .nulls() - .filter(|nulls| nulls.null_count() > 0) - .cloned(); + let present = |nulls: &&NullBuffer| apply && nulls.null_count() > 0; + let left = left.nulls().filter(present).cloned(); + let right = right.nulls().filter(present).cloned(); if left.is_none() && right.is_none() { return Box::new(compare); } @@ -311,6 +362,50 @@ mod tests { Ok(()) } + /// Ignoring the top-level nulls compares a null list or struct by the values under it, and + /// leaves every other slot, and nulls further down, as `spark_comparator` orders them. + #[test] + fn ignoring_nulls_compares_the_values_under_a_null_slot() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Float64, true)); + // The second row of each side is null, over `[-0.0]` and `[1.0]`, and the third holds an + // inner null. + let nulls = || Some(NullBuffer::from(vec![true, false, true])); + let list = |values: Vec>, nulls: Option| -> ArrayRef { + Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths([1, 1, 1]), + Arc::new(Float64Array::from(values)), + nulls, + )) + }; + let left = vec![Some(0.0), Some(-0.0), None]; + let right = vec![Some(-0.0), Some(1.0), Some(f64::NAN)]; + let struct_of = |values: Vec>| -> ArrayRef { + let list = list(values, None); + Arc::new(StructArray::new( + vec![Arc::new(Field::new("v", list.data_type().clone(), true))].into(), + vec![list], + nulls(), + )) + }; + for (l, r) in [ + (list(left.clone(), nulls()), list(right.clone(), nulls())), + (struct_of(left), struct_of(right)), + ] { + let masked = spark_comparator(l.as_ref(), r.as_ref())?; + let unmasked = spark_comparator_ignoring_nulls(l.as_ref(), r.as_ref())?; + let equal = spark_equality_ignoring_nulls(l.as_ref(), r.as_ref())?; + assert_eq!(masked(1, 1), Ordering::Equal); + assert_eq!(unmasked(1, 1), Ordering::Less, "{}", l.data_type()); + assert!(!equal(1, 1)); + for row in [0, 2] { + assert_eq!(unmasked(row, row), masked(row, row), "{}", l.data_type()); + assert_eq!(equal(row, row), masked(row, row).is_eq()); + } + } + Ok(()) + } + #[test] fn struct_fields_without_floats_break_ties() -> Result<()> { let structs = |floats: Vec, ints: Vec| -> ArrayRef { diff --git a/native/spark-expr/src/float_semantics/kernels.rs b/native/spark-expr/src/float_semantics/kernels.rs new file mode 100644 index 00000000000..aa04f7034bb --- /dev/null +++ b/native/spark-expr/src/float_semantics/kernels.rs @@ -0,0 +1,413 @@ +// 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. + +//! Comparison kernels for Float32 and Float64 arrays in Spark's SQL ordering, which read the +//! operands' value buffers as they are instead of normalizing copies of them first. +//! +//! Each operator is one test over a pair of values, written with `&` and `|` rather than `&&` and +//! `||` so that it stays free of branches, and the loops pack 64 results into a word at a time, +//! which the compiler vectorizes. + +use arrow::array::{Array, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::{BooleanBuffer, MutableBuffer, NullBuffer}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use arrow::util::bit_util; +use datafusion::common::{internal_err, Result, ScalarValue}; +use datafusion::logical_expr::Operator; +use num::Float; + +/// `left op right` for two Float32 or two Float64 arrays of the same length, where `op` is one of +/// `=`, `<>`, `<`, `<=`, `>`, `>=`, `IS DISTINCT FROM` and `IS NOT DISTINCT FROM`. Returns `None` +/// for any other pair of types, which the caller compares another way. +pub(crate) fn compare_float_arrays( + op: Operator, + left: &dyn Array, + right: &dyn Array, +) -> Result> { + match (left.data_type(), right.data_type()) { + (DataType::Float32, DataType::Float32) => { + arrays::(op, left.as_primitive(), right.as_primitive()).map(Some) + } + (DataType::Float64, DataType::Float64) => { + arrays::(op, left.as_primitive(), right.as_primitive()).map(Some) + } + _ => Ok(None), + } +} + +/// `left op right` for a Float32 or Float64 array and a scalar of the same type. Returns `None` for +/// any other pair of types. See [`compare_float_arrays`]. +pub(crate) fn compare_float_array_scalar( + op: Operator, + left: &dyn Array, + right: &ScalarValue, +) -> Result> { + match (left.data_type(), right) { + (DataType::Float32, ScalarValue::Float32(right)) => { + array_scalar::(op, left.as_primitive(), *right).map(Some) + } + (DataType::Float64, ScalarValue::Float64(right)) => { + array_scalar::(op, left.as_primitive(), *right).map(Some) + } + _ => Ok(None), + } +} + +fn arrays( + op: Operator, + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result +where + T::Native: Float, +{ + if left.len() != right.len() { + return internal_err!( + "Cannot compare arrays of different lengths, got {} and {}", + left.len(), + right.len() + ); + } + let (l, r) = (left.values().as_ref(), right.values().as_ref()); + // Spark's SQL ordering as single tests, the same as `compare_floats(a, b)` followed by + // `is_eq()`, `is_ne()`, `is_lt()`, `is_le()`, `is_gt()` or `is_ge()`. IEEE 754 comparisons + // already treat `-0.0` as equal to `0.0`; the NaN terms put every NaN above the other values + // and make all NaNs equal. + let values = match op { + Operator::Eq | Operator::IsNotDistinctFrom => { + pack_pairs(l, r, |a, b| (a == b) | (a.is_nan() & b.is_nan())) + } + Operator::NotEq | Operator::IsDistinctFrom => { + pack_pairs(l, r, |a, b| (a != b) & !(a.is_nan() & b.is_nan())) + } + Operator::Lt => pack_pairs(l, r, |a, b| (a < b) | (!a.is_nan() & b.is_nan())), + Operator::LtEq => pack_pairs(l, r, |a, b| (a <= b) | b.is_nan()), + Operator::Gt => pack_pairs(l, r, |a, b| (a > b) | (a.is_nan() & !b.is_nan())), + Operator::GtEq => pack_pairs(l, r, |a, b| (a >= b) | a.is_nan()), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }; + Ok(comparison_with_nulls( + op, + values, + left.nulls(), + right.nulls(), + )) +} + +fn array_scalar( + op: Operator, + left: &PrimitiveArray, + right: Option, +) -> Result +where + T::Native: Float, +{ + let len = left.len(); + let Some(s) = right else { + // Every comparison with a null is null, and only a null is not distinct from one. + return Ok(match op { + Operator::IsNotDistinctFrom => BooleanArray::new( + left.nulls() + .map(|nulls| !nulls.inner()) + .unwrap_or_else(|| BooleanBuffer::new_unset(len)), + None, + ), + Operator::IsDistinctFrom => BooleanArray::new( + left.nulls() + .map(|nulls| nulls.inner().clone()) + .unwrap_or_else(|| BooleanBuffer::new_set(len)), + None, + ), + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq => BooleanArray::new_null(len), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }); + }; + let l = left.values().as_ref(); + // With the scalar fixed, each operator reduces to one or two tests, chosen by whether the + // scalar is NaN. Every value sorts below a NaN or equals it, and a NaN sorts above any other + // value. + let values = match (op, s.is_nan()) { + (Operator::Eq | Operator::IsNotDistinctFrom, true) => pack(l, |a| a.is_nan()), + (Operator::Eq | Operator::IsNotDistinctFrom, false) => pack(l, |a| a == s), + (Operator::NotEq | Operator::IsDistinctFrom, true) => pack(l, |a| !a.is_nan()), + (Operator::NotEq | Operator::IsDistinctFrom, false) => pack(l, |a| a != s), + (Operator::Lt, true) => pack(l, |a| !a.is_nan()), + (Operator::Lt, false) => pack(l, |a| a < s), + (Operator::LtEq, true) => BooleanBuffer::new_set(len), + (Operator::LtEq, false) => pack(l, |a| a <= s), + (Operator::Gt, true) => BooleanBuffer::new_unset(len), + (Operator::Gt, false) => pack(l, |a| (a > s) | a.is_nan()), + (Operator::GtEq, true) => pack(l, |a| a.is_nan()), + (Operator::GtEq, false) => pack(l, |a| (a >= s) | a.is_nan()), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }; + Ok(comparison_with_nulls(op, values, left.nulls(), None)) +} + +/// Applies SQL's null rules to `values`, the comparison of every slot including the null ones. +/// A null operand makes `=`, `<>`, `<`, `<=`, `>` and `>=` null. `IS NOT DISTINCT FROM`, for which +/// `values` holds equality, and `IS DISTINCT FROM`, for which it holds inequality, compare a null +/// as a value instead: equal to another null and distinct from anything else. +pub(crate) fn comparison_with_nulls( + op: Operator, + values: BooleanBuffer, + left: Option<&NullBuffer>, + right: Option<&NullBuffer>, +) -> BooleanArray { + match op { + Operator::IsNotDistinctFrom => { + let values = match (left, right) { + (None, None) => values, + (Some(nulls), None) | (None, Some(nulls)) => &values & nulls.inner(), + (Some(l), Some(r)) => { + let (l, r) = (l.inner(), r.inner()); + &(&(&values & l) & r) | &!&(l | r) + } + }; + BooleanArray::new(values, None) + } + Operator::IsDistinctFrom => { + let values = match (left, right) { + (None, None) => values, + (Some(nulls), None) | (None, Some(nulls)) => &values | &!nulls.inner(), + (Some(l), Some(r)) => { + let (l, r) = (l.inner(), r.inner()); + &(&(&values & l) & r) | &(l ^ r) + } + }; + BooleanArray::new(values, None) + } + _ => BooleanArray::new(values, NullBuffer::union(left, right)), + } +} + +/// Packs `test(value)` for each value into a bitmap, 64 values to a word. +#[inline] +fn pack(values: &[T], test: impl Fn(T) -> bool) -> BooleanBuffer { + let mut buffer = MutableBuffer::new(bit_util::ceil(values.len(), 64) * 8); + let (chunks, remainder) = values.as_chunks::<64>(); + for chunk in chunks { + let mut word = 0u64; + for (bit, &value) in chunk.iter().enumerate() { + word |= (test(value) as u64) << bit; + } + buffer.push(word); + } + if !remainder.is_empty() { + let mut word = 0u64; + for (bit, &value) in remainder.iter().enumerate() { + word |= (test(value) as u64) << bit; + } + buffer.push(word); + } + BooleanBuffer::new(buffer.into(), 0, values.len()) +} + +/// Packs `test(left[i], right[i])` for each pair into a bitmap. See [`pack`]. +#[inline] +fn pack_pairs(left: &[T], right: &[T], test: impl Fn(T, T) -> bool) -> BooleanBuffer { + debug_assert_eq!(left.len(), right.len()); + let mut buffer = MutableBuffer::new(bit_util::ceil(left.len(), 64) * 8); + let (left_chunks, left_remainder) = left.as_chunks::<64>(); + let (right_chunks, right_remainder) = right.as_chunks::<64>(); + for (l, r) in left_chunks.iter().zip(right_chunks) { + let mut word = 0u64; + for bit in 0..64 { + word |= (test(l[bit], r[bit]) as u64) << bit; + } + buffer.push(word); + } + if !left_remainder.is_empty() { + let mut word = 0u64; + for (bit, (&l, &r)) in left_remainder.iter().zip(right_remainder).enumerate() { + word |= (test(l, r) as u64) << bit; + } + buffer.push(word); + } + BooleanBuffer::new(buffer.into(), 0, left.len()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::float_semantics::{compare_floats, EDGE_VALUES}; + use arrow::array::{ArrayRef, Float32Array, Float64Array}; + use std::cmp::Ordering; + use std::sync::Arc; + + const OPERATORS: [Operator; 8] = [ + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + Operator::IsDistinctFrom, + Operator::IsNotDistinctFrom, + ]; + + /// Spark's answer for `left op right` from [`compare_floats`]. + fn expected(op: Operator, left: Option, right: Option) -> Option { + let ordering = match (left, right) { + (Some(l), Some(r)) => Some(compare_floats(l, r)), + _ => None, + }; + match op { + Operator::IsNotDistinctFrom => Some(match (left, right) { + (None, None) => true, + (Some(_), Some(_)) => ordering == Some(Ordering::Equal), + _ => false, + }), + Operator::IsDistinctFrom => { + expected(Operator::IsNotDistinctFrom, left, right).map(|v| !v) + } + _ => ordering.map(|ordering| match op { + Operator::Eq => ordering.is_eq(), + Operator::NotEq => ordering.is_ne(), + Operator::Lt => ordering.is_lt(), + Operator::LtEq => ordering.is_le(), + Operator::Gt => ordering.is_gt(), + Operator::GtEq => ordering.is_ge(), + _ => unreachable!(), + }), + } + } + + /// `values` as a Float64 array, or as a Float32 array holding the same values. + fn array(values: &[Option], float32: bool) -> ArrayRef { + if float32 { + Arc::new(Float32Array::from( + values + .iter() + .map(|v| v.map(|v| v as f32)) + .collect::>(), + )) + } else { + Arc::new(Float64Array::from(values.to_vec())) + } + } + + fn scalar(value: Option, float32: bool) -> ScalarValue { + if float32 { + ScalarValue::Float32(value.map(|v| v as f32)) + } else { + ScalarValue::Float64(value) + } + } + + /// Every pair of edge values, under every operator, as two arrays and as an array and a + /// scalar, with enough rows that the comparison spans whole words and a partial one. + #[test] + fn every_pair_of_edge_values_matches_compare_floats() -> Result<()> { + // Repeat the pairs so that a run covers several 64-value words and a remainder. + let pairs: Vec<(Option, Option)> = EDGE_VALUES + .iter() + .flat_map(|&l| EDGE_VALUES.iter().map(move |&r| (l, r))) + .cycle() + .take(EDGE_VALUES.len() * EDGE_VALUES.len() * 3 + 7) + .collect(); + let left: Vec<_> = pairs.iter().map(|p| p.0).collect(); + let right: Vec<_> = pairs.iter().map(|p| p.1).collect(); + for float32 in [false, true] { + // Float32 holds every edge value exactly, NaNs included, as the same class of value. + let (l, r) = (array(&left, float32), array(&right, float32)); + for op in OPERATORS { + let actual = compare_float_arrays(op, l.as_ref(), r.as_ref())?.unwrap(); + let want: Vec<_> = pairs.iter().map(|&(a, b)| expected(op, a, b)).collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "a {op} b, f32={float32}" + ); + for &value in &EDGE_VALUES { + let s = scalar(value, float32); + let actual = compare_float_array_scalar(op, l.as_ref(), &s)?.unwrap(); + let want: Vec<_> = left.iter().map(|&a| expected(op, a, value)).collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "a {op} {value:?}, f32={float32}" + ); + } + } + } + Ok(()) + } + + /// A sliced array starts its values and its null bitmap at an offset. + #[test] + fn sliced_inputs() -> Result<()> { + let values: Vec> = (0..200) + .map(|i| EDGE_VALUES[(i * 7) % EDGE_VALUES.len()]) + .collect(); + let others: Vec> = (0..200) + .map(|i| EDGE_VALUES[(i * 3 + 1) % EDGE_VALUES.len()]) + .collect(); + let (l, r) = (array(&values, false), array(&others, false)); + for (offset, len) in [(1, 150), (63, 70), (64, 64), (130, 5), (3, 0)] { + let (ls, rs) = (l.slice(offset, len), r.slice(offset + 2, len)); + for op in OPERATORS { + let actual = compare_float_arrays(op, ls.as_ref(), rs.as_ref())?.unwrap(); + let want: Vec<_> = (0..len) + .map(|i| expected(op, values[offset + i], others[offset + 2 + i])) + .collect(); + assert_eq!(actual.iter().collect::>(), want, "{op} at {offset}"); + let actual = + compare_float_array_scalar(op, ls.as_ref(), &ScalarValue::Float64(None))? + .unwrap(); + let want: Vec<_> = (0..len) + .map(|i| expected(op, values[offset + i], None)) + .collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "{op} null at {offset}" + ); + } + } + Ok(()) + } + + #[test] + fn other_types_are_left_to_the_caller() -> Result<()> { + let floats: ArrayRef = Arc::new(Float64Array::from(vec![1.0])); + let narrow: ArrayRef = Arc::new(Float32Array::from(vec![1.0])); + assert!(compare_float_arrays(Operator::Lt, floats.as_ref(), narrow.as_ref())?.is_none()); + assert!(compare_float_array_scalar( + Operator::Lt, + floats.as_ref(), + &ScalarValue::Float32(Some(1.0)) + )? + .is_none()); + assert!( + compare_float_array_scalar(Operator::Lt, floats.as_ref(), &ScalarValue::Null)? + .is_none() + ); + assert!(compare_float_arrays(Operator::Plus, floats.as_ref(), floats.as_ref()).is_err()); + let different_lengths: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + assert!( + compare_float_arrays(Operator::Lt, floats.as_ref(), different_lengths.as_ref()) + .is_err() + ); + Ok(()) + } +} diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs index b295f3ea545..997f8b24648 100644 --- a/native/spark-expr/src/float_semantics/mod.rs +++ b/native/spark-expr/src/float_semantics/mod.rs @@ -39,21 +39,25 @@ //! A native expression must follow the rule of the Spark function it replaces, so build it from //! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes, hashes or //! compares the values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's -//! total order agrees with `compareDoubles`. Comparison operands go through -//! [`normalize_comparison_operand`]. Where Comet compares values itself, or has to return the -//! original bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], -//! [`spark_comparator`] or [`spark_equality`]. Non-canonical NaNs are not a corner case: on x86-64 -//! every NaN that arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. +//! total order agrees with `compareDoubles`. Where Comet compares values itself, or has to return +//! the original bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], +//! [`spark_comparator`] or [`spark_equality`]. The comparison operators compare in place too, with +//! the kernels in `kernels` for flat floats and with [`spark_comparator`] and [`spark_equality`] +//! for lists and structs. Non-canonical NaNs are not a corner case: on x86-64 every NaN that +//! arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. mod compare; +mod kernels; mod normalize; pub use compare::{spark_comparator, spark_equality}; -pub(crate) use normalize::is_nested_with_float_leaf; +pub(crate) use compare::{spark_comparator_ignoring_nulls, spark_equality_ignoring_nulls}; +pub(crate) use kernels::{compare_float_array_scalar, compare_float_arrays, comparison_with_nulls}; pub use normalize::{ has_float_leaf, normalize_comparison_operand, normalize_floats, normalize_nested_floats, NormalizeNaNAndZero, NormalizeNestedFloats, }; +pub(crate) use normalize::{is_nested_with_float_leaf, normalize_float_scalar}; use num::Float; use std::cmp::Ordering; diff --git a/native/spark-expr/src/float_semantics/normalize.rs b/native/spark-expr/src/float_semantics/normalize.rs index ea575aa0454..33b1c27b5ff 100644 --- a/native/spark-expr/src/float_semantics/normalize.rs +++ b/native/spark-expr/src/float_semantics/normalize.rs @@ -257,7 +257,7 @@ pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { } /// Applies [`normalize_float`] to a Float32 or Float64 scalar. Other scalars are returned as is. -fn normalize_float_scalar(value: ScalarValue) -> ScalarValue { +pub(crate) fn normalize_float_scalar(value: ScalarValue) -> ScalarValue { match value { ScalarValue::Float32(v) => ScalarValue::Float32(v.map(normalize_float)), ScalarValue::Float64(v) => ScalarValue::Float64(v.map(normalize_float)), From 00a4b422f0ed6560ef76dce004b94c3697613108 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 9 Oct 2026 05:35:33 +0000 Subject: [PATCH 40/89] test: drain the listener bus before registering Iceberg test listeners (#6809) Co-authored-by: Isaac --- .../org/apache/comet/CometIcebergRewriteActionSuite.scala | 1 + .../test/scala/org/apache/comet/CometIcebergTestBase.scala | 4 ++++ .../scala/org/apache/comet/CometIcebergWriteActionSuite.scala | 4 ++++ .../spark/sql/benchmark/CometIcebergWriteBenchmark.scala | 1 + 4 files changed, 10 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala index 506e9370fc7..0de8ab12c23 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala @@ -373,6 +373,7 @@ class CometIcebergRewriteActionSuite extends CometTestBase with CometIcebergTest case _ => } } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.sparkContext.addSparkListener(listener) try { body diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala index 4fe2803db73..ba34b0f951b 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala @@ -237,6 +237,9 @@ trait CometIcebergTestBase { this: CometTestBase => override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = if (includeFailures) captured += qe.executedPlan } + // Events from earlier queries may still be queued; drain them so they do not reach the + // listener. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action @@ -259,6 +262,7 @@ trait CometIcebergTestBase { this: CometTestBase => override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = captured += qe.executedPlan } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { val error = diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 45a4fe597fd..36a7a235246 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -3411,6 +3411,8 @@ class CometIcebergWriteActionSuite override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = if (taskEnd.reason == Success) JobAbortGate.taskFinished() } + // Drain the setup writes' task-end events first, or they can count toward the gate. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.sparkContext.addSparkListener(listener) try { val run = () => @@ -4309,6 +4311,8 @@ class CometIcebergWriteActionSuite val listener = new IcebergWriteReportListener( new SparkConf() .set(CometConf.COMET_ICEBERG_WRITE_REPORT_DIR.key, reportDir.getAbsolutePath)) + // Drain the setup writes' events first, or they can reach the listener and be reported. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala index fbcde2d46d6..7ab0d9d2161 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala @@ -529,6 +529,7 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { qe: QueryExecution, exception: Exception): Unit = {} } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action From a2ace46c5a0d3235544a810f0f81b2d219c0b01f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 17:18:33 +0000 Subject: [PATCH 41/89] perf: run Spark's null guard for a Scala UDF inside the codegen kernel (#6715) * perf: run Spark's null guard for a Scala UDF inside the codegen kernel Spark's HandleNullInputsForUDF rule wraps a Scala UDF that has a primitive parameter over a nullable input as `if (isnull(a) or ...) null else f(knownnotnull(a), ...)`. The serde sent only the UDF to the codegen dispatcher and ran the If natively, as a CASE that filters the batch for each branch and merges the results, which cost more than the call it protects. CometIf now recognizes that exact shape, including the `false` form that ReplaceNullWithFalseInPredicate leaves in predicates, and sends the whole If to the dispatcher, so the null check is a branch in the kernel's loop. When the dispatcher cannot take the whole guard, the If is converted natively as before and no fallback reason is recorded against it. * fix: keep a null guard native when an expression in it is disabled Converting the guard natively checks spark.comet.expression..enabled on the UDF and on every node of the predicate, and a disabled one takes the projection to Spark. Dispatching the guard whole skipped those checks, so with ScalaUDF, IsNull, Or, or an expression in a guarded argument disabled, the projection stayed in Comet. emitNullGuardDispatch now declines the guard in that case, and the native conversion records the same fallback reason as before. * test: bind the dispatched null guard in the Scala UDF dispatcher benchmark CometScalaUDFDispatchBenchmark, added by #6709, said that Spark's null check around a UDF with a primitive parameter runs natively around the dispatched call, and its JVM layer cases bound the ScalaUDF alone as the subtree the serde would dispatch. The check now runs in the kernel with the call. The scaladoc says so, and the JVM layer cases bind the projected expression, which for a primitive parameter is the check together with the call. * perf: check a repeated argument once when matching a UDF's null guard BooleanSimplification drops a repeated isnull from the predicate that HandleNullInputsForUDF builds, so with a column at two primitive positions, f(a, a) optimizes to if (isnull(a)) null else f(knownnotnull(a), knownnotnull(a)). The predicate rebuilt for the match kept every repeat, so those guards stayed native. It is now rebuilt from an ExpressionSet of the guarded arguments. ExpressionSet keeps every nondeterministic argument, so a guard over one still stays native. The first null guard test also covers addBoth(a, a), addThree(a, b, a) and addBoth(a, 1L), the last with an argument that cannot be null. --- .../apache/comet/serde/CometScalaUDF.scala | 109 ++++++++++--- .../org/apache/comet/serde/conditional.scala | 9 ++ .../org/apache/comet/CometCodegenSuite.scala | 146 +++++++++++++++++- .../CometScalaUDFDispatchBenchmark.scala | 19 ++- 4 files changed, 250 insertions(+), 33 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index c79b4839bb3..e07553e2c14 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.util.control.NonFatal import org.apache.spark.SparkEnv -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, Literal, Or, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf @@ -58,6 +58,64 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = emitJvmCodegenDispatch(expr, inputs, binding) + /** + * Emits Spark's null guard around a Scala UDF, together with the UDF, as one dispatcher call. + * Returns `None` when `expr` is not that guard, when an expression in it is disabled, or when + * the dispatcher cannot take it, and records nothing on `expr` in that case: the caller + * converts `expr` as usual, which sends the UDF to the dispatcher on its own, and the UDF + * records its own fallback reason if it cannot go. + * + * Spark's `HandleNullInputsForUDF` rule wraps a UDF with a primitive parameter over a nullable + * input as `if (isnull(a) or isnull(b)) null else f(knownnotnull(a), knownnotnull(b))`, one + * `IsNull` per such argument. Converted natively, the `If` is a `CASE` that splits the batch by + * the predicate and merges the two branches back together, which costs more than the call it + * protects. In the kernel it is a branch per row. + * + * Only that shape matches: the predicate must be the one the rule builds from the UDF's guarded + * arguments, each checked once. The kernel then runs no expression it would not run for the UDF + * alone, apart from the guard's own `if`, `or` and `isnull`, and it runs Spark's code for all + * of them. In a filter or join condition, and in the predicate of a conditional, + * `ReplaceNullWithFalseInPredicate` replaces the null with `false`, so that form matches too. + */ + def emitNullGuardDispatch(expr: If, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = + if (isNullGuard(expr) && isEnabled(expr)) { + tryJvmCodegenDispatch(expr, inputs, binding).toOption + } else { + None + } + + /** + * Whether converting `guard` natively would pass every `spark.comet.expression..enabled` + * check it makes: on the UDF, and on each node of the predicate and the true branch, the + * guarded arguments included. The dispatcher makes none of these checks, so a guard with a + * disabled expression stays native, where the conversion records the fallback reason. The UDF's + * arguments are not checked, because the dispatcher takes them unchecked with or without the + * guard. + */ + private def isEnabled(guard: If): Boolean = { + def disabled(e: Expression): Boolean = + QueryPlanSerde.exprSerdeMap.get(e.getClass).exists { serde => + !CometConf.isExprEnabled( + serde.asInstanceOf[CometExpressionSerde[Expression]].getExprConfigName(e)) + } + !disabled(guard.falseValue) && !guard.predicate.exists(disabled) && + !guard.trueValue.exists(disabled) + } + + private def isNullGuard(expr: If): Boolean = expr match { + case If(predicate, Literal(null, _) | Literal.FalseLiteral, udf: ScalaUDF) => + val guarded = udf.inputPrimitives.zip(udf.children).collect { + case (true, KnownNotNull(arg)) => arg + } + // The optimizer drops a repeated `isnull`, as in `f(a, a)`, so check each argument once. + // `semanticEquals` is false for a nondeterministic argument, so that guard stays native. + ExpressionSet(guarded).toSeq + .map(IsNull(_)) + .reduceLeftOption[Expression](Or) + .exists(_.semanticEquals(predicate)) + case _ => false + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -73,14 +131,28 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { def emitJvmCodegenDispatch( expr: Expression, inputs: Seq[Attribute], - binding: Boolean): Option[Expr] = { + binding: Boolean): Option[Expr] = + tryJvmCodegenDispatch(expr, inputs, binding) match { + case Right(proto) => Some(proto) + case Left(reason) => + withFallbackReason(expr, reason) + None + } + + /** + * [[emitJvmCodegenDispatch]] without the fallback tagging: returns the reason the dispatcher + * cannot take `expr` rather than recording it on `expr`, for a caller that can still convert + * `expr` another way. + */ + private def tryJvmCodegenDispatch( + expr: Expression, + inputs: Seq[Attribute], + binding: Boolean): Either[String, Expr] = { val exprName = CometExplainInfo.exprDisplayName(expr) if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { - withFallbackReason( - expr, + return Left( s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " + "no native path so the plan falls back to Spark") - return None } // `RuntimeReplaceable` expressions (e.g. Spark 4's `StructsToJson`) have a `doGenCode` that @@ -97,8 +169,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // under a dispatched `map(...)` would run in it too. Check the whole tree. CometInvokeTargets.declineReason(target) match { case Some(reason) => - withFallbackReason(expr, s"$exprName: codegen dispatch: $reason") - return None + return Left(s"$exprName: codegen dispatch: $reason") case None => } @@ -107,12 +178,10 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val attrs = target.collect { case a: AttributeReference => a }.distinct val boundExpr = BindReferences.bindReference(target, AttributeSeq(attrs)) - // Gate at plan time. Surface the reason via withFallbackReason rather than crashing Janino - // at execute. + // Gate at plan time. Surface the reason as a fallback rather than crashing Janino at execute. CometBatchKernelCodegen.canHandle(boundExpr) match { case Some(reason) => - withFallbackReason(expr, s"$exprName: $reason") - return None + return Left(s"$exprName: $reason") case None => } @@ -139,36 +208,28 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // unserializable object graph as one of several exception types depending on where in // the graph it trips, and a custom `writeObject` can throw anything. case NonFatal(e) => - withFallbackReason( - expr, + return Left( s"$exprName: codegen dispatch: expression could not be closure-serialized " + s"(${e.getClass.getSimpleName}: ${e.getMessage})") - return None } // Arg 0 is a digest of the bytes. The dispatcher finds each batch's kernel by it, so it reads // the bytes only to compile on a cache miss instead of copying and hashing them for every // batch. val exprArgs = Seq(CometScalaUDFCodegen.digest(bytes), bytes).map { payload => exprToProtoInternal(Literal(payload, BinaryType), inputs, binding).getOrElse { - withFallbackReason( - expr, + return Left( s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + "expression payload") - return None } } val dataArgs = attrs.map { a => exprToProtoInternal(a, inputs, binding).getOrElse { - withFallbackReason(expr, s"$exprName: codegen dispatch: could not serialize data arg $a") - return None + return Left(s"$exprName: codegen dispatch: could not serialize data arg $a") } } val returnTypeProto = serializeDataType(expr.dataType).getOrElse { - withFallbackReason( - expr, - s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") - return None + return Left(s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") } val udfBuilder = ExprOuterClass.JvmScalarUdf @@ -199,7 +260,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val _ = withCodegenDispatchExpr(expr, CometExplainInfo.exprDisplayName(node)) case _ => } - Some( + Right( ExprOuterClass.Expr .newBuilder() .setJvmScalarUdf(udfBuilder.build()) diff --git a/spark/src/main/scala/org/apache/comet/serde/conditional.scala b/spark/src/main/scala/org/apache/comet/serde/conditional.scala index df94be8c93c..6728c9fcf07 100644 --- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala +++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala @@ -27,6 +27,15 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal object CometIf extends CometExpressionSerde[If] { override def convert( + expr: If, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + // Spark's null guard around a Scala UDF goes to the dispatcher with the UDF when it can. + CometScalaUDF + .emitNullGuardDispatch(expr, inputs, binding) + .orElse(convertNative(expr, inputs, binding)) + + private def convertNative( expr: If, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 7af83a5494c..248e903fbf7 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, If, IsNull, LessThan, Literal, MapConcat, Or, Rand, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.util.GenericArrayData @@ -1028,6 +1028,150 @@ class CometCodegenSuite } } + test("a primitive ScalaUDF's null guard runs in the UDF's kernel (#6704)") { + // Spark's `HandleNullInputsForUDF` wraps each call below as + // `if (isnull(a) or ...) null else f(knownnotnull(a), ...)`. Run natively, the `if` is a + // `CASE` that splits each batch on the predicate and merges the two halves back, so the guard + // goes to the kernel with the call, and neither `if` nor `isnull` is native. + spark.udf.register("plusOne", (x: Long) => x + 1) + spark.udf.register("combine", (a: Long, b: Int, s: String) => s"$s:${a + b}") + spark.udf.register("addBoth", (x: Long, y: Long) => x + y) + spark.udf.register("addThree", (x: Long, y: Int, z: Long) => x + y + z) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b INT, s STRING) USING parquet") + sql( + "INSERT INTO t VALUES (1, 10, 'x'), (NULL, 20, 'y'), (3, NULL, 'z'), " + + "(NULL, NULL, NULL), (-5, 7, NULL)") + val guard = Seq("if", "isnull") + checkSparkAnswerAndImpl(sql("SELECT plusOne(a) FROM t"), dispatched = "plusone" +: guard) + checkSparkAnswerAndImpl( + sql("SELECT max(plusOne(a)), count(plusOne(a)) FROM t"), + dispatched = "plusone" +: guard) + // An argument that is an expression is checked as it is. + checkSparkAnswerAndImpl( + sql("SELECT plusOne(a * 2) FROM t"), + dispatched = "plusone" +: guard) + // One `isnull` per primitive parameter, joined by `or`. The `String` one is not checked. + checkSparkAnswerAndImpl( + sql("SELECT combine(a, b, s) FROM t"), + dispatched = Seq("combine", "or") ++ guard) + // The optimizer drops a repeated `isnull`, so a column passed twice is checked once. + checkSparkAnswerAndImpl(sql("SELECT addBoth(a, a) FROM t"), dispatched = "addboth" +: guard) + checkSparkAnswerAndImpl( + sql("SELECT addThree(a, b, a) FROM t"), + dispatched = Seq("addthree", "or") ++ guard) + // An argument that cannot be null gets no `isnull`. + checkSparkAnswerAndImpl( + sql("SELECT addBoth(a, 1L) FROM t"), + dispatched = "addboth" +: guard) + } + } + + test("a boolean ScalaUDF's null guard runs in the kernel with false for null (#6704)") { + // In a filter, and in the predicate of a conditional, `ReplaceNullWithFalseInPredicate` + // rewrites the guard to `if (isnull(a)) false else f(knownnotnull(a))`. + spark.udf.register("isPositive", (x: Long) => x > 0) + withTypedCol("BIGINT", "1", "NULL", "-3", "0", "7") { + checkSparkAnswerAndImpl( + sql("SELECT c FROM t WHERE isPositive(c)"), + dispatched = Seq("ispositive", "if", "isnull")) + // The query's own `if` is native, so only the guard's `isnull` is checked. + checkSparkAnswerAndImpl( + sql("SELECT IF(isPositive(c), 'yes', 'no') FROM t"), + dispatched = Seq("ispositive", "isnull")) + } + } + + test("a nondeterministic ScalaUDF under its guard is called only for non-null rows (#6704)") { + // Each result counts the calls made so far in the task, so a call on a null row would shift + // every result after it. + import org.apache.spark.sql.functions.udf + var calls = 0L + val countCalls = udf { (x: Long) => + calls += 1 + x * 1000 + calls + } + spark.udf.register("countCalls", countCalls.asNondeterministic()) + withTypedCol("BIGINT", "1", "NULL", "2", "NULL", "NULL", "3") { + checkSparkAnswerAndImpl( + sql("SELECT countCalls(c) FROM t"), + dispatched = Seq("countcalls", "if", "isnull")) + } + } + + test("a null guard the dispatcher cannot take whole stays native around its UDF (#6704)") { + spark.udf.register("plusOne", (x: Long) => x + 1) + withTypedCol("BIGINT", "1", "NULL", "-3") { + // `canHandle` counts each reference to a column against `maxFields`, and the guard reads + // `c` a second time: the UDF alone counts its output and one input, 2, and the guard 3. At + // 2 the guard is refused while the UDF is not, so the guard has to stay a native `if` + // around the dispatched UDF rather than take the projection to Spark. + withSQLConf("spark.sql.codegen.maxFields" -> "2") { + checkSparkAnswerAndImpl( + sql("SELECT plusOne(c) FROM t"), + native = Seq("if", "isnull"), + dispatched = Seq("plusone")) + } + // With the dispatcher off, the fallback reason is the UDF's own, and none is recorded for + // a guard that the query does not contain. + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + sql("SELECT plusOne(c) FROM t"), + s"plusone: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") + val reasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + assert(!reasons.exists(_.startsWith("if:")), s"unexpected fallback reasons: $reasons") + } + } + } + + test("a disabled expression in a null guard keeps the guard native (#6704)") { + // Converting the guard natively checks `spark.comet.expression..enabled` on the UDF and + // on each node of the predicate, the guarded argument's included, and a disabled one takes + // the projection to Spark. Dispatching the guard has to make the same checks. + spark.udf.register("plusOne", (x: Long) => x + 1) + spark.udf.register("addBoth", (a: Long, b: Long) => a + b) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b BIGINT) USING parquet") + sql("INSERT INTO t VALUES (1, 10), (NULL, 20), (-3, NULL)") + Seq( + "ScalaUDF" -> "SELECT plusOne(a) FROM t", + "IsNull" -> "SELECT plusOne(a) FROM t", + "Or" -> "SELECT addBoth(a, b) FROM t", + "Multiply" -> "SELECT plusOne(a * 2) FROM t").foreach { case (name, query) => + val key = CometConf.getExprEnabledConfigKey(name) + withSQLConf(key -> "false") { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(sql(query), s"Set $key=true") + assert( + collect(cometPlan) { case p: CometProjectExec => p }.isEmpty, + s"$query stayed in Comet with $key=false:\n$cometPlan") + } + } + } + } + + test("only the null guard Spark builds goes to the kernel with its UDF (#6704)") { + spark.udf.register("plusOne", (x: Long) => x + 1) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b BIGINT) USING parquet") + val plan = sql("SELECT plusOne(a), b FROM t").queryExecution.optimizedPlan + val guard = plan.expressions.flatMap(_.collect { case g: If => g }).head + val Seq(a, b) = plan.collectLeaves().head.output + def root(expr: Expression): ExprStructCase = + QueryPlanSerde.exprToProto(expr, Seq(a, b)).get.getExprStructCase + assert(root(guard) == ExprStructCase.JVM_SCALAR_UDF) + // A predicate that checks a column the call does not read, or one besides its arguments. + assert(root(guard.copy(predicate = IsNull(b))) == ExprStructCase.IF) + assert(root(guard.copy(predicate = Or(guard.predicate, IsNull(b)))) == ExprStructCase.IF) + // A branch that is not null. + assert(root(guard.copy(trueValue = Literal(0L))) == ExprStructCase.IF) + // A nondeterministic argument, which the guard and the call would each evaluate. + val random = guard.transformUp { case r: AttributeReference => + If(LessThan(Rand(Literal(1L)), Literal(2.0)), r, Literal(null, LongType)) + } + assert(root(random) == ExprStructCase.IF) + } + } + // Arrow Java ignores ArrowArray.offset on import, so the bridge has to zero a sliced boolean's // offset before handing it over, at the top level and inside a struct. // https://github.com/apache/datafusion-comet/issues/6288 diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala index 91a9337f26a..3bb3076016f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala @@ -22,7 +22,7 @@ package org.apache.spark.sql.benchmark import org.apache.arrow.vector.{BigIntVector, ValueVector, VarBinaryVector} import org.apache.spark.SparkEnv import org.apache.spark.benchmark.Benchmark -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSeq, BindReferences, Expression, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, AttributeSeq, BindReferences, Expression} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.comet.{CometArrowAllocator, CometConf} @@ -37,9 +37,9 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * Part 1 runs `SELECT max(f(c))` over a Parquet `bigint` column end to end at several batch * sizes. The total of a cost paid once per batch grows as batches shrink, and the total of a cost * paid per row does not. A tenth of the rows are null. Spark wraps a UDF with a primitive - * parameter in a null check, `if(isnull(c), null, f(knownnotnull(c)))`, which runs natively - * around the dispatched call; a UDF with a boxed parameter gets no such wrapper and converts its - * input through Spark's encoder instead. + * parameter in a null check, `if(isnull(c), null, f(knownnotnull(c)))`, which runs in the + * dispatcher's kernel with the call; a UDF with a boxed parameter gets no such wrapper and + * converts its input through Spark's encoder instead. * * Part 2 calls the dispatcher directly on the same in-memory batch over and over, with no Spark * query and no native code, and does the same with the generated kernel it runs. The difference @@ -124,12 +124,15 @@ object CometScalaUDFDispatchBenchmark extends CometBenchmarkBase { benchmark.run() } - /** The bound UDF subtree the serde would dispatch, and its closure-serialized bytes. */ + /** + * The bound subtree the serde would dispatch for `f`, and its closure-serialized bytes. For a + * primitive parameter that is Spark's null check together with the call. + */ private def boundUdf(f: String): (Expression, Array[Byte]) = { val plan = spark.sql(s"SELECT $f FROM parquetV1Table").queryExecution.optimizedPlan - val udf = plan.expressions.flatMap(_.collect { case u: ScalaUDF => u }).head - val attrs = udf.collect { case a: AttributeReference => a }.distinct - val bound = BindReferences.bindReference(udf: Expression, AttributeSeq(attrs)) + val dispatched = plan.expressions.collectFirst { case Alias(child, _) => child }.get + val attrs = dispatched.collect { case a: AttributeReference => a }.distinct + val bound = BindReferences.bindReference(dispatched, AttributeSeq(attrs)) val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(bound) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) From 6960a48ae8d302ed6b836f3213367ebecebb9515 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 17:46:47 +0000 Subject: [PATCH 42/89] fix: truncate per-row date_trunc formats with the literal-format kernel (#6354) * fix: truncate per-row date_trunc formats with the literal-format kernel * fix: pass the SECOND/MILLISECOND overflow policy through the format-column path * fix: handle a NULL format and a literal timestamp with a format column in date_trunc --- .../expression-audits/datetime_funcs.md | 2 +- native/proto/src/proto/expr.proto | 4 +- native/spark-expr/benches/timestamp_trunc.rs | 61 +- .../src/datetime_funcs/timestamp_trunc.rs | 86 ++- native/spark-expr/src/kernels/temporal.rs | 583 ++++++++++++------ .../org/apache/comet/serde/datetime.scala | 6 +- .../trunc_timestamp_dst_format_column.sql | 74 +++ .../trunc_timestamp_fine_overflow.sql | 17 +- .../trunc_timestamp_fine_overflow_spark42.sql | 24 +- 9 files changed, 646 insertions(+), 211 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql diff --git a/docs/source/contributor-guide/expression-audits/datetime_funcs.md b/docs/source/contributor-guide/expression-audits/datetime_funcs.md index d75661a6502..ba8aedd0b3e 100644 --- a/docs/source/contributor-guide/expression-audits/datetime_funcs.md +++ b/docs/source/contributor-guide/expression-audits/datetime_funcs.md @@ -35,7 +35,7 @@ ## date_trunc -- Spark 4.2.0 raises `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound. +- Spark 4.2.0 and later raise `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound. - Performance (tuned 2026-10-06, PR [#5956](https://github.com/apache/datafusion-comet/pull/5956)): scalar timestamp truncation skips dictionary key masking for infallible fine units and safe UTC/NTZ values with substantial repetition. A values-only range check retains masking near the lower timestamp bound; non-UTC calendar paths and high-cardinality coarse batches retain masking to avoid errors and unnecessary calendar work on unused entries. About 19x faster for MICROSECOND with 65,536 keys and 32 distinct values. Matched Criterion coverage includes sparse/dense/all-NULL keys, high-cardinality YEAR batches, and used/unused overflow values. Benchmark: `benches/timestamp_trunc.rs`. ## dayname diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index adcf2d64dfe..d1bdd9c80a5 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -598,8 +598,8 @@ message TruncTimestamp { Expr child = 2; string timezone = 3; // Whether SECOND and MILLISECOND truncation wraps when the result would fall below the - // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 (SPARK-56663) - // raises `long overflow` like every other unit. This tracks the Spark version Comet is + // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 and later + // (SPARK-56663) raise `long overflow` like every other unit. This tracks the Spark version Comet is // running against. bool wrap_second_millisecond_overflow = 4; } diff --git a/native/spark-expr/benches/timestamp_trunc.rs b/native/spark-expr/benches/timestamp_trunc.rs index 654e30e6393..bc6db2de32b 100644 --- a/native/spark-expr/benches/timestamp_trunc.rs +++ b/native/spark-expr/benches/timestamp_trunc.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, DictionaryArray, Int32Array, RecordBatch, TimestampMicrosecondArray}; +use arrow::array::{ + Array, DictionaryArray, Int32Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_expr::expressions::{lit, Column}; @@ -153,6 +155,63 @@ fn criterion_benchmark(c: &mut Criterion) { } } group.finish(); + + // Formats from a column: one format for every row, and every unit mixed across the rows. + // Daytime instants (16:00 to 22:00 UTC) avoid DST transitions, so older kernels can be + // measured on the same input. + let mut group = c.benchmark_group("timestamp_trunc_format_column"); + let rows = 8_192; + let units = [ + "YEAR", + "QUARTER", + "MONTH", + "WEEK", + "DAY", + "HOUR", + "MINUTE", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ]; + for timezone in ["UTC", "America/Los_Angeles"] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "a", + DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.into())), + true, + ), + Field::new("fmt", DataType::Utf8, true), + ])); + let timestamps = TimestampMicrosecondArray::from_iter_values((0..rows).map(|i| { + BASE_MICROS + + (i % 366) as i64 * MICROS_PER_DAY + + 57_600_000_000 + + (i as i64 * 1_234_567) % 21_600_000_000 + })) + .with_timezone(timezone); + for (tag, unit) in [("YEAR", Some(0)), ("HOUR", Some(5)), ("mixed", None)] { + let formats = StringArray::from_iter_values( + (0..rows).map(|i| units[unit.unwrap_or(i % units.len())]), + ); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(timestamps.clone()), Arc::new(formats)], + ) + .unwrap(); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + Arc::new(Column::new("fmt", 1)), + timezone.to_string(), + true, + ); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{timezone}/{tag}/{rows}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + } + group.finish(); } criterion_group!(benches, criterion_benchmark); diff --git a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs index 2835fab6bd2..e5c07ebbb3b 100644 --- a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs +++ b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs @@ -16,7 +16,7 @@ // under the License. use crate::utils::array_with_timezone; -use arrow::array::ArrayRef; +use arrow::array::{new_null_array, ArrayRef}; use arrow::compute::cast; use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond}; use arrow::record_batch::RecordBatch; @@ -42,7 +42,7 @@ pub struct TimestampTruncExpr { /// reallocating, and parsed once into a `chrono::TimeZone` per batch. timezone: Arc, /// Whether SECOND and MILLISECOND truncation wraps below the smallest timestamp, as Spark - /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 does (SPARK-56663). + /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 and later do (SPARK-56663). wrap_second_millisecond_overflow: bool, } @@ -133,7 +133,7 @@ impl PhysicalExpr for TimestampTruncExpr { Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Array(ts), ColumnarValue::Array(formats)) => { - let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats)?; + let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?; Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => { @@ -142,6 +142,18 @@ impl PhysicalExpr for TimestampTruncExpr { let scalar = ScalarValue::try_from_array(&relabel(result)?, 0)?; Ok(ColumnarValue::Scalar(scalar)) } + (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Array(formats)) => { + let ts = ts_scalar.to_array_of_size(formats.len())?; + let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?; + Ok(ColumnarValue::Array(relabel(result)?)) + } + // A NULL format gives NULL, as in Spark. + (ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(None))) => { + Ok(ColumnarValue::Array(new_null_array(&output_type, ts.len()))) + } + (ColumnarValue::Scalar(_), ColumnarValue::Scalar(Utf8(None))) => { + Ok(ColumnarValue::Scalar(ScalarValue::try_from(&output_type)?)) + } _ => Err(DataFusionError::Execution( "Invalid input to function TimestampTrunc. \ Expected (Timestamp, Utf8)" @@ -184,7 +196,9 @@ fn output_type(input: &DataType) -> DataType { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, AsArray, DictionaryArray, Int32Array, TimestampMicrosecondArray}; + use arrow::array::{ + Array, AsArray, DictionaryArray, Int32Array, StringArray, TimestampMicrosecondArray, + }; use arrow::datatypes::{Field, Int32Type, TimestampMicrosecondType}; use datafusion::physical_expr::expressions::{Column, Literal}; @@ -241,4 +255,68 @@ mod tests { HOUR_IN_KOLKATA ); } + + /// Evaluates `date_trunc` in America/Los_Angeles with the given timestamp and format inputs. + fn trunc_in_los_angeles( + timestamp: Arc, + format: Arc, + batch: &RecordBatch, + ) -> ColumnarValue { + TimestampTruncExpr::new(timestamp, format, "America/Los_Angeles".to_string(), false) + .evaluate(batch) + .unwrap() + } + + #[test] + fn literal_timestamp_with_a_format_column() { + // 2024-11-03 01:30 PDT, inside the overlap. DAY keeps the input's offset, YEAR does not + // depend on it, and a NULL format gives NULL. + let micros = 1_730_622_600_000_000; + let formats = StringArray::from(vec![Some("DAY"), Some("YEAR"), None]); + let schema = Schema::new(vec![Field::new("fmt", DataType::Utf8, true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(formats)]).unwrap(); + let timestamp = Literal::new(ScalarValue::TimestampMicrosecond( + Some(micros), + Some("UTC".into()), + )); + let ColumnarValue::Array(result) = + trunc_in_los_angeles(Arc::new(timestamp), Arc::new(Column::new("fmt", 0)), &batch) + else { + panic!("expected an array"); + }; + assert_eq!(result.data_type(), &utc_timestamp()); + let result = result.as_primitive::(); + // 2024-11-03 00:00 PDT and 2024-01-01 00:00 PST. + assert_eq!(result.value(0), 1_730_617_200_000_000); + assert_eq!(result.value(1), 1_704_096_000_000_000); + assert!(result.is_null(2)); + } + + #[test] + fn null_literal_format_gives_null() { + let input = TimestampMicrosecondArray::from(vec![Some(MICROS), None]).with_timezone("UTC"); + let schema = Schema::new(vec![Field::new("ts", utc_timestamp(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(input)]).unwrap(); + let null_format = || Arc::new(Literal::new(Utf8(None))); + let ColumnarValue::Array(result) = + trunc_in_los_angeles(Arc::new(Column::new("ts", 0)), null_format(), &batch) + else { + panic!("expected an array"); + }; + assert_eq!(result.data_type(), &utc_timestamp()); + assert_eq!(result.null_count(), 2); + + let timestamp = Arc::new(Literal::new(ScalarValue::TimestampMicrosecond( + Some(MICROS), + Some("UTC".into()), + ))); + let ColumnarValue::Scalar(result) = trunc_in_los_angeles(timestamp, null_format(), &batch) + else { + panic!("expected a scalar"); + }; + assert_eq!( + result, + ScalarValue::TimestampMicrosecond(None, Some("UTC".into())) + ); + } } diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index 60bf0e735ce..ce6ad9de5a9 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -177,30 +177,6 @@ where Ok(builder.finish().with_timezone(tz_str)) } -fn as_timestamp_tz_with_op_single( - value: Option, - builder: &mut PrimitiveBuilder, - tz: &Tz, - op: F, -) -> Result<(), SparkError> -where - F: Fn(DateTime) -> i64, - i64: From, -{ - match value { - Some(value) => match as_datetime_with_timezone::(value.into(), *tz) { - Some(time) => builder.append_value(op(time)), - _ => { - return Err(SparkError::Internal( - "Unable to read value as datetime".to_string(), - )); - } - }, - None => builder.append_null(), - } - Ok(()) -} - // Apply the Tz to the Naive Date Time, convert to UTC, and return as microseconds in Unix epoch. // After truncation the carried UTC offset may be wrong if the truncated time falls in a different // DST period than the original (e.g., truncating a December/PST timestamp to QUARTER yields @@ -668,7 +644,7 @@ where /// format is a scalar string specifying the format to apply to the timestamp value. /// /// wrap_second_millisecond_overflow selects how SECOND and MILLISECOND truncation treats a -/// result below the smallest timestamp: Spark before 4.2 wraps, 4.2 raises. +/// result below the smallest timestamp: Spark before 4.2 wraps, 4.2 and later raise. pub(crate) fn timestamp_trunc_dyn( array: &dyn Array, format: String, @@ -920,9 +896,10 @@ where Ok(result) } -/// The zone-aware scalar-format fallback for HOUR/DAY outside DataFusion 55.1's internal -/// TimestampNanosecond range. UTC/NTZ coarse units use integer calendar arithmetic instead. -/// Row-format paths continue to call the same underlying helpers. +/// The zone-aware fallback for HOUR/DAY outside DataFusion 55.1's internal TimestampNanosecond +/// range, for literal formats and format columns alike. UTC/NTZ coarse units use integer calendar +/// arithmetic instead. Like every native local-time conversion, it stops applying DST after +/// chrono-tz's last transition, around 2099, while Spark keeps applying the zone's rules (#6816). fn timestamp_trunc_legacy( array: &TimestampMicrosecondArray, format: &str, @@ -1114,7 +1091,8 @@ fn timestamp_trunc_fine_boundary( micros.wrapping_sub(remainder) } else { // MINUTE/HOUR/DAY use Spark's exact instant-to-microseconds conversion, and - // Spark 4.2 (SPARK-56663) checks the SECOND/MILLISECOND subtraction too. + // Spark 4.2 and later (SPARK-56663) check the SECOND/MILLISECOND + // subtraction too. micros.checked_sub(remainder).ok_or_else(|| { SparkError::Internal(format!( "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" @@ -1241,29 +1219,6 @@ fn timestamp_trunc_upstream( Ok(builder.finish().with_timezone_opt(array.timezone())) } -/// Truncate a single NTZ value and append to builder -fn timestamp_trunc_ntz_single( - value: Option, - builder: &mut PrimitiveBuilder, - op: F, -) -> Result<(), SparkError> -where - F: Fn(NaiveDateTime) -> Option, -{ - match value { - Some(micros) => match micros_to_naive(micros).and_then(op) { - Some(truncated) => builder.append_value(naive_to_micros(truncated)), - None => { - return Err(SparkError::Internal( - "Unable to truncate NTZ timestamp".to_string(), - )) - } - }, - None => builder.append_null(), - } - Ok(()) -} - pub(crate) fn timestamp_trunc( array: &PrimitiveArray, format: String, @@ -1298,158 +1253,126 @@ where /// /// format is an array of strings specifying the format to apply to the corresponding timestamp /// value. The array may be a dictionary array. +/// +/// wrap_second_millisecond_overflow is passed to the literal-format kernel, as in +/// `timestamp_trunc_dyn`. pub(crate) fn timestamp_trunc_array_fmt_dyn( array: &dyn Array, formats: &dyn Array, + wrap_second_millisecond_overflow: bool, ) -> Result { - match (array.data_type().clone(), formats.data_type().clone()) { - (DataType::Dictionary(_, _), DataType::Dictionary(_, _)) => { - downcast_dictionary_array!( - formats => { - downcast_dictionary_array!( - array => { - timestamp_trunc_array_fmt_dict_dict( - &array.downcast_dict::().unwrap(), - &formats.downcast_dict::().unwrap()) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt) - ) - } - fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), - ) - } - (DataType::Dictionary(_, _), DataType::Utf8) => { - downcast_dictionary_array!( - array => { - timestamp_trunc_array_fmt_dict_plain( - &array.downcast_dict::>().unwrap(), - formats.as_any().downcast_ref::() - .expect("Unexpected value type in formats")) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt), - ) + let array = unpack_dictionary(array)?; + let formats = unpack_dictionary(formats)?; + let array = match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, _) => { + array.as_primitive::() } - (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Dictionary(_, _)) => { - downcast_dictionary_array!( - formats => { - downcast_temporal_array!(array => { - timestamp_trunc_array_fmt_plain_dict( - array, - &formats.downcast_dict::().unwrap()) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt), - ) - } - fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), - ) - } - (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Utf8) => { - downcast_temporal_array!( - array => { - timestamp_trunc_array_fmt_plain_plain(array, - formats.as_any().downcast_ref::().expect("Unexpected value type in formats")) - .map(|a| Arc::new(a) as ArrayRef) - }, - dt => return_compute_error_with!("timestamp_trunc does not support", dt), - ) + dt => { + return_compute_error_with!("Unsupported input type for function 'timestamp_trunc'", dt) } - (dt, fmt) => Err(SparkError::Internal(format!( - "Unsupported datatype: {dt:}, format: {fmt:?} for function 'timestamp_trunc'" - ))), - } + }; + let formats = match formats.data_type() { + DataType::Utf8 => formats.as_string::(), + fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), + }; + Ok(Arc::new(timestamp_trunc_by_row_format( + array, + formats, + wrap_second_millisecond_overflow, + )?)) } -macro_rules! timestamp_trunc_array_fmt_helper { - ($array: ident, $formats: ident, $datatype: ident) => {{ - let mut builder = TimestampMicrosecondBuilder::with_capacity($array.len()); - let iter = $array.into_iter(); - assert_eq!( - $array.len(), - $formats.len(), - "lengths of values array and format array must be the same" - ); - match $datatype { - DataType::Timestamp(TimeUnit::Microsecond, None) => { - // TimestampNTZ: operate directly on naive microsecond values - for (index, val) in iter.enumerate() { - let micros_val = val.map(|v| i64::from(v)); - let trunc_fn = ntz_trunc_fn_for_format($formats.value(index))?; - timestamp_trunc_ntz_single(micros_val, &mut builder, trunc_fn)?; - } - Ok(builder.finish()) - } - DataType::Timestamp(TimeUnit::Microsecond, Some(tz_str)) => { - let tz: Tz = tz_str.parse()?; - for (index, val) in iter.enumerate() { - let trunc_fn = tz_trunc_fn_for_format($formats.value(index))?; - as_timestamp_tz_with_op_single::(val, &mut builder, &tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - })?; - } - Ok(builder.finish().with_timezone(tz_str.as_ref())) - } - dt => { - return_compute_error_with!( - "Unsupported input type '{:?}' for function 'timestamp_trunc'", - dt - ) - } - } - }}; +fn unpack_dictionary(array: &dyn Array) -> Result { + match array.data_type() { + DataType::Dictionary(_, value_type) => arrow::compute::cast(array, value_type) + .map_err(|error| SparkError::Internal(error.to_string())), + _ => Ok(make_array(array.to_data())), + } } -fn timestamp_trunc_array_fmt_plain_plain( - array: &PrimitiveArray, +/// Truncates each row with the format in the same row. +/// +/// Rows are grouped by format, and each group goes through the literal-format kernel with the +/// other rows set to NULL. A row is therefore truncated by exactly the rules a literal format +/// applies, including the DST handling, and a value in one group cannot make another group fail. +/// A NULL format gives a NULL result, as in Spark. +fn timestamp_trunc_by_row_format( + array: &TimestampMicrosecondArray, formats: &StringArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - let data_type = array.data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} -fn timestamp_trunc_array_fmt_plain_dict( - array: &PrimitiveArray, - formats: &TypedDictionaryArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, -{ - let data_type = array.data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} + wrap_second_millisecond_overflow: bool, +) -> Result { + if array.len() != formats.len() { + return Err(SparkError::Internal(format!( + "timestamp_trunc has {} values but {} formats", + array.len(), + formats.len() + ))); + } -fn timestamp_trunc_array_fmt_dict_plain( - array: &TypedDictionaryArray>, - formats: &StringArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, -{ - let data_type = array.values().data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} + let mut granularities: Vec<&'static str> = Vec::new(); + // A batch rarely holds more than a few spellings, so parse each one once. + let mut spellings: Vec<(&str, usize)> = Vec::new(); + let mut row_groups: Vec> = Vec::with_capacity(array.len()); + let mut null_format_hides_value = false; + for index in 0..array.len() { + if array.is_null(index) { + row_groups.push(None); + continue; + } + if formats.is_null(index) { + null_format_hides_value = true; + row_groups.push(None); + continue; + } + let spelling = formats.value(index); + let group = match spellings.iter().find(|(seen, _)| *seen == spelling) { + Some((_, group)) => *group, + None => { + let granularity = normalize_timestamp_trunc_format(spelling)?; + let group = match granularities.iter().position(|g| *g == granularity) { + Some(group) => group, + None => { + granularities.push(granularity); + granularities.len() - 1 + } + }; + spellings.push((spelling, group)); + group + } + }; + row_groups.push(Some(group)); + } -fn timestamp_trunc_array_fmt_dict_dict( - array: &TypedDictionaryArray>, - formats: &TypedDictionaryArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, - F: ArrowDictionaryKeyType, -{ - let data_type = array.values().data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) + if granularities.len() == 1 && !null_format_hides_value { + return timestamp_trunc_upstream(array, granularities[0], wrap_second_millisecond_overflow); + } + + let truncated = granularities + .iter() + .enumerate() + .map(|(group, granularity)| { + let outside_group: BooleanArray = row_groups + .iter() + .map(|row_group| Some(*row_group != Some(group))) + .collect(); + let input = arrow::compute::nullif(array, &outside_group) + .map_err(|error| SparkError::Internal(error.to_string()))?; + timestamp_trunc_upstream( + input.as_primitive(), + granularity, + wrap_second_millisecond_overflow, + ) + }) + .collect::, _>>()?; + + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for (index, row_group) in row_groups.iter().enumerate() { + match row_group.map(|group| &truncated[group]) { + Some(result) if result.is_valid(index) => builder.append_value(result.value(index)), + _ => builder.append_null(), + } + } + Ok(builder.finish().with_timezone_opt(array.timezone())) } #[cfg(test)] @@ -2243,7 +2166,7 @@ mod tests { #[test] fn test_timestamp_trunc_dictionary_fine_unused_extremes() { // Unused MIN entries must not raise, whether SECOND/MILLISECOND wrap (Spark before 4.2) - // or fail (Spark 4.2). + // or fail (Spark 4.2 and later). for wrap in [true, false] { for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { let values = TimestampMicrosecondArray::from(vec![ @@ -2388,8 +2311,8 @@ mod tests { #[test] fn test_timestamp_trunc_fine_spark42_overflow() { - // Spark 4.2 (SPARK-56663) checks the SECOND/MILLISECOND subtraction like every other - // unit. First representable boundary for each unit. + // Spark 4.2 and later (SPARK-56663) check the SECOND/MILLISECOND subtraction like every + // other unit. First representable boundary for each unit. for (format, boundary) in [ ("SECOND", -9_223_372_036_854_000_000), ("MILLISECOND", -9_223_372_036_854_775_000), @@ -2577,6 +2500,249 @@ mod tests { ); } + /// Instants around DST and offset transitions, keyed by the session timezone the kernel sees. + /// `None` is a TIMESTAMP_NTZ array. + fn transition_instants() -> Vec<(Option<&'static str>, Vec)> { + let at = |instants: &[&str]| { + instants + .iter() + .map(|i| instant_micros(i)) + .collect::>() + }; + vec![ + // Toronto skipped 1919-03-30 23:30 to 1919-03-31 00:30. + ( + Some("America/Toronto"), + at(&[ + "1919-03-31T04:30:00Z", + "1919-03-31T04:45:00Z", + "1919-04-02T16:00:00Z", + ]), + ), + // Havana repeated midnight on 2020-11-01. + ( + Some("America/Havana"), + at(&[ + "2020-11-01T04:30:00Z", + "2020-11-01T05:30:00Z", + "2020-11-15T12:00:00Z", + ]), + ), + // Sao Paulo skipped midnight on 2018-11-04 and repeated 23:00 on 2019-02-16. + ( + Some("America/Sao_Paulo"), + at(&[ + "2018-11-04T03:30:00Z", + "2018-11-04T12:00:00Z", + "2019-02-17T02:30:00Z", + ]), + ), + ( + Some("America/Los_Angeles"), + at(&[ + "2024-03-10T10:30:00Z", + "2024-03-10T11:15:30Z", + "2024-11-03T08:30:00Z", + "2024-11-03T09:30:00Z", + "1883-06-15T10:30:45Z", + ]), + ), + // Monrovia used -00:44:30 until 1972. + ( + Some("Africa/Monrovia"), + at(&[ + "1960-06-15T11:15:45Z", + "1972-01-07T00:44:29Z", + "1972-01-07T00:44:31Z", + ]), + ), + // Asuncion skipped midnight on 2023-10-01, a MONTH and QUARTER boundary. + ( + Some("America/Asuncion"), + at(&[ + "2023-10-01T03:30:00Z", + "2023-10-01T04:30:00Z", + "2023-10-15T12:00:00Z", + ]), + ), + // Apia skipped 2011-12-30 entirely. + ( + Some("Pacific/Apia"), + at(&[ + "2011-12-30T09:59:59Z", + "2011-12-30T10:00:00Z", + "2011-12-31T12:00:00Z", + ]), + ), + ( + Some("UTC"), + at(&["1500-06-15T12:34:56.123456Z", "3333-05-17T12:34:56.123456Z"]), + ), + ( + None, + vec![ + i64::MAX, + i64::MIN + 400 * MICROS_PER_DAY, + instant_micros("1969-12-31T23:59:59.999999Z"), + instant_micros("2024-05-17T12:34:56.123456Z"), + ], + ), + ] + } + + #[test] + fn row_formats_truncate_like_literal_formats() { + let spellings: Vec<&str> = TIMESTAMP_TRUNC_ALIASES + .iter() + .map(|(name, _)| *name) + .collect(); + for ((timezone, instants), wrap) in transition_instants() + .into_iter() + .flat_map(|zone| [(zone.clone(), true), (zone, false)]) + { + let input = TimestampMicrosecondArray::from(instants).with_timezone_opt(timezone); + let by_spelling: Vec = spellings + .iter() + .map(|spelling| timestamp_trunc(&input, spelling.to_string(), wrap).unwrap()) + .collect(); + // Rotate the spellings across the rows so every row meets every format in a column + // that mixes formats, which exercises the grouping and masking. + for shift in 0..spellings.len() { + let pick = |row: usize| (row + shift) % spellings.len(); + let formats = StringArray::from( + (0..input.len()) + .map(|row| spellings[pick(row)]) + .collect::>(), + ); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let result = result + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.data_type(), input.data_type()); + for row in 0..input.len() { + assert_eq!( + result.value(row), + by_spelling[pick(row)].value(row), + "{timezone:?} {} at {}", + spellings[pick(row)], + input.value(row) + ); + } + } + } + } + + #[test] + fn row_format_week_resolves_a_midnight_gap_to_its_end() { + // Toronto's gap ran from 1919-03-30 23:30 to 1919-03-31 00:30, so the Monday that WEEK + // lands on starts at 00:30. DAY moves its nonexistent midnight forward by the gap to 01:00. + let input = TimestampMicrosecondArray::from(vec![ + instant_micros("1919-04-02T16:00:00Z"), + instant_micros("1919-03-31T04:45:00Z"), + ]) + .with_timezone("America/Toronto"); + let formats = StringArray::from(vec!["WEEK", "DAY"]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let expected = TimestampMicrosecondArray::from(vec![ + instant_micros("1919-03-31T04:30:00Z"), + instant_micros("1919-03-31T05:00:00Z"), + ]) + .with_timezone("America/Toronto"); + assert_eq!( + result + .as_any() + .downcast_ref::() + .unwrap(), + &expected + ); + } + + #[test] + fn row_formats_handle_nulls_like_spark() { + let input = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + Some(instant_micros("2024-05-17T12:34:56Z")), + ]) + .with_timezone("America/Los_Angeles"); + // A NULL format gives NULL, and so does a NULL value, whatever its format says. + let formats = StringArray::from(vec![Some("YEAR"), None, Some("not_a_unit"), Some("HOUR")]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let expected = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-01-01T08:00:00Z")), + None, + None, + Some(instant_micros("2024-05-17T12:00:00Z")), + ]) + .with_timezone("America/Los_Angeles"); + assert_eq!( + result + .as_any() + .downcast_ref::() + .unwrap(), + &expected + ); + + let formats = StringArray::from(vec!["YEAR", "not_a_unit", "YEAR", "YEAR"]); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_err()); + } + + #[test] + fn row_format_groups_do_not_fail_each_other() { + // SECOND wraps at the lower bound like Spark before 4.2, while YEAR would overflow there. + // The YEAR group only sees its own row. + let input = + TimestampMicrosecondArray::from(vec![i64::MIN, instant_micros("2024-05-17T12:34:56Z")]); + let formats = StringArray::from(vec!["SECOND", "YEAR"]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let literal = timestamp_trunc(&input.slice(0, 1), "SECOND".to_string(), true).unwrap(); + let result = result + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), literal.value(0)); + assert_eq!(result.value(1), instant_micros("2024-01-01T00:00:00Z")); + + let formats = StringArray::from(vec!["YEAR", "YEAR"]); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_err()); + } + + #[test] + fn row_formats_accept_dictionaries() { + let instants = [ + instant_micros("2024-03-10T10:30:00Z"), + instant_micros("2024-11-03T09:30:00Z"), + instant_micros("2024-11-03T09:30:00Z"), + ]; + let input = + TimestampMicrosecondArray::from(instants.to_vec()).with_timezone("America/Los_Angeles"); + let formats = StringArray::from(vec!["DAY", "HOUR", "MONTH"]); + let expected = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + + let keys = Int32Array::from(vec![0, 1, 1]); + let values = TimestampMicrosecondArray::from(vec![instants[0], instants[1]]) + .with_timezone("America/Los_Angeles"); + let input_dict = DictionaryArray::try_new(keys, Arc::new(values)).unwrap(); + let mut formats_builder = StringDictionaryBuilder::::new(); + for format in ["DAY", "HOUR", "MONTH"] { + formats_builder.append(format).unwrap(); + } + let formats_dict = formats_builder.finish(); + + for (array, formats) in [ + (&input_dict as &dyn Array, &formats as &dyn Array), + (&input as &dyn Array, &formats_dict as &dyn Array), + (&input_dict as &dyn Array, &formats_dict as &dyn Array), + ] { + assert_eq!( + ×tamp_trunc_array_fmt_dyn(array, formats, true).unwrap(), + &expected + ); + } + } + #[test] // test takes too long with miri #[cfg_attr(miri, ignore)] @@ -2673,7 +2839,7 @@ mod tests { } // test cases - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_array) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_array, true) { for i in 0..array.len() { assert!( array.value(i) @@ -2686,7 +2852,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_array) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_array, true) { for i in 0..array.len() { assert!( array.value(i) @@ -2699,7 +2865,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_dict) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_dict, true) { for i in 0..array.len() { assert!( array.value(i) @@ -2712,7 +2878,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_dict) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_dict, true) { for i in 0..array.len() { assert!( array.value(i) @@ -2727,6 +2893,27 @@ mod tests { } } + #[test] + fn row_formats_follow_the_overflow_policy() { + // Spark before 4.2 wraps SECOND and MILLISECOND below the smallest timestamp, and 4.2 and + // later raise. Cover a column with one format and one that mixes formats. + let input = TimestampMicrosecondArray::from(vec![ + i64::MIN, + i64::MIN, + instant_micros("2024-05-17T12:34:56Z"), + ]) + .with_timezone("UTC"); + for formats in [ + vec!["SECOND", "SECOND", "SECOND"], + vec!["SECOND", "MILLISECOND", "YEAR"], + ] { + let formats = StringArray::from(formats); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_ok()); + let error = timestamp_trunc_array_fmt_dyn(&input, &formats, false).unwrap_err(); + assert!(error.to_string().contains("long overflow"), "{error}"); + } + } + /// Truncating a November timestamp in `America/Denver` to QUARTER must land on the start of /// Q4, which is October 1 — and October 1 is still MDT (UTC-6), not MST (UTC-7). The /// pre-fix kernel reused the input's MST offset for the truncated date, producing a result diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index eaecffac875..0317132ef00 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -644,9 +644,9 @@ object CometTruncTimestamp builder.setChild(childExpr.get) builder.setFormat(formatExpr.get) builder.setTimezone(timeZone.get) - // Spark 4.2.0 (SPARK-56663) raises `long overflow` when SECOND or MILLISECOND truncation - // falls below the smallest timestamp; earlier versions wrap the Long subtraction. The - // native side has to match whichever Spark we are running against. + // Spark 4.2.0 and later (SPARK-56663) raise `long overflow` when SECOND or MILLISECOND + // truncation falls below the smallest timestamp; earlier versions wrap the Long + // subtraction. The native side has to match whichever Spark we are running against. builder.setWrapSecondMillisecondOverflow(!isSpark42Plus) Some( diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql new file mode 100644 index 00000000000..e9533ae3cc4 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.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. + + +-- date_trunc with the format in a column, around DST overlaps, gaps and historical offsets. Each +-- row is truncated with the rules a literal format applies (see trunc_timestamp_dst_ambiguous.sql +-- for what each zone covers), and a NULL format gives NULL. A non-literal format is Incompatible, +-- because an invalid format throws instead of returning NULL. + +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles,America/New_York,America/Sao_Paulo,Africa/Monrovia,Asia/Aden,America/Havana,America/Toronto,America/Asuncion + +statement +CREATE TABLE test_trunc_fmt_ts(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_fmt_ts VALUES + (timestamp('1919-04-02T16:00:00Z')), + (timestamp('1919-03-31T04:30:00Z')), + (timestamp('1919-03-31T04:45:00Z')), + (timestamp('2023-10-15T12:00:00Z')), + (TIMESTAMP '1960-06-15 10:30:45'), + (timestamp('1947-03-13T20:53:30.123Z')), + (timestamp('1972-01-07T00:44:45Z')), + (timestamp('2018-11-04T02:30:15.123456Z')), + (timestamp('2018-11-04T03:30:15.123456Z')), + (timestamp('2024-03-10T09:30:15.123456Z')), + (timestamp('2024-03-10T10:30:15.123456Z')), + (timestamp('2024-11-03T08:30:15.123456Z')), + (timestamp('2024-11-03T09:30:15.123456Z')), + (timestamp('2020-11-01T04:30:00Z')), + (timestamp('2020-11-01T05:30:00Z')), + (timestamp('2020-11-15T12:00:00Z')), + (NULL) + +statement +CREATE TABLE test_trunc_fmt_unit(fmt string) USING parquet + +statement +INSERT INTO test_trunc_fmt_unit VALUES + ('YEAR'), ('yyyy'), ('YY'), ('QUARTER'), ('MONTH'), ('mon'), ('MM'), ('WEEK'), ('DAY'), ('dd'), + ('HOUR'), ('MINUTE'), ('SECOND'), ('MILLISECOND'), ('MICROSECOND'), (NULL) + +statement +CREATE TABLE test_trunc_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_fmt SELECT ts, fmt FROM test_trunc_fmt_ts CROSS JOIN test_trunc_fmt_unit + +query expect_native(date_trunc) +SELECT ts, fmt, date_trunc(fmt, ts) FROM test_trunc_fmt ORDER BY ts, fmt + +-- A NULL literal format gives NULL. +query expect_native(date_trunc) +SELECT ts, date_trunc(NULL, ts) FROM test_trunc_fmt_ts ORDER BY ts + +-- A literal timestamp with the format in a column. 2024-11-03 01:30 is in the US fall-back overlap. +query expect_native(date_trunc) +SELECT fmt, date_trunc(fmt, TIMESTAMP '2024-11-03 01:30:00') FROM test_trunc_fmt_unit ORDER BY fmt diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql index dc8a5478d18..4ef39e12a95 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql @@ -16,12 +16,14 @@ -- under the License. -- Before 4.2, Spark wraps the Long subtraction when SECOND or MILLISECOND truncation falls below --- Long.MinValue. Spark 4.2 (SPARK-56663) raises instead, which +-- Long.MinValue. Spark 4.2 and later (SPARK-56663) raise instead, which -- trunc_timestamp_fine_overflow_spark42.sql covers. -- MaxSparkVersion: 4.1 -- Use microsecond storage to preserve the extreme value without an INT96 conversion. -- Config: spark.sql.session.timeZone=UTC +-- A format column is Incompatible, because an invalid format throws instead of returning NULL. +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true -- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED -- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED -- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS @@ -39,3 +41,16 @@ INSERT INTO test_trunc_ts_fine_overflow VALUES query expect_native(date_trunc) SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) FROM test_trunc_ts_fine_overflow + +statement +CREATE TABLE test_trunc_ts_fine_overflow_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_ts_fine_overflow_fmt VALUES + (timestamp_micros(-9223372036854775808), 'SECOND'), + (timestamp_micros(-9223372036854775808), 'MILLISECOND'), + (timestamp('2024-05-17 12:34:56.123456'), 'SECOND') + +-- The same truncation with the format in a column. +query expect_native(date_trunc) +SELECT fmt, date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt ORDER BY fmt, ts diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql index 60f917b8ef7..d6fadeeb288 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql @@ -15,13 +15,15 @@ -- specific language governing permissions and limitations -- under the License. --- Spark 4.2 (SPARK-56663) raises when SECOND or MILLISECOND truncation falls below +-- Spark 4.2 and later (SPARK-56663) raise when SECOND or MILLISECOND truncation falls below -- Long.MinValue, like every other unit. trunc_timestamp_fine_overflow.sql covers 4.1 and -- earlier, where Spark wraps the Long subtraction instead. -- MinSparkVersion: 4.2 -- Use microsecond storage to preserve the extreme value without an INT96 conversion. -- Config: spark.sql.session.timeZone=UTC +-- A format column is Incompatible, because an invalid format throws instead of returning NULL. +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true -- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED -- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED -- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS @@ -47,3 +49,23 @@ SELECT date_trunc('SECOND', ts) FROM test_trunc_ts_fine_overflow query expect_error(long overflow) SELECT date_trunc('MILLISECOND', ts) FROM test_trunc_ts_fine_overflow + +statement +CREATE TABLE test_trunc_ts_fine_overflow_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_ts_fine_overflow_fmt VALUES + (timestamp_micros(-9223372036854775808), 'SECOND'), + (timestamp_micros(-9223372036854775808), 'MILLISECOND'), + (timestamp('2024-05-17 12:34:56.123456'), 'SECOND') + +-- With the format in a column. The first query shows it runs natively. +query expect_native(date_trunc) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt +WHERE ts > timestamp_micros(-9223372036854775808) + +query expect_error(long overflow) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt + +query expect_error(long overflow) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt WHERE fmt = 'MILLISECOND' From 1e2d78345325ff60c27171f47b060815b71b8378 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 19:53:21 +0000 Subject: [PATCH 43/89] feat: run nullable nested float sort keys natively in strict floating-point mode (#6735) * feat: run nullable nested float sort keys natively in strict floating-point mode Strict floating-point mode declined a sort or window key that nests floats in an array or struct whose type can hold a null element or field, because the native sort placed such a null by the key's null order (#6476) and a RANGE frame ordered it above every other value (#6477). #6683 and #6684 now decline exactly those shapes in every mode: ASC NULLS LAST or DESC NULLS FIRST on such a key, and a RANGE frame that has to find a row's peers over one. What strict mode still declined beyond them, the default null orders, TopK, ranks, rank limits, ROWS frames and RANGE frames bounded only by the partition, already matches Spark, so the strict-only decline goes. The strict fixture now runs those shapes natively under strict mode, with a null element, a -0.0 and a NaN with the sign bit set among the keys, and keeps expect_fallback for the shapes the general guards decline. CometWindowExecSuite expects a native rank limit in strict mode, and the nested strict sort test in CometExpressionSuite adds a key that can hold a null. The floating-point compatibility page, the operator tuning page and the spark.comet.exec.strictFloatingPoint doc no longer describe a strict exception. * test: expect native strict sort over nullable nested float keys * test: sort a real null element in the strict nested float sort test The IF key in the strict nested sort test never produced a null, because id runs from 0 to 4. It now nulls row 4's element, so the test checks that a null element sorts first in strict mode, ahead of the signed zeros and NaN payloads, with that row's bits check skipped. Also update two comments that still described the strict-only decline: nested_float_order_keys.sql and CometWindowGroupLimitExec now say that strict mode runs these keys natively, and that CometSortOrder declines ASC NULLS LAST and DESC NULLS FIRST on them in every mode. --- .../latest/compatibility/floating-point.md | 17 +++-- .../user-guide/latest/tuning/operators.md | 10 +-- .../scala/org/apache/comet/CometConf.scala | 4 +- .../apache/comet/serde/CometSortOrder.scala | 26 ++------ .../sql/comet/CometWindowGroupLimitExec.scala | 4 +- .../windows/nested_float_order_keys.sql | 4 +- .../nested_float_order_keys_strict.sql | 53 ++++++++++++---- .../apache/comet/CometExpressionSuite.scala | 63 ++++++++++--------- .../comet/exec/CometWindowExecSuite.scala | 10 +-- 9 files changed, 102 insertions(+), 89 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index adf8518d42d..39e1de9b431 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -82,15 +82,14 @@ Because those comparison keys match Spark, `spark.comet.exec.strictFloatingPoint force a fallback for them: sort keys, window and rank order keys, and range partitioning keys all stay native under strict mode, whether the floats in them are scalar or nested. -The exception is a key that nests floats in an array or struct whose type can hold a null element -or field. Spark orders such a null below every other value, whatever the key's `NULLS FIRST` or -`NULLS LAST`. The native sort places it by that null order, so `ASC NULLS LAST` and -`DESC NULLS FIRST` can differ from Spark, and a `RANGE` window frame orders it above every other -value, so a running aggregate can span the whole partition -([#6476](https://github.com/apache/datafusion-comet/issues/6476), -[#6477](https://github.com/apache/datafusion-comet/issues/6477)). Strict mode makes those keys fall -back to Spark. A key whose type cannot hold a null, such as `array(coalesce(x, 0.0D))`, stays -native. +That includes a key that nests floats in an array or struct whose type can hold a null element or +field. Spark orders such a null below every other value, whatever the key's `NULLS FIRST` or +`NULLS LAST`, so Comet falls back in every mode, not only in strict mode, for the two shapes where +the native sort or window frame would place it differently: `ASC NULLS LAST` or `DESC NULLS FIRST` +on such a key ([#6476](https://github.com/apache/datafusion-comet/issues/6476)), and a `RANGE` +window frame that has to find a row's peers over it +([#6477](https://github.com/apache/datafusion-comet/issues/6477)). The +[operator compatibility notes](operators.md) describe both. `array_min` and `array_max` use Spark-compatible native comparisons in both strict and non-strict floating-point modes. Signed zeros compare equal, and all NaN representations compare equal and diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 97baabfccad..cb538fb63d4 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -162,12 +162,12 @@ nested in arrays and structs, so `ORDER BY`, window ordering and range partition even with `spark.comet.exec.strictFloatingPoint=true`. Only the comparison key is normalized; returned values keep their original NaN representation and zero sign. -The exception is a key that nests floating-point values in an array or struct whose type can hold a null element or -field. Spark orders such a null below every other value, and the native sort and `RANGE` window frames do not +A key that nests values in an array or struct whose type can hold a null element or field falls back in every mode under +`ASC NULLS LAST` or `DESC NULLS FIRST`, and in a `RANGE` window frame that has to find a row's peers, because Spark orders +such a null below every other value and the native sort and window frames would not ([#6476](https://github.com/apache/datafusion-comet/issues/6476), -[#6477](https://github.com/apache/datafusion-comet/issues/6477)), so -`spark.comet.exec.strictFloatingPoint=true` makes those keys fall back to Spark. They can be forced back onto the -native path with `spark.comet.expression.SortOrder.allowIncompatible=true`. +[#6477](https://github.com/apache/datafusion-comet/issues/6477)). The default null orders run natively, and +`spark.comet.expression.SortOrder.allowIncompatible=true` forces the other null orders onto the native path. `sort_array` sorts array elements rather than ordering rows. It follows Spark's floating-point ordering as well, so it also stays native with `spark.comet.exec.strictFloatingPoint=true`. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a187f8fa5f8..290a3b05358 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1223,9 +1223,7 @@ object CometConf extends ShimCometConf { "When enabled, fall back to Spark for floating-point operations that may differ from " + "Spark, such as comparing -0.0 and 0.0. `ORDER BY`, window ordering and range " + "partitioning keys are unaffected, including floating-point values nested in arrays " + - "and structs, because Comet normalizes those comparison keys to match Spark. The " + - "exception is a nested key whose type can hold a null element or field, which falls " + - "back until the native sort and window frames order those nulls as Spark does. " + + "and structs, because Comet normalizes those comparison keys to match Spark. " + "`sort_array` is unaffected too, because it follows Spark's ordering. " + s"$COMPAT_GUIDE.") .booleanConf diff --git a/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala b/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala index 1241c0bad38..c70d34dd202 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Descending, NullsFirst, NullsLast, SortOrder} import org.apache.spark.sql.types.{ArrayType, DataType, StructType} -import org.apache.comet.CometConf import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal /** @@ -41,9 +40,8 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal * null orders, `ASC NULLS FIRST` and `DESC NULLS LAST`, place it where Spark does. * https://github.com/apache/datafusion-comet/issues/6476 * - * A `RANGE` window frame orders such a null above every other value, so its frames can differ - * from Spark too. Strict mode keeps the fallback it applied to every nested float key before they - * were normalized, for the keys that can hit this. + * A `RANGE` window frame orders such a null above every other value, so `CometWindowExec` + * declines a `RANGE` frame that has to find a row's peers over such a key. * https://github.com/apache/datafusion-comet/issues/6477 */ object CometSortOrder extends CometExpressionSerde[SortOrder] { @@ -62,26 +60,14 @@ object CometSortOrder extends CometExpressionSerde[SortOrder] { "places it by the null order " + "([#6476](https://github.com/apache/datafusion-comet/issues/6476))." - private val nullableNestedFloatingPointSort = - "Sorting on floating-point values nested in an array or struct that can hold a null " + - "element or field" - - override def getIncompatibleReasons(): Seq[String] = Seq( - nestedNullOrderReason, - s"$nullableNestedFloatingPointSort is not 100% compatible with Spark when " + - s"`${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true`") + override def getIncompatibleReasons(): Seq[String] = Seq(nestedNullOrderReason) override def getSupportLevel(expr: SortOrder): SupportLevel = { - val dataType = expr.child.dataType - if (!canHoldNestedNull(dataType)) { - Compatible() - } else if (expr.nullOrdering != expr.direction.defaultNullOrdering) { + if (canHoldNestedNull(expr.child.dataType) && + expr.nullOrdering != expr.direction.defaultNullOrdering) { Incompatible(Some(nestedNullOrderReason)) } else { - SupportLevel - .strictFloatingPointReason(dataType, nullableNestedFloatingPointSort) - .map(reason => Incompatible(Some(reason))) - .getOrElse(Compatible()) + Compatible() } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala index ee3cbd61c63..e4c24583337 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -102,8 +102,8 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { // The same byte equality decides RANK and DENSE_RANK ties, but it runs on keys that the native // planner has normalized: a `FLOAT` or `DOUBLE` order key, scalar or nested in an array or // struct, has its NaNs folded together and its signed zeros tied first, as Spark does, so no - // float key needs a guard here. Strict floating-point mode declines the nested keys that can - // hold a null element or field, through `CometSortOrder`. + // float key needs a guard here. `CometSortOrder` declines `ASC NULLS LAST` and + // `DESC NULLS FIRST` on a nested key that can hold a null element or field, in every mode. val childOutput = op.children.head.output val partitionProtos = fields.partitionSpec.map(e => e -> exprToProto(e, childOutput)) diff --git a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql index 711cf63bc47..92e7d29f058 100644 --- a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql +++ b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql @@ -24,8 +24,8 @@ -- total order sorts below -Infinity. Each ORDER BY ends with a unique tiebreaker, so peers come -- out in the tiebreaker's order. --- Strict floating-point mode declines these keys, because their types can hold a null element or --- field: see nested_float_order_keys_strict.sql. +-- Strict floating-point mode runs these keys natively too. nested_float_order_keys_strict.sql covers +-- that, along with the null orders and RANGE frames that fall back in every mode. statement CREATE TABLE nested_float_keys(id INT, g INT, d DOUBLE, f FLOAT, s BOOLEAN) USING parquet diff --git a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql index 995470b2273..3d83601fa8a 100644 --- a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql +++ b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql @@ -16,11 +16,11 @@ -- under the License. -- Strict floating-point mode and sort or window keys that nest floats in arrays and structs. --- nested_float_order_keys.sql checks that Comet normalizes those floats to match Spark. Strict --- mode still declines a key whose type can hold a null element or field, because Spark orders --- such a null below every value whatever the key's null order, the native sort places it by the --- null order (#6476), and a RANGE window frame orders it above every value (#6477). A key whose --- type cannot hold a null stays native. +-- nested_float_order_keys.sql checks that Comet normalizes those floats to match Spark, so strict +-- mode runs these keys natively too. Spark orders a null element or field below every other +-- value, whatever the key's null order. Comet falls back, in every mode, for the two shapes where +-- that differs: `ASC NULLS LAST` or `DESC NULLS FIRST` on a key whose type can hold such a null +-- (#6476), and a RANGE window frame that has to find a row's peers over such a key (#6477). -- -- The test harness admits incompatible sort orders by default, so turn that off to check the -- shipped policy. @@ -42,10 +42,11 @@ INSERT INTO nested_float_strict VALUES (8, 1, NULL, NULL, false), (9, 2, double('-Infinity'), float('-Infinity'), false) --- Keys over the nullable columns can hold a null element or field, whether or not any row does, --- so they fall back, and Spark's answers come back. Row 8 holds one. Native, ORDER BY array(d) --- NULLS LAST puts row 8 last, where Spark puts it first, and the running sum spans the whole --- partition on row 8 and every row after it. +-- Keys over the nullable columns can hold a null element or field, whether or not any row does. +-- Row 8 holds one. Under a non-default null order, or in a RANGE frame that finds a row's peers, +-- they fall back, and Spark's answers come back. Native, ORDER BY array(d) NULLS LAST would put +-- row 8 last, where Spark puts it first, and the running sum would span the whole partition on +-- row 8 and every row after it. query expect_fallback(can hold a null element or field) SELECT id FROM nested_float_strict ORDER BY array(d) NULLS LAST, id @@ -56,25 +57,51 @@ query expect_fallback(can hold a null element or field) SELECT id FROM nested_float_strict ORDER BY named_struct('x', f) DESC NULLS FIRST, id query expect_fallback(can hold a null element or field) +SELECT id, RANK() OVER (ORDER BY array(IF(s, -d, d)) DESC NULLS FIRST, id) AS r +FROM nested_float_strict + +-- The default null orders place a null element or field where Spark does, so the same keys run +-- natively in a Sort, a TopK, ranking functions, a rank limit, a ROWS frame and a RANGE frame +-- bounded only by the partition. Row 8's null element sorts first ascending and last descending. +-- IF(s, ...) makes row 1 a -0.0 and row 5 a NaN with the sign bit set. +query +SELECT id FROM nested_float_strict ORDER BY array(IF(s, -d, d)), id + +query +SELECT id FROM nested_float_strict ORDER BY named_struct('x', IF(s, -f, f)) DESC, id + +query SELECT id FROM nested_float_strict ORDER BY array(IF(s, -d, d)), id LIMIT 4 -query expect_fallback(can hold a null element or field) +query SELECT id, RANK() OVER (PARTITION BY g ORDER BY named_struct('x', IF(s, -d, d))) AS r, DENSE_RANK() OVER (ORDER BY array(IF(s, -f, f)) DESC) AS dr FROM nested_float_strict -query expect_fallback(can hold a null element or field) +query SELECT id FROM ( SELECT id, RANK() OVER (ORDER BY array(IF(s, -d, d))) AS r FROM nested_float_strict ) WHERE r <= 3 +query +SELECT id, + SUM(id) OVER (ORDER BY array(IF(s, -d, d)), id + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS running, + SUM(id) OVER (PARTITION BY g ORDER BY named_struct('x', IF(s, -f, f)) + RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS total +FROM nested_float_strict + -- coalesce with a literal is never null, so an array or struct of it cannot hold a null, and --- these keys stay native. Row 8 becomes a third zero. Each ORDER BY ends with a unique --- tiebreaker, so peers come out in its order, and the running sums need no null row to leave out. +-- these keys stay native under any null order and in RANGE frames. Row 8 becomes a third zero. +-- Each ORDER BY ends with a unique tiebreaker, so peers come out in its order, and the running +-- sums need no null row to leave out. query SELECT id FROM nested_float_strict ORDER BY array(coalesce(IF(s, -d, d), 0.0D)), id DESC +query +SELECT id FROM nested_float_strict ORDER BY array(coalesce(IF(s, -d, d), 0.0D)) NULLS LAST, id + query SELECT id FROM nested_float_strict ORDER BY named_struct('x', coalesce(IF(s, -f, f), 0.0F)) DESC, id diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 5d2ca1d91de..6e1d77b3619 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -242,11 +242,10 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - // Floats nested in array and struct sort keys are normalized natively as well, but strict mode - // declines a key whose type can hold a null element or field: Spark orders that null below - // every value whatever the key's null order, and the native sort places it by the null order - // (#6476). Every field the generator makes is nullable, so these sorts fall back. A unique `id` - // sorted last makes the ordering total, as above. + // Floats nested in array and struct sort keys are normalized natively, in strict mode too. A key + // whose type can hold a null element or field sorts natively under the default null order, which + // agrees with Spark's. Every field the generator makes is nullable. A unique `id` sorted last + // makes the ordering total, as above. private def checkStrictNestedFloatingPointSort(schema: StructType): Unit = { val df = FuzzDataGenerator.generateDataFrame( new Random(42), @@ -260,21 +259,15 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { df.withColumn("id", monotonically_increasing_id()).write.parquet(path) spark.read.parquet(path).createOrReplaceTempView("tbl") - withSQLConf( - CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false", - CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndFallbackReason( - sql("select * from tbl order by 1, 2, 3"), - "can hold a null element or field") - } - - // The default null order agrees with Spark's, so opting in keeps the sort native and right. - withSQLConf( - CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "true", - CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndOperator( - sql("select * from tbl order by 1, 2, 3"), - Seq(classOf[CometSortExec])) + // The sort stays native whether or not incompatible expressions are allowed. + Seq("false", "true").foreach { allowIncompat => + withSQLConf( + CometConf.getExprAllowIncompatConfigKey("SortOrder") -> allowIncompat, + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { + checkSparkAnswerAndOperator( + sql("select * from tbl order by 1, 2, 3"), + Seq(classOf[CometSortExec])) + } } } } @@ -301,8 +294,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { test("strict floating point: nested sort keeps NaN payloads and zero signs unchanged") { // As in the scalar test, a local relation keeps the raw bits that Parquet would canonicalize. // `d` is a primitive `Double`, so it is not nullable, and neither are the element of - // `array(d)` and the field of `named_struct('v', d)`: these keys cannot hold a null, and - // strict mode admits them. + // `array(d)` and the field of `named_struct('v', d)`. The `IF` key's element is null on row + // 4, and strict mode admits it too, because its default null order sorts a null element first, + // where Spark does. val negNan = java.lang.Double.longBitsToDouble(0xfff8000000000002L) val posNan = java.lang.Double.longBitsToDouble(0x7ff8000000000002L) val rows = Seq((0, negNan), (1, posNan), (2, -0.0d), (3, 0.0d), (4, 1.0d)) @@ -315,9 +309,14 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { withSQLConf( CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false", CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - for ((key, value) <- Seq[(String, Row => Double)]( - "array(d)" -> (_.getSeq[Double](1).head), - "named_struct('v', d)" -> (_.getStruct(1).getDouble(0)))) { + // The zeros are peers, and so are the two NaNs, which sort last. + for ((key, value, order) <- Seq[(String, Row => Option[Double], Seq[Int])]( + ("array(d)", r => Some(r.getSeq[Double](1).head), Seq(2, 3, 4, 0, 1)), + ("named_struct('v', d)", r => Some(r.getStruct(1).getDouble(0)), Seq(2, 3, 4, 0, 1)), + ( + "array(IF(id = 4, CAST(NULL AS DOUBLE), d))", + r => Option(r.getSeq[java.lang.Double](1).head).map(_.doubleValue), + Seq(4, 2, 3, 0, 1)))) { val query = s"SELECT id, $key AS k FROM strict_fp_nested_bits ORDER BY k, id" checkSparkAnswerAndOperator( sql(query), @@ -326,12 +325,16 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { val actual = sql(query).collect().toSeq actual.foreach { row => - assert( - java.lang.Double.doubleToRawLongBits(value(row)) == expected(row.getInt(0)), - s"row ${row.getInt(0)} had its floating-point bits rewritten by the sort on $key") + val id = row.getInt(0) + value(row) match { + case Some(v) => + assert( + java.lang.Double.doubleToRawLongBits(v) == expected(id), + s"row $id had its floating-point bits rewritten by the sort on $key") + case None => assert(id == 4, s"row $id lost its value in the sort on $key") + } } - // The zeros are peers, and so are the two NaNs, which sort last. - assert(actual.map(_.getInt(0)) == Seq(2, 3, 4, 0, 1), s"sort on $key") + assert(actual.map(_.getInt(0)) == order, s"sort on $key") } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 700748b6edc..51a918606ce 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -214,10 +214,10 @@ class CometWindowExecSuite extends CometTestBase { assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") // The native RANK and DENSE_RANK find ties by byte equality, on order keys that the native // planner has normalized, nested floats included. So [-0.0] and [0.0] are peers, and the - // cutoff keeps both rows, as Spark does (#5507). Strict floating-point mode declines a nested - // key whose type can hold a null element or field (#6476, #6477), and these columns, read - // back from Parquet, can. ROW_NUMBER never compares peers, and Spark normalizes nested - // floating-point partition keys itself. + // cutoff keeps both rows, as Spark does (#5507). These keys, read back from Parquet, can hold + // a null element, but their default null order places it where Spark does (#6476), so they + // stay native in strict floating-point mode too. ROW_NUMBER never compares peers, and Spark + // normalizes nested floating-point partition keys itself. withTempDir { dir => val path = new Path(dir.toString, "nested_float_order").toString Seq((1, -0.0f, -0.0d), (2, 0.0f, 0.0d), (3, 1.0f, 1.0d)) @@ -248,7 +248,7 @@ class CometWindowExecSuite extends CometTestBase { assert(rankLimit().collect().map(_.getInt(0)).sorted.toSeq == Seq(1, 2)) withSQLConf(CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndFallbackReason(rankLimit(), "can hold a null element or field") + checkSparkAnswerAndOperator(rankLimit(), Seq(classOf[CometWindowGroupLimitExec])) assert(rankLimit().collect().map(_.getInt(0)).sorted.toSeq == Seq(1, 2)) } } From c118b3590bbd8937efe076d006344f2d2e136e95 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 19:54:08 +0000 Subject: [PATCH 44/89] feat: enable Comet's in-memory cache by default (#5634) * perf: project cached batches by buffer selection, prune on collated strings Follow-up to #5051, applying items from #5487. Replace the per-column Arrow IPC stream layout of `CometCachedBatch` with a single encapsulated IPC record batch message per cached batch, carrying no Schema message and no end-of-stream marker. The reader rebuilds the schema from the cached relation's attributes, so a wide relation no longer repeats the same schema bytes once per cached batch. Compression moves from a whole-payload Spark codec to Arrow's per-buffer IPC compression. That is what makes projection cheap: the message metadata records every buffer's offset and length in the body, so `CachedBatchIpc.readProjected` copies out only the byte ranges of the columns a scan selected and decompresses just those. This subsumes the separate "drop the schema message" item, since there is no longer a per-column stream to frame. Dictionary-encoded columns are decoded before being stored: a payload with no schema message cannot describe a dictionary encoding. The codec defaults to zstd, and lz4 is deliberately not offered. Arrow's lz4 is commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`. Over a 200k-row six-column relation it measured 205s to write against 347ms for zstd, while also producing larger output, so no workload prefers it. zstd also beats storing batches uncompressed on both axes (347ms and 2 MiB against 1743ms and 13 MiB), because the bytes it saves cost more to copy and store than compressing them costs. Decompression is done here rather than left to `VectorLoader`, which leaks: `VectorLoader.loadBuffers` collects a field's decompressed buffers into a local list and releases them only after the whole field loads, so a buffer that fails to decompress strands every buffer of that field decompressed before it. A string column reaches this, its offsets buffer decompressing before its data buffer throws. Also track statistics bounds for collated string columns, comparing with the collation's own ordering through a new `CometTypeShim.compareStrings`. Matching the bare `StringType` object excluded collated columns, which then got null bounds and no pruning. Benchmark over a 5M-row six-column relation, keeping the cached scan native against falling back to a Spark cache scan and converting: 1.3x on a repeated scan, 1.3x on a narrow projection and 2.3x on a full projection. * refactor: use Spark's interpreted ordering for bounds, hoist projection layout Cleanup pass over the cache format change. No behaviour change. Drop the `compareStrings` shim in favour of `TypeUtils.getInterpretedOrdering`. That method is public with the same signature on every supported Spark version, and on Spark 4 it resolves a `StringType` through `CollationFactory.fetchCollation(collationId).comparator` -- the comparison the shim was reaching for. So the collation awareness comes from Spark itself and the shim, its Spark 3.x stub and the hand-rolled per-type `compare` all go. The ordering is now resolved once per column per partition rather than being re-dispatched on the `DataType` twice per row. Build the projection's index layout once per partition instead of per batch. The node, buffer and variadic index arithmetic is a pure function of the cached schema and the selected columns, but it walks every field of the relation, so recomputing it per batch made the bookkeeping O(total columns) against O(selected columns) of useful work -- worst in the wide-relation, narrow-projection case the format exists for. `CachedBatchIpc.Projection` now holds that layout and the projected schema, and owns the whole decode; `ProjectedBatch` is left with ownership only. This also puts the projected schema next to the code that packs buffers in the same order, an invariant that previously spanned two files unstated. Smaller cleanups: use Arrow's `DataSizeRoundingUtil.roundUpTo8Multiple` rather than open-coding IPC body alignment; size the serialization buffer from the record batch's known body length instead of growing from 32 bytes; resolve decompressors once instead of per batch; share the dictionary lookup guard between `Utils.combineDictionaryProviders` and the cache writer; read the codec config through one helper carrying the driver-vs-executor rationale; and collapse the duplicated compressed-buffer predicate and scramble loop in the test helper. Corrects two `Utils` scaladocs that still described the per-column stream format this change replaced. Benchmark and codec figures in the docs re-measured against the current code. * fix: relocate the arrow-compression service file when shading arrow-compression ships META-INF/services/org.apache.arrow.vector.compression.CompressionCodec$Factory. The shade plugin copies it verbatim without a ServicesResourceTransformer, so the jar declared a provider for Spark's own unshaded Arrow interface while naming a class that exists here only under the relocated package. Every ServiceLoader lookup Spark's Arrow made then failed with a ServiceConfigurationError, which took CompressionCodec.Factory's static initializer down with it and broke unrelated Arrow IPC reads, including mapInArrow. Add ServicesResourceTransformer so the service file name and its contents are both relocated. arrow-compression is the only bundled artifact that ships one. Also drop an unused NonFatal import that scalafix flagged. * test: drop the cache leak test that depends on zstd corruption detection "releases its vectors when a column fails part way through" zeroed the last 16 bytes of a compressed buffer and required the read to fail. Whether that fails is a property of the zstd runtime, not of Comet: the cached payload is byte-identical across Spark versions, but Comet takes zstd-jni from Spark rather than from arrow-compression, and 1.5.5 (Spark 3.4, 3.5) decodes that frame while 1.5.7 (Spark 4.x) reports it corrupt. So the test passed on 4.x and failed on 3.4 and 3.5. The scenario it claimed to cover is also unreachable: CachedBatchIpc decompresses every selected buffer before VectorLoader runs, so no content corruption can fail part way through the load. The two remaining leak tests corrupt a frame from its header onwards, which every zstd release rejects, and already cover a failure at a column's first buffer and a failure after an earlier buffer of the same column decoded. Records the constraint on scramble so a future test does not reach for a tail-only corruption again, and drops the now unused truncateColumn helper and the dictionary fixture's payload argument. * feat: enable Comet's in-memory cache by default Flip spark.comet.exec.inMemoryCache.enabled to true so cached tables are stored and scanned in Comet's Arrow format without an opt-in. CometDriverPlugin.maybeSetCacheSerializer read the config out of SparkConf with a hardcoded false default, so flipping the ConfigEntry alone would have left the serializer uninstalled unless the user set the key explicitly. It now falls back to the entry's own default, matching how the plugin reads spark.comet.metrics.enabled. Stacked on #5543. * test: cover nested columns in the cached-batch projection tests and benchmark Addresses review feedback asking whether nested data should be tested and benchmarked. Nested columns were already round-tripped, but only under a full projection, which cannot see the part of the format that is nontrivial for them. A flat column always owns one field node and two or three buffers; a nested one owns a run as long as its subtree, and selecting every column covers the whole sequence however it is partitioned. So the buffer-span arithmetic was only exercised in the one shape where getting it wrong does not show. Adds two tests over a six-column relation whose middle four columns are a struct, an array, a map and a struct wrapping an array: - Each column takes its turn as the sole projection while the other five are corrupted, so a run computed short or long is caught by reaching into a corrupted neighbour. - Values are compared against the uncached query across single-column, paired and out-of-order projections. Row counts cannot catch a window that is misaligned but still decompresses, and out-of-order is the case a full projection cannot stand in for. The per-column statistics test now runs over the nested relation too, since a nested column's recorded size is the sum of its whole subtree. Both new tests fail if fieldNodeCount stops recursing into children. In the benchmark, adds the three projection widths over a relation of struct columns, and asserts the width each case claims. That assertion caught the existing "full projection (6 of 6 columns)" case reading three: count() over a non-nullable column is rewritten to count(1) by NullPropagation, which prunes the column out of the scan, and only k, s1 and s2 were nullable -- and those only incidentally, because Remainder can divide by zero. Every column of both relations is now nullable so count(c) genuinely reads c, and the documented numbers are regenerated. Array and map columns are left out of the benchmark deliberately: the baseline arm needs Spark's cache scan to bridge into Comet operators, and CometSparkToColumnarExec declines ArrayType and MapType, so for those the arm does not exist and the two cases stop measuring the same boundary. The docs say so rather than leaving it to be rediscovered. * fix: drop a redundant string interpolator flagged by scalafix RedundantSyntax * review: check the cached layout, and address the rest of the review Reader-side: a cached payload carries no schema, so `Projection` derived every node and buffer window from `Utils.toArrowSchema(cacheAttributes)` with nothing checking the writer had produced that layout. `load` now compares `nodesLength()`/`buffersLength()` against the totals `selectedRange` already computes, before any unchecked `batch.buffers(j)`. Writer-side: `isArrowBacked` accepts a `FixedSizeBinaryVector` for a `BinaryType` column, which is two buffers where the reader rebuilds three, and it answers for the top-level vector only -- so a struct of large strings passes it and is stored with 64-bit offsets. `matchesReaderLayout` compares the batch's Arrow types against the reader's recursively, and a batch that disagrees takes the conversion path instead. A dictionary column's field carries the index type, so the dictionary's field is what is compared. Also: an unrecognized body-compression byte is rejected rather than read as plain bytes, `fieldVariadicCount` and the variadic plumbing are gone (the length check covers view vectors, which the counts would not have), `columnSizes` no longer re-walks each column's subtree, the write codec is a case class rather than a bare tuple, the per-partition `Projection` is lazy so a row-count-only read never builds it, `hydrateDictionaries` is `decodeDictionaries`, `Projection` takes an `IndexedSeq`, and the stale `readProjected` links and some over-long comments are fixed. Tests: the two projection tests become one parameterized over both relations, caching once and restoring the payload between columns instead of re-caching; the two leak tests become one with two corruption points. New tests cover the reader's layout check and the writer declining a fixed-size-binary batch. * review: own the write-side compression buffers, fix the activation example Compressing through VectorUnloader leaks on the failure path: appendNodes retains each input buffer and accumulates the compressed ones into a list local to getRecordBatch, so a buffer that fails to compress strands that retain and leaves every buffer compressed before it reachable from nothing. Closing the input batch afterwards undoes neither. Unload plain and compress in CachedBatchIpc.compressed instead, mirroring what decompressed already does on the read side, so every allocation stays reachable from an error path that owns it. The docs enabled the cache with spark.conf.set, which cannot work: the driver plugin picks spark.sql.cache.serializer while the SparkContext is initializing. Show it as a startup --conf. Also drops a redundant s interpolator that the scalafix lint rejected. * fix: write cached batches to the schema width, not the batch width ArrowWriter.writeColumns drove its loop from the input ColumnarBatch's width while indexing the writer's fields, which come from the schema the batch is written under. That assumed every producer hands over a batch exactly as wide as the schema. Iceberg's vectorized reader does not. BatchDeleteFilter.filterBatch reads with the delete filter's requiredSchema, which carries _pos after the projected columns when a data file has position deletes, and trims the extras back only when the file also has equality deletes. A merge-on-read UPDATE writes position deletes and no equality deletes, so the extra column survives into the batch, and caching such a relation failed with ArrayIndexOutOfBoundsException inside the write loop. Drive the loop from the writer's fields instead, which writes exactly the columns the schema describes: the extras are trailing, the same prefix Iceberg keeps when it does trim. A batch narrower than the schema is a genuine contract violation and is now refused with a message naming both widths. Closes #6087. * fix: report a cached relation's decoded size to the planner, not its compressed size Comet's cache format reported each batch's sizeInBytes as its stored, compressed payload. Spark's planner reads the sum of those as the size of a materialized cached relation, for the broadcast threshold and the shuffled hash join build side among others, and both of Spark's own cache formats report decoded sizes there. With the default zstd codec a cached relation could look several times smaller than in Spark's format and be broadcast where Spark would shuffle it. Record each column's decoded Arrow size, measured before compression as Spark's ArrowCachedBatchSerializer does, and let CometCachedBatch inherit sizeInBytes from SimpleMetricsCachedBatch as the sum of those. * fix: keep Spark's cache scan for a relation whose cached plan records observed metrics Spark collects Dataset.observe metrics after a query with CollectMetricsExec.collect, which reaches the metrics recorded inside a cached plan only through an InMemoryTableScanExec over it. With Comet's native cache scan in its place those metrics were lost: QueryExecution.observedMetrics came back empty, Observation.get returned an empty map on Spark 3.5 and later, and on Spark 3.4 it never returned. Keep Spark's scan for such a relation, with a fallback reason. The data stays in Comet's format and is read through the existing fallback path. The check walks the cached plan the way CollectMetricsExec.collect does, through subqueries, adaptive plans and nested caches. * test: install Comet's cache serializer in the Spark SQL test sessions The Spark SQL diffs turn Comet on through SharedSparkSession and TestHive rather than by loading CometPlugin, and the plugin is what installs Comet's cache serializer. So even with spark.comet.exec.inMemoryCache.enabled on by default, every Spark SQL suite cached in Spark's own format and none of them exercised Comet's. Set spark.sql.cache.serializer in both places when Comet is enabled, as the plugin would, and adapt the tests that assume Spark's cache scan: - plan checks that only need a cache scan to be present also accept CometInMemoryTableScanExec, reading its originalPlan where the test uses the scan's relation; - PartitionBatchPruningSuite checks its answers as before, but reads the test-only accumulators only from Spark's scan, which Comet's lacks; - CacheTableInKryoSuite registers Comet's classes with CometKryoRegistrator, as Comet asks of any application that sets spark.kryo.registrationRequired; - tests of Spark internals that Comet's scan replaces are tagged IgnoreComet: exact cache size estimates, the union's columnar support, subquery reuse through the scan's predicates, and AQE coalescing under a union (#6454). Each diff was regenerated from a clone of its Spark tag with the existing diff applied, after checking that it round-tripped unchanged. * fix: coalesce shuffle partitions under Comet unions the way Spark does Spark's CoalesceShufflePartitions coalesces each child of a UnionExec as its own group, and from Spark 4.0 each child of a CartesianProductExec, BroadcastHashJoinExec or BroadcastNestedLoopJoinExec too. It matches those classes, so a CometUnionExec fell through to the case that coalesces only when every leaf below it is an exchange stage. A union with a scan or a table-cache stage in one branch kept every partition of the shuffles in the others. Add CometCoalesceShufflePartitions, a query-stage optimizer rule that runs after Spark's. For a Comet operator whose Spark original is one of those, and whose shuffle stages no AQE rule has touched, it rebuilds the Spark operators over the Comet children, runs Spark's own rule over that, and swaps the Comet operators back in. It extends AQEShuffleReadRule, so AQE gates and validates it the way it does Spark's rule. Spark 3.4 has no hook for query-stage optimizer rules, so the fix applies from Spark 3.5. Closes #6454. * test: fix three Spark SQL test adaptations for Comet's cache format - CachedBatchSerializerNoUnwrapSuite sets its own cache serializer, but Spark keeps the first one it loads for the rest of the JVM, so the suite ran with Comet's. Clear it before and after the suite, as CachedBatchSerializerSuite does (4.1.3, 4.2.0). - The SPARK-19993 helper counted the subqueries in a Comet cache scan's original plan, which subquery reuse never reaches, on top of the same subqueries in the filter above the scan. - SPARK-35332's AQE check reads the cached plan from the SQL UI's plan tree, which shows it only under Spark's own cache scan (#6463). Skip it under Comet (4.0.4 and later). * fix: keep Spark's cache format when Kryo would reject Comet's or Comet disables itself CometDriverPlugin installed ArrowCachedBatchSerializer whenever Comet, its native execution and the cache config were enabled at startup. Two more startup settings decide whether that format can work: - With spark.kryo.registrationRequired=true and no CometKryoRegistrator, Kryo rejects CometCachedBatch, so caching failed with "Class is not registered" as soon as Spark serialized a cached block, where Spark's own format works. The plugin only warned. It now keeps Spark's format, and the warning says so. - With Comet shuffle enabled but neither of Comet's shuffle managers configured, Comet disables itself, so every cache was stored in Comet's format with only Spark operators to read it. The plugin now keeps Spark's format there too. The plugin's boolean config reads also honor a deprecated alternative key, such as spark.comet.exec.shuffle.enabled, as a session does. * docs: describe the in-memory cache default in the upgrade guide Add a 1.2.0 section to the upgrade guide for the new default of spark.comet.exec.inMemoryCache.enabled: what changes, how to keep Spark's cache format, and when Comet keeps it without being asked. The operator and compatibility pages still called the feature disabled by default, and the cache guide still said an application that started with the default kept Spark's format. * test: benchmark the in-memory cache against Spark's format with AQE on The published comparisons either read the same Comet-written cache both ways, with spark.comet.sparkToColumnar.enabled turned on, or read both formats with Comet off. Neither is what an application gets from the feature's default. Add cases with Comet and AQE on and Comet's other settings at their defaults, reading a relation cached in Spark's format and in Comet's, with Comet operators above the cache scan and with a Spark operator above it. * feat: record why Spark scans a relation cached in Comet's format when Comet is off spark.sql.cache.serializer is static, so a relation cached in Comet's format keeps it for the life of the application. A session that turns Comet or its native execution off then reads it with Spark's InMemoryTableScanExec, which is slower than reading Spark's own format (#5485). CometExecRule records a fallback reason for the other ways Spark ends up scanning a cache, but it returned before looking at the plan in these two cases, so nothing explained them. Record a reason on each such scan, naming the cause and the startup setting that keeps caches in Spark's format. * docs: say an application keeps the cache format it started with The sentence described a session started with the default, which keeps Spark's format only while the default is off. * test: format the in-memory cache benchmark * fix: bind the fallback reason result for the strict Scala warnings build -Ywarn-value-discard rejects the node that withFallbackReason returns, discarded as the last expression of a case. * docs: compare the in-memory cache with Spark's format as an application gets it Add the benchmark's adaptive results to the cache guide: with Comet and AQE on and Comet's other settings at their defaults, Comet's format is as fast or faster than Spark's in every shape but a read of three long columns, which zstd decompression makes about 10% slower. That includes a Spark operator above the native scan. The guide said that, without the feature, a CometSparkColumnarToColumnar converts each batch of Spark's cache scan for Comet. That takes spark.comet.sparkToColumnar.enabled, which is off by default, so by default the operators above Spark's scan run on Spark. The limitation that Spark operators read Comet's format slowly applies to Spark's own cache scan, not to Spark operators above Comet's. * fix: count Kryo registrations however the application made them The gate kept Spark's cache format whenever spark.kryo.registrator did not list CometKryoRegistrator, even where the application had registered Comet's cached batch another way, such as spark.kryo.classesToRegister. Comet's format works there, while Spark's does not before 4.1, which registers DefaultCachedBatch itself only from then on, so the gate broke caches that worked. Ask a Kryo instance built from the application's conf which of Comet's classes it has registered, and keep Spark's format only if Comet's cached batch is not among them. If Kryo cannot be built, fall back to looking for CometKryoRegistrator in spark.kryo.registrator. The startup warning uses the same check and names the missing classes. * docs: describe the Kryo condition by registration in the upgrade guide The plugin now keeps Spark's cache format when Kryo has not registered Comet's cached batch, however the application registers classes, rather than when spark.kryo.registrator does not list CometKryoRegistrator. * docs: name spark.comet.convert.inMemoryCache.enabled for converting Spark's cache scan spark.comet.sparkToColumnar.enabled is deprecated for in-memory cached tables since #6602, which gave the conversion its own config. * test: accept Comet's cache scan in BroadcastJoinSuite's cached-relation tests Since #6415, suites that build their own SparkSession run Comet, and with the cache on by default BroadcastJoinSuite's cached relations are read by CometInMemoryTableScanExec. SPARK-23214 and SPARK-37742 count Spark's cache scan and broadcast hash join, so count Comet's as well. * fix: coalesce with the whole stage in view and keep partitionings above Comet unions Address review of the CometCoalesceShufflePartitions rule: - Stand in only for Comet operators whose output partitioning is unknown. From Spark 4.1 a union of co-partitioned children reports their partitioning, and an aggregate above reads it without a shuffle. Coalescing one branch made Comet's aggregate, which states no required distribution, return each key twice. - Hand Spark's rule the whole stage rather than the subtree below each Comet operator, so it sees a Cartesian product or nested loop join above a union and uses its smaller target size there, and divides the minimum partition count over every group. Leave a stage that already has a read over a shuffle, which Spark's rule would coalesce again. - Copy the Spark original when withNewChildren hands it back, so a union that AQE plans again over its stages is coalesced too. - Drop the CartesianProductExec case, which has no Comet counterpart, and document the rule in CometSparkSessionExtensions. * test: run Spark's SPARK-42101 union coalescing test with Comet again The test failed with Comet's cache format because AQE did not coalesce the shuffled branch of a union that Comet ran (#6454), which #6459 fixes. * fix: coalesce Comet union branches in stages that already hold an AQE shuffle read CometCoalesceShufflePartitions returned a stage unchanged when it held any AQEShuffleReadExec, because Spark's rule asserts on reads that are not skew splits. Below a Cartesian product, Spark's rule coalesces the shuffle on the other side by itself first, so the read it leaves there kept the shuffled branch of a Comet union at every partition. Hide each read behind a leaf while Spark's rule runs, and put it back afterwards. Spark's rule then leaves alone the shuffles it would coalesce together with a read one, and coalesces the rest of the stage. * docs: say Spark's cache scan can read Comet's format more slowly The fused reader from #5859 reads Comet's format about as fast as Spark's own, and faster for the narrowest reads, so the upgrade note should not say Spark's scan is always slower. * test: leave Comet's cache scan out of the DPP suite's subquery check Since #6577, CometInMemoryTableScanExec exposes Spark's own InMemoryTableScanExec as its one subquery, so that Spark's UI draws the cached plan below it. With the cache on by default, the DPP suite's 'filtering ratio policy fallback' caches its dimension table, and checkPartitionPruningPredicate requires every subquery of an adaptive plan to contain an AdaptiveSparkPlanExec, which Spark's scan does not. Skip it there in every diff, since the query never runs it as a subquery. * feat: keep Comet's cache format opt-in on Spark 3.4 Spark 3.4 has no hook for query-stage optimizer rules, so CometCoalesceShufflePartitions cannot run there, and AQE cannot coalesce the shuffled branch of a Comet union whose other branch reads a cache in Comet's format (#6454). spark.comet.exec.inMemoryCache.enabled now defaults to true only from Spark 3.5, and the plugin, which follows the entry's default, keeps Spark's format on 3.4 unless it is set. The 3.4 Spark SQL diff goes back to main's: its sessions install Comet's serializer only as the plugin would, which on 3.4 it no longer does by default, so its cache test adaptations are not needed there. * docs: rewrap the Spark 3.4 cache default sentences * test: run Spark's SPARK-35332 AQE cache test with Comet again The diffs skipped CachedTableSuite's 'SPARK-35332: Make cache plan disable configs configurable - check AQE' on Spark 4.0 and later, because SparkPlanInfo drew the cached plan only below Spark's own InMemoryTableScanExec. Since #6577 it is drawn below CometInMemoryTableScanExec too, so the test finds the cached plan. The test's check that the cached plan was coalesced, which the diffs already adapt, looked for AQEShuffleRead only below a columnar-to-row transition, where it is when Comet's plan is cached in Spark's format. In Comet's columnar format the cached plan ends in the read itself, as in Spark, so the check accepts it there too. The test passes on Spark 4.0, 4.1 and 4.2 with either format, and with Comet off. * docs: stop describing the in-memory cache as experimental It is enabled by default from Spark 3.5. The operator table now marks InMemoryTableScanExec as supported, and the plan-node table no longer says the cache scan is disabled by default. * test: skip only the format-dependent checks in the Spark SQL diffs SPARK-36120, SPARK-33687, SPARK-22673 and SPARK-43376 now run with Comet's cache format and skip only the size or subquery-reuse assertions that depend on it. The union test in CoalesceShufflePartitionsSuite runs its coalescing check again on Spark 3.5 and later, where CometCoalesceShufflePartitions coalesces the shuffles on the union's join side. From Spark 3.5, CometTestSettings also makes Comet's cache serializer the default of the sql and hive test JVMs, next to the shuffle manager. Spark keeps the first cache serializer it loads for the rest of the JVM, so a suite that caches first through a session of its own left the suites after it in Spark's format. * test: time building the cache in Spark's format in the cache benchmark runBuildBenchmark caches the same relation in Spark's format and in Comet's, with each codec, from Comet operators and from Spark operators, and reports the footprint of each. The cache guide has the results. * test: skip only the columnar checks of SPARK-37371 under Comet Comet runs the union of two cached relations as a CometUnion over its own cache scans, so the checks on Spark's scans and union are skipped there. The union with a local relation stays Spark's non-columnar UnionExec, and both answers are checked, under Comet as well. --------- Co-authored-by: test --- dev/diffs/3.5.9.diff | 501 ++++++++++++-- dev/diffs/4.0.4.diff | 548 ++++++++++++++-- dev/diffs/4.1.3.diff | 618 ++++++++++++++++-- dev/diffs/4.2.0.diff | 618 ++++++++++++++++-- .../contributor-guide/spark-sql-tests.md | 4 +- .../latest/compatibility/operators.md | 11 +- .../user-guide/latest/in-memory-cache.md | 78 ++- .../user-guide/latest/migration-guide.md | 28 + docs/source/user-guide/latest/operators.md | 2 +- .../latest/understanding-comet-plans.md | 2 +- .../scala/org/apache/comet/CometConf.scala | 12 +- .../comet/exec/CometInMemoryCacheSuite.scala | 6 + .../CometInMemoryCacheBenchmark.scala | 266 +++++++- 13 files changed, 2395 insertions(+), 299 deletions(-) diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index 801199ddb8e..4ee26ebe93c 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -38,7 +38,7 @@ index 49eca0f7555..b9e5d69a6c6 100644 diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index daa3ebac992..2b6138c4280 100644 +index daa3ebac992..2c721d85961 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -459,9 +459,11 @@ object SparkBuild extends PomBuild { @@ -53,7 +53,7 @@ index daa3ebac992..2b6138c4280 100644 enable(SparkConnectCommon.settings)(connectCommon) enable(SparkConnect.settings)(connect) -@@ -1198,6 +1200,25 @@ object SQL { +@@ -1198,6 +1200,35 @@ object SQL { ) } @@ -61,16 +61,26 @@ index daa3ebac992..2b6138c4280 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -260,10 +270,14 @@ index 0efe0877e9b..423d3b3d76d 100644 -- SELECT_HAVING -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index e5494726695..00937f025c2 100644 +index e5494726695..9ceb28855db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -35,10 +35,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -272,7 +286,28 @@ index e5494726695..00937f025c2 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -519,7 +519,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -113,6 +114,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -393,7 +398,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -519,7 +525,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -282,6 +317,18 @@ index e5494726695..00937f025c2 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { +@@ -1581,7 +1588,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 6f3090d8908..52523429c87 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -348,6 +395,30 @@ index 56e9520fdab..917932336df 100644 spark.range(50).write.saveAsTable(s"$dbName.$table1Name") spark.range(100).write.saveAsTable(s"$dbName.$table2Name") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +index 747f43fa2a7..1b76251f0bd 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +@@ -1417,11 +1417,14 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index 7ee18df3756..d09f70e5d99 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -471,6 +542,83 @@ index a1d5d579338..8825683ebcd 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index bda8c7f2608..10a70745d07 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -35,6 +37,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -42,7 +48,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -252,7 +258,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -285,14 +291,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -300,7 +306,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -308,7 +314,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index c4fb4fa943c..a04b23870a8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -494,10 +642,10 @@ index c4fb4fa943c..a04b23870a8 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index f33432ddb6f..b375e285dde 100644 +index f33432ddb6f..4e8522d3aab 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -505,7 +653,11 @@ index f33432ddb6f..b375e285dde 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -513,7 +665,19 @@ index f33432ddb6f..b375e285dde 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.index @@ -522,7 +686,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -535,7 +699,7 @@ index f33432ddb6f..b375e285dde 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -554,7 +718,7 @@ index f33432ddb6f..b375e285dde 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1577,6 +1597,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -562,7 +726,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1729,6 +1750,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1054,10 +1218,18 @@ index c26757c9cff..d55775f09d7 100644 protected val baseResourcePath = { // use the same way as `SQLQueryTestSuite` to get the resource path diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 3cf2bfd17ab..5bcf9478e9b 100644 +index 3cf2bfd17ab..11d3ca4d5a9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1067,7 +1239,7 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1979,8 +1980,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1979,8 +1981,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1084,7 +1256,15 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -3750,7 +3758,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3289,6 +3298,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -3750,7 +3760,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark } } @@ -1126,7 +1306,7 @@ index 71af1fd69c3..81a04c93c9c 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index 8b4ac474f87..5218e9125c6 100644 +index 8b4ac474f87..b097a14a972 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -167,7 +167,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1145,7 +1325,15 @@ index 8b4ac474f87..5218e9125c6 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -261,6 +269,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -210,6 +218,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -261,6 +270,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1153,7 +1341,7 @@ index 8b4ac474f87..5218e9125c6 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -316,6 +325,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -316,6 +326,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1161,7 +1349,7 @@ index 8b4ac474f87..5218e9125c6 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -484,6 +494,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -484,6 +495,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1169,6 +1357,33 @@ index 8b4ac474f87..5218e9125c6 100644 def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = { assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index e827396009d..efbee3975e3 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -829,14 +829,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index 04702201f82..4d38d8d6e51 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -1406,7 +1621,7 @@ index ae1c0a86a14..1d3b914fd64 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index e11191da6a9..005fd4fa657 100644 +index e11191da6a9..0988faf4fcd 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -51,6 +51,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite { @@ -1466,23 +1681,6 @@ index e11191da6a9..005fd4fa657 100644 } } } -@@ -473,10 +484,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite { - // aggregate on the other side of the union. - val finalPlan = resultDf.queryExecution.executedPlan - .asInstanceOf[AdaptiveSparkPlanExec].executedPlan -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -1749,7 +1947,7 @@ index 5a413c77754..207b66e1d7b 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 2f8e401e743..f7923d2b9a7 100644 +index 2f8e401e743..fc4ce35be4b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._ @@ -2155,7 +2353,18 @@ index 2f8e401e743..f7923d2b9a7 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite +@@ -2932,13 +2975,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -2165,7 +2374,7 @@ index 2f8e401e743..f7923d2b9a7 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -2980,7 +3024,9 @@ class AdaptiveQueryExecSuite +@@ -2980,7 +3028,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2176,6 +2385,147 @@ index 2f8e401e743..f7923d2b9a7 100644 plan.inputPlan.output.zip(plan.finalPhysicalPlan.output).foreach { case (o1, o2) => assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index de04938f247..f3351a86946 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -543,7 +550,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 885286843a1..f0f805a6cd9 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd52d038ca6..154c800be67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -2618,10 +2968,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 5479be86e9f..07c81e4b830 100644 +index 5479be86e9f..103dab4657c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,12 +21,13 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -2630,7 +2980,28 @@ index 5479be86e9f..07c81e4b830 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -2640,6 +3011,21 @@ index 5479be86e9f..07c81e4b830 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -655,12 +659,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 5cdbdc27b32..307fba16578 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -3275,10 +3661,10 @@ index e937173a590..263934fbe7b 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index c23bf4204f7..07d215aad2b 100644 +index c23bf4204f7..d922f0b3d23 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -97,6 +97,19 @@ trait SharedSparkSessionBase +@@ -97,6 +97,24 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3294,6 +3680,11 @@ index c23bf4204f7..07d215aad2b 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -3379,10 +3770,10 @@ index 6160c3e5f6c..bfc0c618a9b 100644 test("SPARK-4963 DataFrame sample on mutable row return wrong result") { diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index 1d646f40b3e..c8192f52f98 100644 +index 1d646f40b3e..b8a043c6248 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,25 +53,41 @@ object TestHive +@@ -53,25 +53,45 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -3434,6 +3825,10 @@ index 1d646f40b3e..c8192f52f98 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + conf diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 21ef06bdfec..895f6bd3db8 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -78,7 +78,7 @@ index 046f357ff79..7c6fd199e41 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 9c2312ed256..b151f6a630c 100644 +index 9c2312ed256..09026e8dc8a 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -414,9 +414,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 9c2312ed256..b151f6a630c 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1183,6 +1185,25 @@ object SQL { +@@ -1183,6 +1185,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 9c2312ed256..b151f6a630c 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -375,10 +385,14 @@ index 21a3ce1e122..f4762ab98f0 100644 -- In COMPENSATION views get invalidated if the type can't cast diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 0f42502f1d9..e9ff802141f 100644 +index 0f42502f1d9..98479690165 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -36,10 +36,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -387,7 +401,28 @@ index 0f42502f1d9..e9ff802141f 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -520,7 +520,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -114,6 +115,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -394,7 +399,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -520,7 +526,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -397,7 +432,19 @@ index 0f42502f1d9..e9ff802141f 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1659,9 +1660,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1588,7 +1595,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1659,9 +1669,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -405,20 +452,29 @@ index 0f42502f1d9..e9ff802141f 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1750,6 +1770,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 9db406ff12f..e59a931c06a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -485,6 +541,30 @@ index ed182322aec..1ae6afa686a 100644 spark.range(50).write.saveAsTable(s"$dbName.$table1Name") spark.range(100).write.saveAsTable(s"$dbName.$table2Name") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +index 332be4c7bbc..ec66595e873 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +@@ -1493,11 +1493,14 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index d9ce3000a0c..f2d044ed6b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -608,6 +688,97 @@ index 552e2b2e274..17a5ae20f0f 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 9d8aaf8d90e..41afdbdafc1 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -36,6 +38,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -43,7 +49,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -253,7 +259,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -286,14 +292,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -301,7 +307,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -309,7 +315,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -321,11 +327,11 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 81713c777bc..b5f92ed9742 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -631,10 +802,10 @@ index 81713c777bc..b5f92ed9742 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index 2c24cc7d570..12d897866da 100644 +index 2c24cc7d570..505157acd91 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -642,7 +813,11 @@ index 2c24cc7d570..12d897866da 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -650,7 +825,19 @@ index 2c24cc7d570..12d897866da 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -659,7 +846,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -672,7 +859,7 @@ index 2c24cc7d570..12d897866da 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -691,7 +878,7 @@ index 2c24cc7d570..12d897866da 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -699,7 +886,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1201,10 +1388,18 @@ index ad424b3a7cc..4ece0117a34 100644 protected val baseResourcePath = { // use the same way as `SQLQueryTestSuite` to get the resource path diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index f294ff81021..62ca9a9bcd5 100644 +index f294ff81021..f019c6de8e9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1524,7 +1524,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1524,7 +1525,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1214,7 +1409,7 @@ index f294ff81021..62ca9a9bcd5 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1985,8 +1986,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1985,8 +1987,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1231,7 +1426,15 @@ index f294ff81021..62ca9a9bcd5 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -4744,7 +4752,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3278,6 +3287,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4744,7 +4754,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark ) } @@ -1273,7 +1476,7 @@ index 575a4ae69d1..129d9f27232 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index c1c041509c3..d068241d585 100644 +index c1c041509c3..6fc92c40ae7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -179,7 +179,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1292,7 +1495,15 @@ index c1c041509c3..d068241d585 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -273,6 +281,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -222,6 +230,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -273,6 +282,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1300,7 +1511,7 @@ index c1c041509c3..d068241d585 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -328,6 +337,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -328,6 +338,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1308,7 +1519,7 @@ index c1c041509c3..d068241d585 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -496,6 +506,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -496,6 +507,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1331,6 +1542,33 @@ index 5ba69c8f9d9..ac1256afe88 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 0df7f806272..9cdfe8b8f46 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1530,7 +1768,7 @@ index a40e34d94d0..abc1f035d15 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 11e9547dfc5..ba340c4ebcf 100644 +index 11e9547dfc5..327e0b1bc5b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -24,6 +24,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1538,7 +1776,7 @@ index 11e9547dfc5..ba340c4ebcf 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1597,6 +1835,15 @@ index 11e9547dfc5..ba340c4ebcf 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1897,7 +1909,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index 3eeed2e4175..9f21d547c1c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1779,7 +2026,7 @@ index 04d33ecd3d5..450df347297 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index da43b0cfc58..ca1f6f47a9f 100644 +index da43b0cfc58..653cd7036d1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -54,6 +54,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -1839,23 +2086,6 @@ index da43b0cfc58..ca1f6f47a9f 100644 } } } -@@ -468,10 +479,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2406,7 +2636,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index fb8fab6a80f..c5768935805 100644 +index fb8fab6a80f..cfcff43d557 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -28,12 +28,14 @@ import org.apache.spark.SparkException @@ -2891,7 +3121,18 @@ index fb8fab6a80f..c5768935805 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3048,7 +3101,8 @@ class AdaptiveQueryExecSuite +@@ -3042,13 +3095,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -2901,7 +3142,7 @@ index fb8fab6a80f..c5768935805 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3138,7 +3192,8 @@ class AdaptiveQueryExecSuite +@@ -3138,7 +3196,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2911,6 +3152,147 @@ index fb8fab6a80f..c5768935805 100644 plan.inputPlan.output.zip(plan.finalPhysicalPlan.output).foreach { case (o1, o2) => assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..9610c703610 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -544,7 +551,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3353,10 +3735,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 69dd04e07d5..781018ecc66 100644 +index 69dd04e07d5..a78dd593985 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3365,7 +3747,29 @@ index 69dd04e07d5..781018ecc66 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3375,6 +3779,21 @@ index 69dd04e07d5..781018ecc66 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 0dd90925d3c..7d53ec845ef 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4077,10 +4496,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4097,6 +4516,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4219,10 +4643,10 @@ index b67370f6eb9..746b3974b29 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4267,6 +4691,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index d034d20b990..0b814aa1b10 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -78,7 +78,7 @@ index 370bfa4397f..ccb4398de58 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 01aa71047fe..cde29aabc83 100644 +index 01aa71047fe..06913c0656b 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -443,9 +443,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 01aa71047fe..cde29aabc83 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1347,6 +1349,25 @@ object SQL { +@@ -1347,6 +1349,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 01aa71047fe..cde29aabc83 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -400,11 +410,41 @@ index 21a3ce1e122..f4762ab98f0 100644 SET spark.sql.ansi.enabled = false; -- In COMPENSATION views get invalidated if the type can't cast +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 26d8f750f6e..c888f8e0844 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -32,9 +32,16 @@ class CacheTableInKryoSuite extends QueryTest + with SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 0d807aeae4d..6d7744e771b 100644 +index 0d807aeae4d..b5a00b75580 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp import org.apache.spark.sql.execution.columnar._ import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation @@ -413,7 +453,28 @@ index 0d807aeae4d..6d7744e771b 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -534,7 +534,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -128,6 +129,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -408,7 +413,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -534,7 +540,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -423,7 +484,19 @@ index 0d807aeae4d..6d7744e771b 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1673,9 +1674,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1602,7 +1609,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1673,9 +1683,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -431,20 +504,45 @@ index 0d807aeae4d..6d7744e771b 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1764,6 +1784,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2630,6 +2651,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2657,6 +2679,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index bfe15b33768..df8fb4c0a35 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -512,7 +610,7 @@ index ed182322aec..1ae6afa686a 100644 spark.range(100).write.saveAsTable(s"$dbName.$table2Name") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala -index 93ff7becaec..87537a25b3b 100644 +index 93ff7becaec..78b02628d42 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -23,10 +23,11 @@ import java.util.Locale @@ -528,7 +626,24 @@ index 93ff7becaec..87537a25b3b 100644 import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData} -@@ -1519,11 +1520,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1497,9 +1498,13 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported, 2) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported, 1) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported, 2) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[UnionExec], supported, 1) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) +@@ -1519,11 +1524,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"a")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -542,7 +657,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1554,11 +1556,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1554,11 +1560,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -556,7 +671,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1573,10 +1576,10 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1573,10 +1580,10 @@ class DataFrameSetOperationsSuite extends QueryTest // Avoid unnecessary shuffle if union output partitioning is enabled val shuffledUnion = union.repartition($"a") val shuffleNumBefore = union.queryExecution.executedPlan.collect { @@ -569,7 +684,7 @@ index 93ff7becaec..87537a25b3b 100644 } if (enabled) { -@@ -1605,6 +1608,7 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1605,6 +1612,7 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -710,6 +825,99 @@ index 4a070becfa6..61d515d127e 100644 val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value") val window = Window.partitionBy($"key").orderBy($"value") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 627811eaecf..7269aeab949 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -39,6 +40,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -46,7 +51,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -256,7 +261,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -289,14 +294,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -304,7 +309,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -312,7 +317,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -324,12 +329,12 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 6df8d66ee7f..35e270c7241 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -733,10 +941,10 @@ index 6df8d66ee7f..35e270c7241 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index e1a2fd33c7c..632f4b695df 100644 +index e1a2fd33c7c..ac0682e0e1a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -744,7 +952,11 @@ index e1a2fd33c7c..632f4b695df 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -752,7 +964,19 @@ index e1a2fd33c7c..632f4b695df 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -761,7 +985,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -774,7 +998,7 @@ index e1a2fd33c7c..632f4b695df 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -793,7 +1017,7 @@ index e1a2fd33c7c..632f4b695df 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1579,6 +1595,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1579,6 +1599,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -801,7 +1025,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1731,6 +1748,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1731,6 +1752,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1322,10 +1546,18 @@ index cb9d0909554..084d6515e8b 100644 } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 74cdee49e55..1a7bc4567b4 100644 +index 74cdee49e55..da1e9bd966a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1335,7 +1567,7 @@ index 74cdee49e55..1a7bc4567b4 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1982,8 +1983,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1982,8 +1984,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1352,7 +1584,15 @@ index 74cdee49e55..1a7bc4567b4 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -4741,7 +4749,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3275,6 +3284,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4741,7 +4751,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark ) } @@ -1407,7 +1647,7 @@ index 23f0144dcec..40d536bb23a 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index 66826a9ca76..a330d142eb1 100644 +index 66826a9ca76..db56c426dc0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -196,7 +196,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1426,7 +1666,15 @@ index 66826a9ca76..a330d142eb1 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -290,6 +298,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -239,6 +247,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -290,6 +299,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1434,7 +1682,7 @@ index 66826a9ca76..a330d142eb1 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -345,6 +354,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -345,6 +355,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1442,7 +1690,7 @@ index 66826a9ca76..a330d142eb1 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -513,6 +523,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -513,6 +524,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1465,6 +1713,33 @@ index d7b2511eac2..d5f5b940b94 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 7bfc8cf4fa6..4bd387801db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1664,7 +1939,7 @@ index 8a0e2c29653..d276a51cbc6 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 8f7a68bcbe6..88dbe1793c9 100644 +index 8f7a68bcbe6..c09c5d74309 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -26,6 +26,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1672,7 +1947,7 @@ index 8f7a68bcbe6..88dbe1793c9 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1731,6 +2006,15 @@ index 8f7a68bcbe6..88dbe1793c9 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1948,7 +1960,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index a09b7e0827c..ffc29f764bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1913,7 +2197,7 @@ index fcecaf25d4c..e5a511022cc 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index 28762f01d7a..6cdc3953bf3 100644 +index 28762f01d7a..d19f7ff9d34 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -56,6 +56,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -1973,23 +2257,6 @@ index 28762f01d7a..6cdc3953bf3 100644 } } } -@@ -475,10 +486,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2540,7 +2807,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 188a28ff1c0..2717d45f3d0 100644 +index 188a28ff1c0..d1d8714cd32 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,12 +27,14 @@ import org.apache.spark.SparkException @@ -3025,7 +3292,18 @@ index 188a28ff1c0..2717d45f3d0 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3229,7 +3282,8 @@ class AdaptiveQueryExecSuite +@@ -3223,13 +3276,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -3035,7 +3313,7 @@ index 188a28ff1c0..2717d45f3d0 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3319,7 +3373,8 @@ class AdaptiveQueryExecSuite +@@ -3319,7 +3377,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3046,10 +3324,29 @@ index 188a28ff1c0..2717d45f3d0 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 47b935a2880..3fdeab3113c 100644 +index 47b935a2880..65ee66c1975 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -@@ -230,9 +230,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest +@@ -212,6 +212,18 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -230,9 +242,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3074,6 +3371,147 @@ index 47b935a2880..3fdeab3113c 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..9610c703610 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -544,7 +551,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3561,10 +3999,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 9bd858608cb..2682ba53513 100644 +index 9bd858608cb..896d33a8f0e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3573,7 +4011,29 @@ index 9bd858608cb..2682ba53513 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3583,6 +4043,21 @@ index 9bd858608cb..2682ba53513 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index f2e9121d566..2c9f517034f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4390,10 +4865,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4410,6 +4885,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4532,10 +5012,10 @@ index b67370f6eb9..746b3974b29 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4580,6 +5060,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index e53ec96c3f5..f58f210fbef 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -78,7 +78,7 @@ index 46558134f41..7fa990be776 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 7c5e49f3a69..28f83381fa0 100644 +index 7c5e49f3a69..1091462f8ef 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -455,9 +455,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 7c5e49f3a69..28f83381fa0 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1426,6 +1428,25 @@ object SQL { +@@ -1426,6 +1428,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 7c5e49f3a69..28f83381fa0 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -418,11 +428,41 @@ index 21a3ce1e122..f4762ab98f0 100644 SET spark.sql.ansi.enabled = false; -- In COMPENSATION views get invalidated if the type can't cast +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 72a2da16054..54146036f8a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -30,9 +30,16 @@ import org.apache.spark.storage.StorageLevel + class CacheTableInKryoSuite extends SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 085dbcd8046..3090d321b6c 100644 +index 085dbcd8046..2569c53815f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp import org.apache.spark.sql.execution.columnar._ import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation @@ -431,7 +471,28 @@ index 085dbcd8046..3090d321b6c 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -564,7 +564,8 @@ class CachedTableSuite extends SharedSparkSession +@@ -127,6 +128,10 @@ class CachedTableSuite extends SharedSparkSession + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -407,7 +412,8 @@ class CachedTableSuite extends SharedSparkSession + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -564,7 +570,8 @@ class CachedTableSuite extends SharedSparkSession df.collect() } assert( @@ -441,7 +502,19 @@ index 085dbcd8046..3090d321b6c 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1703,9 +1704,18 @@ class CachedTableSuite extends SharedSparkSession +@@ -1632,7 +1639,10 @@ class CachedTableSuite extends SharedSparkSession + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1703,9 +1713,19 @@ class CachedTableSuite extends SharedSparkSession _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -449,20 +522,45 @@ index 085dbcd8046..3090d321b6c 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1794,6 +1814,7 @@ class CachedTableSuite extends SharedSparkSession + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2660,6 +2681,7 @@ class CachedTableSuite extends SharedSparkSession + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2687,6 +2709,7 @@ class CachedTableSuite extends SharedSparkSession + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 5b8154d2900..d42374b127f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -530,7 +628,7 @@ index 9733d51a91c..395a108abc8 100644 spark.range(100).write.saveAsTable(s"$dbName.$table2Name") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala -index d838ba4c234..cb0573d56d0 100644 +index d838ba4c234..95ba4076c4d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -23,10 +23,11 @@ import java.util.Locale @@ -546,7 +644,24 @@ index d838ba4c234..cb0573d56d0 100644 import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData} -@@ -1518,11 +1519,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1496,9 +1497,13 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported, 2) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported, 1) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported, 2) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[UnionExec], supported, 1) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) +@@ -1518,11 +1523,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartition($"a").union(df2.repartition($"a")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -560,7 +675,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1553,11 +1555,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1553,11 +1559,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartition($"a").union(df2.repartition($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -574,7 +689,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1572,10 +1575,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1572,10 +1579,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP // Avoid unnecessary shuffle if union output partitioning is enabled val shuffledUnion = union.repartition($"a") val shuffleNumBefore = union.queryExecution.executedPlan.collect { @@ -587,7 +702,7 @@ index d838ba4c234..cb0573d56d0 100644 } if (enabled) { -@@ -1604,6 +1607,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1604,6 +1611,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -728,6 +843,99 @@ index f79824de8ff..5432984960f 100644 val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value") val window = Window.partitionBy($"key").orderBy($"value") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 0d1b0e1d981..e12c11e7b59 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -38,6 +39,10 @@ class DatasetCacheSuite extends SharedSparkSession + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -45,7 +50,7 @@ class DatasetCacheSuite extends SharedSparkSession + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -255,7 +260,7 @@ class DatasetCacheSuite extends SharedSparkSession + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -288,14 +293,14 @@ class DatasetCacheSuite extends SharedSparkSession + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -303,7 +308,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -311,7 +316,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -323,12 +328,12 @@ class DatasetCacheSuite extends SharedSparkSession + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 879569045b6..f3ff89067d2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -751,10 +959,10 @@ index 879569045b6..f3ff89067d2 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index c68f64d52f2..32489f2cb15 100644 +index c68f64d52f2..59f9f9bb092 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -762,7 +970,11 @@ index c68f64d52f2..32489f2cb15 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -192,6 +193,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -192,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -770,7 +982,19 @@ index c68f64d52f2..32489f2cb15 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -246,6 +248,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -228,7 +231,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -246,6 +252,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -779,7 +1003,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(buf.distinct.size == n) } -@@ -261,6 +265,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -261,6 +269,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -792,7 +1016,7 @@ index c68f64d52f2..32489f2cb15 100644 case _ => Nil } } -@@ -1203,10 +1213,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1203,10 +1217,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -811,7 +1035,7 @@ index c68f64d52f2..32489f2cb15 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -819,7 +1043,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1371,10 +1595,18 @@ index 291aa7cab72..7783c37683e 100644 super.test(testName, testTags: _*) { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index da6f6aca2ad..12f48c4bb5d 100644 +index da6f6aca2ad..64c981ad3d7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1529,7 +1529,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper +@@ -40,6 +40,7 @@ import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1529,7 +1530,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1384,7 +1616,7 @@ index da6f6aca2ad..12f48c4bb5d 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1990,8 +1991,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper +@@ -1990,8 +1992,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper countAcc.add(1) x }) @@ -1401,7 +1633,15 @@ index da6f6aca2ad..12f48c4bb5d 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -4755,7 +4763,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper +@@ -3283,6 +3292,7 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4755,7 +4765,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper ) } @@ -1477,7 +1717,7 @@ index 395cb67f441..33ac6ed19af 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index bfcf583a705..6c177bff0e3 100644 +index bfcf583a705..749e6d7f1a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -217,7 +217,15 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp @@ -1496,7 +1736,15 @@ index bfcf583a705..6c177bff0e3 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -311,6 +319,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -260,6 +268,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -311,6 +320,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1504,7 +1752,7 @@ index bfcf583a705..6c177bff0e3 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -366,6 +375,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -366,6 +376,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("reset column vectors") { @@ -1512,7 +1760,7 @@ index bfcf583a705..6c177bff0e3 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -553,6 +563,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -553,6 +564,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1535,6 +1783,33 @@ index d7b2511eac2..d5f5b940b94 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 63589472854..f8c07a9b037 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1734,7 +2009,7 @@ index 2d26356890d..2c5994f5fbc 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 37684c7fce3..f3574dec867 100644 +index 37684c7fce3..a0a2f72b6c9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -26,6 +26,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1742,7 +2017,7 @@ index 37684c7fce3..f3574dec867 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1801,6 +2076,15 @@ index 37684c7fce3..f3574dec867 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -2005,7 +2017,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index 5ae23bc3338..5c2c3fff284 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1969,7 +2253,7 @@ index ce549da03b4..47adf7cf94d 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index 3ef22ccb77e..54f67b45387 100644 +index 3ef22ccb77e..9aeb967617f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -56,6 +56,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -2029,23 +2313,6 @@ index 3ef22ccb77e..54f67b45387 100644 } } } -@@ -530,10 +541,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 60bf3a21e96..7693107f107 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2620,7 +2887,7 @@ index d70bd715879..074a9fa29d9 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index d6d19d21e65..85e2f441f81 100644 +index d6d19d21e65..7b3942d936b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,13 +27,15 @@ import org.apache.spark.SparkException @@ -3096,7 +3363,18 @@ index d6d19d21e65..85e2f441f81 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3278,7 +3330,8 @@ class AdaptiveQueryExecSuite +@@ -3272,13 +3324,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -3106,7 +3384,7 @@ index d6d19d21e65..85e2f441f81 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3368,7 +3421,8 @@ class AdaptiveQueryExecSuite +@@ -3368,7 +3425,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3117,10 +3395,29 @@ index d6d19d21e65..85e2f441f81 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 88be4adb6a4..f8fe831744e 100644 +index 88be4adb6a4..23ec5374f8a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -@@ -228,9 +228,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv +@@ -210,6 +210,18 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -228,9 +240,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3145,6 +3442,147 @@ index 88be4adb6a4..f8fe831744e 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 57da12e8797..396d0745758 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -543,7 +550,10 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd8d1308e99..0e1f80045a3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3631,10 +4069,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 52746720eba..15df5e78440 100644 +index 52746720eba..98266bf06f5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3643,7 +4081,29 @@ index 52746720eba..15df5e78440 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest } } @@ -3653,6 +4113,21 @@ index 52746720eba..15df5e78440 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -654,12 +658,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala index 2ae2ea6339f..30c8959a057 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala @@ -4526,10 +5001,10 @@ index e2c74533e7f..a12d55848ea 100644 val tblTargetName = "tbl_target" val tblSourceQualified = s"default.$tblSourceName" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index fb26d3311eb..13dc8b89a91 100644 +index fb26d3311eb..188b4d439a2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -102,6 +102,20 @@ trait SharedSparkSessionBase +@@ -102,6 +102,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4546,6 +5021,11 @@ index fb26d3311eb..13dc8b89a91 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4624,10 +5104,10 @@ index 59022deaed7..f9aeacb5a9b 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index 8e7ff526a95..ea1072a7195 100644 +index 8e7ff526a95..518e1bae289 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -54,24 +54,34 @@ object TestHive +@@ -54,24 +54,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4672,6 +5152,10 @@ index 8e7ff526a95..ea1072a7195 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/docs/source/contributor-guide/spark-sql-tests.md b/docs/source/contributor-guide/spark-sql-tests.md index c4c39764a37..3918f69ffa5 100644 --- a/docs/source/contributor-guide/spark-sql-tests.md +++ b/docs/source/contributor-guide/spark-sql-tests.md @@ -32,7 +32,9 @@ Here is an overview of the changes that we need to make to Spark: - Modify `project/SparkBuild.scala` so that, when Comet is enabled, the `sql` and `hive` test JVMs use `CometShuffleManager` as the default shuffle manager. Comet disables itself when its shuffle manager is not registered, so without this default the suites that build their own `SparkSession` or `SparkContext` would - run without Comet + run without Comet. From Spark 3.5, where Comet's cache format is the default, the same JVMs also default + `spark.sql.cache.serializer` to Comet's `ArrowCachedBatchSerializer`, because Spark keeps the first cache + serializer it loads for the rest of the JVM Here are the steps involved in running the Spark SQL tests with Comet, using Spark 3.4.3 for this example. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 40732ced79f..afab52bfce8 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -37,11 +37,12 @@ infer the schema from. ## In-Memory Cache Comet can store cached relations (`df.cache()`, `CACHE TABLE`) in Arrow format and scan them -natively. This is experimental and disabled by default; see [In-Memory Cache](../in-memory-cache.md) -for how to enable it. Comet does not replace a `spark.sql.cache.serializer` that the application -has already set. Relations whose schema Comet's Arrow writer does not support are cached in -Spark's default format, and their scans fall back to Spark. Reads that feed Spark operators rather -than Comet operators can be slower than Spark's cache. +natively. This is enabled by default from Spark 3.5 but not on Spark 3.4; see +[In-Memory Cache](../in-memory-cache.md) for how to turn it off or on. Comet does not replace a +`spark.sql.cache.serializer` that the application has already set. Relations whose schema Comet's +Arrow writer does not support are cached in Spark's default format, and their scans fall back to +Spark. Reads that feed Spark operators rather than Comet operators can be slower than Spark's +cache. With Kryo and `spark.kryo.registrationRequired=true`, Comet needs its Kryo registrator whether or not the cache is enabled; see [Kryo serialization](../installation.md#kryo-serialization). diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 20091ac3172..d6867defeab 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -21,18 +21,24 @@ Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, `df.persist()`) in an Arrow format that Comet operators read directly. Without it, a cached table is stored in Spark's own -format and every scan of it has to convert each batch before Comet can continue, which shows up in -the plan as a `CometSparkColumnarToColumnar` above the cache scan. +format, which Comet operators cannot read. Under Comet's default settings the operators above the +cache scan then run on Spark. With `spark.comet.convert.inMemoryCache.enabled`, a +`CometSparkColumnarToColumnar` above the scan converts each batch for Comet operators instead. -This feature is **experimental and disabled by default**. Turn it on at startup, alongside the rest -of Comet's configuration: +This feature is **enabled by default** from Spark 3.5. To turn it off, set the config at startup, +alongside the rest of Comet's configuration: ```shell $SPARK_HOME/bin/spark-shell \ ... \ - --conf spark.comet.exec.inMemoryCache.enabled=true + --conf spark.comet.exec.inMemoryCache.enabled=false ``` +On Spark 3.4 it is disabled by default, and setting the config to `true` turns it on. Spark 3.4 has +no hook for the rule through which Comet lets AQE coalesce the shuffle partitions of a union that +Comet runs, so a union of a shuffle and a relation cached in Comet's format keeps every shuffle +partition there. + It has to be set before the `SparkContext` starts. Comet's driver plugin chooses `spark.sql.cache.serializer` once, while the context is initializing, so an application keeps the cache format it started with however the config is set afterwards. The plugin installs Comet's @@ -112,7 +118,7 @@ nowhere to record either that a column is dictionary encoded or the dictionary i | Config | Default | Description | | ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.exec.inMemoryCache.enabled` | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | +| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | | `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. | | `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | @@ -161,10 +167,57 @@ back to Spark row execution above the scan and the two columns stop measuring th Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs on Comet either way and only the cache-scan boundary moves: on the left, Spark's `InMemoryTableScanExec` feeds those same Comet operators through a `CometSparkColumnarToColumnar` -bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both columns read the same +bridge, which the benchmark turns on with `spark.comet.convert.inMemoryCache.enabled`; on the +right, `CometInMemoryTableScan` feeds them directly. Both columns read the same Comet-written `CometCachedBatch`. These numbers are therefore "keep the cached scan native" against "fall back to a Spark cache scan and convert", not Comet against Spark execution, and not a -comparison with Spark's own cache format. That comparison is under [Limitations](#limitations). +comparison with Spark's own cache format, which follows. + +### Against Spark's cache format + +What turning the feature on changes for a query that Comet runs is measured against Spark's own +cache format by the benchmark's adaptive cases. Comet and AQE are on, Comet's other settings are at +their defaults, and the same 5M-row relation is cached in each format. The defaults leave +`spark.comet.convert.inMemoryCache.enabled` off, so Comet operators cannot read Spark's cache scan, +and with Spark's format the operators directly above the scan run on Spark. Measured on an AMD +Ryzen 9 7950X3D (JDK 17, Spark 4.1, release build): + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 29 ms | 24 ms | 1.2x | +| Narrow projection (1 of 6) | 52 ms | 34 ms | 1.5x | +| 3 of 6 columns | 102 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 299 ms | 224 ms | 1.3x | + +A Spark operator above the cache scan, standing in for any operator Comet does not support, is +measured the same way, with Comet's aggregate turned off. With Comet's format, the native scan feeds +that operator through a columnar-to-row transition: + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 39 ms | 16 ms | 2.4x | +| Narrow projection (1 of 6) | 50 ms | 27 ms | 1.8x | +| 3 of 6 columns | 97 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 303 ms | 299 ms | 1.0x | + +Comet's format is as fast or faster in every shape but one: the read of three of the six columns, +all of them longs, is about 10% slower under either kind of operator. That cost is `zstd` +decompression. With the `none` codec, the same read is 2.7x faster than Spark's format with Comet +operators above the scan, and 1.6x faster with a Spark operator above it. + +Building the cache is measured the same way, from the same source, with its rows produced either by +Comet operators or by Spark operators. Comet's format is written straight from the Arrow batches of +Comet operators, while Spark's is always built from rows, so above Comet operators it first converts +their batches. Measured on an Apple M3 Max (JDK 17, Spark 4.1, release build): + +| Rows produced by | Spark's cache format | Comet's cache format, `zstd` | Comet's cache format, `none` | +| ---------------- | -------------------: | ---------------------------: | ---------------------------: | +| Comet operators | 3285 ms | 944 ms | 601 ms | +| Spark operators | 3680 ms | 2079 ms | 1710 ms | + +Spark's format holds the relation in 217 MiB, and Comet's in 51 to 55 MiB with `zstd` and 315 MiB +with `none`. So with the default codec, Comet's format builds 1.8x to 3.5x faster than Spark's and +takes about a quarter of the memory. ## Kryo @@ -220,6 +273,9 @@ operator reads a relation cached in Comet's format depends on the scan below it: and operators that do not take part in code generation, such as exchanges and limits, or a query that returns the cached rows as they are. +Whenever Spark's scan reads Comet's format, which also happens for a relation whose cached plan +records `Dataset.observe` metrics, Comet records a fallback reason on the scan. + Measured by the same benchmark over the same 5M-row relation, with native execution off so that Spark operators consume the cached data, Comet disabled for the row reader and enabled for the fused reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): @@ -234,9 +290,9 @@ reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): The fused reader is faster than Spark's own format for the narrowest reads and within 20% of it for the others. The row reader takes up to 1.6 times as long, and it is the only reader for relations wider than `spark.sql.codegen.maxFields`: reading every column of relations of 100, 200 and 1500 -nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. These gaps are why -the feature is still off by default; -[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks them. +nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. A Spark operator +above Comet's native cache scan does not pay this; see [Performance](#performance). +[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks these gaps. Comet's serializer exists because Spark's own Arrow cache format ([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index 3ce3e88d1da..c7ceea0ee60 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -59,6 +59,34 @@ key is removed. ## Upgrading to Comet 1.2.0 +Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below +need none either, but check whether any of them applies to your deployment. + +### In-Memory Cache Enabled by Default + +`spark.comet.exec.inMemoryCache.enabled` now defaults to `true` on Spark 3.5 and later. There, an +application that loads `CometPlugin` now stores what it caches with `CACHE TABLE`, `df.cache()` or +`df.persist()` in Comet's Arrow format instead of Spark's, and Comet scans it natively. On Spark 3.4 +the default stays `false`; see [In-Memory Cache](in-memory-cache.md) for why. The format does not +change query results, but it can change performance. Spark's own cache scan can read Comet's format +more slowly than Spark's, which matters when a session turns Comet or its native execution off +after caching, and Comet records a fallback reason on such a scan. See +[In-Memory Cache](in-memory-cache.md#limitations). + +The format is chosen once, when the application starts. To keep Spark's format, set +`spark.comet.exec.inMemoryCache.enabled=false` then. Comet also keeps Spark's format without that +setting when the application: + +- starts with `spark.comet.enabled` or `spark.comet.exec.enabled` set to `false`. +- leaves Comet shuffle enabled without one of Comet's shuffle managers, so that Comet disables + itself. +- uses Kryo with `spark.kryo.registrationRequired=true` and has not registered Comet's cached + batch, because Kryo would reject it. To use Comet's format, register Comet's classes with + `spark.kryo.registrator=org.apache.comet.CometKryoRegistrator`; see + [Kryo](in-memory-cache.md#kryo). + +An application that sets `spark.sql.cache.serializer` itself keeps the serializer it chose. + ### Deprecated and Removed Settings Using `spark.comet.sparkToColumnar.enabled` and `spark.comet.sparkToColumnar.supportedOperatorList` diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index acefd4f0890..c0054d2213b 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -66,7 +66,7 @@ omitted from the tables below and may be reconsidered based on demand: | `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). | | `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support, including native Parquet writes. | | `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. | -| `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | +| `InMemoryTableScanExec` | ✅ | Enabled by default from Spark 3.5. Comet installs its Arrow cache serializer as the application starts if `spark.comet.exec.inMemoryCache.enabled` is true, its default except on Spark 3.4. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | ## Projection and filtering diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 13733f0b510..4791f1abf5a 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -298,7 +298,7 @@ by role. Names match what is shown in the plan output. | `CometNativeScan` | Parquet scan that runs entirely in Rust via DataFusion. | | `CometIcebergNativeScan` | Iceberg Parquet scan that runs entirely in Rust via DataFusion. | | `CometCsvNativeScan` | CSV scan that runs entirely in Rust via DataFusion (experimental). | -| `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (experimental, disabled by default). | +| `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (enabled by default from Spark 3.5). | | `CometLocalTableScan` | JVM-side scan of a `LocalTableScanExec` (disabled by default). | | `CometEmptyRelation` | Empty native input that replaces `EmptyRelationExec` (Spark 4.0 and later). | | `CometRange` | Generates the rows of `spark.range` and SQL `range()` in native code, replacing `RangeExec` (disabled by default). | diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 290a3b05358..4aab7c72a11 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -375,9 +375,11 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "Whether to enable Comet native scans and fused Spark reads of in-memory cached tables. " + - "Requires spark.comet.enabled=true. At startup, this setting also decides whether " + - "CometDriverPlugin installs Comet's cache serializer, which stores cached data in " + - "Arrow format. The plugin installs it only if spark.comet.enabled and " + + "Requires spark.comet.enabled=true. Defaults to true from Spark 3.5 and to false on " + + "Spark 3.4, where AQE cannot coalesce the shuffle partitions of a union that Comet " + + "runs with a cached relation in one branch. At startup, this setting also decides " + + "whether CometDriverPlugin installs Comet's cache serializer, which stores cached " + + "data in Arrow format. The plugin installs it only if spark.comet.enabled and " + "spark.comet.exec.enabled are also enabled at startup, and only with one of Comet's " + "shuffle managers while Comet shuffle is enabled. " + "Because spark.sql.cache.serializer is a " + @@ -398,7 +400,9 @@ object CometConf extends ShimCometConf { "soon as it is serialized, including the disk half of the default " + "MEMORY_AND_DISK storage level.") .booleanConf - .createWithDefault(false) + // CometCoalesceShufflePartitions needs a query-stage optimizer rule hook, which Spark 3.4 + // lacks. + .createWithDefault(CometSparkSessionExtensions.isSpark35Plus) val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = conf("spark.comet.exec.inMemoryCache.compression.codec") diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index e161466c4fb..aba97d80219 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1437,6 +1437,12 @@ class CometInMemoryCacheSuite extends CometTestBase { extraConfs.containsKey(serializerKey) } + test("Comet's cache format is the default from Spark 3.5 and opt-in on Spark 3.4") { + // Spark 3.4 has no hook for CometCoalesceShufflePartitions, so AQE could not coalesce the + // shuffled branch of a Comet union whose other branch reads a cache in Comet's format. + assert(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.defaultValue.get == isSpark35Plus) + } + test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { val cometOn = CometConf.COMET_ENABLED.key -> "true" val execOn = CometConf.COMET_EXEC_ENABLED.key -> "true" diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index 046fe3b822c..230126a7025 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -28,7 +28,9 @@ import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.execution.ColumnarToRowExec -import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.aggregate.HashAggregateExec +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types.{DataType, LongType, StringType} @@ -236,7 +238,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { scanned = 6) runCodecBenchmark(flatRelation) + runBuildBenchmark(flatRelation) runSparkOperatorBenchmark(flatRelation) + runAdaptiveBenchmark(flatRelation) } runWideSparkOperatorBenchmark() } @@ -344,17 +348,101 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * Building the cache in each format, and the footprint each format leaves, against Spark's own + * cache format. Turning the feature on changes how every cached relation is written, and a job + * that caches a relation to read it only once or twice pays more for the build than for the + * reads. + * + * What the build costs depends on what produces the rows being cached. Comet's format is + * written straight from the Arrow batches of a plan that runs natively, and from rows + * otherwise, while Spark's format is always built from rows, so above Comet operators it pays a + * columnar-to-row transition first. The relation's source stands in for both kinds of input: + * with spark.comet.exec.range.enabled Comet generates and projects its rows natively, as above + * a native scan, and without it the source runs on Spark. Comet and AQE are on and Comet's + * other settings are at their defaults, as in runAdaptiveBenchmark. + * + * Only one copy is cached at a time, as in runCodecBenchmark. + */ + private def runBuildBenchmark(relation: CachedRelation): Unit = { + val view = s"${relation.table}_build" + val formats = Seq( + ("Spark's cache format", sparkSerializer, Seq.empty[(String, String)]), + ("Comet's cache format, zstd", cometSerializer, codecConf("zstd")), + ("Comet's cache format, none", cometSerializer, codecConf("none"))) + val inputs = Seq( + ("Comet operators", true, Seq(CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true")), + ("Spark operators", false, Seq.empty[(String, String)])) + + spark.catalog.clearCache() + withTempTable(view) { + spark + .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") + .createOrReplaceTempView(view) + + inputs.foreach { case (input, nativeInput, inputConf) => + val benchmark = + new Benchmark( + s"in-memory cache build from $input", + relation.rows.toLong, + output = output) + val footprints = new Array[Long](formats.length) + formats.zipWithIndex.foreach { case ((name, serializer, formatConf), i) => + // Timed around the caching alone: dropping the previous copy is setup, as in + // runCodecBenchmark. + benchmark.addTimerCase(name) { timer => + spark.catalog.uncacheTable(view) + withCacheSerializer(serializer) { + withSQLConf(adaptiveConf ++ inputConf ++ formatConf: _*) { + timer.startTiming() + spark.catalog.cacheTable(view) + spark.table(view).count() + timer.stopTiming() + } + } + // On the case's first call, which is a warmup iteration. + if (footprints(i) == 0L) { + verifyBuild(view, serializer, nativeInput) + footprints(i) = cachedBytes(view) + } + } + } + benchmark.run() + + // As in runCodecBenchmark, footprint has no column in a Benchmark table. + formats.zip(footprints).foreach { case ((name, _, _), bytes) => + benchmark.out.println(f"Cached footprint ($name): ${bytes / (1024.0 * 1024.0)}%.1f MiB") + } + } + + spark.catalog.uncacheTable(view) + } + } + + // Pins what a build case claims: the named serializer cached the relation, from a plan that ran + // natively exactly when the case says its input came from Comet operators. + private def verifyBuild(view: String, serializer: String, nativeInput: Boolean): Unit = { + val relation = cachedRelation(view) + val actual = relation.cacheBuilder.serializer.getClass.getName + assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + val plan = relation.cacheBuilder.cachedPlan + val cometOperators = collect(plan) { case p if p.nodeName.startsWith("Comet") => p } + assert( + cometOperators.nonEmpty == nativeInput, + s"Expected the cached plan to run ${if (nativeInput) "natively" else "on Spark"}:\n$plan") + } + /** * Reads that feed Spark operators rather than Comet ones, against Spark's own cache format. * * Native execution is off in every case, so this measures Spark consuming the cached data: the - * shape where Comet's format has something to lose, and the reason the feature is off by - * default. Comet's format is read two ways. With Comet off, every Spark operator reads rows - * from the cache scan's row reader. With Comet on, a generated Spark operator instead reads the - * cached vectors through a ColumnarToRowExec fused into its generated code, which is how these - * reads run when Comet is enabled without native execution. With native execution on, Spark - * operators above the cache read CometInMemoryTableScan's batches through CometColumnarToRow - * instead, which this does not measure. + * shape where Comet's format has the most to lose. Comet's format is read two ways. With Comet + * off, every Spark operator reads rows from the cache scan's row reader. With Comet on, a + * generated Spark operator instead reads the cached vectors through a ColumnarToRowExec fused + * into its generated code, which is how these reads run when Comet is enabled without native + * execution. With native execution on, Spark operators above the cache read + * CometInMemoryTableScan's batches through CometColumnarToRow instead, which + * runAdaptiveBenchmark measures. * * Both formats are cached from the same relation, one copy at a time as in runCodecBenchmark, * and each case checks which serializer cached the relation it reads and which reader it uses. @@ -377,14 +465,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .createOrReplaceTempView(view) val cache = new OneCachedCopy(view) - Seq( - ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), - ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), - ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), - ( - "full projection (6 of 6 columns)", - s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", - 6)).foreach { case (label, query, scanned) => + readShapes(view).foreach { case (label, query, scanned) => val benchmark = new Benchmark( s"in-memory cache read by Spark operators, $label", relation.rows.toLong, @@ -446,17 +527,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { fused: Boolean = false) /** - * Holds one cached copy of `view` at a time. Two copies could not coexist anyway: the cache - * manager keys on the plan rather than the name, so a second one would find the first. + * Holds one cached copy of `view` at a time, cached under `conf`. Two copies could not coexist + * anyway: the cache manager keys on the plan rather than the name, so a second one would find + * the first. */ - private class OneCachedCopy(view: String) { + private class OneCachedCopy(view: String, conf: Seq[(String, String)] = sparkOperatorConf) { private var cachedBy: String = _ def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { spark.catalog.uncacheTable(view) cachedBy = null withCacheSerializer(serializer) { - withSQLConf(sparkOperatorConf: _*) { + withSQLConf(conf: _*) { spark.catalog.cacheTable(view) spark.table(view).count() } @@ -488,6 +570,116 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * What the feature changes for a query that runs with Comet, against Spark's own cache format, + * with AQE on and Comet's other settings at their defaults. This is the comparison an + * application gets from turning the feature on or off. Both formats are cached from the same + * relation, one copy at a time as in runSparkOperatorBenchmark. + * + * Two shapes of plan read the cache. With Comet operators above the cache scan, Comet's format + * runs the whole query natively, while Spark's leaves the operators directly above its scan on + * Spark: Comet reads Spark's cache scan only through spark.comet.convert.inMemoryCache.enabled, + * which is off by default. With a Spark operator above the scan, Comet's format is read by the + * native scan and converted to rows for that operator, where Spark's is read by Spark's own + * scan. The Spark operator is the aggregate, with Comet's turned off, standing in for any + * operator Comet does not support. + */ + private def runAdaptiveBenchmark(relation: CachedRelation): Unit = { + val view = s"${relation.table}_adaptive" + val formats = + Seq("Spark's cache format" -> sparkSerializer, "Comet's cache format" -> cometSerializer) + val operatorsAbove = Seq( + "Comet operators" -> Seq.empty[(String, String)], + "a Spark operator" -> Seq(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false")) + + spark.catalog.clearCache() + withTempTable(view) { + spark + .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") + .createOrReplaceTempView(view) + + val cache = new OneCachedCopy(view, adaptiveConf) + for { + (operators, operatorConf) <- operatorsAbove + (label, query, scanned) <- readShapes(view) + } { + val benchmark = new Benchmark( + s"in-memory cache with AQE, $operators above the scan, $label", + relation.rows.toLong, + output = output) + formats.foreach { case (name, serializer) => + var verified = false + // Re-caching in this case's format is setup, so it is outside the timer, and it only + // happens on the case's first call, which is a warmup iteration. + benchmark.addTimerCase(name) { timer => + cache.cacheBy(serializer) + withSQLConf(adaptiveConf ++ operatorConf: _*) { + if (!verified) { + verifyAdaptiveRead(query, scanned, serializer, operatorConf.nonEmpty) + verified = true + } + timer.startTiming() + spark.sql(query).noop() + timer.stopTiming() + } + } + } + benchmark.run() + } + + spark.catalog.uncacheTable(view) + } + } + + // The reads runSparkOperatorBenchmark and runAdaptiveBenchmark measure: no columns of the flat + // relation, one, three, and all six. + private def readShapes(view: String): Seq[(String, String, Int)] = Seq( + ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), + ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), + ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), + ( + "full projection (6 of 6 columns)", + s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", + 6)) + + // Pins what an adaptive case claims, in the plan AQE settles on, which it does only by running + // the query: one cache scan, native exactly when it reads Comet's format, reading the columns its + // label counts from a relation the named serializer cached. Spark aggregates run only where the + // case puts them: above Spark's scan, which nothing bridges into Comet, or wherever Comet's + // aggregate is turned off. + private def verifyAdaptiveRead( + query: String, + scanned: Int, + serializer: String, + sparkOperator: Boolean): Unit = { + val df = spark.sql(query) + df.collect() + val executed = df.queryExecution.executedPlan + val plan = executed.toString() + assert(executed.isInstanceOf[AdaptiveSparkPlanExec], s"Expected an adaptive plan:\n$plan") + + val nativeScans = collect(executed) { case s: CometInMemoryTableScanExec => s } + val sparkScans = collect(executed) { case s: InMemoryTableScanExec => s } + assert( + nativeScans.length + sparkScans.length == 1, + s"Expected exactly one cache scan:\n$plan") + val cometFormat = serializer == cometSerializer + assert( + nativeScans.nonEmpty == cometFormat, + s"Expected a native scan exactly for Comet's format:\n$plan") + val (relation, columns) = nativeScans.headOption + .map(s => (s.originalPlan.relation, s.scanOutput.length)) + .getOrElse((sparkScans.head.relation, sparkScans.head.attributes.length)) + assert(columns == scanned, s"Expected the scan to read $scanned columns:\n$plan") + val actual = relation.cacheBuilder.serializer.getClass.getName + assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + + val sparkAggregates = collect(executed) { case a: HashAggregateExec => a } + assert( + sparkAggregates.nonEmpty == (sparkOperator || !cometFormat), + s"Expected Spark aggregates only where this case puts them:\n$plan") + } + // spark.sql.cache.serializer is static, and InMemoryRelation memoizes the serializer it names // for the life of the JVM. It looks the name up in the active session's conf when a relation is // cached, though, so setting it there directly and clearing the memoized instance around one @@ -535,20 +727,26 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { /** What the cached relation behind `view` occupies, summed over its batches as written. */ private def cachedBytes(view: String): Long = { - val relation = spark - .table(view) - .queryExecution - .optimizedPlan - .collectFirst { case r: InMemoryRelation => r } - .getOrElse(sys.error(s"$view is not cached")) + val relation = cachedRelation(view) // The stored payloads rather than computeStats, which reports the relation's decoded size and // so is the same for every codec. assert(relation.cacheBuilder.isCachedColumnBuffersLoaded, s"$view is not materialized") relation.cacheBuilder.cachedColumnBuffers - .map(CometCachedBatchHelper.payloadSize) + .map { + case batch: DefaultCachedBatch => batch.buffers.map(_.length.toLong).sum + case batch => CometCachedBatchHelper.payloadSize(batch) + } .fold(0L)(_ + _) } + private def cachedRelation(view: String): InMemoryRelation = + spark + .table(view) + .queryExecution + .optimizedPlan + .collectFirst { case r: InMemoryRelation => r } + .getOrElse(sys.error(s"$view is not cached")) + private def runStatsBenchmark(): Unit = { val batchSize = 10000 val types: Seq[DataType] = Seq.fill(3)(LongType) ++ Seq.fill(3)(StringType) @@ -629,8 +827,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // scan and convert" -- which is the overhead this feature exists to remove. // // Neither case is a baseline for Spark's own cache format: both read the same Comet-written - // CometCachedBatch. Spark's format is only measured by runSparkOperatorBenchmark, with Comet - // off, since that is the only comparison it answers. + // CometCachedBatch. Spark's format is measured by runSparkOperatorBenchmark, with Comet off, + // and by runAdaptiveBenchmark, with Comet on. withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { spark .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") @@ -713,6 +911,16 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { "spark.comet.exec.onHeap.enabled" -> "true", "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + // Comet and AQE on, and Comet's other settings at their defaults, unlike cacheConf. The batch + // size matches the other confs, and on-heap mode is what lets Comet run in this session. + private val adaptiveConf: Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + private val sparkSerializer = classOf[DefaultCachedBatchSerializer].getName private val cometSerializer = classOf[ArrowCachedBatchSerializer].getName } From bfacc43a50aba13da4fa8a828db52309541191b4 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Fri, 9 Oct 2026 19:55:38 +0000 Subject: [PATCH 45/89] =?UTF-8?q?feat:=20run=20array=5Fcontains=20on=20flo?= =?UTF-8?q?at=20elements=20natively=20with=20Spark's=20equa=E2=80=A6=20(#6?= =?UTF-8?q?599)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: run array_contains on float elements natively with Spark's equality A Comet kernel compares float elements as Spark's genEqual does, with -0.0 equal to 0.0 and all NaNs equal, at any depth, so these calls no longer go through the codegen dispatcher. Closes #6520 * fix: keep nested floats and eagerly evaluated values on the dispatcher Route only flat FLOAT and DOUBLE arrays to spark_array_contains. Nested float elements, and a value other than a literal or column read over a nullable array, stay on the codegen dispatcher, which is faster for nested elements, keeps string collations, and skips the value for a null array as Spark does. * test: import EDGE_VALUES from float_semantics in array_contains tests Main moved the constant there, so the crate's unit tests failed to compile. --------- Co-authored-by: Andy Grove --- .../expression-audits/array_funcs.md | 2 +- .../latest/compatibility/floating-point.md | 9 +- docs/source/user-guide/latest/expressions.md | 2 +- native/spark-expr/benches/float_arrays.rs | 30 +- .../src/array_funcs/array_contains.rs | 447 ++++++++++++++++++ .../src/array_funcs/array_remove.rs | 19 +- native/spark-expr/src/array_funcs/mod.rs | 26 +- native/spark-expr/src/comet_scalar_funcs.rs | 8 +- .../scala/org/apache/comet/serde/arrays.scala | 69 ++- .../expressions/array/array_contains.sql | 6 +- .../array/array_contains_collation.sql | 32 ++ .../array/array_contains_eager_eval.sql | 34 ++ .../array/array_contains_floating_point.sql | 85 ++++ .../array/routing_arrays_disabled.sql | 3 +- .../array/routing_arrays_enabled.sql | 3 +- .../comet/CometMapExpressionSuite.scala | 14 +- 16 files changed, 734 insertions(+), 55 deletions(-) create mode 100644 native/spark-expr/src/array_funcs/array_contains.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index 2eed573d7f0..50ac45b852d 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -48,7 +48,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayContains(left, right) extends BinaryExpression with NullIntolerant with Predicate`; `inputTypes` uses `findWiderTypeWithoutStringPromotionForTwo`. Wired as `CometScalarFunction("array_contains")`. - Spark 4.0.1 (audited 2026-05-27): `NullIntolerant` trait replaced by `nullIntolerant: Boolean`; `checkInputDataTypes` adopts `DataTypeUtils.sameType` (collation-aware in 4.x). - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Float/double arrays containing NaN and signed zero match Spark; DataFusion canonicalizes them the same way as Spark's `SQLOrderingUtil`. +- Spark 3.4, 3.5, 4.0, 4.1, 4.2 (2026-10-02, [#6520](https://github.com/apache/datafusion-comet/issues/6520)): datafusion-spark's `array_contains` compares floats by their bits, where Spark's `genEqual` treats `-0.0` as `0.0` and all NaNs as equal. Flat `FLOAT` and `DOUBLE` arrays now run through Comet's `spark_array_contains`, which applies Spark's equality, instead of the codegen dispatcher. Nested float elements stay on the dispatcher, where Spark's generated code is faster and keeps string collations, as does a value other than a literal or column read over a nullable array, which Spark skips for a null array. Tracked upstream as [apache/datafusion#25981](https://github.com/apache/datafusion/issues/25981). ## array_distinct diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index 39e1de9b431..1324dbf87fd 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -146,7 +146,12 @@ reports 4.2.0 but includes SPARK-59602 still runs natively; set `spark.comet.expression.ArrayDistinct.enabled=false` and `spark.comet.expression.ArrayUnion.enabled=false` on such a build. -## `array_remove` and `sort_array` +## `array_contains`, `array_remove` and `sort_array` + +`array_contains` compares `FLOAT` and `DOUBLE` elements as Spark does: `-0.0` equals `0.0`, and all +NaN representations are equal, inside nested arrays and structs too. The result keeps Spark's +three-valued form: null when nothing matches and the array holds a null element. Flat arrays run a +native kernel; nested float elements go through the codegen dispatcher. `array_remove` compares `FLOAT` and `DOUBLE` elements as Spark does: `-0.0` equals `0.0`, and all NaN representations are equal, inside nested arrays too. The elements it keeps retain their @@ -156,4 +161,4 @@ original NaN representations and zero signs. representations tie, and keeps equal elements in their original order. In an array whose elements can be null, such as one built from nullable columns, `-0.0` and `0.0` therefore tie. Sorting an array whose elements cannot be null in ascending order, Spark's generated code puts `-0.0` before -`0.0`, and so does Comet. Both expressions run natively in strict floating-point mode. +`0.0`, and so does Comet. All three expressions run natively in strict floating-point mode. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index dffd6b49884..2fd8333a79b 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -155,7 +155,7 @@ The tables below list every Spark built-in expression with its current status. | `array` | ✅ | Native | | | `array_append` | ✅ | Native | | | `array_compact` | ✅ | — | | -| `array_contains` | ✅ | Hybrid | Float/double element arrays and non-default string collations route through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible | +| `array_contains` | ✅ | Hybrid | Flat float/double element arrays use a native kernel with Spark's equality; nested float elements, non-default string collations, and a value that is not a literal or column over a nullable array, route through the JVM codegen dispatcher | | `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)). Non-default string collations fall back | | `array_except` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | | `array_insert` | ✅ | Native | | diff --git a/native/spark-expr/benches/float_arrays.rs b/native/spark-expr/benches/float_arrays.rs index 44a4d78a2fc..87f0dd27e80 100644 --- a/native/spark-expr/benches/float_arrays.rs +++ b/native/spark-expr/benches/float_arrays.rs @@ -28,7 +28,8 @@ use datafusion::common::ScalarValue; use datafusion::functions_nested::remove::array_remove_all_udf; use datafusion::functions_nested::sort::array_sort_udf; use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; -use datafusion_comet_spark_expr::{SparkArrayRemove, SparkSortArray}; +use datafusion_comet_spark_expr::{SparkArrayRemove, SparkFloatArrayContains, SparkSortArray}; +use datafusion_spark::function::array::array_contains::SparkArrayContains; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -76,6 +77,9 @@ fn criterion_benchmark(c: &mut Criterion) { group.measurement_time(Duration::from_secs(1)); let comet_sort = ScalarUDF::new_from_impl(SparkSortArray::default()); let comet_remove = ScalarUDF::new_from_impl(SparkArrayRemove::default()); + let comet_contains = ScalarUDF::new_from_impl(SparkFloatArrayContains::default()); + // The native path other element types take: array_has plus Spark's null semantics. + let datafusion_contains = Arc::new(ScalarUDF::new_from_impl(SparkArrayContains::default())); let (datafusion_sort, datafusion_remove) = (array_sort_udf(), array_remove_all_udf()); for len in [8, 50] { for nulls in [false, true] { @@ -107,6 +111,16 @@ fn criterion_benchmark(c: &mut Criterion) { let row_values = Float64Array::from_iter_values((0..ROWS).map(|row| value(row * len + 1))); let remove_column_args = args( + vec![column.clone(), ColumnarValue::Array(Arc::new(row_values))], + &list, + ); + let contains_args = args( + vec![column.clone(), scalar(ScalarValue::Float64(Some(26.0)))], + &list, + ); + let row_values = + Float64Array::from_iter_values((0..ROWS).map(|row| value(row * len + 1))); + let contains_column_args = args( vec![column, ColumnarValue::Array(Arc::new(row_values))], &list, ); @@ -132,6 +146,20 @@ fn criterion_benchmark(c: &mut Criterion) { &remove_column_args, &remove_column_args, ), + ( + "array_contains", + &comet_contains, + &datafusion_contains, + &contains_args, + &contains_args, + ), + ( + "array_contains_column", + &comet_contains, + &datafusion_contains, + &contains_column_args, + &contains_column_args, + ), ]; for (op, comet, datafusion, comet_args, datafusion_args) in cases { let evaluate = |udf: &ScalarUDF, args: &ScalarFunctionArgs| { diff --git a/native/spark-expr/src/array_funcs/array_contains.rs b/native/spark-expr/src/array_funcs/array_contains.rs new file mode 100644 index 00000000000..652e8ae24e2 --- /dev/null +++ b/native/spark-expr/src/array_funcs/array_contains.rs @@ -0,0 +1,447 @@ +// 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 std::sync::Arc; + +use super::set_bits_before; +use crate::float_semantics::{compare_floats, spark_equality}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use num::Float; + +/// Spark's `array_contains` for arrays whose elements hold floats at any depth. +/// +/// Spark tests each element against the value with `genEqual`, in which `-0.0` equals `0.0` and +/// all NaNs are equal, at any depth of an array or struct element. The result is true when an +/// element matches, null when none does and the array holds a null element, and false otherwise; +/// a null array or value gives null. datafusion-spark's `array_contains`, which Comet uses for +/// other element types, compares the bits instead, so `0.0` misses a `-0.0` and a NaN misses a +/// NaN whose bits differ. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkFloatArrayContains { + signature: Signature, +} + +impl Default for SparkFloatArrayContains { + fn default() -> Self { + Self { + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkFloatArrayContains { + fn name(&self) -> &str { + "spark_array_contains" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [array, value] = take_function_args(self.name(), &args.args)?; + let all_scalars = matches!( + (array, value), + (ColumnarValue::Scalar(_), ColumnarValue::Scalar(_)) + ); + let rows = if all_scalars { 1 } else { args.number_rows }; + let array = array.to_array(rows)?; + // Spark arrays use Arrow's 32-bit List layout. + let Some(list) = array.as_list_opt::() else { + return exec_err!( + "spark_array_contains takes a list, got {}", + array.data_type() + ); + }; + let result = match value { + ColumnarValue::Scalar(needle) => contains_constant(list, needle)?, + ColumnarValue::Array(value) => array_contains(list, value)?, + }; + if all_scalars { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean( + result.is_valid(0).then(|| result.value(0)), + ))) + } else { + Ok(ColumnarValue::Array(Arc::new(result))) + } + } +} + +fn array_contains(array: &ListArray, value: &ArrayRef) -> Result { + // Spark casts the value to the element type. Anything else fails in `spark_equality`. + match (array.value_type(), value.data_type()) { + (DataType::Float32, DataType::Float32) => Ok(contains_floats::(array, value)), + (DataType::Float64, DataType::Float64) => Ok(contains_floats::(array, value)), + _ => { + let equal = spark_equality(array.values().as_ref(), value.as_ref())?; + contains_where(array, value, equal) + } + } +} + +/// [`array_contains`] with the same value for every row, which a float array answers in one pass +/// over all the values. +fn contains_constant(array: &ListArray, needle: &ScalarValue) -> Result { + match (array.value_type(), needle) { + // A null value gives null for every row. + (_, needle) if needle.is_null() => Ok(BooleanArray::new_null(array.len())), + (DataType::Float32, ScalarValue::Float32(Some(needle))) => { + Ok(contains_float_constant::(array, *needle)) + } + (DataType::Float64, ScalarValue::Float64(Some(needle))) => { + Ok(contains_float_constant::(array, *needle)) + } + _ => array_contains(array, &needle.to_array_of_size(array.len())?), + } +} + +fn contains_float_constant(array: &ListArray, needle: T::Native) -> BooleanArray +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let offsets = array.offsets(); + let first = offsets[0] as usize; + let values = array + .values() + .slice(first, offsets[array.len()] as usize - first); + let floats = values.as_primitive::(); + let buffer = floats.values(); + // For a fixed value, Spark's genEqual is one test per element: IEEE `==`, under which the two + // zeros are equal and a NaN differs from every number, or `is_nan` for a NaN value. + let matches = if needle.is_nan() { + BooleanBuffer::collect_bool(buffer.len(), |index| buffer[index].is_nan()) + } else { + BooleanBuffer::collect_bool(buffer.len(), |index| buffer[index] == needle) + }; + // A null element never matches. + let matches = match floats.nulls() { + Some(nulls) => &matches & nulls.inner(), + None => matches, + }; + // One pass over each bitmap gives every row's match count and valid count. + let matched = set_bits_before(&matches, offsets); + let valid = floats + .nulls() + .map(|nulls| set_bits_before(nulls.inner(), offsets)); + let mut values = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in offsets.windows(2).enumerate() { + let found = matched[row + 1] > matched[row]; + let has_null = valid + .as_ref() + .is_some_and(|valid| valid[row + 1] - valid[row] < bounds[1] - bounds[0]); + values.append(found); + validity.append(array.is_valid(row) && (found || !has_null)); + } + BooleanArray::new(values.finish(), Some(NullBuffer::new(validity.finish()))) +} + +/// Tests each row's floats against the row's value in Spark's `genEqual`, straight from the +/// flattened values. A row is null when the array or the value is. +fn contains_floats(array: &ListArray, value: &ArrayRef) -> BooleanArray +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let values = array.values().as_primitive::(); + let needles = value.as_primitive::().values(); + let (buffer, nulls) = (values.values(), values.nulls()); + let row_nulls = NullBuffer::union(array.nulls(), value.nulls()); + let mut found_rows = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in array.offsets().windows(2).enumerate() { + let (mut found, mut has_null) = (false, false); + let row_valid = row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)); + if row_valid { + let needle = needles[row]; + for index in bounds[0] as usize..bounds[1] as usize { + if nulls.is_some_and(|nulls| nulls.is_null(index)) { + has_null = true; + } else if compare_floats(buffer[index], needle).is_eq() { + found = true; + break; + } + } + } + found_rows.append(found); + validity.append(row_valid && (found || !has_null)); + } + BooleanArray::new( + found_rows.finish(), + Some(NullBuffer::new(validity.finish())), + ) +} + +/// Tests each row's non-null elements with `equal(element, row)`. A row is null when the array or +/// the value is. +fn contains_where(array: &ListArray, value: &ArrayRef, equal: F) -> Result +where + F: Fn(usize, usize) -> bool, +{ + let element_nulls = array.values().logical_nulls(); + let row_nulls = NullBuffer::union(array.nulls(), value.nulls()); + let mut found_rows = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in array.offsets().windows(2).enumerate() { + let (mut found, mut has_null) = (false, false); + let row_valid = row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)); + if row_valid { + for element in bounds[0] as usize..bounds[1] as usize { + if element_nulls + .as_ref() + .is_some_and(|nulls| nulls.is_null(element)) + { + has_null = true; + } else if equal(element, row) { + found = true; + break; + } + } + } + found_rows.append(found); + validity.append(row_valid && (found || !has_null)); + } + Ok(BooleanArray::new( + found_rows.finish(), + Some(NullBuffer::new(validity.finish())), + )) +} + +#[cfg(test)] +mod tests { + use super::super::test_util::list; + use super::*; + use crate::float_semantics::EDGE_VALUES; + use arrow::array::{Float64Array, Int32Array, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion::config::ConfigOptions; + + fn invoke(array: ColumnarValue, value: ColumnarValue, rows: usize) -> Result { + let return_field = Arc::new(Field::new("result", DataType::Boolean, true)); + SparkFloatArrayContains::default().invoke_with_args(ScalarFunctionArgs { + args: vec![array, value], + arg_fields: vec![], + number_rows: rows, + return_field, + config_options: Arc::new(ConfigOptions::default()), + }) + } + + fn booleans(result: ColumnarValue) -> Vec> { + let ColumnarValue::Array(array) = result else { + panic!("expected an array"); + }; + array.as_boolean().iter().collect() + } + + /// Spark's result for one row: true if a non-null element `genEqual`s the value, else null if + /// the row holds a null, else false. + fn spark_contains(row: &[Option], value: f64) -> Option { + if row + .iter() + .any(|element| element.is_some_and(|e| compare_floats(e, value).is_eq())) + { + Some(true) + } else if row.contains(&None) { + None + } else { + Some(false) + } + } + + /// Every edge value looked up in rows of edge values, so each zero and each NaN meets every + /// other, through the constant path and the per-row path. + #[test] + fn floats_follow_spark_equality() -> Result<()> { + let rows: Vec>>> = EDGE_VALUES + .iter() + .map(|skipped| { + Some( + EDGE_VALUES + .iter() + .filter(|v| v != &skipped) + .copied() + .collect(), + ) + }) + .collect(); + let array = list(&rows); + for value in EDGE_VALUES.into_iter().flatten() { + let expected: Vec> = rows + .iter() + .map(|row| spark_contains(row.as_ref().unwrap(), value)) + .collect(); + let constant = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(Some(value))), + rows.len(), + )?; + assert_eq!(booleans(constant), expected, "constant {value:?}"); + let per_row = Float64Array::from(vec![Some(value); rows.len()]); + let column = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Array(Arc::new(per_row)), + rows.len(), + )?; + assert_eq!(booleans(column), expected, "per row {value:?}"); + } + Ok(()) + } + + #[test] + fn null_rows_null_values_and_empty_arrays() -> Result<()> { + let array = list(&[ + None, + Some(vec![]), + Some(vec![None]), + Some(vec![Some(1.0), None]), + ]); + let constant = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))), + 4, + )?; + assert_eq!( + booleans(constant), + vec![None, Some(false), None, Some(true)] + ); + let null_value = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(None)), + 4, + )?; + assert_eq!(booleans(null_value), vec![None; 4]); + let per_row = Float64Array::from(vec![Some(1.0), Some(1.0), None, Some(2.0)]); + let column = invoke( + ColumnarValue::Array(array), + ColumnarValue::Array(Arc::new(per_row)), + 4, + )?; + assert_eq!(booleans(column), vec![None, Some(false), None, None]); + Ok(()) + } + + #[test] + fn all_scalars() -> Result<()> { + let row = ScalarValue::List(Arc::new( + list(&[Some(vec![Some(-0.0), None])]) + .as_list::() + .clone(), + )); + let result = invoke( + ColumnarValue::Scalar(row.clone()), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))), + 1, + )?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) + )); + let result = invoke( + ColumnarValue::Scalar(row), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))), + 1, + )?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(None)) + )); + Ok(()) + } + + /// Elements that are themselves arrays go through `spark_equality`, which compares their + /// floats the same way. + #[test] + fn nested_elements() -> Result<()> { + // Two rows of two arrays each: [[-0.0], [1.0]] and [[NaN, 1.0], [2.0]]. + let inner = list(&[ + Some(vec![Some(-0.0)]), + Some(vec![Some(1.0)]), + Some(vec![Some(f64::NAN), Some(1.0)]), + Some(vec![Some(2.0)]), + ]); + let array: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 4].into()), + inner, + None, + )); + // Looked-up values: [0.0] for the first row, [-NaN, 1.0] for the second. + let values = list(&[ + Some(vec![Some(0.0)]), + Some(vec![Some(-f64::NAN), Some(1.0)]), + ]); + let result = invoke(ColumnarValue::Array(array), ColumnarValue::Array(values), 2)?; + assert_eq!(booleans(result), vec![Some(true), Some(true)]); + Ok(()) + } + + /// Struct elements compare their float fields the same way; another field still has to match. + #[test] + fn struct_elements() -> Result<()> { + let fields = Fields::from(vec![ + Field::new("x", DataType::Float64, true), + Field::new("y", DataType::Int32, true), + ]); + let elements: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(Float64Array::from(vec![-0.0, f64::NAN, 2.0])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ], + None, + )); + // One row holding all three structs. + let array: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", elements.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 3].into()), + elements, + None, + )); + let lookup = |x: f64, y: i32| -> Result { + let value: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(Float64Array::from(vec![x])), + Arc::new(Int32Array::from(vec![y])), + ], + None, + )); + invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Array(value), + 1, + ) + }; + assert_eq!(booleans(lookup(0.0, 1)?), vec![Some(true)]); + assert_eq!(booleans(lookup(-f64::NAN, 1)?), vec![Some(true)]); + assert_eq!(booleans(lookup(0.0, 2)?), vec![Some(false)]); + Ok(()) + } +} diff --git a/native/spark-expr/src/array_funcs/array_remove.rs b/native/spark-expr/src/array_funcs/array_remove.rs index 7688435e433..9eda990c3ce 100644 --- a/native/spark-expr/src/array_funcs/array_remove.rs +++ b/native/spark-expr/src/array_funcs/array_remove.rs @@ -17,7 +17,7 @@ use std::sync::Arc; -use super::with_values; +use super::{set_bits_before, with_values}; use crate::float_semantics::{compare_floats, spark_equality}; use arrow::array::{ Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray, PrimitiveArray, @@ -152,22 +152,9 @@ where } /// The offsets of the elements that `keep` keeps, for a `keep` that starts at `offsets[0]`: the -/// number of its set bits before each offset, counted a word at a time in one pass. +/// number of its set bits before each offset. fn kept_offsets(keep: &BooleanBuffer, offsets: &OffsetBuffer) -> OffsetBuffer { - let chunks = keep.inner().bit_chunks(keep.offset(), keep.len()); - let mut words = chunks.iter_padded(); - let (mut word, mut word_start, mut before_word) = (words.next().unwrap_or(0), 0, 0); - let kept = offsets.iter().map(|offset| { - let end = (offset - offsets[0]) as usize; - while end >= word_start + 64 { - before_word += word.count_ones() as usize; - word = words.next().unwrap_or(0); - word_start += 64; - } - let below_end = word & ((1u64 << (end - word_start)) - 1); - (before_word + below_end.count_ones() as usize) as i32 - }); - OffsetBuffer::new(kept.collect::>().into()) + OffsetBuffer::new(set_bits_before(keep, offsets).into()) } /// Copies the floats of each row that are null or differ from the row's value in Spark's diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 18d35d7ac0a..ff0d55f0bdf 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod array_contains; mod array_extrema; mod array_insert; mod array_position; @@ -31,6 +32,7 @@ mod sequence; mod size; mod sort_array; +pub use array_contains::SparkFloatArrayContains; pub use array_extrema::SparkArrayExtrema; pub use array_insert::ArrayInsert; pub use array_position::SparkArrayPositionFunc; @@ -48,9 +50,31 @@ pub use size::{spark_size, SparkSizeFunc}; pub use sort_array::SparkSortArray; use arrow::array::{ArrayRef, ListArray}; -use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; use std::sync::Arc; +/// The number of set bits of `bits` before each offset, for `bits` that starts at `offsets[0]`, +/// counted a word at a time in one pass. The difference between two consecutive entries is the +/// count within that row. +fn set_bits_before(bits: &BooleanBuffer, offsets: &OffsetBuffer) -> Vec { + let chunks = bits.inner().bit_chunks(bits.offset(), bits.len()); + let mut words = chunks.iter_padded(); + let (mut word, mut word_start, mut before_word) = (words.next().unwrap_or(0), 0, 0); + offsets + .iter() + .map(|offset| { + let end = (offset - offsets[0]) as usize; + while end >= word_start + 64 { + before_word += word.count_ones() as usize; + word = words.next().unwrap_or(0); + word_start += 64; + } + let below_end = word & ((1u64 << (end - word_start)) - 1); + (before_word + below_end.count_ones() as usize) as i32 + }) + .collect() +} + /// A list with the field of `array` and the given row nulls, holding `values` at `offsets`. fn with_values( array: &ListArray, diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 0e4a3824e96..9cc1cadbd8d 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -34,9 +34,10 @@ use crate::{ spark_round, spark_rpad, spark_sequence, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, - SparkDayOfWeek, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, - SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMapExtract, - SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkSortArray, SparkWeekDay, + SparkDayOfWeek, SparkFlatten, SparkFloatArrayContains, SparkIcebergBucket, + SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, + SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, + SparkSortArray, SparkWeekDay, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -342,6 +343,7 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkArrayExtrema::new(false))), Arc::new(ScalarUDF::new_from_impl(SparkArrayPositionFunc::default())), Arc::new(ScalarUDF::new_from_impl(SparkArrayRemove::default())), + Arc::new(ScalarUDF::new_from_impl(SparkFloatArrayContains::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraySlice::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraysOverlap::default())), Arc::new(ScalarUDF::new_from_impl(SparkSortArray::default())), diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 5aeb6824346..67ab955296d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -120,29 +120,53 @@ object CometArrayContains override def hasConditionalNativeDefault: Boolean = true - private val floatingPointReason: String = + private val nestedFloatReason: String = "Spark compares array elements with ordering.equiv, so -0.0 matches +0.0 and all NaNs match " + - "each other; Comet's native array_contains compares the raw Arrow values bitwise" + "each other; Comet's native array_contains applies that only to flat FLOAT and DOUBLE " + + "elements, and nested float elements stay on Spark's comparison" + + private val eagerEvalReason: String = + "array_contains evaluates its value eagerly, while Spark skips it when the array is null" private val collationReason: String = ArrayElementEqualitySupport.collationReason("array_contains") - override def getIncompatibleReasons(): Seq[String] = Seq(floatingPointReason, collationReason) + override def getIncompatibleReasons(): Seq[String] = + Seq(nestedFloatReason, eagerEvalReason, collationReason) - override def getSupportLevel(expr: ArrayContains): SupportLevel = expr.left.dataType match { - // See ArrayElementEqualitySupport: bytewise string comparison ignores the collation. - case ArrayType(elementType, _) if hasNonDefaultStringCollation(elementType) => + private def isFlatFloat(elementType: DataType): Boolean = elementType match { + case FloatType | DoubleType => true + case _ => false + } + + // Whether evaluating `expr` earlier than Spark would is unobservable: a literal or column read + // cannot throw. See CometArrayJoin.orderInsensitive. + private def orderInsensitive(expr: Expression): Boolean = expr match { + case _: Literal | _: Attribute | _: BoundReference => true + case _ => false + } + + override def getSupportLevel(expr: ArrayContains): SupportLevel = { + val elementType = expr.left.dataType.asInstanceOf[ArrayType].elementType + if (hasNonDefaultStringCollation(elementType)) { + // See ArrayElementEqualitySupport: bytewise string comparison ignores the collation. Incompatible(Some(collationReason)) - // Native array_contains compares floating-point elements bitwise, disagreeing with Spark for - // -0.0/+0.0 and NaN. Report Incompatible (not Unsupported) for float/double element types (at - // any nesting level) so the expression routes through the JVM codegen dispatcher (Spark's own - // doGenCode) and stays native + Spark-exact under the default config, while non-float arrays - // keep the fast native kernel. Under allowIncompatible=true the native kernel is used - // as before. - case ArrayType(elementType, _) - if SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType]) => - Incompatible(Some(floatingPointReason)) - case _ => Compatible() + } else if (isFlatFloat(elementType)) { + // Spark skips the value when the array is null, so a value that can throw (an ANSI cast, + // say) must not run for a null array. The native kernel evaluates both arguments first. + if (expr.left.nullable && !orderInsensitive(expr.right)) { + Incompatible(Some(eagerEvalReason)) + } else { + Compatible() + } + } else if (SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])) { + // Float leaves inside arrays or structs stay on the codegen dispatcher: the per-element + // comparator is slower than Spark's generated code there, and it compares any string field + // bytewise, ignoring its collation. + Incompatible(Some(nestedFloatReason)) + } else { + Compatible() + } } override def convert( @@ -151,8 +175,17 @@ object CometArrayContains binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) val keyExprProto = exprToProtoInternal(expr.children(1), inputs, binding) - - scalarFunctionExprToProto("array_contains", arrayExprProto, keyExprProto) + // datafusion-spark's array_contains compares floats by their bits. spark_array_contains + // compares them as Spark does, with -0.0 equal to 0.0 and all NaNs equal, at any depth, so + // it also serves nested float elements under allowIncompatible. + val elementType = expr.left.dataType.asInstanceOf[ArrayType].elementType + val function = + if (SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])) { + "spark_array_contains" + } else { + "array_contains" + } + scalarFunctionExprToProto(function, arrayExprProto, keyExprProto) } } diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql index e0fb8709b93..6a38f7efdd1 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql @@ -59,8 +59,8 @@ SELECT array_contains(cast(NULL as array), val) FROM test_array_contains -- ============================================================ -- Floating-point elements: Spark compares with ordering.equiv, so -0.0 == +0.0 and NaN == NaN. --- Native array_contains compares raw Arrow values bitwise, so float/double element types route --- through the JVM codegen dispatcher (Spark's own doGenCode) to stay native and Spark-exact. +-- Flat float/double arrays run Comet's spark_array_contains, which compares the same way; nested +-- float elements route through the JVM codegen dispatcher (Spark's own doGenCode). -- ============================================================ statement @@ -87,6 +87,6 @@ SELECT array_contains(array(-0.0D, 1.0D), 0.0D), array_contains(array(cast('-0.0' as float), 1.0F), cast('0.0' as float)), array_contains(array(cast('NaN' as float), 1.0F), cast('NaN' as float)) --- nested float element: the element type is array, so the float decline must recurse +-- nested float element: the element type is array, so the dispatch decision must recurse query SELECT array_contains(array(array(0.0D), array(1.0D)), array(-0.0D)) diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql new file mode 100644 index 00000000000..0a4d4306bc1 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql @@ -0,0 +1,32 @@ +-- 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. + +-- A collated string field inside a struct element compares by its collation in Spark. Nested +-- elements stay on the codegen dispatcher, which keeps that. Collation syntax requires Spark 4.0+. +-- MinSparkVersion: 4.0 + +statement +CREATE TABLE ac_collation(d DOUBLE, s STRING) USING parquet + +statement +INSERT INTO ac_collation VALUES (0.0D, 'a'), (1.0D, 'b'), (NULL, NULL) + +query expect_dispatch(array_contains) +SELECT array_contains( + array(named_struct('d', d, 's', CAST(s AS STRING COLLATE UTF8_LCASE))), + named_struct('d', d, 's', CAST(upper(s) AS STRING COLLATE UTF8_LCASE))) +FROM ac_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql new file mode 100644 index 00000000000..359f4b83803 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.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. + +-- Spark skips array_contains' value when the array is null, so a value that throws is not +-- evaluated for that row. The native kernel evaluates both arguments first, so a value other than +-- a literal or a column read stays on the codegen dispatcher when the array can be null. +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE ac_eager(a ARRAY, s STRING) USING parquet + +statement +INSERT INTO ac_eager VALUES (NULL, 'bad'), (array(1.0D), '1') + +query expect_dispatch(array_contains) +SELECT array_contains(a, CAST(s AS DOUBLE)) FROM ac_eager + +-- A column read is safe to evaluate early, so it stays native +query expect_native(array_contains) +SELECT array_contains(a, 1.0D), array_contains(a, CAST(NULL AS DOUBLE)) FROM ac_eager diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql new file mode 100644 index 00000000000..010608cecfb --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql @@ -0,0 +1,85 @@ +-- 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. + +-- array_contains compares floats the way Spark's generated code does: -0.0 equals 0.0, and every +-- NaN equals every other NaN, at any depth of an array or struct element. Flat FLOAT and DOUBLE +-- arrays run natively; nested float elements go through the codegen dispatcher. The result is true when +-- an element matches, null when none does and the array holds a null element, and false +-- otherwise; a null array or value gives null. +-- +-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every +-- platform. That is the NaN that arithmetic produces on x86-64. + +statement +CREATE TABLE ac_float(id INT, d DOUBLE, f FLOAT) USING parquet + +statement +INSERT INTO ac_float VALUES + (1, 0.0D, float('0.0')), + (2, double('-0.0'), float('-0.0')), + (3, double('NaN'), float('NaN')), + (4, 1.0D, float('1.0')), + (5, NULL, NULL) + +-- A constant value: either zero finds both zeros, and a NaN finds both NaNs +query expect_native(array_contains) +SELECT id, array_contains(array(d, -d), 0.0D), array_contains(array(d, -d), -0.0D), + array_contains(array(-d, 1.0D), double('NaN')), array_contains(array(d, 1.0D), 2.0D) +FROM ac_float + +query expect_native(array_contains) +SELECT id, array_contains(array(f, -f), 0.0F), array_contains(array(-f, 2.0F), float('NaN')) +FROM ac_float + +-- A value per row +query expect_native(array_contains) +SELECT id, array_contains(array(0.0D, 1.0D), d), array_contains(array(double('NaN'), 1.0D), -d), + array_contains(array(float('-0.0'), 2.0F), f) +FROM ac_float + +-- Null elements: true when another element matches, null when nothing matches +query expect_native(array_contains) +SELECT id, array_contains(array(d, NULL), 0.0D), array_contains(array(2.0D, NULL), d), + array_contains(array(NULL, -d), double('NaN')) +FROM ac_float + +-- A null array or a null value gives null; an empty array gives false +query expect_native(array_contains) +SELECT id, array_contains(IF(d IS NULL, NULL, array(d, -d)), 0.0D), + array_contains(array(d, -d), CAST(NULL AS DOUBLE)), array_contains(array(), d) +FROM ac_float + +-- Arrays as elements compare their floats the same way. Nested float elements stay on the codegen +-- dispatcher. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(array(-d), array(1.0D)), array(d)), + array_contains(array(array(d, 1.0D), array(2.0D)), array(-d, 1.0D)), + array_contains(array(array(d), NULL), array(2.0D)) +FROM ac_float + +-- Struct elements too +query expect_dispatch(array_contains) +SELECT id, array_contains(array(named_struct('x', -d, 'y', 1), named_struct('x', 2.0D, 'y', 1)), + named_struct('x', d, 'y', 1)) +FROM ac_float + +-- Only literals, which the suite does not constant fold. `-0.0D` is a literal with the sign bit +-- set, while `double('NaN')` is a cast. +query expect_native(array_contains) +SELECT array_contains(array(0.0D, 1.0D), -0.0D), array_contains(array(-0.0D), 0.0D), + array_contains(array(double('NaN'), 1.0D), double('NaN')), + array_contains(array(1.0D, NULL), -0.0D) diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql index a6e9499c047..b3991c1a1dc 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql @@ -44,7 +44,8 @@ SELECT array_join(words, '-') FROM routing_arrays query expect_native(reverse) SELECT reverse(a) FROM routing_arrays -query expect_fallback(array_contains: spark.comet.exec.scalaUDF.codegen.enabled=false) +-- array_contains on floats follows Spark natively, so it needs no dispatcher. +query expect_native(array_contains) SELECT array_contains(f, 1.0D) FROM routing_arrays -- sort_array on floats follows Spark natively, even in strict floating-point mode. diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql index 1e688857a9e..c401a040a9c 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql @@ -44,7 +44,8 @@ SELECT array_join(words, '-') FROM routing_arrays query expect_native(reverse) SELECT reverse(a) FROM routing_arrays -query expect_dispatch(array_contains) +-- array_contains on floats follows Spark natively, even in strict floating-point mode. +query expect_native(array_contains) SELECT array_contains(f, 1.0D) FROM routing_arrays -- sort_array on floats follows Spark natively, even in strict floating-point mode. diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 091f777dbe9..da07c05f673 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -442,11 +442,11 @@ class CometMapExpressionSuite extends CometTestBase { } } - // `map_contains_key` lowers to `array_contains(map_keys(...), key)`. `CometArrayContains` reports - // Incompatible for floating-point element types, because native `array_contains` compares - // -0.0/+0.0 and NaN bitwise unlike Spark. So under the default config it codegen-dispatches to - // Spark and the query stays native with Spark-exact results and no fallback. Forcing the native - // kernel with `ArrayContains.allowIncompatible=true` makes it serialize its children instead, so + // `map_contains_key` lowers to `array_contains(map_keys(...), key)`. The key here is a cast, + // which `CometArrayContains` keeps on the codegen dispatcher over a nullable array, because Spark + // skips the value for a null array while the native kernel evaluates it first. So under the + // default config the query stays native with Spark-exact results and no fallback. Forcing the + // native kernel with `ArrayContains.allowIncompatible=true` makes it serialize its children, so // the folded double-keyed map literal reaches the literal-expansion path. That map is nested // inside the INT-keyed outer map, past the outer `element_at` key guard which only sees the INT // key. The nested-key walk in the expansion (`mapKeyTypesExpandable`) is what has to decline the @@ -455,8 +455,8 @@ class CometMapExpressionSuite extends CometTestBase { test("map_contains_key over nested floating-point map keys falls back (multirow)") { withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { val negZero = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)" - // allowIncompatible flips array_contains from codegen dispatch to the native kernel, so the - // nested-map literal is actually serialized and the expansion guard runs. + // allowIncompatible lifts array_contains' eager-evaluation guard, so the nested-map literal + // is actually serialized and the expansion guard runs. withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[ArrayContains]) -> "true") { checkSparkAnswerAndFallbackReason( "SELECT _1 AS id, map_contains_key(" + From a3b297c4d1b0aa26776213342716aef9cb4ec34d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 10 Oct 2026 12:08:23 +0000 Subject: [PATCH 46/89] fix: don't let S3A committer, delete and read defaults block native Iceberg writes (#6808) #6441 makes a native Iceberg write to S3 fall back when the Hadoop configuration has an fs.s3a.* setting the native writer doesn't support. Spark 4.1 sets fs.s3a.committer.magic.enabled and fs.s3a.committer.name for every application when spark-hadoop-cloud is on the classpath (SPARK-47618), and many deployments set fs.s3a.experimental.input.fadvise, fs.s3a.bulk.delete.page.size or fs.s3a.committer.threads cluster-wide. Any of them makes every native S3 write fall back. None of these settings affects an Iceberg data-file write: the native writer doesn't go through S3A, fadvise is a read hint, and Iceberg commits through table metadata rather than a Hadoop output committer. Ignore them like the Spark-seeded read settings already in IgnoredHadoopS3Keys. --- .../operator/CometIcebergNativeWrite.scala | 11 ++++++++++- .../comet/CometIcebergWriteDetectionSuite.scala | 17 +++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index 34c07ff8dd5..9c708e2d0d1 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -158,10 +158,19 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // from spark.hadoop.*. They do not alter an Iceberg data-file write request, so they must not // make every otherwise-clean S3 write ineligible. This also permits explicit overrides, which // are harmless on the write path for the same reason. + // + // Some Spark distributions also set S3A delete, read and committer defaults cluster-wide. The + // native writer does not go through S3A, `fadvise` is a read hint, and Iceberg commits through + // table metadata rather than a Hadoop output committer, so these are harmless here as well. private val IgnoredHadoopS3Keys: Set[String] = Set( "fs.s3a.downgrade.syncable.exceptions", "fs.s3a.vectored.read.max.merged.size", - "fs.s3a.vectored.read.min.seek.size") + "fs.s3a.vectored.read.min.seek.size", + "fs.s3a.bulk.delete.page.size", + "fs.s3a.experimental.input.fadvise", + "fs.s3a.committer.magic.enabled", + "fs.s3a.committer.name", + "fs.s3a.committer.threads") // Audited against the pinned iceberg-rust S3 parser // (`iceberg/src/io/storage/config/s3.rs` and `storage/opendal/src/s3.rs`). Do not broaden this diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 922b5d6073f..7caffd0e7ee 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -772,6 +772,23 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("cluster-wide S3A delete, read and committer defaults do not block a native write") { + val hadoopConf = new Configuration(false) + Seq( + "fs.s3a.bulk.delete.page.size" -> "1000", + "fs.s3a.experimental.input.fadvise" -> "random", + "fs.s3a.committer.magic.enabled" -> "true", + "fs.s3a.committer.name" -> "magic", + "fs.s3a.committer.threads" -> "8").foreach { case (k, v) => hadoopConf.set(k, v) } + assert( + CometIcebergNativeWrite.unsupportedHadoopS3Settings(hadoopConf, Some("target")).isEmpty) + + hadoopConf.set("fs.s3a.encryption.algorithm", "SSE-KMS") + assert( + CometIcebergNativeWrite.unsupportedHadoopS3Settings(hadoopConf, Some("target")) == Seq( + "fs.s3a.encryption.algorithm")) + } + test("Hadoop S3A built-in defaults are ignored but custom resources are effective") { def xml(key: String, value: String): java.io.ByteArrayInputStream = new java.io.ByteArrayInputStream(s""" From dc880d6d265f8391b6202efa43574cfeea71e75a Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 10 Oct 2026 12:46:12 +0000 Subject: [PATCH 47/89] fix: run InsertIntoHadoopFsRelationCommand's post-write steps in the Spark 3.x native writer (#6821) * fix: run InsertIntoHadoopFsRelationCommand's post-write steps in the Spark 3.x native writer On Spark 3.x CometNativeWriteExec replaces the whole InsertIntoHadoopFsRelationCommand, but only ran the job/task commit lifecycle. After the job commits it now refreshes the table's file index, recaches cached plans that read the output path, and updates or clears the table stats, as the command does. It also rejects duplicate output column names up front, which the command checks before writing. Co-authored-by: Isaac * test: assert the duplicate-column write is native Spark's own writer rejects the same columns, so a fallback would keep the test green without exercising the native writer's check. Capture the plan of the failed write as well, and assert it carried a native write. Co-authored-by: Isaac --------- Co-authored-by: Isaac --- .../sql/comet/CometNativeWriteExec.scala | 36 ++++- .../parquet/CometParquetWriterSuite.scala | 125 +++++++++++++++++- .../parquet/CometParquetWriterTestBase.scala | 43 ++++-- 3 files changed, 189 insertions(+), 15 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index 200c7831321..ec903fc8937 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -37,9 +37,10 @@ import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} -import org.apache.spark.sql.execution.command.DataWritingCommandExec -import org.apache.spark.sql.execution.datasources.{OutputWriterFactory, WriteFilesExec} +import org.apache.spark.sql.execution.command.{CommandUtils, DataWritingCommandExec} +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, OutputWriterFactory, WriteFilesExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.util.SchemaUtils import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.{SerializableConfiguration, Utils} @@ -106,7 +107,21 @@ case class CometNativeWriteExec( sparkContext.emptyRDD[ColumnarBatch] } + /** The command this exec replaces, which holds the table this write updates, if any. */ + private def command: InsertIntoHadoopFsRelationCommand = originalPlan.cmd match { + case cmd: InsertIntoHadoopFsRelationCommand => cmd + case other => + throw new IllegalStateException( + s"${getClass.getSimpleName} replaced an unexpected write command: ${other.getClass}") + } + private def executeWriteAndCommit(): Unit = { + val cmd = command + // Like InsertIntoHadoopFsRelationCommand.run, reject duplicate columns before anything else. + SchemaUtils.checkColumnNameDuplication( + cmd.outputColumnNames, + session.sessionState.conf.caseSensitiveAnalysis) + if (!prepareOutputPathForMode()) { logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") return @@ -124,6 +139,23 @@ case class CometNativeWriteExec( s"${metrics("files_written").value} files, " + s"${metrics("bytes_written").value} bytes, ${metrics("rows_written").value} rows") })(catchBlock = committer.abortJob(job)) + // Outside the block above: the job has committed, so a failure here must not abort it. + refreshAfterWrite(cmd) + } + + /** + * Repeats what InsertIntoHadoopFsRelationCommand.run does once the write has committed, since + * this exec replaces that command. Partition metadata updates are left out because partitioned + * writes never reach this exec. + */ + private def refreshAfterWrite(cmd: InsertIntoHadoopFsRelationCommand): Unit = { + val fs = cmd.outputPath.getFileSystem(serializableHadoopConf.value) + // Refresh the files the table's relation has listed. + cmd.fileIndex.foreach(_.refresh()) + // Recache any cached plan that reads the output path. + session.sharedState.cacheManager.recacheByPath(session, cmd.outputPath, fs) + // Update or clear the table stats, depending on spark.sql.statistics.size.autoUpdate. + cmd.catalogTable.foreach(CommandUtils.updateTableStats(session, _)) } private def runNativeWriteJob( diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 466f6cad9e6..2feb4e12e2d 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -34,7 +34,7 @@ import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometSparkToColumnarExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, FileSourceScanExec, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.command.DataWritingCommandExec @@ -1318,9 +1318,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("INSERT INTO ... SELECT is visible to subsequent reads") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3521 - reads returned no rows because the - // bespoke write path never refreshed the catalog cache. Spark's command does that itself. + // bespoke write path never refreshed the catalog cache. On Spark 4.0+ Spark's command does + // that itself; on 3.x CometNativeWriteExec repeats it after the job commits. withTable("comet_write_target", "comet_write_source") { withNativeWriter { sql("CREATE TABLE comet_write_source(id bigint, name string) USING parquet") @@ -1340,6 +1340,125 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a native overwrite of a cached path recaches the data") { + withTempPath { dir => + val target = new File(dir, "target").getAbsolutePath + val source = new File(dir, "source").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 5).write.parquet(target) + spark.range(100, 103).write.parquet(source) + } + val cached = spark.read.parquet(target).cache() + try { + assert(cached.count() == 5) + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + spark.read.parquet(source).write.mode(SaveMode.Overwrite).parquet(target))) + } + // Spark recaches every cached plan that reads the output path once the write commits. + checkAnswer(cached, (100L until 103L).map(Row(_))) + } finally { + cached.unpersist() + } + } + } + + test("a native write rejects case-insensitively duplicate column names") { + withTempPath { dir => + val source = new File(dir, "source").getAbsolutePath + val target = new File(dir, "target") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(3).selectExpr("id AS a", "id AS b").write.parquet(source) + } + withNativeWriter { + val (plan, e) = captureFailedWritePlan( + spark.read + .parquet(source) + .selectExpr("a", "b AS A") + .write + .mode(SaveMode.Overwrite) + .parquet(target.getAbsolutePath)) + // Spark's own writer rejects these columns too, so a fallback would also pass the checks + // below without exercising the native writer's check. + assertHasCometNativeWriteExec(plan) + assert(e.isInstanceOf[AnalysisException], e) + assert(e.getMessage.contains("COLUMN_ALREADY_EXISTS"), e.getMessage) + } + assert(!target.exists(), "nothing may be written when the columns are rejected") + } + } + + test("a native INSERT into a cached table recaches the table") { + withTable("comet_write_target", "comet_write_source") { + sql("CREATE TABLE comet_write_source(id bigint) USING parquet") + sql("CREATE TABLE comet_write_target(id bigint) USING parquet") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("INSERT INTO comet_write_source VALUES (1), (2)") + sql("INSERT INTO comet_write_target VALUES (0)") + } + sql("CACHE TABLE comet_write_target") + try { + checkAnswer(spark.table("comet_write_target"), Row(0L)) + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + sql("INSERT INTO comet_write_target SELECT id FROM comet_write_source"))) + } + checkAnswer(spark.table("comet_write_target"), Seq(Row(0L), Row(1L), Row(2L))) + } finally { + sql("UNCACHE TABLE IF EXISTS comet_write_target") + } + } + } + + for { + autoUpdate <- Seq(true, false) + overwrite <- Seq(false, true) + } { + val insert = if (overwrite) "INSERT OVERWRITE" else "INSERT INTO" + test(s"native $insert maintains table stats with size autoUpdate=$autoUpdate") { + withTable("comet_write_target", "comet_write_source") { + withSQLConf(SQLConf.AUTO_SIZE_UPDATE_ENABLED.key -> autoUpdate.toString) { + sql("CREATE TABLE comet_write_source(id bigint) USING parquet") + sql("CREATE TABLE comet_write_target(id bigint) USING parquet") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("INSERT INTO comet_write_source SELECT id FROM range(100)") + sql("INSERT INTO comet_write_target VALUES (0)") + } + sql("ANALYZE TABLE comet_write_target COMPUTE STATISTICS NOSCAN") + val target = TableIdentifier("comet_write_target") + val catalog = spark.sessionState.catalog + val statsBefore = catalog.getTableMetadata(target).stats + assert(statsBefore.isDefined) + + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + sql(s"$insert comet_write_target SELECT id FROM comet_write_source"))) + } + + // Mirrors CommandUtils.updateTableStats, which Spark runs after the write commits. + val statsAfter = catalog.getTableMetadata(target).stats + if (autoUpdate) { + val location = new Path(catalog.getTableMetadata(target).location) + val fs = location.getFileSystem(spark.sessionState.newHadoopConf()) + val dataSize = fs + .listStatus(location) + .filterNot(f => + f.getPath.getName.startsWith("_") || f.getPath.getName.startsWith(".")) + .map(_.getLen) + .sum + assert(statsAfter.map(_.sizeInBytes) == Some(BigInt(dataSize))) + assert(statsAfter != statsBefore) + } else { + assert(statsAfter.isEmpty, s"stale stats were left in place: $statsAfter") + } + } + } + } + } + test("dynamic partition overwrite falls back to Spark") { // A dynamic overwrite is a partitioned write, which both native writers decline - but the // consequence of getting it wrong is silent data loss across untouched partitions, so assert diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index 20caba36b1f..605dcbbb3ee 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -21,6 +21,8 @@ package org.apache.comet.parquet import java.util.concurrent.atomic.AtomicReference +import scala.util.control.NonFatal + import org.apache.spark.CometListenerBusUtils import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometNativeWriteExec, CometWriteFilesExec} @@ -64,19 +66,37 @@ abstract class CometParquetWriterTestBase extends CometTestBase { /** As above, for a write that names its own target (an `INSERT INTO`, for example). */ protected def captureWritePlan(writeOp: => Unit): SparkPlan = { + val (plan, failure) = captureWriteOutcome(writeOp) + failure.foreach(throw _) + plan.getOrElse(fail("Listener was not called - no execution plan captured")) + } + + /** + * As above, for a write that is expected to fail. Returns the plan of the failed write, so a + * test can check that the failure came from the native writer rather than from Spark's. + */ + protected def captureFailedWritePlan(writeOp: => Unit): (SparkPlan, Throwable) = { + val (plan, failure) = captureWriteOutcome(writeOp) + val error = failure.getOrElse(fail("Expected the write to fail, but it succeeded")) + (plan.getOrElse(fail("Listener was not called - no execution plan captured")), error) + } + + private def captureWriteOutcome(writeOp: => Unit): (Option[SparkPlan], Option[Throwable]) = { val capturedPlan = new AtomicReference[QueryExecution]() + def isWrite(funcName: String): Boolean = funcName == "save" || funcName.contains("command") val listener = new org.apache.spark.sql.util.QueryExecutionListener { override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { - if (funcName == "save" || funcName.contains("command")) { + if (isWrite(funcName)) { capturedPlan.set(qe) } } - override def onFailure( - funcName: String, - qe: QueryExecution, - exception: Exception): Unit = {} + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + if (isWrite(funcName)) { + capturedPlan.set(qe) + } + } } // Listener events are delivered asynchronously, so drain the bus before registering: an @@ -85,12 +105,15 @@ abstract class CometParquetWriterTestBase extends CometTestBase { spark.listenerManager.register(listener) try { - writeOp + val failure = + try { + writeOp + None + } catch { + case NonFatal(e) => Some(e) + } CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - - val plan = capturedPlan.get() - assert(plan != null, "Listener was not called - no execution plan captured") - stripAQEPlan(plan.executedPlan) + (Option(capturedPlan.get()).map(qe => stripAQEPlan(qe.executedPlan)), failure) } finally { spark.listenerManager.unregister(listener) } From b94d52f43c9edaf341f5c2cc00fa2e351ba9222b Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 10 Oct 2026 14:30:39 +0000 Subject: [PATCH 48/89] refactor: move the plan entry point logic into core functions (#6695) * refactor: move the plan entry point logic into core functions Separate the native logic of createPlan, setShufflePartitionPusher, executePlan and releasePlan from their JNI argument handling: - create_execution_context builds the execution context from plain values (PlanInputs) and a memory pool factory. prepare_plan_settings parses the Spark configs and initializes the Tokio runtime. - register_shuffle_partition_pusher checks the plan state before building the pusher. - execute_plan returns the plan's next batch; the export still exports it. - release_plan stops and drops the plan. These plans call back into the JVM, so the JVM objects stay JNI global references in the execution context (PlanJvmRefs), and the steps that need an Env are passed in as callbacks: the memory pool factory, the pusher factory, and PublishMetrics. ExecutionContext::metrics becomes an Option, which is always set by createPlan, so a plan can be created and run without a JVM in unit tests. Exception classes and messages are unchanged. * Test that release_plan returns the memory of a plan holding a reservation The lifecycle test ran a RangeScan, which never reserves memory, so its check that nothing stays reserved after release_plan could not fail. Run a descending sort over the range instead, release it after its first batch while it still holds its buffered input, and check that the plan holds a reservation before the release and none after. Do the same when the final metrics publish fails. Draining a plan to the end stays covered by a separate test. --- native/core/src/execution/jni_api.rs | 921 +++++++++++++++++++-------- 1 file changed, 651 insertions(+), 270 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 98baf952ef3..0812b9f7087 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -24,7 +24,10 @@ //! //! `createPlan`, `setShufflePartitionPusher`, `executePlan` and `releasePlan` also depend on //! upcalls into the JVM (input iterators, the task memory manager, metrics, UDFs and scalar -//! subqueries), so their core logic keeps holding JNI references. +//! subqueries). Their core functions (`create_execution_context`, +//! `register_shuffle_partition_pusher`, `execute_plan` and `release_plan`) take the JVM objects +//! as JNI global references held by the `ExecutionContext` (`PlanJvmRefs`), and the steps that +//! need an `Env` as callbacks: the memory pool factory, the pusher factory, and `PublishMetrics`. use super::{serde, utils::SparkArrowConvert}; use crate::{ @@ -115,7 +118,7 @@ use tokio::sync::mpsc; use tokio::task::JoinHandle; use crate::execution::memory_pools::{ - create_memory_pool, overcommit, parse_memory_pool_config, PlanMemoryPool, + create_memory_pool, overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, }; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ @@ -539,8 +542,9 @@ struct ExecutionContext { batch_producer: Option, /// The pool every reservation the plan makes goes through plan_memory: Arc, - /// Native metrics - pub metrics: Arc>>, + /// The JVM `CometMetricNode` the plan's metrics are published to. `None` only without a + /// JVM (unit tests). + pub metrics: Option>>>, // The interval in milliseconds to update metrics pub metrics_update_interval: Option, // The last update time of metrics @@ -607,33 +611,16 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( class_loader_obj: JObject, ) -> jlong { try_unwrap_or_throw(&e, |env| { - // Deserialize Spark configs let bytes = env.convert_byte_array(serialized_spark_configs)?; - let spark_configs = serde::deserialize_config(bytes.as_slice())?; - let spark_config: HashMap = spark_configs.entries.into_iter().collect(); - - // Initialize the tokio runtime with spark.executor.cores as the default - // worker thread count, falling back to 1 if not set. - let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); - init_runtime(executor_cores); - - // Access Comet configs - let debug_native = spark_config.get_bool(COMET_DEBUG_ENABLED); - let explain_native = spark_config.get_bool(COMET_EXPLAIN_NATIVE_ENABLED); - let tracing_enabled = spark_config.get_bool(COMET_TRACING_ENABLED); - let max_temp_directory_size = - spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024); - let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY); - - with_trace("createPlan", tracing_enabled, || { + let settings = prepare_plan_settings(&bytes)?; + + with_trace("createPlan", settings.tracing_enabled, || { // Init JVM classes JVMClasses::init(env); let start = Instant::now(); - // Deserialize query plan - let bytes = env.convert_byte_array(serialized_query)?; - let spark_plan = serde::deserialize_op(bytes.as_slice())?; + let serialized_plan = env.convert_byte_array(serialized_query)?; let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?); @@ -646,40 +633,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources.push(input_source); } - // Create DataFusion memory pool let task_memory_manager = Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); let memory_pool_type = memory_pool_type.try_to_string(env)?; - let memory_pool_config = parse_memory_pool_config( - off_heap_mode != JNI_FALSE, - memory_pool_type, - memory_limit, - )?; - let memory_pool = - create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); - - // Register the shared base pool before wrapping it for per-plan debug logging. The - // guard removes the entry if any later plan setup step fails. - // - // Registration is not conditional on this plan's tracing setting. `tracing.enabled` is - // a session config, so an executor can run a traced and an untraced plan at once, - // while the allocation counter a trace compares against is process-wide. Registering - // only traced plans would leave the untraced plan's reservation out of the total and - // report it as allocation held outside any pool. - let rust_thread_id = get_thread_id(); - let memory_pool_registration = Some(ThreadMemoryPoolRegistration::new( - rust_thread_id, - id, - Arc::clone(&memory_pool), - )); - - let memory_pool = if logging_memory_pool { - Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) - } else { - memory_pool - }; - let plan_memory = Arc::new(PlanMemoryPool::new(memory_pool)); // Get local directories for storing spill files let num_local_dirs = local_dirs.len(env)?; @@ -691,46 +648,13 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( local_dirs_vec.push(local_dir); } - // We need to keep the session context alive. Some session state like temporary - // dictionaries are stored in session context. If it is dropped, the temporary - // dictionaries will be dropped as well. - let session = prepare_datafusion_session_context( - batch_size as usize, - Arc::clone(&plan_memory) as Arc, - local_dirs_vec, - max_temp_directory_size, - task_cpus as usize, - &spark_config, - &spark_plan, - )?; - - let plan_creation_time = start.elapsed(); - - let metrics_update_interval = if metrics_update_interval > 0 { - Some(Duration::from_millis(metrics_update_interval as u64)) + // Key unwrapper for encrypted files + let key_unwrapper = if !key_unwrapper_obj.is_null() { + Some(Arc::new(jni_new_global_ref!(env, key_unwrapper_obj)?)) } else { None }; - // Handle key unwrapper for encrypted files - if !key_unwrapper_obj.is_null() { - let encryption_factory = CometEncryptionFactory { - key_unwrapper: Arc::new(jni_new_global_ref!(env, key_unwrapper_obj)?), - }; - session.runtime_env().register_parquet_encryption_factory( - ENCRYPTION_FACTORY_ID, - Arc::new(encryption_factory), - ); - } - - let session = Arc::new(session); - - let tracing_event_name = if tracing_enabled { - build_tracing_event_name(&spark_plan) - } else { - String::new() - }; - // Capture the driving Spark task's TaskContext and context ClassLoader as JNI global // references when non-null. The `Arc>` releases its global ref on // drop, so cleanup is automatic when the ExecutionContext drops. @@ -746,41 +670,225 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( None }; - let exec_context = Box::new(ExecutionContext { + let inputs = PlanInputs { id, - spark_plan, + settings, + serialized_plan, partition_count: partition_count as usize, - root_op: None, - scans: vec![], - shuffle_scans: vec![], - input_sources, - stream: None, - batch_producer: None, - plan_memory, - metrics, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, metrics_update_interval, - metrics_last_update_time: Instant::now(), - plan_creation_time, - session_ctx: session, - debug_native, - explain_native, - tracing_enabled, - rust_thread_id, - tracing_memory_metric_name: format!( - "thread_{rust_thread_id}_comet_memory_reserved" - ), - tracing_event_name, - task_context, - class_loader, - shuffle_partition_pusher: None, - memory_pool_registration, - }); + start, + jvm: PlanJvmRefs { + metrics: Some(metrics), + input_sources, + key_unwrapper, + task_context, + class_loader, + }, + }; + // The pool acquires memory from Spark's task memory manager. + let exec_context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; Ok(Box::into_raw(exec_context) as i64) }) }) } +/// What `createPlan` takes from the Spark configs. +struct PlanSettings { + spark_config: HashMap, + debug_native: bool, + explain_native: bool, + tracing_enabled: bool, + max_temp_directory_size: u64, + logging_memory_pool: bool, +} + +/// Deserializes the Spark configs of a plan and initializes the Tokio runtime from them. The +/// first step of `Native.createPlan`, before its trace span opens. +fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult { + let spark_configs = serde::deserialize_config(serialized_spark_configs)?; + let spark_config: HashMap = spark_configs.entries.into_iter().collect(); + + // Initialize the tokio runtime with spark.executor.cores as the default + // worker thread count, falling back to 1 if not set. + let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); + init_runtime(executor_cores); + + // Access Comet configs + Ok(PlanSettings { + debug_native: spark_config.get_bool(COMET_DEBUG_ENABLED), + explain_native: spark_config.get_bool(COMET_EXPLAIN_NATIVE_ENABLED), + tracing_enabled: spark_config.get_bool(COMET_TRACING_ENABLED), + max_temp_directory_size: spark_config + .get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024), + logging_memory_pool: spark_config.get_bool(COMET_DEBUG_MEMORY), + spark_config, + }) +} + +/// The arguments of `Native.createPlan`, converted from JNI. +struct PlanInputs { + id: i64, + settings: PlanSettings, + serialized_plan: Vec, + partition_count: usize, + batch_size: usize, + off_heap_mode: bool, + memory_pool_type: String, + memory_limit: i64, + task_attempt_id: i64, + task_cpus: usize, + local_dirs: Vec, + /// In milliseconds; zero or less publishes metrics only when the plan is released. + metrics_update_interval: i64, + /// When plan creation started, the start of `ExecutionContext::plan_creation_time`. + start: Instant, + jvm: PlanJvmRefs, +} + +/// The JVM objects a plan calls back into, as JNI global references. Each is `None` or empty +/// only without a JVM (unit tests) or when Spark passes none. +struct PlanJvmRefs { + metrics: Option>>>, + input_sources: Vec>>>, + key_unwrapper: Option>>>, + task_context: Option>>>, + class_loader: Option>>>, +} + +/// Creates the execution context of a native plan: deserializes the plan, sets up its memory +/// pool and DataFusion session, and registers the pool for tracing. The plan's operators are +/// created by the first `execute_plan`. Core of `Native.createPlan`. +/// +/// `memory_pool_for` builds the pool every reservation of the plan goes through from its +/// config; `Native.createPlan` builds one that acquires memory from Spark. +fn create_execution_context( + inputs: PlanInputs, + memory_pool_for: impl FnOnce(&MemoryPoolConfig) -> Arc, +) -> CometResult> { + let PlanInputs { + id, + settings, + serialized_plan, + partition_count, + batch_size, + off_heap_mode, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus, + local_dirs, + metrics_update_interval, + start, + jvm, + } = inputs; + + // Deserialize query plan + let spark_plan = serde::deserialize_op(serialized_plan.as_slice())?; + + // Create DataFusion memory pool + let memory_pool_config = + parse_memory_pool_config(off_heap_mode, memory_pool_type, memory_limit)?; + let memory_pool = memory_pool_for(&memory_pool_config); + + // Register the shared base pool before wrapping it for per-plan debug logging. The + // guard removes the entry if any later plan setup step fails. + // + // Registration is not conditional on this plan's tracing setting. `tracing.enabled` is + // a session config, so an executor can run a traced and an untraced plan at once, + // while the allocation counter a trace compares against is process-wide. Registering + // only traced plans would leave the untraced plan's reservation out of the total and + // report it as allocation held outside any pool. + let rust_thread_id = get_thread_id(); + let memory_pool_registration = Some(ThreadMemoryPoolRegistration::new( + rust_thread_id, + id, + Arc::clone(&memory_pool), + )); + + let memory_pool = if settings.logging_memory_pool { + Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) + } else { + memory_pool + }; + let plan_memory = Arc::new(PlanMemoryPool::new(memory_pool)); + + // We need to keep the session context alive. Some session state like temporary + // dictionaries are stored in session context. If it is dropped, the temporary + // dictionaries will be dropped as well. + let session = prepare_datafusion_session_context( + batch_size, + Arc::clone(&plan_memory) as Arc, + local_dirs, + settings.max_temp_directory_size, + task_cpus, + &settings.spark_config, + &spark_plan, + )?; + + let plan_creation_time = start.elapsed(); + + let metrics_update_interval = if metrics_update_interval > 0 { + Some(Duration::from_millis(metrics_update_interval as u64)) + } else { + None + }; + + // Handle key unwrapper for encrypted files + if let Some(key_unwrapper) = jvm.key_unwrapper { + let encryption_factory = CometEncryptionFactory { key_unwrapper }; + session.runtime_env().register_parquet_encryption_factory( + ENCRYPTION_FACTORY_ID, + Arc::new(encryption_factory), + ); + } + + let session = Arc::new(session); + + let tracing_event_name = if settings.tracing_enabled { + build_tracing_event_name(&spark_plan) + } else { + String::new() + }; + + Ok(Box::new(ExecutionContext { + id, + spark_plan, + partition_count, + root_op: None, + scans: vec![], + shuffle_scans: vec![], + input_sources: jvm.input_sources, + stream: None, + batch_producer: None, + plan_memory, + metrics: jvm.metrics, + metrics_update_interval, + metrics_last_update_time: Instant::now(), + plan_creation_time, + session_ctx: session, + debug_native: settings.debug_native, + explain_native: settings.explain_native, + tracing_enabled: settings.tracing_enabled, + rust_thread_id, + tracing_memory_metric_name: format!("thread_{rust_thread_id}_comet_memory_reserved"), + tracing_event_name, + task_context: jvm.task_context, + class_loader: jvm.class_loader, + shuffle_partition_pusher: None, + memory_pool_registration, + })) +} + /// Binds one task-owned shuffle callback before native execution is initialized. /// /// Keeping callback registration separate preserves the existing `createPlan` JNI ABI for @@ -800,23 +908,36 @@ pub extern "system" fn Java_org_apache_comet_Native_setShufflePartitionPusher( } let exec_context = get_execution_context(exec_context); - if exec_context.root_op.is_some() { - return Err(CometError::Internal( - "Remote shuffle callback cannot be registered after native execution starts" - .to_string(), - )); - } + register_shuffle_partition_pusher(exec_context, || { + Ok(Arc::new(JavaShufflePartitionPusher::try_new( + env, &callback, + )?)) + }) + }) +} - if exec_context.shuffle_partition_pusher.is_some() { - return Err(CometError::Internal( - "Remote shuffle callback has already been registered for this task".to_string(), - )); - } +/// Registers the plan's remote shuffle callback, built by `make_pusher` only once the plan +/// is known to accept one: before execution starts, and once. Core of +/// `Native.setShufflePartitionPusher`, whose pusher calls back into the JVM. +fn register_shuffle_partition_pusher( + exec_context: &mut ExecutionContext, + make_pusher: impl FnOnce() -> CometResult>, +) -> CometResult<()> { + if exec_context.root_op.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback cannot be registered after native execution starts" + .to_string(), + )); + } - let pusher = JavaShufflePartitionPusher::try_new(env, &callback)?; - exec_context.shuffle_partition_pusher = Some(Arc::new(pusher)); - Ok(()) - }) + if exec_context.shuffle_partition_pusher.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback has already been registered for this task".to_string(), + )); + } + + exec_context.shuffle_partition_pusher = Some(make_pusher()?); + Ok(()) } /// Skipping is opt-in (`spark.comet.exec.aggregate.skipPartial.enabled`): once DataFusion's probe @@ -1278,93 +1399,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( }; let result = with_trace(tracing_label, tracing_enabled, || { - let exec_context_id = exec_context.id; - - // Initialize the execution stream. - // Because we don't know if input arrays are dictionary-encoded when we create - // query plan, we need to defer stream initialization to first time execution. - if exec_context.root_op.is_none() { - let start = Instant::now(); - let planner = - PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) - .with_exec_id(exec_context_id) - .with_sql_text_pool(&exec_context.spark_plan) - .with_task_context(exec_context.task_context.clone()) - .with_class_loader(exec_context.class_loader.clone()) - .with_shuffle_partition_pusher( - exec_context.shuffle_partition_pusher.clone(), - ); - let (scans, shuffle_scans, root_op) = planner.create_plan( - &exec_context.spark_plan, - &mut exec_context.input_sources.clone(), - exec_context.partition_count, - )?; - let physical_plan_time = start.elapsed(); - - exec_context.plan_creation_time += physical_plan_time; - exec_context.scans = scans; - exec_context.shuffle_scans = shuffle_scans; - - if exec_context.explain_native { - let formatted_plan_str = - DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true); - info!("Comet native query plan:\n{formatted_plan_str:}"); - } - - let task_ctx = exec_context.session_ctx.task_ctx(); - // Each Comet native execution corresponds to a single Spark partition, - // so we should always execute partition 0. - let stream = root_op.native_plan.execute(0, task_ctx)?; - - if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() { - // No JVM data sources — spawn onto tokio so the executor - // thread parks in blocking_recv instead of busy-polling. - exec_context.batch_producer = - Some(BatchProducer::spawn(&get_runtime(), stream)); - } else { - exec_context.stream = Some(stream); - } - exec_context.root_op = Some(root_op); - } else { - // Pull input batches - pull_input_batches(exec_context)?; - } - - if let Some(producer) = &mut exec_context.batch_producer { - match producer.next_batch()? { - Some(batch) => { - // Publish on the configured interval, as the ScanExec path below does, - // since each publish walks the whole metric tree and calls into the JVM. - // `releasePlan` publishes the final values. - update_metrics_on_interval(env, exec_context)?; - return prepare_output( - env, - array_addrs, - schema_addrs, - batch, - exec_context.debug_native, - ); - } - None => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); - } - } - } - - // ScanExec path: JVM-fed scans return `Pending` until `pull_input_batches` refills - // them and wakes the stream. A poll that is still pending, with nothing having woken - // the stream by the end of the pull, waits on native I/O, and `next_batch` parks - // until it completes. - let mut stream = exec_context.stream.take().unwrap(); - let next = get_runtime().block_on(next_batch(&mut stream, || { - pull_input_batches(exec_context)?; - update_metrics_on_interval(env, exec_context) - })); - exec_context.stream = Some(stream); - let next = next?; - update_metrics_on_interval(env, exec_context)?; - match next { + let mut publish_metrics = |ctx: &ExecutionContext| update_metrics(env, ctx); + match execute_plan(exec_context, stage_id, partition, &mut publish_metrics)? { Some(batch) => prepare_output( env, array_addrs, @@ -1372,10 +1408,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( batch, exec_context.debug_native, ), - None => { - log_plan_metrics(exec_context, stage_id, partition); - Ok(-1) - } + None => Ok(-1), } }); @@ -1400,6 +1433,103 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( }) } +/// Publishes a plan's metrics to the JVM. The plan entry points pass one that calls into the JVM +/// through their `Env`, so the core functions that drive a plan take no JNI types. +type PublishMetrics<'a> = dyn FnMut(&ExecutionContext) -> CometResult<()> + 'a; + +/// Returns the plan's next output batch, or `None` once the plan has produced all of its output. +/// The first call creates the plan's operators and starts its stream; a plan with no JVM input +/// then runs on a Tokio task. Metrics are published through `publish_metrics` on the configured +/// interval. Core of `Native.executePlan`; the JVM-fed scans call back into the JVM for their +/// input. +fn execute_plan( + exec_context: &mut ExecutionContext, + stage_id: i32, + partition: i32, + publish_metrics: &mut PublishMetrics, +) -> CometResult> { + let exec_context_id = exec_context.id; + + // Initialize the execution stream. + // Because we don't know if input arrays are dictionary-encoded when we create + // query plan, we need to defer stream initialization to first time execution. + if exec_context.root_op.is_none() { + let start = Instant::now(); + let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) + .with_exec_id(exec_context_id) + .with_sql_text_pool(&exec_context.spark_plan) + .with_task_context(exec_context.task_context.clone()) + .with_class_loader(exec_context.class_loader.clone()) + .with_shuffle_partition_pusher(exec_context.shuffle_partition_pusher.clone()); + let (scans, shuffle_scans, root_op) = planner.create_plan( + &exec_context.spark_plan, + &mut exec_context.input_sources.clone(), + exec_context.partition_count, + )?; + let physical_plan_time = start.elapsed(); + + exec_context.plan_creation_time += physical_plan_time; + exec_context.scans = scans; + exec_context.shuffle_scans = shuffle_scans; + + if exec_context.explain_native { + let formatted_plan_str = + DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true); + info!("Comet native query plan:\n{formatted_plan_str:}"); + } + + let task_ctx = exec_context.session_ctx.task_ctx(); + // Each Comet native execution corresponds to a single Spark partition, + // so we should always execute partition 0. + let stream = root_op.native_plan.execute(0, task_ctx)?; + + if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() { + // No JVM data sources — spawn onto tokio so the executor + // thread parks in blocking_recv instead of busy-polling. + exec_context.batch_producer = Some(BatchProducer::spawn(&get_runtime(), stream)); + } else { + exec_context.stream = Some(stream); + } + exec_context.root_op = Some(root_op); + } else { + // Pull input batches + pull_input_batches(exec_context)?; + } + + if let Some(producer) = &mut exec_context.batch_producer { + return match producer.next_batch()? { + Some(batch) => { + // Publish on the configured interval, as the ScanExec path below does, + // since each publish walks the whole metric tree and calls into the JVM. + // `releasePlan` publishes the final values. + update_metrics_on_interval(exec_context, publish_metrics)?; + Ok(Some(batch)) + } + None => { + log_plan_metrics(exec_context, stage_id, partition); + Ok(None) + } + }; + } + + // ScanExec path: JVM-fed scans return `Pending` until `pull_input_batches` refills + // them and wakes the stream. A poll that is still pending, with nothing having woken + // the stream by the end of the pull, waits on native I/O, and `next_batch` parks + // until it completes. + let mut stream = exec_context.stream.take().unwrap(); + let next = get_runtime().block_on(next_batch(&mut stream, || { + pull_input_batches(exec_context)?; + update_metrics_on_interval(exec_context, publish_metrics) + })); + exec_context.stream = Some(stream); + let next = next?; + update_metrics_on_interval(exec_context, publish_metrics)?; + if next.is_none() { + log_plan_metrics(exec_context, stage_id, partition); + } + Ok(next) +} + #[no_mangle] /// Drop the native query plan object and context object. pub extern "system" fn Java_org_apache_comet_Native_releasePlan( @@ -1419,48 +1549,58 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( // Reclaim ownership of the context up front so that it is always freed, even if a step // below fails. Dropping it releases the memory pool and every JNI global ref the context // holds. - let mut execution_context: Box = + let execution_context: Box = Box::from_raw(exec_context as *mut ExecutionContext); - // A plan with no JVM input runs on a Tokio task. Stop it and drop its stream first, so the - // stream's memory is returned before the Spark task can end, and the metrics flushed below - // include what the stream records as it is dropped. - let producer_stopped = execution_context - .batch_producer - .take() - .map_or(Ok(()), BatchProducer::stop); - - // Unregister this context's pool and, when tracing, emit the remaining total for the - // thread. Every context registers, but only a traced one writes counters, so the - // reservations are read only then. - drop(execution_context.memory_pool_registration.take()); - if execution_context.tracing_enabled { - log_memory_usage( - &execution_context.tracing_memory_metric_name, - total_reserved_for_thread(execution_context.rust_thread_id) as u64, - ); - } + release_plan(execution_context, &mut |ctx| update_metrics(env, ctx)) + }) +} - let metrics_flushed = update_metrics(env, &mut execution_context); - - // Once this returns the Spark task can end, and Spark hands whatever the task still holds - // to other tasks. Dropping the plan aborts the tasks its operators spawned, and they give - // back what they hold the next time they yield, so wait for them. - let id = execution_context.id; - let plan_memory = Arc::clone(&execution_context.plan_memory); - drop(execution_context); - let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT); - if held > 0 { - warn!( - "Native plan {id} still holds {held} bytes {PLAN_MEMORY_RELEASE_TIMEOUT:?} after \ - it was released. They go back to Spark only when whatever holds them drops them, \ - which can be after the Spark task has ended and its memory has been given to \ - another task." - ); - } +/// Releases a plan: stops it, publishes its final metrics through `publish_metrics`, drops it, +/// and waits briefly for its memory reservations to be returned. The context is freed even if a +/// step fails. Core of `Native.releasePlan`. +fn release_plan( + mut execution_context: Box, + publish_metrics: &mut PublishMetrics, +) -> CometResult<()> { + // A plan with no JVM input runs on a Tokio task. Stop it and drop its stream first, so the + // stream's memory is returned before the Spark task can end, and the metrics flushed below + // include what the stream records as it is dropped. + let producer_stopped = execution_context + .batch_producer + .take() + .map_or(Ok(()), BatchProducer::stop); + + // Unregister this context's pool and, when tracing, emit the remaining total for the + // thread. Every context registers, but only a traced one writes counters, so the + // reservations are read only then. + drop(execution_context.memory_pool_registration.take()); + if execution_context.tracing_enabled { + log_memory_usage( + &execution_context.tracing_memory_metric_name, + total_reserved_for_thread(execution_context.rust_thread_id) as u64, + ); + } - producer_stopped.and(metrics_flushed) - }) + let metrics_flushed = publish_metrics(&execution_context); + + // Once this returns the Spark task can end, and Spark hands whatever the task still holds + // to other tasks. Dropping the plan aborts the tasks its operators spawned, and they give + // back what they hold the next time they yield, so wait for them. + let id = execution_context.id; + let plan_memory = Arc::clone(&execution_context.plan_memory); + drop(execution_context); + let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT); + if held > 0 { + warn!( + "Native plan {id} still holds {held} bytes {PLAN_MEMORY_RELEASE_TIMEOUT:?} after \ + it was released. They go back to Spark only when whatever holds them drops them, \ + which can be after the Spark task has ended and its memory has been given to \ + another task." + ); + } + + producer_stopped.and(metrics_flushed) } /// How long `releasePlan` waits for a released plan's reservations to be returned. The tasks @@ -1469,11 +1609,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( /// Spark task. const PLAN_MEMORY_RELEASE_TIMEOUT: Duration = Duration::from_secs(1); -/// Runs `update_metrics` once the configured interval has passed and, with tracing on, samples +/// Runs `publish_metrics` once the configured interval has passed and, with tracing on, samples /// this thread's pool reservation at the same cadence. fn update_metrics_on_interval( - env: &mut Env, exec_context: &mut ExecutionContext, + publish_metrics: &mut PublishMetrics, ) -> CometResult<()> { let Some(interval) = exec_context.metrics_update_interval else { return Ok(()); @@ -1482,7 +1622,7 @@ fn update_metrics_on_interval( if now - exec_context.metrics_last_update_time < interval { return Ok(()); } - update_metrics(env, exec_context)?; + publish_metrics(exec_context)?; exec_context.metrics_last_update_time = now; if exec_context.tracing_enabled { log_memory_usage( @@ -1493,12 +1633,13 @@ fn update_metrics_on_interval( Ok(()) } -fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { - if let Some(native_query) = &exec_context.root_op { - let metrics = exec_context.metrics.as_obj(); - update_comet_metric(env, metrics, native_query) - } else { - Ok(()) +/// Publishes the plan's metrics to its JVM `CometMetricNode`, once the plan has been created. +fn update_metrics(env: &mut Env, exec_context: &ExecutionContext) -> CometResult<()> { + match (&exec_context.root_op, &exec_context.metrics) { + (Some(native_query), Some(metrics)) => { + update_comet_metric(env, metrics.as_obj(), native_query) + } + _ => Ok(()), } } @@ -3400,4 +3541,244 @@ mod tests { "{error}" ); } + + /// A range of `num_elements` longs from 0, which needs no JVM input and reserves no memory. + fn range_operator(num_elements: i64) -> Operator { + use datafusion_comet_proto::spark_operator::RangeScan; + + Operator { + plan_id: 1, + op_struct: Some(OpStruct::RangeScan(RangeScan { + start: 0, + step: 1, + num_elements, + num_slices: 1, + })), + ..Default::default() + } + } + + fn range_plan(num_elements: i64) -> Vec { + use prost::Message; + range_operator(num_elements).encode_to_vec() + } + + /// A range of `num_elements` longs sorted in descending order. The sort buffers its whole + /// input, so the plan holds a memory reservation while it emits its output. + fn sorted_range_plan(num_elements: i64) -> Vec { + use datafusion_comet_proto::spark_expression::{ + self, expr::ExprStruct, NullOrdering, SortDirection, + }; + use datafusion_comet_proto::spark_operator::Sort; + use prost::Message; + + let column = Expr { + expr_struct: Some(ExprStruct::Bound(spark_expression::BoundReference { + index: 0, + // A Spark `long`. + datatype: Some(spark_expression::DataType { + type_id: 4, + type_info: None, + }), + })), + ..Default::default() + }; + let descending = Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(column)), + direction: SortDirection::Descending as i32, + null_ordering: NullOrdering::NullsLast as i32, + }, + ))), + ..Default::default() + }; + Operator { + plan_id: 2, + children: vec![range_operator(num_elements)], + op_struct: Some(OpStruct::Sort(Sort { + sort_orders: vec![descending], + ..Default::default() + })), + ..Default::default() + } + .encode_to_vec() + } + + fn long_values(batch: &RecordBatch) -> Vec { + arrow::array::AsArray::as_primitive::(batch.column(0)) + .values() + .to_vec() + } + + /// Creates a plan without a JVM: no metrics node or JVM input, and a memory pool backed by a + /// fake Spark. + fn create_test_plan(serialized_plan: Vec) -> CometResult> { + use crate::execution::memory_pools::create_memory_pool_with_fake_spark; + use datafusion_comet_proto::spark_config::ConfigMap; + use prost::Message; + + let configs = ConfigMap { + entries: HashMap::from([(SPARK_EXECUTOR_CORES.to_string(), "2".to_string())]), + }; + let inputs = PlanInputs { + id: -6201, + settings: prepare_plan_settings(&configs.encode_to_vec())?, + serialized_plan, + partition_count: 1, + batch_size: 4, + off_heap_mode: true, + memory_pool_type: "greedy_unified".to_string(), + memory_limit: 1 << 30, + task_attempt_id: -6201, + task_cpus: 1, + local_dirs: vec![std::env::temp_dir().to_str().unwrap().to_string()], + metrics_update_interval: 0, + start: Instant::now(), + jvm: PlanJvmRefs { + metrics: None, + input_sources: vec![], + key_unwrapper: None, + task_context: None, + class_loader: None, + }, + }; + create_execution_context(inputs, |config| { + create_memory_pool_with_fake_spark(config, -6201, 1 << 30) + }) + } + + #[test] + fn execute_plan_drains_a_plan() { + let _guard = serial(); + let mut exec_context = create_test_plan(range_plan(10)).unwrap(); + assert!( + exec_context.root_op.is_none(), + "operators are created on first execution" + ); + + let mut values = Vec::new(); + while let Some(batch) = execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap() { + assert!( + batch.num_rows() <= 4, + "batches follow the configured batch size" + ); + values.extend(long_values(&batch)); + } + assert_eq!(values, (0..10).collect::>()); + assert!( + exec_context.batch_producer.is_some(), + "a plan without JVM input runs on a Tokio task" + ); + release_plan(exec_context, &mut |_| Ok(())).unwrap(); + } + + /// The sort holds a reservation while it emits its 10,000 sorted rows, 4 at a time. + const SORTED_ROWS: i64 = 10_000; + + #[test] + fn release_plan_returns_the_plans_memory() { + let _guard = serial(); + let mut exec_context = create_test_plan(sorted_range_plan(SORTED_ROWS)).unwrap(); + + // Without an update interval, metrics are published only when the plan is released. + let mut publishes = 0; + let batch = execute_plan(&mut exec_context, 0, 0, &mut |_| { + publishes += 1; + Ok(()) + }) + .unwrap() + .expect("the sort emits a first batch"); + assert_eq!(long_values(&batch), [9999, 9998, 9997, 9996]); + assert_eq!(publishes, 0); + + // Released mid-stream, as when Spark stops reading early, with the sort still holding + // its buffered input. + let plan_memory = Arc::clone(&exec_context.plan_memory); + assert!(plan_memory.reserved() > 0, "the sort holds a reservation"); + release_plan(exec_context, &mut |ctx| { + assert!( + ctx.root_op.is_some(), + "the final metrics are of the executed plan" + ); + publishes += 1; + Ok(()) + }) + .unwrap(); + assert_eq!(publishes, 1); + assert_eq!(plan_memory.reserved(), 0); + } + + #[test] + fn release_plan_reports_a_failed_metrics_publish() { + let _guard = serial(); + let mut exec_context = create_test_plan(sorted_range_plan(SORTED_ROWS)).unwrap(); + execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap(); + let plan_memory = Arc::clone(&exec_context.plan_memory); + assert!(plan_memory.reserved() > 0, "the sort holds a reservation"); + + let error = release_plan(exec_context, &mut |_| { + Err(CometError::Internal("metrics node is gone".to_string())) + }) + .unwrap_err(); + assert!( + error.to_string().contains("metrics node is gone"), + "{error}" + ); + // The plan is still released. + assert_eq!(plan_memory.reserved(), 0); + } + + #[test] + fn create_execution_context_rejects_bad_input() { + let _guard = serial(); + let error = create_test_plan(vec![0xff, 0xff, 0xff]) + .err() + .expect("a malformed plan is rejected"); + assert!(error.to_string().contains("deserialize"), "{error}"); + } + + struct NoopPusher; + + impl ShufflePartitionPusher for NoopPusher { + fn push_partition_data( + &self, + _partition_id: i32, + _data: &[u8], + ) -> datafusion::common::Result<()> { + Ok(()) + } + } + + #[test] + fn register_shuffle_partition_pusher_once_before_execution() { + let _guard = serial(); + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + register_shuffle_partition_pusher(&mut exec_context, || Ok(Arc::new(NoopPusher))).unwrap(); + assert!(exec_context.shuffle_partition_pusher.is_some()); + + // The pusher is built only once the plan is known to accept it. + let error = register_shuffle_partition_pusher(&mut exec_context, || { + panic!("a rejected pusher must not be built") + }) + .unwrap_err(); + assert!( + error.to_string().contains("already been registered"), + "{error}" + ); + + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap(); + let error = register_shuffle_partition_pusher(&mut exec_context, || { + panic!("a rejected pusher must not be built") + }) + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot be registered after native execution starts"), + "{error}" + ); + release_plan(exec_context, &mut |_| Ok(())).unwrap(); + } } From e7e71987198ca04de646dea68ef30f80831fed98 Mon Sep 17 00:00:00 2001 From: Yu Hsun Hung <75130467+sam-1112@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:38:19 +0000 Subject: [PATCH 49/89] refactor: remove unused cast compatibility flag (#6135) * refactor: remove unused cast compatibility flag * refactor: drop unused parquet allow_incompat option The flag has no remaining reader after the cast-option cleanup, and the unused CometConf import fails scalafix. * fix: preserve DPP filters when reverting native scans * Fix SparkCastOptions constructor calls * refactor: drop the DPP scan reversion from the cast cleanup * fix: drop the removed cast-option argument from case-when calls * fix: drop the removed parquet-option argument from the page-index test * fix: drop the removed cast-option argument from the cast_nested bench #6648 added this bench after the last merge of main into this branch, and it still passed the third argument that this PR removes from SparkCastOptions::new. --------- Co-authored-by: Cursor Co-authored-by: Andy Grove --- .../expression-audits/conversion_funcs.md | 2 +- .../benches/parquet_timestamp_conversion.rs | 2 +- .../schema_adapter/tests/resolution.rs | 2 +- .../src/execution/operators/iceberg_scan.rs | 4 +- native/core/src/execution/planner.rs | 8 +-- native/core/src/parquet/cast_column.rs | 4 +- .../eager_page_index_reader_factory.rs | 2 +- native/core/src/parquet/parquet_exec.rs | 3 +- native/core/src/parquet/parquet_support.rs | 24 +++---- native/core/src/parquet/schema_adapter.rs | 47 +++++++------- native/proto/src/proto/expr.proto | 3 +- .../benches/cast_binary_to_string.rs | 2 +- .../benches/cast_decimal_to_boolean.rs | 2 +- .../benches/cast_decimal_to_string.rs | 2 +- .../benches/cast_float_to_decimal.rs | 2 +- .../benches/cast_float_to_string.rs | 2 +- .../spark-expr/benches/cast_from_boolean.rs | 2 +- native/spark-expr/benches/cast_from_string.rs | 8 +-- .../spark-expr/benches/cast_int_to_decimal.rs | 2 +- .../benches/cast_int_to_timestamp.rs | 2 +- native/spark-expr/benches/cast_nested.rs | 2 +- .../benches/cast_non_int_numeric_timestamp.rs | 2 +- native/spark-expr/benches/cast_numeric.rs | 4 +- .../spark-expr/benches/cast_string_to_date.rs | 2 +- .../benches/cast_string_to_timestamp.rs | 4 +- native/spark-expr/benches/conditional.rs | 2 +- native/spark-expr/benches/to_csv.rs | 2 +- native/spark-expr/benches/wide_decimal.rs | 2 +- .../src/conditional_funcs/case_when.rs | 4 +- .../src/conversion_funcs/boolean.rs | 2 +- .../spark-expr/src/conversion_funcs/cast.rs | 63 ++++++++----------- .../src/conversion_funcs/numeric.rs | 2 +- .../spark-expr/src/conversion_funcs/string.rs | 4 +- .../src/conversion_funcs/temporal.rs | 12 ++-- native/spark-expr/src/csv_funcs/to_csv.rs | 2 +- native/spark-expr/src/json_funcs/to_json.rs | 2 +- .../internal/decimal_rescale_check.rs | 2 +- .../spark-expr/src/math_funcs/modulo_expr.rs | 6 +- .../apache/comet/expressions/CometCast.scala | 5 -- .../comet/expressions/CometEvalMode.scala | 8 --- .../org/apache/comet/serde/datetime.scala | 3 - 41 files changed, 112 insertions(+), 148 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/conversion_funcs.md b/docs/source/contributor-guide/expression-audits/conversion_funcs.md index e31d5c92d9b..0b52c96d014 100644 --- a/docs/source/contributor-guide/expression-audits/conversion_funcs.md +++ b/docs/source/contributor-guide/expression-audits/conversion_funcs.md @@ -24,7 +24,7 @@ ## cast - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8 modulo `Cast.canUpCast` refactored to delegate to `UpCastRule.canUpCast`. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Cast(child, dataType, timeZoneId, evalMode)`; eval modes are `LEGACY`, `ANSI`, `TRY`. The legacy `Cast.canCast` matrix and the `Cast.canAnsiCast` matrix decide acceptance per type pair. Comet routes via `CometCast` (`spark/src/main/scala/org/apache/comet/expressions/CometCast.scala`) using a per-source-type support matrix that returns `Compatible`, `Incompatible(reason)`, or `Unsupported(reason)`; literal children are short-circuited to `Compatible()` so `CometLiteral` validates them. The serialized `Cast` proto carries `datatype`, `evalMode`, `timezone` (default `UTC`), `allowIncompat` (from `spark.comet.expression.Cast.allowIncompatible`), and `isSpark4Plus`. The native side (`native/spark-expr/src/conversion_funcs/cast.rs`) implements explicit per-eval-mode branches for narrowing numeric casts that match Spark's overflow exceptions, and falls through to DataFusion `cast_with_options(safe = !ANSI)` for the rest. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Cast(child, dataType, timeZoneId, evalMode)`; eval modes are `LEGACY`, `ANSI`, `TRY`. The legacy `Cast.canCast` matrix and the `Cast.canAnsiCast` matrix decide acceptance per type pair. Comet routes via `CometCast` (`spark/src/main/scala/org/apache/comet/expressions/CometCast.scala`) using a per-source-type support matrix that returns `Compatible`, `Incompatible(reason)`, or `Unsupported(reason)`; literal children are short-circuited to `Compatible()` so `CometLiteral` validates them. The serialized `Cast` proto carries `datatype`, `evalMode`, `timezone` (default `UTC`), and `isSpark4Plus`; `allowIncompatible` is enforced only by the JVM planner and is not sent to the native cast kernel. The native side (`native/spark-expr/src/conversion_funcs/cast.rs`) implements explicit per-eval-mode branches for narrowing numeric casts that match Spark's overflow exceptions, and falls through to DataFusion `cast_with_options(safe = !ANSI)` for the rest. - Spark 4.0.1 (audited 2026-05-27): `VariantType` added; `StringType` literals replaced with `_: StringType` to accommodate collated strings. `(TimestampType, ByteType|ShortType|IntegerType)` added to `canAnsiCast`. `NullIntolerant` -> `nullIntolerant: Boolean` refactor. New `ToPrettyString.BinaryFormatter` semantics for `Binary -> String` are replicated natively via `spark_binary_formatter`. Numeric-to-numeric matrix unchanged. - Spark 4.1.1 (audited 2026-05-27): `TimeType` added; many `TimeType` arms in `canCast`/`canAnsiCast`. Geospatial `GeographyType` / `GeometryType` types added with their own conversion rules. Numeric-to-numeric matrix unchanged. - Known divergences and gaps: diff --git a/native/core/benches/parquet_timestamp_conversion.rs b/native/core/benches/parquet_timestamp_conversion.rs index d090fc79691..6d6e6209637 100644 --- a/native/core/benches/parquet_timestamp_conversion.rs +++ b/native/core/benches/parquet_timestamp_conversion.rs @@ -119,7 +119,7 @@ fn timestamp_containers() -> [(ArrayRef, DataType); 2] { } fn benchmark(c: &mut Criterion) { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let mut group = c.benchmark_group("parquet_timestamp_conversion"); for width in [8, 1024] { let (array, target) = array_sibling(width); diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs index 21b7bcc6315..706f75f7f2e 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs @@ -27,7 +27,7 @@ fn spark_factory( use_field_id: bool, defaults: Option>, ) -> Arc { - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.use_field_id = use_field_id; Arc::new(SparkPhysicalExprAdapterFactory::new(options, defaults)) } diff --git a/native/core/src/execution/operators/iceberg_scan.rs b/native/core/src/execution/operators/iceberg_scan.rs index 9f465329ca0..ad69cac4f03 100644 --- a/native/core/src/execution/operators/iceberg_scan.rs +++ b/native/core/src/execution/operators/iceberg_scan.rs @@ -230,7 +230,7 @@ impl IcebergScanExec { let scan_metrics = scan_result.metrics().clone(); let stream = scan_result.stream(); - let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let adapter_factory = SparkPhysicalExprAdapterFactory::new(spark_options, None); let adapted_stream = @@ -630,7 +630,7 @@ mod tests { Field::new("a", DataType::Int64, false), Field::new("b", DataType::Int64, false), ])); - let mut options = super::SparkParquetOptions::new(super::EvalMode::Legacy, "UTC", false); + let mut options = super::SparkParquetOptions::new(super::EvalMode::Legacy, "UTC"); options.case_sensitive = true; let factory = super::SparkPhysicalExprAdapterFactory::new(options, None); for name in ["a", "b"] { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index b88b477931d..96802ca805d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -710,7 +710,6 @@ impl PhysicalPlanner { SparkCastOptions::new_with_version( eval_mode, &expr.timezone, - expr.allow_incompat, expr.is_spark4_plus, ), spark_expr.expr_id, @@ -803,7 +802,7 @@ impl PhysicalPlanner { "md5" => Ok(Arc::new(Cast::new( func?, DataType::Utf8, - SparkCastOptions::new_without_timezone(EvalMode::Try, true), + SparkCastOptions::new_without_timezone(EvalMode::Try), None, None, ))), @@ -925,8 +924,7 @@ impl PhysicalPlanner { ))) } ExprStruct::ToPrettyString(expr) => { - let mut spark_cast_options = - SparkCastOptions::new(EvalMode::Try, &expr.timezone, true); + let mut spark_cast_options = SparkCastOptions::new(EvalMode::Try, &expr.timezone); let null_string = "NULL"; spark_cast_options.null_string = null_string.to_string(); spark_cast_options.binary_output_style = @@ -6907,7 +6905,7 @@ mod tests { .with_table_parquet_options(TableParquetOptions::new()), ) as Arc; - let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let expr_adapter_factory: Arc = Arc::new( SparkPhysicalExprAdapterFactory::new(spark_parquet_options, None), diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index ad7b909768b..3a76e89916f 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -393,7 +393,7 @@ mod tests { let schema = Schema::new(vec![Arc::clone(&input_field)]); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_arr)]).unwrap(); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let col_expr: Arc = Arc::new(Column::new("s", 0)); @@ -466,7 +466,7 @@ mod tests { let expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::try_new(expr, input_field, target_field, None) .unwrap() - .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC", false)); + .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC")); let input = TimestampMillisecondArray::from(vec![Some(1_234), Some(-1_234), None]) .with_timezone_opt(source_tz.clone()); diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index dbc7d0f3b80..15bafdd2d63 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -1341,7 +1341,7 @@ mod tests { false, )])); let mut options = - SparkParquetOptions::new(datafusion_comet_spark_expr::EvalMode::Legacy, "UTC", false); + SparkParquetOptions::new(datafusion_comet_spark_expr::EvalMode::Legacy, "UTC"); options.case_sensitive = false; let factory = EagerPageIndexReaderFactory::new( store, diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 93cffbca67a..fba2d6fc13e 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -351,8 +351,7 @@ fn get_options( // distinguish INT96-derived TimestampLTZ from a true TimestampNTZ source // and apply the pre-Spark-4 SPARK-36182 rejection (#4219). table_parquet_options.global.coerce_int96_tz = Some("UTC".to_string()); - let mut spark_parquet_options = - SparkParquetOptions::new(EvalMode::Legacy, session_timezone, false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, session_timezone); spark_parquet_options.allow_cast_unsigned_ints = true; spark_parquet_options.case_sensitive = case_sensitive; spark_parquet_options.return_null_struct_if_all_fields_missing = diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index e8987a79da6..cc17429a72b 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -87,8 +87,6 @@ pub struct SparkParquetOptions { /// session local timezone by an analyzer in Spark. // TODO we should change timezone to Tz to avoid repeated parsing pub timezone: String, - /// Allow casts that are supported but not guaranteed to be 100% compatible - pub allow_incompat: bool, /// Support casting unsigned ints to signed ints (used by Parquet SchemaAdapter) pub allow_cast_unsigned_ints: bool, /// Whether to read dates/timestamps that were written in the legacy hybrid Julian + Gregorian calendar as it is. If false, throw exceptions instead. If the spark type is TimestampNTZ, this should be true. @@ -121,11 +119,10 @@ pub struct SparkParquetOptions { } impl SparkParquetOptions { - pub fn new(eval_mode: EvalMode, timezone: &str, allow_incompat: bool) -> Self { + pub fn new(eval_mode: EvalMode, timezone: &str) -> Self { Self { eval_mode, timezone: timezone.to_string(), - allow_incompat, allow_cast_unsigned_ints: false, use_legacy_date_timestamp_or_ntz: false, case_sensitive: false, @@ -137,11 +134,10 @@ impl SparkParquetOptions { } } - pub fn new_without_timezone(eval_mode: EvalMode, allow_incompat: bool) -> Self { + pub fn new_without_timezone(eval_mode: EvalMode) -> Self { Self { eval_mode, timezone: "".to_string(), - allow_incompat, allow_cast_unsigned_ints: false, use_legacy_date_timestamp_or_ntz: false, case_sensitive: false, @@ -1405,7 +1401,7 @@ mod tests { let err = spark_parquet_convert( ColumnarValue::Array(Arc::new(array)), &to_type, - &SparkParquetOptions::new(EvalMode::Legacy, "UTC", false), + &SparkParquetOptions::new(EvalMode::Legacy, "UTC"), ) .expect_err("array -> int must be an error"); assert!( @@ -1513,7 +1509,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let overflow_millis = 9_223_372_036_854_776_i64; let millis: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ Some(overflow_millis), @@ -1633,7 +1629,7 @@ mod tests { )); let target = DataType::Struct(target_fields.into()); let input: ArrayRef = Arc::new(input); - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let output = parquet_convert_array(Arc::clone(&input), &target, &options).unwrap(); assert_eq!(output.data_type(), &target); assert!(output.is_null(0)); @@ -1682,7 +1678,7 @@ mod tests { for timezone in [None::>, Some(Arc::from("UTC"))] { for overflow in [i64::MAX, i64::MIN] { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let millis: ArrayRef = Arc::new( TimestampMillisecondArray::from(vec![overflow, 7, overflow]) .with_timezone_opt(timezone.clone()), @@ -1822,7 +1818,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); for timezone in [None::>, Some(Arc::from("UTC"))] { for overflow in [i64::MIN, i64::MAX] { let values: ArrayRef = Arc::new( @@ -1918,7 +1914,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let values: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ i64::MAX, 7, @@ -2125,7 +2121,7 @@ mod tests { )); let to_type = struct_type_with_field_id(vec![("f", DataType::Int32, 1)]); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let err = parquet_convert_array(from, &to_type, &opts).unwrap_err(); @@ -2165,7 +2161,7 @@ mod tests { )); let to_type = struct_type_with_field_id(vec![("f", DataType::Int32, 2)]); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let result = parquet_convert_array(from, &to_type, &opts).unwrap(); diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index 6bcc8d3518b..495eeb545ec 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -1475,7 +1475,6 @@ impl SparkPhysicalExprAdapter { let mut cast_options = SparkCastOptions::new( self.parquet_options.eval_mode, &self.parquet_options.timezone, - self.parquet_options.allow_incompat, ); cast_options.allow_cast_unsigned_ints = self.parquet_options.allow_cast_unsigned_ints; cast_options.is_adapting_schema = true; @@ -2202,7 +2201,7 @@ pub(crate) mod test { let required_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_type_promotion = false; let expr_adapter_factory: Arc = Arc::new( @@ -2247,7 +2246,7 @@ pub(crate) mod test { let required_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_type_promotion = false; let expr_adapter_factory: Arc = Arc::new( @@ -2375,7 +2374,7 @@ pub(crate) mod test { batch: &RecordBatch, required_schema: SchemaRef, ) -> Result { - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_cast_unsigned_ints = true; let mut stream = scan_parquet(batch, required_schema, spark_parquet_options)?; stream.next().await.unwrap() @@ -2890,7 +2889,7 @@ pub(crate) mod test { #[test] fn nested_dictionary_containers_are_checked_by_value_type() -> Result<(), DataFusionError> { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let physical_struct = DataType::Struct(Fields::from(vec![ Field::new("x", DataType::Int64, true), Field::new("unused", DataType::Int32, true), @@ -2925,7 +2924,7 @@ pub(crate) mod test { #[test] fn nested_map_shape_mismatch_is_rejected() -> Result<(), DataFusionError> { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let valid = map_type(DataType::Int64); let DataType::Map(entries, _) = &valid else { unreachable!() @@ -3048,7 +3047,7 @@ pub(crate) mod test { let required_schema = struct_schema(vec![ Field::new("b", DataType::Int32, true).with_metadata(id_meta("1")) ]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.use_field_id = true; let mut stream = scan_parquet(&batch, required_schema, options)?; let err = stream @@ -3075,7 +3074,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(vec![1, 2, 3])), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = false; let mut stream = scan_parquet(&batch, required_schema, options)?; let err = stream @@ -3101,7 +3100,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(Vec::::new())), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = false; let mut stream = scan_parquet(&batch, required_schema, options)?; while let Some(batch) = stream.next().await { @@ -3486,7 +3485,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(vec![1, 2, 3])), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = true; let mut stream = scan_parquet(&batch, required_schema, options)?; let result = stream.next().await.unwrap()?; @@ -3579,7 +3578,7 @@ pub(crate) mod test { Field::new("a", DataType::Int64, nullable), Field::new("b", DataType::Int64, false), ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(logical, physical) @@ -3613,7 +3612,7 @@ pub(crate) mod test { Arc::new(Field::new("dup", DataType::Int64, true)), Arc::new(Field::new("dup", DataType::Int64, true)), ]; - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let error = super::match_struct_fields(&fields, &fields[..1], &options) .expect_err("selected duplicate child must fail") @@ -3648,7 +3647,7 @@ pub(crate) mod test { default.rewrite(Arc::clone(&column)).is_err(), "fixture must reach the default-adapter fallback" ); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(Arc::clone(&logical), Arc::clone(&physical)) @@ -3714,7 +3713,7 @@ pub(crate) mod test { if mode == "missing" { requested.push(Field::new("missing", DataType::Int64, true)); } - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let mut stream = scan_parquet(&batch, struct_schema(requested), options).unwrap(); let error = stream @@ -3816,7 +3815,7 @@ pub(crate) mod test { reader.schema().field(0).data_type(), required.field(0).data_type() ); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let mut stream = scan_parquet(&batch, required, options).unwrap(); let error = stream @@ -3859,7 +3858,7 @@ pub(crate) mod test { // Read with case-insensitive mode, requesting column "b" which matches both "B" and "b" let required_schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.case_sensitive = false; let expr_adapter_factory: Arc = Arc::new( @@ -3919,7 +3918,7 @@ pub(crate) mod test { None, )?)); let defaults = HashMap::from([(Column::new("missing", 0), default)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = false; let adapter = SparkPhysicalExprAdapterFactory::new(options, Some(defaults)) .create(logical, physical)?; @@ -4014,7 +4013,7 @@ pub(crate) mod test { Field::new("ω", DataType::Int32, true).with_metadata(id_meta("2")), ])); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.case_sensitive = false; opts.use_field_id = true; let adapter = SparkPhysicalExprAdapterFactory::new(opts, None) @@ -4045,7 +4044,7 @@ pub(crate) mod test { let physical = Arc::new(Schema::new(vec![ Field::new("MÜNCHEN", storage, true).with_extension_type(VariantType) ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = false; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(Arc::clone(&logical), Arc::clone(&physical)) @@ -4104,7 +4103,7 @@ pub(crate) mod test { .is_some()); } let adapter = SparkPhysicalExprAdapterFactory::new( - SparkParquetOptions::new(EvalMode::Legacy, "UTC", false), + SparkParquetOptions::new(EvalMode::Legacy, "UTC"), None, ) .create(logical, physical) @@ -4140,7 +4139,7 @@ pub(crate) mod test { Field::new("other", storage.clone(), true).with_metadata(id_meta("1")), Field::new("__comet_unmatched_field_id_1", DataType::Binary, true), ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = case_sensitive; options.use_field_id = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) @@ -4210,7 +4209,7 @@ pub(crate) mod test { )); let target_type = DataType::List(Arc::clone(&to_item_field)); - let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let comet_result = spark_parquet_convert( ColumnarValue::Array(Arc::clone(&list_array)), &target_type, @@ -4321,7 +4320,7 @@ pub(crate) mod test { } fn default_options() -> SparkParquetOptions { - SparkParquetOptions::new(EvalMode::Legacy, "UTC", false) + SparkParquetOptions::new(EvalMode::Legacy, "UTC") } /// Dropping a struct field by exact name, including through nested struct-in-struct and @@ -4683,7 +4682,7 @@ pub(crate) mod test { )) as Arc; let batch = RecordBatch::try_new(file_schema, vec![col]).unwrap(); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let err = match scan_parquet(&batch, required_schema, opts) { diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index d1bdd9c80a5..b85f08f6531 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -434,7 +434,8 @@ message Cast { DataType datatype = 2; string timezone = 3; EvalMode eval_mode = 4; - bool allow_incompat = 5; + reserved 5; + reserved "allow_incompat"; // True when running against Spark 4.0+. Controls version-specific cast behaviour // such as the handling of leading whitespace before T-prefixed time-only strings. bool is_spark4_plus = 6; diff --git a/native/spark-expr/benches/cast_binary_to_string.rs b/native/spark-expr/benches/cast_binary_to_string.rs index 8ae9b617615..7c3e5d3a9c5 100644 --- a/native/spark-expr/benches/cast_binary_to_string.rs +++ b/native/spark-expr/benches/cast_binary_to_string.rs @@ -46,7 +46,7 @@ fn build(size: usize, width: usize, null_every: usize) -> ArrayRef { } fn options(style: Option) -> SparkCastOptions { - let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); options.binary_output_style = style; options } diff --git a/native/spark-expr/benches/cast_decimal_to_boolean.rs b/native/spark-expr/benches/cast_decimal_to_boolean.rs index 0d1bfe255c3..2e5af4809e2 100644 --- a/native/spark-expr/benches/cast_decimal_to_boolean.rs +++ b/native/spark-expr/benches/cast_decimal_to_boolean.rs @@ -59,7 +59,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast_to_bool = Cast::new( expr, DataType::Boolean, - SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + SparkCastOptions::new(EvalMode::Legacy, "UTC"), None, None, ); diff --git a/native/spark-expr/benches/cast_decimal_to_string.rs b/native/spark-expr/benches/cast_decimal_to_string.rs index d0629300a97..5fd76e44fd4 100644 --- a/native/spark-expr/benches/cast_decimal_to_string.rs +++ b/native/spark-expr/benches/cast_decimal_to_string.rs @@ -43,7 +43,7 @@ fn cast_to_utf8() -> Cast { Cast::new( Arc::new(Column::new("a", 0)), DataType::Utf8, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ) diff --git a/native/spark-expr/benches/cast_float_to_decimal.rs b/native/spark-expr/benches/cast_float_to_decimal.rs index fab3c3d1122..1bb071092d3 100644 --- a/native/spark-expr/benches/cast_float_to_decimal.rs +++ b/native/spark-expr/benches/cast_float_to_decimal.rs @@ -52,7 +52,7 @@ fn cast(to: DataType, mode: EvalMode) -> Cast { Cast::new( Arc::new(Column::new("a", 0)), to, - SparkCastOptions::new_without_timezone(mode, false), + SparkCastOptions::new_without_timezone(mode), None, None, ) diff --git a/native/spark-expr/benches/cast_float_to_string.rs b/native/spark-expr/benches/cast_float_to_string.rs index deb24a3e04c..71af3074c8b 100644 --- a/native/spark-expr/benches/cast_float_to_string.rs +++ b/native/spark-expr/benches/cast_float_to_string.rs @@ -67,7 +67,7 @@ fn criterion_benchmark(c: &mut Criterion) { let size = 8192; let f64_array = create_f64_array(size); let f32_array = create_f32_array(size); - let cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let mut group = c.benchmark_group("cast_float_to_string"); group.bench_function("cast_f64_to_utf8", |b| { diff --git a/native/spark-expr/benches/cast_from_boolean.rs b/native/spark-expr/benches/cast_from_boolean.rs index caccd67e26d..48f080568aa 100644 --- a/native/spark-expr/benches/cast_from_boolean.rs +++ b/native/spark-expr/benches/cast_from_boolean.rs @@ -26,7 +26,7 @@ use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { let expr = Arc::new(Column::new("a", 0)); let boolean_batch = create_boolean_batch(); - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let cast_to_i8 = Cast::new( expr.clone(), DataType::Int8, diff --git a/native/spark-expr/benches/cast_from_string.rs b/native/spark-expr/benches/cast_from_string.rs index 31cf11a84ba..cd344670127 100644 --- a/native/spark-expr/benches/cast_from_string.rs +++ b/native/spark-expr/benches/cast_from_string.rs @@ -41,7 +41,7 @@ fn criterion_benchmark(c: &mut Criterion) { let expr = Arc::new(Column::new("a", 0)); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let cast_to_i8 = Cast::new( expr.clone(), DataType::Int8, @@ -88,7 +88,7 @@ fn criterion_benchmark(c: &mut Criterion) { } // Benchmark decimal truncation (Legacy mode only) - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, ""); let cast_to_i32 = Cast::new( expr.clone(), DataType::Int32, @@ -116,7 +116,7 @@ fn criterion_benchmark(c: &mut Criterion) { // str -> decimal benchmark let decimal_string_batch = create_decimal_cast_string_batch(); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let mut group = c.benchmark_group(format!("cast_string_to_decimal/{}", mode_name)); for (data_type, name) in [ (DataType::Decimal128(38, 10), "decimal_38_10"), @@ -143,7 +143,7 @@ fn criterion_benchmark(c: &mut Criterion) { let float_batch = create_float_string_batch(false); let float_padded_batch = create_float_string_batch(true); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let mut group = c.benchmark_group(format!("cast_string_to_bool_and_float/{}", mode_name)); for (data_type, name, batch) in [ (DataType::Boolean, "boolean", &bool_batch), diff --git a/native/spark-expr/benches/cast_int_to_decimal.rs b/native/spark-expr/benches/cast_int_to_decimal.rs index 8949c5c3276..a27b5d277b0 100644 --- a/native/spark-expr/benches/cast_int_to_decimal.rs +++ b/native/spark-expr/benches/cast_int_to_decimal.rs @@ -56,7 +56,7 @@ fn cast(col: &str, to: DataType, mode: EvalMode) -> Cast { Cast::new( Arc::new(Column::new(col, 0)), to, - SparkCastOptions::new_without_timezone(mode, false), + SparkCastOptions::new_without_timezone(mode), None, None, ) diff --git a/native/spark-expr/benches/cast_int_to_timestamp.rs b/native/spark-expr/benches/cast_int_to_timestamp.rs index 4479627ae66..04a62280d5a 100644 --- a/native/spark-expr/benches/cast_int_to_timestamp.rs +++ b/native/spark-expr/benches/cast_int_to_timestamp.rs @@ -27,7 +27,7 @@ const BATCH_SIZE: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { // Test with UTC timezone - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let timestamp_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); let mut group = c.benchmark_group("cast_int_to_timestamp"); diff --git a/native/spark-expr/benches/cast_nested.rs b/native/spark-expr/benches/cast_nested.rs index 5c6efd5cb47..0474f2b9af2 100644 --- a/native/spark-expr/benches/cast_nested.rs +++ b/native/spark-expr/benches/cast_nested.rs @@ -93,7 +93,7 @@ fn list_input() -> ArrayRef { } fn criterion_benchmark(c: &mut Criterion) { - let options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let col: Arc = Arc::new(Column::new("c", 0)); let mut group = c.benchmark_group("cast_nested"); diff --git a/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs b/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs index 18ed3041eb6..9c071aae4b6 100644 --- a/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs +++ b/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs @@ -26,7 +26,7 @@ use std::sync::Arc; const BATCH_SIZE: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let timestamp_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); let mut group = c.benchmark_group("cast_non_int_numeric_to_timestamp"); diff --git a/native/spark-expr/benches/cast_numeric.rs b/native/spark-expr/benches/cast_numeric.rs index 5153fb7d011..41fd6006d83 100644 --- a/native/spark-expr/benches/cast_numeric.rs +++ b/native/spark-expr/benches/cast_numeric.rs @@ -28,7 +28,7 @@ const NUM_ROWS: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { let batch = create_int32_batch(); let expr = Arc::new(Column::new("a", 0)); - let spark_cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let spark_cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let cast_i32_to_i8 = Cast::new( expr.clone(), DataType::Int8, @@ -61,7 +61,7 @@ fn criterion_benchmark(c: &mut Criterion) { Cast::new( Arc::new(Column::new("a", 0)), data_type, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ) diff --git a/native/spark-expr/benches/cast_string_to_date.rs b/native/spark-expr/benches/cast_string_to_date.rs index aee0fc0b03a..b2a9d7dabc5 100644 --- a/native/spark-expr/benches/cast_string_to_date.rs +++ b/native/spark-expr/benches/cast_string_to_date.rs @@ -32,7 +32,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast_to_date = Cast::new( expr, DataType::Date32, - SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + SparkCastOptions::new(EvalMode::Legacy, "UTC"), None, None, ); diff --git a/native/spark-expr/benches/cast_string_to_timestamp.rs b/native/spark-expr/benches/cast_string_to_timestamp.rs index d499cec96ae..117c00592e5 100644 --- a/native/spark-expr/benches/cast_string_to_timestamp.rs +++ b/native/spark-expr/benches/cast_string_to_timestamp.rs @@ -218,7 +218,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast = Cast::new( Arc::clone(&expr), to_type.clone(), - SparkCastOptions::new(mode, timezone, false), + SparkCastOptions::new(mode, timezone), None, None, ); @@ -238,7 +238,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast = Cast::new( Arc::clone(&expr), DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), - SparkCastOptions::new_with_version(EvalMode::Legacy, "UTC", false, true), + SparkCastOptions::new_with_version(EvalMode::Legacy, "UTC", true), None, None, ); diff --git a/native/spark-expr/benches/conditional.rs b/native/spark-expr/benches/conditional.rs index 8af04c2d52f..2febcddee99 100644 --- a/native/spark-expr/benches/conditional.rs +++ b/native/spark-expr/benches/conditional.rs @@ -122,7 +122,7 @@ fn spark_cast(child: Expr, data_type: DataType) -> Expr { Arc::new(Cast::new( child, data_type, - SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles"), None, None, )) diff --git a/native/spark-expr/benches/to_csv.rs b/native/spark-expr/benches/to_csv.rs index 8620dd0f160..f51dde2a5d0 100644 --- a/native/spark-expr/benches/to_csv.rs +++ b/native/spark-expr/benches/to_csv.rs @@ -86,7 +86,7 @@ fn criterion_benchmark(c: &mut Criterion) { let default_null_value = ""; let default_quote = "\""; let default_escape = "\\"; - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, timezone, false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, timezone); cast_options.null_string = default_null_value.to_string(); let csv_write_options = CsvWriteOptions::new( default_delimiter.to_string(), diff --git a/native/spark-expr/benches/wide_decimal.rs b/native/spark-expr/benches/wide_decimal.rs index b38ddbfaa06..6de9fb5e83b 100644 --- a/native/spark-expr/benches/wide_decimal.rs +++ b/native/spark-expr/benches/wide_decimal.rs @@ -85,7 +85,7 @@ fn build_old_expr( ) -> Arc { let left_col: Arc = Arc::new(Column::new("left", 0)); let right_col: Arc = Arc::new(Column::new("right", 1)); - let cast_opts = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let cast_opts = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let left_cast = Arc::new(Cast::new( left_col, DataType::Decimal256(p1, s1), diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 7dec89c61c0..74df26c977c 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -152,7 +152,7 @@ fn coerce_branch( if data_type == common_type { return expr; } - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); Arc::new(Cast::new( expr, common_type.clone(), @@ -1164,7 +1164,7 @@ mod tests { Arc::new(Cast::new( e, to, - SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + SparkCastOptions::new_without_timezone(EvalMode::Ansi), None, None, )) diff --git a/native/spark-expr/src/conversion_funcs/boolean.rs b/native/spark-expr/src/conversion_funcs/boolean.rs index 1db2746ce24..7e9dffd3307 100644 --- a/native/spark-expr/src/conversion_funcs/boolean.rs +++ b/native/spark-expr/src/conversion_funcs/boolean.rs @@ -67,7 +67,7 @@ mod tests { } fn test_input_spark_opts() -> SparkCastOptions { - SparkCastOptions::new(EvalMode::Legacy, "Asia/Kolkata", false) + SparkCastOptions::new(EvalMode::Legacy, "Asia/Kolkata") } #[test] diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 26659f7489f..30b3fe997a7 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -149,8 +149,6 @@ pub struct SparkCastOptions { /// session local timezone by an analyzer in Spark. // TODO we should change timezone to Tz to avoid repeated parsing pub timezone: String, - /// Allow casts that are supported but not guaranteed to be 100% compatible - pub allow_incompat: bool, /// True when running against Spark 4.0+. Enables version-specific cast behaviour /// such as the handling of leading whitespace before T-prefixed time-only strings. pub is_spark4_plus: bool, @@ -166,11 +164,10 @@ pub struct SparkCastOptions { } impl SparkCastOptions { - pub fn new(eval_mode: EvalMode, timezone: &str, allow_incompat: bool) -> Self { + pub fn new(eval_mode: EvalMode, timezone: &str) -> Self { Self { eval_mode, timezone: timezone.to_string(), - allow_incompat, is_spark4_plus: false, allow_cast_unsigned_ints: false, is_adapting_schema: false, @@ -179,11 +176,10 @@ impl SparkCastOptions { } } - pub fn new_without_timezone(eval_mode: EvalMode, allow_incompat: bool) -> Self { + pub fn new_without_timezone(eval_mode: EvalMode) -> Self { Self { eval_mode, timezone: "".to_string(), - allow_incompat, is_spark4_plus: false, allow_cast_unsigned_ints: false, is_adapting_schema: false, @@ -192,15 +188,10 @@ impl SparkCastOptions { } } - pub fn new_with_version( - eval_mode: EvalMode, - timezone: &str, - allow_incompat: bool, - is_spark4_plus: bool, - ) -> Self { + pub fn new_with_version(eval_mode: EvalMode, timezone: &str, is_spark4_plus: bool) -> Self { Self { is_spark4_plus, - ..Self::new(eval_mode, timezone, allow_incompat) + ..Self::new(eval_mode, timezone) } } } @@ -941,7 +932,7 @@ mod tests { let error = cast_array( Arc::new(StringArray::from(vec!["a"])), &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap_err(); @@ -979,7 +970,7 @@ mod tests { let error = cast_array( input, &output_type, - &SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + &SparkCastOptions::new_without_timezone(EvalMode::Ansi), ) .unwrap_err(); @@ -1014,7 +1005,7 @@ mod tests { Some(&[0xEDu8, 0xA0, 0x80][..]), ]); // binary_output_style defaults to None, i.e. the plain (non-ToPrettyString) cast path. - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = cast_binary_to_string::(&input, &cast_options).unwrap(); @@ -1038,7 +1029,7 @@ mod tests { Some("abc".as_bytes()), Some(&[0xEDu8, 0xA0, 0x80][..]), ]); - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); cast_options.binary_output_style = Some(BinaryOutputStyle::Utf8); let result = cast_binary_to_string::(&input, &cast_options).unwrap(); @@ -1060,7 +1051,7 @@ mod tests { Some(b"hi".as_slice()), ])); let cast = |style: Option| { - let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); options.binary_output_style = style; let result = spark_cast( ColumnarValue::Array(Arc::clone(&input)), @@ -1124,7 +1115,7 @@ mod tests { None, Some("héllo".as_bytes()), ])); - let options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = spark_cast( ColumnarValue::Array(Arc::clone(&input)), &DataType::Utf8, @@ -1143,7 +1134,7 @@ mod tests { fn test_cast_unsupported_timestamp_to_date() { // Since datafusion uses chrono::Datetime internally not all dates representable by TimestampMicrosecondType are supported let timestamps: PrimitiveArray = vec![i64::MAX].into(); - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = cast_array( Arc::new(timestamps.with_timezone("Europe/Copenhagen")), &DataType::Date32, @@ -1155,7 +1146,7 @@ mod tests { #[test] fn test_cast_invalid_timezone() { let timestamps: PrimitiveArray = vec![i64::MAX].into(); - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "Not a valid timezone", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "Not a valid timezone"); let result = cast_array( Arc::new(timestamps.with_timezone("Europe/Copenhagen")), &DataType::Date32, @@ -1181,7 +1172,7 @@ mod tests { let string_array = cast_array( c, &DataType::Utf8, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); let string_array = string_array.as_string::(); @@ -1215,7 +1206,7 @@ mod tests { let cast_array = spark_cast( ColumnarValue::Array(c), &DataType::Struct(fields), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); if let ColumnarValue::Array(cast_array) = cast_array { @@ -1253,7 +1244,7 @@ mod tests { let result = spark_cast( ColumnarValue::Array(outer), &DataType::Struct(to_fields), - &SparkCastOptions::new(EvalMode::Ansi, "UTC", false), + &SparkCastOptions::new(EvalMode::Ansi, "UTC"), ); assert!(result.is_err()); @@ -1278,7 +1269,7 @@ mod tests { let cast_array = spark_cast( ColumnarValue::Array(c), &DataType::Struct(fields), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); if let ColumnarValue::Array(cast_array) = cast_array { @@ -1304,11 +1295,9 @@ mod tests { Arc::new(values_array), None, )); - let string_array = cast_array_to_string( - &list_array, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), - ) - .unwrap(); + let string_array = + cast_array_to_string(&list_array, &SparkCastOptions::new(EvalMode::Legacy, "UTC")) + .unwrap(); let string_array = string_array.as_string::(); assert_eq!(r#"[a, b, c]"#, string_array.value(0)); assert_eq!(r#"[a, null]"#, string_array.value(1)); @@ -1327,11 +1316,9 @@ mod tests { Arc::new(values_array), None, )); - let string_array = cast_array_to_string( - &list_array, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), - ) - .unwrap(); + let string_array = + cast_array_to_string(&list_array, &SparkCastOptions::new(EvalMode::Legacy, "UTC")) + .unwrap(); let string_array = string_array.as_string::(); assert_eq!(r#"[1, 2, 3]"#, string_array.value(0)); assert_eq!(r#"[1, null]"#, string_array.value(1)); @@ -1354,7 +1341,7 @@ mod tests { let to_array = cast_array( from_array, &to_type, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); @@ -1368,7 +1355,7 @@ mod tests { assert!(values.iter().all(|value| value.is_none())); } fn legacy_opts() -> SparkCastOptions { - SparkCastOptions::new(EvalMode::Legacy, "UTC", false) + SparkCastOptions::new(EvalMode::Legacy, "UTC") } /// Build a `Map` MapArray (Parquet-style "key_value" field names). @@ -1532,7 +1519,7 @@ mod tests { let casted = cast_array( map_array, &to_type, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index f7fbae54f98..aad9006937c 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -2291,7 +2291,7 @@ mod tests { ); for eval_mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { - let options = SparkCastOptions::new(eval_mode, "UTC", false); + let options = SparkCastOptions::new(eval_mode, "UTC"); let doubles = cast_array(Arc::clone(&decimals_38_18), &DataType::Float64, &options).unwrap(); diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index adedd3ba0d2..f278bafda68 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -2491,7 +2491,7 @@ mod tests { for (position, input, expect_value) in cases { for eval_mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { let array: ArrayRef = Arc::new(StringArray::from(vec![Some(input.as_str())])); - let options = SparkCastOptions::new(eval_mode, "UTC", false); + let options = SparkCastOptions::new(eval_mode, "UTC"); let result = cast_array(array, to_type, &options); let context = format!("cast {input:?} ({position}) to {to_type} in {eval_mode:?}"); if expect_value { @@ -2736,7 +2736,7 @@ mod tests { let timezone = "UTC".to_string(); // test casting string dictionary array to timestamp array - let cast_options = SparkCastOptions::new(EvalMode::Legacy, &timezone, false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, &timezone); let result = cast_array( dict_array, &DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.clone().into())), diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index 31e63644fd5..15926047f73 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -121,7 +121,7 @@ mod tests { let target_tz: Option> = Some("UTC".into()); let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), &target_tz, ) .unwrap(); @@ -134,7 +134,7 @@ mod tests { // validate LA timezone (follows Daylight savings) let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles"), &target_tz, ) .unwrap(); @@ -148,7 +148,7 @@ mod tests { // Phoenix timezone (does not follow Daylight savings) let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Phoenix", false), + &SparkCastOptions::new(EvalMode::Legacy, "America/Phoenix"), &target_tz, ) .unwrap(); @@ -189,7 +189,7 @@ mod tests { ] { let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, tz, false), + &SparkCastOptions::new(EvalMode::Legacy, tz), &ntz_target, ) .unwrap(); @@ -215,7 +215,7 @@ mod tests { assert!( cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), &ntz_target, ) .is_err(), @@ -240,7 +240,7 @@ mod tests { Some(-106_751_992), ]; for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { - let options = SparkCastOptions::new(mode, "America/Los_Angeles", false); + let options = SparkCastOptions::new(mode, "America/Los_Angeles"); let result = spark_cast( ColumnarValue::Array(Arc::new(Date32Array::from(days.clone()))), &target, diff --git a/native/spark-expr/src/csv_funcs/to_csv.rs b/native/spark-expr/src/csv_funcs/to_csv.rs index 01fdc901cb7..d79b5cbca3d 100644 --- a/native/spark-expr/src/csv_funcs/to_csv.rs +++ b/native/spark-expr/src/csv_funcs/to_csv.rs @@ -86,7 +86,7 @@ impl PhysicalExpr for ToCsv { fn evaluate(&self, batch: &RecordBatch) -> Result { let input_array = self.expr.evaluate(batch)?.into_array(batch.num_rows())?; - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, &self.timezone, false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, &self.timezone); cast_options.null_string = self.csv_write_options.null_value.clone(); let struct_array = as_struct_array(&input_array); diff --git a/native/spark-expr/src/json_funcs/to_json.rs b/native/spark-expr/src/json_funcs/to_json.rs index caf87514fad..94f29e3f458 100644 --- a/native/spark-expr/src/json_funcs/to_json.rs +++ b/native/spark-expr/src/json_funcs/to_json.rs @@ -139,7 +139,7 @@ fn array_to_json_string( spark_cast( ColumnarValue::Array(Arc::clone(arr)), &DataType::Utf8, - &SparkCastOptions::new(EvalMode::Legacy, timezone, false), + &SparkCastOptions::new(EvalMode::Legacy, timezone), )? .into_array(arr.len()) } diff --git a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs b/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs index 190d7b060c7..580a01c9afd 100644 --- a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs +++ b/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs @@ -195,7 +195,7 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow { Err(_) if self.fail_on_error => spark_cast( arg, &target_type, - &SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + &SparkCastOptions::new_without_timezone(EvalMode::Ansi), ), Err(error) => Err(error), } diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index 5c0652e9a01..09d006c9e49 100644 --- a/native/spark-expr/src/math_funcs/modulo_expr.rs +++ b/native/spark-expr/src/math_funcs/modulo_expr.rs @@ -163,14 +163,14 @@ pub fn create_modulo_expr( let left_256 = Arc::new(Cast::new( left, DataType::Decimal256(p1, s1), - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, )); let right_256 = Arc::new(Cast::new( right_non_ansi_safe, DataType::Decimal256(p2, s2), - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, )); @@ -192,7 +192,7 @@ pub fn create_modulo_expr( Ok(Arc::new(Cast::new( modulo_scalar_func, data_type, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ))) diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index e14fb086b12..b6fcdf32806 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -23,7 +23,6 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, L import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, MapType, NullType, StructType, TimestampNTZType, TimestampType} -import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason} import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, CometTimeZone, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported} @@ -167,10 +166,6 @@ object CometCast castBuilder.setChild(childExpr) castBuilder.setDatatype(dataType) castBuilder.setEvalMode(evalModeToProto(evalMode)) - castBuilder.setAllowIncompat( - SQLConf.get - .getConfString(CometConf.getExprAllowIncompatConfigKey(classOf[Cast]), "false") - .toBoolean) castBuilder.setTimezone(timeZone) castBuilder.setIsSpark4Plus(isSpark40Plus) Some( diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala index 59e9c89a6b5..d61f0eb3c7b 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala @@ -31,12 +31,4 @@ package org.apache.comet.expressions */ object CometEvalMode extends Enumeration { val LEGACY, ANSI, TRY = Value - - def fromBoolean(ansiEnabled: Boolean): Value = if (ansiEnabled) { - ANSI - } else { - LEGACY - } - - def fromString(str: String): CometEvalMode.Value = CometEvalMode.withName(str) } diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index 0317132ef00..8b3ddab5085 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -69,7 +69,6 @@ trait CometExprGetDateField[T <: GetDateField] { .setChild(e) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() }) @@ -503,7 +502,6 @@ object CometUnixDate extends CometExpressionSerde[UnixDate] { .setChild(child) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() } @@ -872,7 +870,6 @@ object CometDays extends CometExpressionSerde[Days] { .setChild(dateExpr) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() } From c1e00e7b92b0ada119b140b016dbefdd198bbc35 Mon Sep 17 00:00:00 2001 From: dustin Date: Sat, 10 Oct 2026 17:53:52 +0000 Subject: [PATCH 50/89] fix: retry a native acquire that Spark failed after dropping the task's entry (#6310) Spark removes a task's execution pool entry when its balance reaches zero, and an acquire that was parked in the pool then throws NoSuchElementException when it wakes up. Releases take no task monitor, so another native thread or a JVM consumer of the same task can empty the balance while a native acquire waits. The error is not a memory refusal, so the operator could not spill and the task failed. CometTaskMemoryManager now recognises that one exception and calls Spark again. A missing entry means the task held nothing from Spark, so the failed call has no grant to reconcile, and the retry registers the task again and waits for its share as the first call would have. After three attempts it returns a zero grant, which the native side treats as a refusal and spills. Closes #6224. --- .../contributor-guide/memory_management.md | 9 + .../apache/spark/CometTaskMemoryManager.java | 57 +++++- .../spark/CometTaskMemoryManagerSuite.scala | 168 +++++++++++++++++- 3 files changed, 231 insertions(+), 3 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 196e0d6608f..af84e4a2786 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -214,6 +214,15 @@ the call never returns. Nothing releases them until Spark's final task cleanup, nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile that grant, and the only figure available for doing so is the task-wide one above. +**A parked acquire can wake up to a missing task entry.** Spark removes the task's `memoryForTask` +entry when its balance reaches zero, and an acquire that was waiting in `lock.wait()` and wakes +afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant. +`CometTaskMemoryManager` retries that one case. A missing entry means the task held nothing from +Spark at that moment, so the failed call has no partial grant to reconcile, and the retry registers +the task again and waits for its share as the first call would have. After a few attempts it returns +a zero grant, which the native side treats as a refusal and spills. JVM consumers that call +`allocatePage` directly are not covered; that gap is tracked in #6304. + **A consumer whose `spill` returns zero takes budget it can never give back.** `NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim nothing from it; it is only ever a spill trigger. The bytes it holds are real, so other consumers in diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index afdaafd828b..a468446bae7 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -20,6 +20,7 @@ package org.apache.spark; import java.io.IOException; +import java.util.NoSuchElementException; import java.util.concurrent.atomic.AtomicLong; import org.slf4j.Logger; @@ -38,6 +39,9 @@ public class CometTaskMemoryManager { private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class); + /** Attempts at an acquire whose task entry Spark keeps losing, see acquireFromSpark. */ + private static final int MAX_ACQUIRE_ATTEMPTS = 3; + /** The id of the native plan this memory manager was created for, the first in its task. */ private final long id; @@ -70,7 +74,7 @@ public long acquireMemory(long size) { if (logger.isTraceEnabled()) { logger.trace("Task {} requested {} bytes", taskAttemptId, size); } - long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer); + long acquired = acquireFromSpark(size); long newUsed = used.addAndGet(acquired); // A partial grant is routine, not an error: the native pool either refuses the reservation, // which tells the operator to spill, or carries the shortfall as overcommit. A refusal that @@ -94,6 +98,57 @@ public long acquireMemory(long size) { return acquired; } + /** + * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry. Spark + * removes a task's entry from its execution pool when the task's balance reaches zero, and an + * acquire that was waiting for memory then fails when it wakes up with a NoSuchElementException + * ("key not found: " and the task id). Releases take no task monitor, so another native thread or + * a JVM consumer of the same task can empty the balance while an acquire waits. The failed call + * was granted nothing, because Spark only drops the entry at a zero balance, which cannot hold + * while this call has a partial grant. So trying again is safe: it registers the task again and + * waits for its share as Spark would have. After a few failed attempts the grant is refused. The + * helper holds no lock of its own across attempts, so a retry parks in Spark exactly as the first + * attempt did. + */ + private long acquireFromSpark(long size) { + for (int attempt = 1; ; attempt++) { + try { + return internal.acquireExecutionMemory(size, nativeMemoryConsumer); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e)) { + throw e; + } + if (attempt >= MAX_ACQUIRE_ATTEMPTS) { + logger.warn( + "Task {} lost its execution memory entry in Spark on {} attempts to acquire {} " + + "bytes, refusing the request", + taskAttemptId, + attempt, + size); + return 0; + } + logger.info( + "Task {} lost its execution memory entry in Spark while waiting to acquire {} bytes, " + + "trying again", + taskAttemptId, + size); + } + } + } + + /** Whether {@code e} is Spark's execution pool failing to find this task's entry. */ + private boolean isMissingTaskEntry(NoSuchElementException e) { + if (!("key not found: " + taskAttemptId).equals(e.getMessage())) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; + } + // Called by Comet native through JNI public void releaseMemory(long size) { if (logger.isTraceEnabled()) { diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index c9b6d3de7fc..efd899252ac 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -20,11 +20,13 @@ package org.apache.spark import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference} import org.apache.logging.log4j.{Level, LogManager} import org.apache.logging.log4j.core.{LogEvent, LoggerContext} import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} class CometTaskMemoryManagerSuite extends SparkFunSuite { @@ -152,10 +154,166 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { appender.loggingEvents.toSeq } + test("an acquire waiting in Spark survives a release that empties the task's balance") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(otherTask.acquireMemory(90L) == 90L) + assert(manager.acquireMemory(10L) == 10L) + + // Another native thread releases the task's last 10 bytes while the acquire waits. + acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { + manager.releaseMemory(10L) + } + } + } + + test("an acquire waiting in Spark survives a sibling consumer freeing the task's last bytes") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + val sibling = new OffHeapConsumer(taskMemoryManager) + assert(otherTask.acquireMemory(90L) == 90L) + assert(sibling.acquireMemory(10L) == 10L) + + // A JVM consumer of the same task frees with freeMemory, which takes no task monitor. + acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { + sibling.freeMemory(10L) + } + } + } + + test("a NoSuchElementException other than Spark's missing task entry is rethrown") { + val wrongTask = missingEntryError("key not found: 7", fromExecutionMemoryPool = true) + val notFromPool = missingEntryError("key not found: 0", fromExecutionMemoryPool = false) + for (error <- Seq(wrongTask, notFromPool)) { + val taskMemoryManager = new FailingTaskMemoryManager(error) + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + val thrown = intercept[NoSuchElementException](manager.acquireMemory(10L)) + assert(thrown eq error) + assert(taskMemoryManager.calls.get == 1) + assert(manager.getUsed == 0L) + } + } + } + + test("an acquire that keeps losing its task entry gives up with a zero grant") { + val error = missingEntryError("key not found: 0", fromExecutionMemoryPool = true) + val taskMemoryManager = new FailingTaskMemoryManager(error) + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(manager.acquireMemory(10L) == 0L) + assert(taskMemoryManager.calls.get == maxAcquireAttempts) + assert(manager.getUsed == 0L) + assert(nativeMemoryConsumer(manager).getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + + private val TimeoutSeconds = 10L + + /** + * The task holds 10 bytes of a 100 byte pool and another task holds 90. An acquire of 20 bytes + * waits in Spark below its minimum share of 25 while `release` takes the task's balance to + * zero, which removes the task's entry from Spark's pool. Once the other task frees its memory + * the acquire must be granted all 20 bytes. + */ + private def acquireWhileBalanceEmpties( + manager: CometTaskMemoryManager, + taskMemoryManager: TaskMemoryManager, + otherTask: OffHeapConsumer)(release: => Unit): Unit = { + val granted = new AtomicLong(-1L) + val failure = new AtomicReference[Throwable]() + val acquire = new Thread(() => + try granted.set(manager.acquireMemory(20L)) + catch { case t: Throwable => failure.set(t) }) + acquire.setDaemon(true) + + try { + acquire.start() + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(acquire) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(acquire), s"the acquire is ${acquire.getState}") + release + } finally { + // Free the other task's memory so that the acquire does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + + assert(!acquire.isAlive, s"the acquire is ${acquire.getState}") + assert(failure.get == null, s"the acquire failed: ${failure.get}") + assert(granted.get == 20L) + assert(manager.getUsed == 20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) + manager.releaseMemory(20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + + private def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** A 100 byte off-heap execution pool. */ + private def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** A task memory manager whose every acquire throws `error`. */ + private class FailingTaskMemoryManager(error: Throwable) + extends TaskMemoryManager(new TestMemoryManager(new SparkConf()), 0L) { + val calls = new AtomicInteger() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (calls.incrementAndGet() > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + throw error + } + } + + private def missingEntryError( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = { val memoryManager = new TestMemoryManager(new SparkConf()) memoryManager.limit(1024) - val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L) + withTaskContext(new TaskMemoryManager(memoryManager, 0L))(f) + } + + private def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { val taskContext = new TaskContextImpl( stageId = 0, stageAttemptNumber = 0, @@ -182,6 +340,12 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } } + private def maxAcquireAttempts: Int = { + val field = classOf[CometTaskMemoryManager].getDeclaredField("MAX_ACQUIRE_ATTEMPTS") + field.setAccessible(true) + field.getInt(null) + } + private def nativeMemoryConsumer(manager: CometTaskMemoryManager): MemoryConsumer = { val field = classOf[CometTaskMemoryManager].getDeclaredField("nativeMemoryConsumer") field.setAccessible(true) From e84684750cc76227ee000fbf0b4206bd0c39ba1e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 10 Oct 2026 23:26:04 +0000 Subject: [PATCH 51/89] perf: bulk-append long nullable primitive arrays in SparkUnsafeArray conversion (#6824) * Builder appending refactoring * Formatting, small refactoring and adding docstrings * Added docstring regarding alignment * Fixed the fallback in an unaligned but NULLABLE case, that was failing the CI * Changes fixed * Fixed cargo formatting * Remade the building process using arrow buffers with no hot-spots in Vec<> loops * Formatting fix * Added docstring about runtime check of alignment * Addressed the review comments * Fixed formatting * Added docstring regarding endianess * test: cover a timestamp array with a timezone in the append round trip The round-trip test built every builder with its type's default data type, so the timestamp case never had a timezone. append_array panics when the array's data type, timezone included, differs from the builder's. The test now builds the builder with the column's data type and adds a timestamp column with a timezone, which fails if the bitmap path drops the timezone. * perf: append a nullable array with nulls in bulk only from 64 elements Building the validity and value buffers takes four allocations per array, which cost as much as appending 40 to 55 elements one by one, so the bulk path made lists of shorter arrays with a null up to 4.8 times slower than the per-element loop. MIN_BULK_NULLABLE_APPEND_ELEMENTS keeps those arrays on the loop. No array length is slower than on main now, and arrays from 64 elements keep the speedup. - The timestamp appender takes the timezone by reference, so the list path doesn't clone it for every row. - process_primitive_lists is main's single arm again, with an optional argument list. - The SAFETY comments are back, and the alignment comment points at SparkUnsafeObject. - The partial-byte, all-null and all-valid tests run at lengths on both sides of the cutoff, from an 8-byte-aligned copy. The round trip covers 63 and 64 elements with nulls, and its "no nulls and long enough" case now has 8 elements. A new test appends several arrays into one builder, so a validity buffer lands at an offset that isn't a whole byte. * test: cover the list path's timezone and make the bulk cutoff measurable in-tree - timestamp_list_column_keeps_its_timezone drives a Timestamp(Microsecond, Some("UTC")) list column through append_columns, so append_list_column_batch must forward the timezone. With &None in its place, append_array's data type assertion fails. CometColumnarShuffleSuite gets an array case with 70-element arrays that hold a null, which fails the same way end to end. - The impl_append_to_builder doc states that `timezone` must be the builder's, and when a mismatch panics. The has_null() comment covers the validity buffer path too. - array_element_append gets a list_with_one_null group: 8192 list rows of 32, 64 and 128 elements with one null each, for i32 and i64, appended as the list path appends them. Its doc comment shows how to compare the bulk path with the per-element loop, and MIN_BULK_NULLABLE_APPEND_ELEMENTS cites it instead of an out-of-tree harness. * test: read the bench's array words with as_chunks Clippy 1.99 flags chunks_exact with a constant chunk size (chunks_exact_to_as_chunks), so the rust-test job failed under -D warnings before any Rust test ran. The buffer's length is a multiple of 8, so as_chunks::<8>() leaves no remainder. --------- Co-authored-by: sandudb Co-authored-by: Matt Butrovich --- native/core/benches/array_element_append.rs | 88 ++++- native/shuffle/src/spark_unsafe/list.rs | 300 ++++++++++++++---- native/shuffle/src/spark_unsafe/row.rs | 91 +++++- .../exec/CometColumnarShuffleSuite.scala | 19 ++ 4 files changed, 425 insertions(+), 73 deletions(-) diff --git a/native/core/benches/array_element_append.rs b/native/core/benches/array_element_append.rs index 2c46f9ba117..0d4f2850af5 100644 --- a/native/core/benches/array_element_append.rs +++ b/native/core/benches/array_element_append.rs @@ -22,14 +22,24 @@ //! used when processing List/Array columns in JVM shuffle. use arrow::array::builder::{ - Date32Builder, Float64Builder, Int32Builder, Int64Builder, TimestampMicrosecondBuilder, + make_builder, Date32Builder, Float64Builder, Int32Builder, Int64Builder, ListBuilder, + TimestampMicrosecondBuilder, }; use arrow::datatypes::{DataType, TimeUnit}; -use comet::execution::shuffle::spark_unsafe::list::{append_to_builder, SparkUnsafeArray}; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use comet::execution::shuffle::spark_unsafe::list::{ + append_list_element, append_to_builder, SparkUnsafeArray, +}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; const NUM_ELEMENTS: usize = 10000; +/// Rows in each batch of `benchmark_list_with_one_null`, as in a default Comet batch. +const LIST_ROWS: usize = 8192; + +/// Distinct arrays `benchmark_list_with_one_null` cycles through, so it does not read one array +/// from cache for every row. +const DISTINCT_ARRAYS: usize = 64; + /// Create a SparkUnsafeArray in memory with i32 elements. /// Layout: /// - 8 bytes: num_elements (i64) @@ -137,6 +147,29 @@ fn create_spark_unsafe_array_f64(num_elements: usize, with_nulls: bool) -> Vec Vec { + let header_size = 8 + num_elements.div_ceil(64) * 8; + let mut buffer = vec![0u8; header_size + (num_elements * width).div_ceil(8) * 8]; + buffer[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + buffer[8 + null_idx / 8] |= 1 << (null_idx % 8); + for i in 0..num_elements { + let offset = header_size + i * width; + buffer[offset..offset + width].copy_from_slice(&(i as i64).to_le_bytes()[..width]); + } + buffer + .as_chunks::<8>() + .0 + .iter() + .map(|word| u64::from_ne_bytes(*word)) + .collect() +} + fn benchmark_array_conversion(c: &mut Criterion) { let mut group = c.benchmark_group("spark_unsafe_array_to_arrow"); @@ -260,6 +293,53 @@ fn benchmark_array_conversion(c: &mut Criterion) { group.finish(); } +/// Appends a batch of list rows whose arrays each hold one null, as JVM shuffle appends a list +/// column, for array lengths on both sides of `MIN_BULK_NULLABLE_APPEND_ELEMENTS` in +/// `native/shuffle/src/spark_unsafe/list.rs`. Arrays at least that long are appended with one copy +/// of their values and a validity buffer, shorter ones element by element. To compare the two at +/// every length, save a baseline with that constant set to `usize::MAX`, which keeps every array +/// on the per-element loop, then restore it and compare against the baseline: +/// +/// ```shell +/// cargo bench --bench array_element_append -- list_with_one_null --save-baseline loop +/// cargo bench --bench array_element_append -- list_with_one_null --baseline loop +/// ``` +fn benchmark_list_with_one_null(c: &mut Criterion) { + let mut group = c.benchmark_group("list_with_one_null"); + group.throughput(Throughput::Elements(LIST_ROWS as u64)); + + for (name, element_type, width) in [("i32", DataType::Int32, 4), ("i64", DataType::Int64, 8)] { + for num_elements in [32, 64, 128] { + let buffers: Vec> = (0..DISTINCT_ARRAYS) + .map(|i| { + create_spark_unsafe_array_with_one_null(num_elements, width, i % num_elements) + }) + .collect(); + let arrays: Vec = buffers + .iter() + .map(|buffer| SparkUnsafeArray::new(buffer.as_ptr() as i64)) + .collect(); + + group.bench_with_input( + BenchmarkId::new(name, num_elements), + &arrays, + |b, arrays| { + b.iter(|| { + // The capacity Comet gives a list column's values builder. + let mut builder = ListBuilder::new(make_builder(&element_type, 100)); + for array in arrays.iter().cycle().take(LIST_ROWS) { + append_list_element(&element_type, &mut builder, array).unwrap(); + } + builder.finish() + }); + }, + ); + } + } + + group.finish(); +} + fn config() -> Criterion { Criterion::default() } @@ -267,6 +347,6 @@ fn config() -> Criterion { criterion_group! { name = benches; config = config(); - targets = benchmark_array_conversion + targets = benchmark_array_conversion, benchmark_list_with_one_null } criterion_main!(benches); diff --git a/native/shuffle/src/spark_unsafe/list.rs b/native/shuffle/src/spark_unsafe/list.rs index 3d8e89ebc56..f74fcf95baa 100644 --- a/native/shuffle/src/spark_unsafe/list.rs +++ b/native/shuffle/src/spark_unsafe/list.rs @@ -27,10 +27,15 @@ use arrow::array::{ ListBuilder, NullBuilder, StringBuilder, StructBuilder, Time64NanosecondBuilder, TimestampMicrosecondBuilder, }, - MapBuilder, + MapBuilder, PrimitiveArray, +}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer}; +use arrow::datatypes::{ + DataType, Date32Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, + Time64NanosecondType, TimeUnit, TimestampMicrosecondType, }; -use arrow::datatypes::{DataType, TimeUnit}; use datafusion_comet_jni_bridge::errors::CometError; +use std::sync::Arc; /// Element count from which a nullable array with no null elements is appended with one copy /// instead of element by element. Below it the loop is as fast: checking the null bitset and @@ -39,55 +44,93 @@ use datafusion_comet_jni_bridge::errors::CometError; /// time with the copy, and arrays of 50 about a fifth. const MIN_BULK_APPEND_ELEMENTS: usize = 8; +/// Element count from which an aligned nullable array that holds a null is appended with one copy +/// of its values and a validity buffer built from its null bitset, instead of element by element. +/// Building the two buffers takes four allocations, which cost as much as appending about 50 +/// elements one by one, so shorter arrays stay on the loop. The `list_with_one_null` group of the +/// `array_element_append` benchmark times lengths on both sides of the cutoff, and its doc comment +/// shows how to compare the two paths. On an M3 Max with the system allocator, arrays of 32 +/// elements took 1.4 times as long with the copy, and arrays of 64 took 69% of the time for 4-byte +/// elements and 86% for 8-byte ones. +const MIN_BULK_NULLABLE_APPEND_ELEMENTS: usize = 64; + /// Generates bulk append methods for primitive types in SparkUnsafeArray. /// /// # Safety invariants for all generated methods: /// - `element_offset` points to contiguous element data of length `num_elements` /// - `null_bitset_ptr()` returns a pointer to `ceil(num_elements/64)` i64 words /// - These invariants are guaranteed by the SparkUnsafeArray layout from the JVM +/// +/// A timestamp appender also takes `timezone`, which must be the builder's timezone, because +/// `append_array` requires the array's data type to match the builder's. A mismatch panics in +/// `append_array` for an aligned array of `MIN_BULK_NULLABLE_APPEND_ELEMENTS` or more elements that +/// holds a null, and goes unnoticed for every other array. macro_rules! impl_append_to_builder { - ($method_name:ident, $builder_type:ty, $element_type:ty) => { - pub(crate) fn $method_name(&self, builder: &mut $builder_type) { + ($method_name:ident, $builder_type:ty, $element_type:ty, $arrow_type:ty + $(, $timezone:ident)?) => { + pub(crate) fn $method_name( + &self, + builder: &mut $builder_type, + $($timezone: &Option>,)? + ) { let num_elements = self.num_elements; if num_elements == 0 { return; } + debug_assert!(self.element_offset != 0, "element_offset is null"); + let ptr = self.element_offset as *const $element_type; + // An array can start at an address that is not aligned for its elements (see + // `SparkUnsafeObject`), and only aligned elements can be read as a slice. + let aligned = (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()); // A nullable array without nulls takes the copy below too, once it is long enough. if NULLABLE && (num_elements < MIN_BULK_APPEND_ELEMENTS || self.has_null()) { - let mut ptr = self.element_offset as *const $element_type; let null_words = self.null_bitset_ptr(); - debug_assert!(!null_words.is_null(), "null_bitset_ptr is null"); - debug_assert!(!ptr.is_null(), "element_offset pointer is null"); - for idx in 0..num_elements { - // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { - builder.append_null(); - } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { ptr.read_unaligned() }); - } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; - } - } else { - // SAFETY: element_offset points to contiguous data of length num_elements - debug_assert!(self.element_offset != 0, "element_offset is null"); - let ptr = self.element_offset as *const $element_type; - // Use bulk copy when data is properly aligned, fall back to - // per-element unaligned reads otherwise - if (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()) { - let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; - builder.append_slice(slice); + if aligned && num_elements >= MIN_BULK_NULLABLE_APPEND_ELEMENTS { + // SAFETY: element_offset points to num_elements aligned elements. + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + let null_mask_len = num_elements.div_ceil(8); + // SAFETY: the ceil(num_elements/64) words of the null bitset cover these + // bytes. The words are little-endian, so byte i holds the bits of elements + // 8i to 8i+7, least significant first, as in an Arrow bitmap. + let null_mask = unsafe { + std::slice::from_raw_parts(null_words as *const u8, null_mask_len) + }; + // Spark sets the bit of a null element and Arrow the bit of a valid one. The + // padding bits past the last element turn to ones, outside the buffer's length. + let flipped: Vec = null_mask.iter().map(|byte| !byte).collect(); + let validity = + NullBuffer::new(BooleanBuffer::new(Buffer::from(flipped), 0, num_elements)); + let arr = PrimitiveArray::<$arrow_type>::new( + ScalarBuffer::from(Buffer::from_slice_ref(values)), + Some(validity), + ); + $(let arr = arr.with_timezone_opt($timezone.clone());)? + builder.append_array(&arr); } else { let mut ptr = ptr; - for _ in 0..num_elements { - builder.append_value(unsafe { ptr.read_unaligned() }); + for idx in 0..num_elements { + // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + // SAFETY: ptr is within element data bounds + builder.append_value(unsafe { ptr.read_unaligned() }); + } + // SAFETY: ptr stays within bounds, iterating num_elements times ptr = unsafe { ptr.add(1) }; } } + } else if aligned { + // SAFETY: element_offset points to num_elements aligned elements. + let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + builder.append_slice(slice); + } else { + let mut ptr = ptr; + for _ in 0..num_elements { + builder.append_value(unsafe { ptr.read_unaligned() }); + ptr = unsafe { ptr.add(1) }; + } } } }; @@ -180,7 +223,9 @@ impl SparkUnsafeArray { fn has_null(&self) -> bool { let null_words = self.null_bitset_ptr(); // SAFETY: the null bitset holds ceil(num_elements/64) words. Spark zeroes the bits past - // the last element, and a stray one would only send the array down the per-element path. + // the last element. A stray one would send the array to the per-element loop, or to the + // validity buffer if it is aligned and long enough, and both ignore the bits past the + // last element. (0..self.num_elements.div_ceil(64)) .any(|word| unsafe { null_words.add(word).read_unaligned() } != 0) } @@ -196,19 +241,26 @@ impl SparkUnsafeArray { (null_words.add(word_idx).read_unaligned() & (1i64 << bit_idx)) != 0 } - impl_append_to_builder!(append_ints_to_builder, Int32Builder, i32); - impl_append_to_builder!(append_longs_to_builder, Int64Builder, i64); - impl_append_to_builder!(append_time64s_to_builder, Time64NanosecondBuilder, i64); - impl_append_to_builder!(append_shorts_to_builder, Int16Builder, i16); - impl_append_to_builder!(append_bytes_to_builder, Int8Builder, i8); - impl_append_to_builder!(append_floats_to_builder, Float32Builder, f32); - impl_append_to_builder!(append_doubles_to_builder, Float64Builder, f64); + impl_append_to_builder!(append_ints_to_builder, Int32Builder, i32, Int32Type); + impl_append_to_builder!(append_longs_to_builder, Int64Builder, i64, Int64Type); + impl_append_to_builder!( + append_time64s_to_builder, + Time64NanosecondBuilder, + i64, + Time64NanosecondType + ); + impl_append_to_builder!(append_shorts_to_builder, Int16Builder, i16, Int16Type); + impl_append_to_builder!(append_bytes_to_builder, Int8Builder, i8, Int8Type); + impl_append_to_builder!(append_floats_to_builder, Float32Builder, f32, Float32Type); + impl_append_to_builder!(append_doubles_to_builder, Float64Builder, f64, Float64Type); impl_append_to_builder!( append_timestamps_to_builder, TimestampMicrosecondBuilder, - i64 + i64, + TimestampMicrosecondType, + timezone ); - impl_append_to_builder!(append_dates_to_builder, Date32Builder, i32); + impl_append_to_builder!(append_dates_to_builder, Date32Builder, i32, Date32Type); /// Bulk append boolean values to builder. /// Booleans are stored as 1 byte each in SparkUnsafeArray, requiring special handling. @@ -220,37 +272,28 @@ impl SparkUnsafeArray { if num_elements == 0 { return; } - - let mut ptr = self.element_offset as *const u8; debug_assert!( - !ptr.is_null(), + self.element_offset != 0, "append_booleans: element_offset pointer is null" ); + // SAFETY: element_offset points to num_elements one-byte booleans, which need no + // alignment. + let values = + unsafe { std::slice::from_raw_parts(self.element_offset as *const u8, num_elements) }; if NULLABLE { let null_words = self.null_bitset_ptr(); - debug_assert!( - !null_words.is_null(), - "append_booleans: null_bitset_ptr is null" - ); - for idx in 0..num_elements { + for (idx, &value) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); + builder.append_value(value != 0); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - for _ in 0..num_elements { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); - ptr = unsafe { ptr.add(1) }; + for &value in values { + builder.append_value(value != 0); } } } @@ -307,9 +350,9 @@ pub fn append_to_builder( let builder = downcast_builder_ref!(Float64Builder, builder); array.append_doubles_to_builder::(builder); } - DataType::Timestamp(TimeUnit::Microsecond, _) => { + DataType::Timestamp(TimeUnit::Microsecond, tz) => { let builder = downcast_builder_ref!(TimestampMicrosecondBuilder, builder); - array.append_timestamps_to_builder::(builder); + array.append_timestamps_to_builder::(builder, tz); } DataType::Date32 => { let builder = downcast_builder_ref!(Date32Builder, builder); @@ -415,6 +458,7 @@ pub fn append_list_element( mod test { use super::*; use arrow::array::builder::PrimitiveBuilder; + use arrow::array::Array; use arrow::datatypes::{ ArrowPrimitiveType, Date32Type, Int16Type, Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, @@ -455,7 +499,7 @@ mod test { let base = unsafe { (words.as_mut_ptr() as *mut u8).add(shift) }; unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), base, bytes.len()) }; let array = SparkUnsafeArray::new(base as i64); - let mut builder = PrimitiveBuilder::::new(); + let mut builder = PrimitiveBuilder::::new().with_data_type(data_type.clone()); append_to_builder::(&data_type, &mut builder, &array).unwrap(); builder.finish().iter().map(|v| v.map(Into::into)).collect() } @@ -465,6 +509,11 @@ mod test { let long: Vec> = (0..70).map(Some).collect(); let mut null_in_second_bitset_word = long.clone(); null_in_second_bitset_word[66] = None; + let nulls_at_both_ends = |len: i64| -> Vec> { + (0..len) + .map(|v| (v != 0 && v != len - 1).then_some(-v)) + .collect() + }; let cases = [ vec![], // Shorter than MIN_BULK_APPEND_ELEMENTS: element by element. @@ -472,8 +521,22 @@ mod test { // Holds a null: element by element. vec![Some(1), None, Some(3), Some(-4), Some(5)], // No nulls and long enough: one copy. - vec![Some(-1), Some(2), Some(-3), Some(4)], + vec![ + Some(-1), + Some(2), + Some(-3), + Some(4), + Some(-5), + Some(6), + Some(-7), + Some(8), + ], long, + // Holds a null, one element short of MIN_BULK_NULLABLE_APPEND_ELEMENTS: element by + // element. + nulls_at_both_ends(63), + // Holds a null and long enough: one copy and a validity buffer. + nulls_at_both_ends(64), // The null is only in the second bitset word, which the copy check must read. null_in_second_bitset_word, ]; @@ -508,7 +571,116 @@ mod test { ), *values ); + // The builder carries the column's timezone, which `append_array` checks. + assert_eq!( + round_trip::( + DataType::Timestamp(TimeUnit::Microsecond, Some("America/Denver".into())), + values, + shift + ), + *values + ); } } } + + fn make_i32_array(num_elements: usize, null_indices: &[usize]) -> Vec { + let null_bitset_words = num_elements.div_ceil(64); + let header_size = 8 + null_bitset_words * 8; + let mut buffer = vec![0u8; header_size + num_elements * 4]; + + buffer[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + + for &i in null_indices { + let word_offset = 8 + (i / 64) * 8; + let word = i64::from_le_bytes(buffer[word_offset..word_offset + 8].try_into().unwrap()); + buffer[word_offset..word_offset + 8] + .copy_from_slice(&(word | (1i64 << (i % 64))).to_le_bytes()); + } + + for i in 0..num_elements { + let offset = header_size + i * 4; + buffer[offset..offset + 4].copy_from_slice(&(i as i32).to_le_bytes()); + } + buffer + } + + /// Copies `bytes` to an 8-byte boundary, as inside an UnsafeRow. + fn aligned_copy(bytes: &[u8]) -> Vec { + let mut words = vec![0u64; bytes.len().div_ceil(8)]; + // SAFETY: `words` covers `bytes.len()` bytes. + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr(), + words.as_mut_ptr() as *mut u8, + bytes.len(), + ) + }; + words + } + + fn append_i32(buffer: &[u8]) -> arrow::array::Int32Array { + let words = aligned_copy(buffer); + let array = SparkUnsafeArray::new(words.as_ptr() as i64); + let mut builder = Int32Builder::new(); + append_to_builder::(&DataType::Int32, &mut builder, &array).unwrap(); + builder.finish() + } + + #[test] + fn test_nullable_null_in_last_byte() { + // None of the lengths is a multiple of 8, so the null sits in a partial last byte. The + // first three take the per-element loop, the rest the validity buffer. + for num_elements in [7, 9, 17, 65, 71, 73, 81] { + let last = num_elements - 1; + let arr = append_i32(&make_i32_array(num_elements, &[last])); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), 1); + assert!(arr.is_null(last)); + for i in 0..last { + assert!(arr.is_valid(i)); + assert_eq!(arr.value(i), i as i32); + } + } + } + + #[test] + fn test_all_null() { + for num_elements in [9, 73] { + let indices: Vec = (0..num_elements).collect(); + let arr = append_i32(&make_i32_array(num_elements, &indices)); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), num_elements); + } + } + + #[test] + fn test_all_valid() { + for num_elements in [9, 73] { + let arr = append_i32(&make_i32_array(num_elements, &[])); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), 0); + } + } + + #[test] + fn nullable_arrays_append_into_one_builder() { + // The list path appends every row's array to one values builder, so the validity buffer + // of a long array is appended at an offset that is not a whole byte. + let arrays: [Vec>; 4] = [ + vec![Some(1), None, Some(3)], + (0..70).map(|v| (v % 9 != 4).then_some(v)).collect(), + vec![None; 5], + (0..65).map(|v| (v != 64).then_some(-v)).collect(), + ]; + let mut builder = Int32Builder::new(); + for values in &arrays { + let words = aligned_copy(&unsafe_array_bytes(values, 4)); + let array = SparkUnsafeArray::new(words.as_ptr() as i64); + append_to_builder::(&DataType::Int32, &mut builder, &array).unwrap(); + } + let appended: Vec> = + builder.finish().iter().map(|v| v.map(i64::from)).collect(); + assert_eq!(appended, arrays.concat()); + } } diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 180673c37fb..ce18a398df5 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -622,7 +622,7 @@ fn append_list_column_batch( // Helper macro for primitive element types - gets builder fresh each iteration // to avoid borrow conflicts with list_builder.append() macro_rules! process_primitive_lists { - ($builder_type:ty, $append_fn:ident) => {{ + ($builder_type:ty, $append_fn:ident $(, $extra:expr)*) => {{ for i in row_start..row_end { read_row_at!(row, row_addresses_ptr, row_sizes_ptr, i); @@ -636,7 +636,7 @@ fn append_list_column_batch( .as_any_mut() .downcast_mut::<$builder_type>() .expect(stringify!($builder_type)); - array.$append_fn::(values_builder); + array.$append_fn::(values_builder $(, $extra)*); list_builder.append(true); } } @@ -668,8 +668,12 @@ fn append_list_column_batch( DataType::Date32 => { process_primitive_lists!(Date32Builder, append_dates_to_builder); } - DataType::Timestamp(TimeUnit::Microsecond, _) => { - process_primitive_lists!(TimestampMicrosecondBuilder, append_timestamps_to_builder); + DataType::Timestamp(TimeUnit::Microsecond, tz) => { + process_primitive_lists!( + TimestampMicrosecondBuilder, + append_timestamps_to_builder, + tz + ); } DataType::Time64(TimeUnit::Nanosecond) => { process_primitive_lists!(Time64NanosecondBuilder, append_time64s_to_builder); @@ -1504,7 +1508,8 @@ fn make_batch(arrays: Vec, row_count: usize) -> Result`. assert_eq!(&*row.get_string(0), "\u{FFFD}\u{FFFD}A"); } + + /// Lays out an 8-byte-aligned `UnsafeRow` with one array field of 8-byte elements: the null + /// bitset, the field's offset and size, then the `UnsafeArrayData`. `None` is a null field. + fn unsafe_row_with_array(values: Option<&[Option]>) -> Vec { + let Some(values) = values else { + return vec![1, 0]; + }; + let mut null_bitset = vec![0u64; values.len().div_ceil(64)]; + let mut elements = Vec::with_capacity(values.len()); + for (i, value) in values.iter().enumerate() { + if value.is_none() { + null_bitset[i / 64] |= 1 << (i % 64); + } + elements.push(value.unwrap_or(0) as u64); + } + let array_size = 8 * (1 + null_bitset.len() + elements.len()) as u64; + let mut row = vec![0, (16 << 32) | array_size, values.len() as u64]; + row.extend(null_bitset); + row.extend(elements); + row + } + + #[test] + fn timestamp_list_column_keeps_its_timezone() { + // An array of at least MIN_BULK_NULLABLE_APPEND_ELEMENTS elements that holds a null is + // appended with `append_array`, which panics unless the array has the builder's timezone. + let data_type = DataType::List(Arc::new(Field::new( + "item", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ))); + let with_null: Vec> = (0..70).map(|v| (v != 3).then_some(v << 20)).collect(); + let without_null: Vec> = (0..70).map(|v| Some(-v)).collect(); + let short = vec![Some(1), None]; + let lists = [ + Some(with_null.as_slice()), + None, + Some(short.as_slice()), + Some(without_null.as_slice()), + ]; + let rows: Vec> = lists.iter().map(|l| unsafe_row_with_array(*l)).collect(); + let mut row_addresses: Vec = rows.iter().map(|r| r.as_ptr() as jlong).collect(); + let mut row_sizes: Vec = rows.iter().map(|r| (r.len() * 8) as jint).collect(); + + let schema = [data_type.clone()]; + let mut builder = make_builders(&data_type, rows.len(), 1.0).unwrap(); + append_columns( + row_addresses.as_mut_ptr(), + row_sizes.as_mut_ptr(), + 0, + rows.len(), + &schema, + 0, + &mut builder, + 1.0, + ) + .unwrap(); + let array = builder.finish(); + + assert_eq!(array.data_type(), &data_type); + let appended: Vec>>> = array + .as_list::() + .iter() + .map(|values| { + values.map(|values| { + values + .as_primitive::() + .iter() + .collect() + }) + }) + .collect(); + let expected: Vec>>> = + lists.iter().map(|l| l.map(<[_]>::to_vec)).collect(); + assert_eq!(appended, expected); + } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index c6a6805c72f..a84d4df3e7d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -419,6 +419,25 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar } } + test("columnar shuffle on array [timestamp]") { + // The native row-to-Arrow conversion appends an array of 64 or more elements that holds a + // null in one call, which requires the array to carry the column's timezone. + Seq(10, 201).foreach { numPartitions => + withParquetTable( + (0 until 50).map { i => + val timestamps: Seq[java.sql.Timestamp] = (0 until 70).map { j => + if (j == i) null else new java.sql.Timestamp(i * 1000000L + j) + } + (i, timestamps) + }, + "tbl") { + val df = sql("SELECT * FROM tbl").repartition(numPartitions, $"_1") + + checkShuffleAnswer(df, 1) + } + } + } + test("columnar shuffle on nested array") { Seq("false", "true").foreach { _ => Seq(10, 201).foreach { numPartitions => From a56db5a6c0e47e9ffdb446e01312598db6c54896 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 10 Oct 2026 23:26:54 +0000 Subject: [PATCH 52/89] perf: skip DecimalPrecision.promote's rewrite when there is no decimal arithmetic (#5216) * perf: skip DecimalPrecision.promote's rewrite when there is no decimal arithmetic QueryPlanSerde.exprToProto runs DecimalPrecision.promote over every expression it converts. The rule is a transformUp, so it walks and rebuilds every node even when nothing matches, which is the common case. Guard it with a cheap traversal that looks for arithmetic producing a decimal, and only run the rewrite when one exists. * ci: register DecimalPrecisionSuite in the PR build workflows * docs: describe transformUp's real cost in the promote guard comment --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../spark/sql/comet/DecimalPrecision.scala | 21 +++++ .../sql/comet/DecimalPrecisionSuite.scala | 83 +++++++++++++++++++ 4 files changed, 106 insertions(+) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 2bb8fd363d1..49eafc6e3a1 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -566,6 +566,7 @@ jobs: org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite org.apache.spark.sql.comet.CometDecimalPromotionSuite + org.apache.spark.sql.comet.DecimalPrecisionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 148dde4784e..22bd6251dcf 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -270,6 +270,7 @@ jobs: org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite org.apache.spark.sql.comet.CometDecimalPromotionSuite + org.apache.spark.sql.comet.DecimalPrecisionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala index 08efcfc9c9c..14ab3554776 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala @@ -40,6 +40,27 @@ import org.apache.spark.sql.types.DecimalType */ object DecimalPrecision { def promote(expr: Expression): Expression = { + // `transformUp` walks the whole tree and calls `mapChildren` at every node, which allocates + // through `children.map` / `mapProductIterator` even when no case matches (the specialized + // `mapChildren` overrides do skip the node copy via `fastEquals`, so the cost is the traversal + // and its allocations, not a rebuild). The serde calls this once per expression it converts, + // so skip it when there is nothing to rewrite. + if (containsDecimalArithmetic(expr)) rewrite(expr) else expr + } + + /** + * Whether [[promote]] could rewrite anything in this tree. Deliberately a superset of the rule: + * it ignores the operands, so any arithmetic node the rule handles that produces a decimal + * qualifies. Wrapping a child in `CheckOverflow` does not change that child's data type, so a + * tree without such a node cannot grow one part way through the rewrite. + */ + private def containsDecimalArithmetic(expr: Expression): Boolean = expr.exists { + case e @ (_: Add | _: Subtract | _: Multiply | _: Divide | _: Remainder) => + e.dataType.isInstanceOf[DecimalType] + case _ => false + } + + private def rewrite(expr: Expression): Expression = { expr.transformUp { // Recursive exprToProto calls can re-promote every decimal binary operator below. Collapse // only equivalent wrappers so the shared promotion rule remains idempotent. diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala new file mode 100644 index 00000000000..be4c7db5de8 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala @@ -0,0 +1,83 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, CaseWhen, Cast, CheckOverflow, Coalesce, Divide, Expression, GreaterThan, IsNotNull, Literal, Multiply, Remainder, Subtract} +import org.apache.spark.sql.types.{DecimalType, IntegerType, LongType, StringType} + +class DecimalPrecisionSuite extends SparkFunSuite { + + private val decA = AttributeReference("a", DecimalType(10, 2))() + private val decB = AttributeReference("b", DecimalType(10, 2))() + private val intA = AttributeReference("c", IntegerType)() + private val intB = AttributeReference("d", IntegerType)() + + private def assertUnchanged(expr: Expression): Unit = { + // Reference equality, not `===`: the point of the guard is that no node is rebuilt. + val promoted = DecimalPrecision.promote(expr) + assert(promoted.eq(expr), s"expected promote to leave $expr untouched, got: $promoted") + } + + test("promote leaves trees without decimal arithmetic untouched") { + Seq( + decA, + Literal(1), + Add(intA, intB), + Multiply(intA, Cast(decA, IntegerType)), + GreaterThan(decA, decB), + Cast(Add(intA, intB), DecimalType(20, 4)), + Coalesce(Seq(decA, decB)), + CaseWhen(Seq((IsNotNull(decA), decA)), Some(decB)), + Alias(Add(Cast(decA, LongType), Cast(decB, LongType)), "x")()).foreach(assertUnchanged) + } + + test("promote wraps decimal arithmetic wherever it appears in the tree") { + Seq[Expression]( + Add(decA, decB), + Subtract(decA, decB), + Multiply(decA, decB), + Divide(decA, decB), + Remainder(decA, decB)).foreach { arithmetic => + DecimalPrecision.promote(arithmetic) match { + case CheckOverflow(child, dt, _) => + assert(child.eq(arithmetic)) + assert(dt === arithmetic.dataType) + case other => fail(s"expected $arithmetic to be wrapped in CheckOverflow, got: $other") + } + + // The guard has to inspect the whole tree, not just the root: a decimal arithmetic node + // buried under expressions the rule does not rewrite must still be promoted. + val buried = Cast(Coalesce(Seq(decA, arithmetic)), StringType) + val promoted = DecimalPrecision.promote(buried) + assert(!promoted.eq(buried)) + assert(promoted.collectFirst { case c: CheckOverflow => c }.isDefined) + } + } + + test("promote rewrites only the decimal arithmetic in a mixed tree") { + val intAdd = Add(intA, intB) + val decAdd = Add(decA, decB) + val promoted = DecimalPrecision.promote(GreaterThan(Cast(intAdd, DecimalType(20, 2)), decAdd)) + val overflows = promoted.collect { case c: CheckOverflow => c } + assert(overflows.length === 1) + assert(overflows.head.child.eq(decAdd)) + } +} From 1e5297fc26ed2319eb6ebee15bb398938cf05b2c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 10 Oct 2026 23:29:11 +0000 Subject: [PATCH 53/89] deps: bump to datafusion 55.2.0 (#6766) --- contrib/delta/native/Cargo.lock | 120 ++++++++-------- contrib/delta/native/Cargo.toml | 2 +- native/Cargo.lock | 135 +++++++++--------- native/Cargo.toml | 10 +- native/core/Cargo.toml | 2 +- .../eager_page_index_reader_factory.rs | 5 +- 6 files changed, 138 insertions(+), 136 deletions(-) diff --git a/contrib/delta/native/Cargo.lock b/contrib/delta/native/Cargo.lock index c3a4c3289a2..dbbbf9dfc41 100644 --- a/contrib/delta/native/Cargo.lock +++ b/contrib/delta/native/Cargo.lock @@ -488,9 +488,9 @@ dependencies = [ [[package]] name = "datafusion" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14313430439ef858ed5f52ca9c840e9eea9844da34d3716d05a476449634da5d" +checksum = "4cd5317db77ba6bca25c761a8091b9d294a9f4188e4f9dcd0243e6a57de862ac" dependencies = [ "arrow", "arrow-schema", @@ -534,9 +534,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c0756ebc770bcebae366dfd273dad50935f6103b4d212253cb17e2e53c5427f" +checksum = "f35d36794509b7bfcbf90e11a1b5e4d047123c34610105e2e80a2048351cdb2f" dependencies = [ "arrow", "async-trait", @@ -559,9 +559,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbde2961698705d8402f0f382f08876fc6f1126ffde790ef17833cf0ab87cd59" +checksum = "a6454bfb7f43b83942e2ca2576fd09342eedbd0d323e1eef72c3c8e4d57d7203" dependencies = [ "arrow", "async-trait", @@ -591,9 +591,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e06fe711ff815c8cf3143158547aea599702fd43f76f5866c70d5a5b3693a0" +checksum = "d26626e12b224b83351239465dcd00f72d47b02fb5e5b65a643e2dc8ed37f687" dependencies = [ "arrow", "arrow-ipc", @@ -616,9 +616,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7a5107acff4e3ed50a7732acec06bba40e74ab49df021397ed250d60f04e42a" +checksum = "a106420cbfa980f242c22243ae9a6cce661d4d86ce96cff87dacabf7754a4470" dependencies = [ "futures", "log", @@ -627,9 +627,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f69f45d2d700785be8df9d4b44410e1c8e4a856f10cc728b418a1066c0b594fb" +checksum = "210aa1e9525e9454f2399dbb8b0b95510145bd8c1fdc35de74e9ccbf81bcb699" dependencies = [ "arrow", "async-trait", @@ -657,9 +657,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81650af64894ef34b46d2ef818ef244b40a81d90818f3bc642bf8fbf6758f8e0" +checksum = "65b1b041688cad707fa103c3e7f4c6d89b3134ebf67da0129070dad3381851eb" dependencies = [ "arrow", "arrow-ipc", @@ -681,9 +681,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e2519bed59a6907c79298b9c21fe807bbc34e4ff8414577d699b0741083593" +checksum = "2a7b34e253c4fa5251f174977af910a2f2ec831ec39dbbb56b29ce505cff278b" dependencies = [ "arrow", "async-trait", @@ -704,9 +704,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635e61290e171597f6e50427e0cb486eda468ab47e880fa24d15bbb0fa24ac7e" +checksum = "1663338f43be6a5746473008e2d0fbde282577ddf0bae90d6c7386e988d71b29" dependencies = [ "arrow", "async-trait", @@ -727,9 +727,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cd518f85b3a027d913047b4ab0811111a2b734fdf84a51b78234455f10d8fa" +checksum = "7c521ae91b6d899d7190812881f45fe9463c222c11d90149f5332a2e56555904" dependencies = [ "arrow", "arrow-schema", @@ -759,15 +759,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c5919b560862dd1cb20a5c75b3e48c3d9e49719c91d78da5c38eaad5495c4" +checksum = "0c2569c2e74a34fcae3a369ab3912d61e41a062592510f11e31b1dae2725ecdc" [[package]] name = "datafusion-execution" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7620ad535af8730186bdec05d720e40d9a4584b40c7eb75f72c43aa7978e5484" +checksum = "39cfdd5ac0769d1e74a885450e49fde8f63ca7ab41eb479a068564aaea80b837" dependencies = [ "arrow", "arrow-buffer", @@ -791,9 +791,9 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea229af8b73e566fc6ebb1c8dd42ab250f8855c5304bd173f4b50ead78b64304" +checksum = "86560fd7df4dda924f5474b588c56e7ae8e96dcf158c7327c9cfe9daff448ec5" dependencies = [ "arrow", "arrow-schema", @@ -807,14 +807,15 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap", "itertools 0.15.0", + "serde", "serde_json", ] [[package]] name = "datafusion-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a18baa840f77760dcc9eaad281329bc57536c5c39e65e8649596f2a51216be7" +checksum = "e82e313cf655cf6bc2dce6cbd23bcb0d20236a54a1b09c436615e4bac4d8c49f" dependencies = [ "arrow", "datafusion-common", @@ -824,9 +825,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763bbfeec40b55a8a961d7349c0fcfcec4f3f66cc234d6295250b95371b85e34" +checksum = "939ae2771da8a3fdc23ab4de610ae8fef9515476626b49d96d0d5c782b105771" dependencies = [ "arrow", "arrow-buffer", @@ -852,9 +853,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c0c2b2e856883ecf6a4fa8b23e7c6041183af8a68d0583dc95d1d21e4eaee0" +checksum = "fd6c3aee9f2bc65716787d651e023743a90b4efbc57a2f679f036632e16fb110" dependencies = [ "arrow", "datafusion-common", @@ -873,9 +874,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e35e42c87dfbff06934ea0e99cfed3440ac614c3f3bfb8c33be35ca0587d4140" +checksum = "988b5d3dfb4939d81c2f9fab267063610fe52a843620f2fa6cb2a1598e43369c" dependencies = [ "arrow", "datafusion-common", @@ -885,9 +886,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe2070ecf5638a93c81c8babc7858339ef52165a9f54da1411afd69adf7d85ce" +checksum = "2c767567cce2782d4d0bdfa55fab010aaba0bdbe533509e9e0b9a6a3ee29ed59" dependencies = [ "arrow", "async-trait", @@ -901,9 +902,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9a164ebe1998c5c3f9e15c9a0c4e6b599a0c8739b5471bfb089752441963f1" +checksum = "0a311c1f558d1f792f1f631c3670da12ee72a20fbf5d6ca151325d66191c03c1" dependencies = [ "arrow", "datafusion-common", @@ -918,9 +919,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98965b183f92201c16798994820ae032649c4821e089c4c80e7d0f4485d41df5" +checksum = "fffd54b8acf6ecef284cc0677e01e1b8c71be64fec50612ba00a0ff0d9b41a52" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -928,9 +929,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a711bd5a06ec85d9683e142d28a98ad0679a7220669f5c1b8589e57e9b29e46" +checksum = "14d6df9763c8fb968b2417d797aef4c55ac2e4a895aaef38c72c16a287aa4a40" dependencies = [ "datafusion-doc", "quote", @@ -939,9 +940,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3256b62ff44616ec83c7f04741d21124970833bb377ed56d9abc5ec6f707ec" +checksum = "acfb8cea6e05e2e75bae183340cf7d53cccb3db0449e2f1fbc07654d980fd19f" dependencies = [ "arrow", "chrono", @@ -958,9 +959,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1e37a25b5aa1f3a1db114f0ee3605bf01639b83a831dcde35119ee61f0b11d6" +checksum = "9f002cef849e10c0597a217ec7891c2796c0f8b42d739f5d24abc24496373406" dependencies = [ "arrow", "arrow-schema", @@ -980,9 +981,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a4d12787d42e88b26b0d3ec006192fa23ceb832ef2a9e87439ce820a8e038cc" +checksum = "82761522c4536e7059d3b185cd99f10525c2bf48f1bd26a0a57e73e19eadc0fb" dependencies = [ "arrow", "datafusion-common", @@ -995,9 +996,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d13bf7f50fd902f5afc46d50759f5185b5620c40b24054f80f7c648897d3e5a6" +checksum = "f8c5c67aa252886f4d6b331c8114797f5b140464e774c4da36889269a1bf1bd8" dependencies = [ "arrow", "chrono", @@ -1012,9 +1013,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11fbe7e6665072957d44ed9587676e17a1f9ab3acfadfc287dac23750585f18a" +checksum = "f1a49dd3a76e4879012682c59dc3735282e60640c97483546a4113564a60bf8f" dependencies = [ "arrow", "datafusion-common", @@ -1031,9 +1032,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" +checksum = "b1451d6817e5ac0dd21819b19949d1c420ab51d748f55dd9a96ecb5418a24c97" dependencies = [ "arrow", "arrow-data", @@ -1060,15 +1061,16 @@ dependencies = [ "num-traits", "parking_lot", "pin-project-lite", + "serde", "serde_json", "tokio", ] [[package]] name = "datafusion-pruning" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +checksum = "f11f0dbe2b5a642f6106e44de8bbd6654fa7220ca14f32313d4efd5514a86968" dependencies = [ "arrow", "datafusion-common", @@ -1082,9 +1084,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7616b17b1d2fee3e9b13c651eac91b768d153265e6ba88e06e0b4af7883f69c" +checksum = "62dc5ce37a5ce1ff5362949776829aded2529a11dbaa13d24b0ffa6a059e106a" dependencies = [ "arrow-schema", "async-trait", @@ -2119,6 +2121,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -2147,7 +2150,6 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ - "indexmap", "itoa", "memchr", "serde", diff --git a/contrib/delta/native/Cargo.toml b/contrib/delta/native/Cargo.toml index abe73473c4b..43cc1dda273 100644 --- a/contrib/delta/native/Cargo.toml +++ b/contrib/delta/native/Cargo.toml @@ -47,4 +47,4 @@ crate-type = ["rlib"] datafusion-comet-proto = { path = "../../../native/proto" } # For `ExecutionPlan` (return type), `SchemaRef` (the shim hands these in), and `DataFusionError` # (the not-implemented error). Versions/features mirror core so the shared lockfile resolves cleanly. -datafusion = { version = "55.1.0", default-features = false, features = ["parquet"] } +datafusion = { version = "55.2.0", default-features = false, features = ["parquet"] } diff --git a/native/Cargo.lock b/native/Cargo.lock index c73688fe1d8..e4d2aac7e1d 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1881,9 +1881,9 @@ dependencies = [ [[package]] name = "datafusion" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14313430439ef858ed5f52ca9c840e9eea9844da34d3716d05a476449634da5d" +checksum = "4cd5317db77ba6bca25c761a8091b9d294a9f4188e4f9dcd0243e6a57de862ac" dependencies = [ "arrow", "arrow-schema", @@ -1930,9 +1930,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c0756ebc770bcebae366dfd273dad50935f6103b4d212253cb17e2e53c5427f" +checksum = "f35d36794509b7bfcbf90e11a1b5e4d047123c34610105e2e80a2048351cdb2f" dependencies = [ "arrow", "async-trait", @@ -1955,9 +1955,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbde2961698705d8402f0f382f08876fc6f1126ffde790ef17833cf0ab87cd59" +checksum = "a6454bfb7f43b83942e2ca2576fd09342eedbd0d323e1eef72c3c8e4d57d7203" dependencies = [ "arrow", "async-trait", @@ -2164,9 +2164,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e06fe711ff815c8cf3143158547aea599702fd43f76f5866c70d5a5b3693a0" +checksum = "d26626e12b224b83351239465dcd00f72d47b02fb5e5b65a643e2dc8ed37f687" dependencies = [ "arrow", "arrow-ipc", @@ -2191,9 +2191,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7a5107acff4e3ed50a7732acec06bba40e74ab49df021397ed250d60f04e42a" +checksum = "a106420cbfa980f242c22243ae9a6cce661d4d86ce96cff87dacabf7754a4470" dependencies = [ "futures", "log", @@ -2202,9 +2202,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f69f45d2d700785be8df9d4b44410e1c8e4a856f10cc728b418a1066c0b594fb" +checksum = "210aa1e9525e9454f2399dbb8b0b95510145bd8c1fdc35de74e9ccbf81bcb699" dependencies = [ "arrow", "async-compression", @@ -2238,9 +2238,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81650af64894ef34b46d2ef818ef244b40a81d90818f3bc642bf8fbf6758f8e0" +checksum = "65b1b041688cad707fa103c3e7f4c6d89b3134ebf67da0129070dad3381851eb" dependencies = [ "arrow", "arrow-ipc", @@ -2262,9 +2262,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e2519bed59a6907c79298b9c21fe807bbc34e4ff8414577d699b0741083593" +checksum = "2a7b34e253c4fa5251f174977af910a2f2ec831ec39dbbb56b29ce505cff278b" dependencies = [ "arrow", "async-trait", @@ -2285,9 +2285,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635e61290e171597f6e50427e0cb486eda468ab47e880fa24d15bbb0fa24ac7e" +checksum = "1663338f43be6a5746473008e2d0fbde282577ddf0bae90d6c7386e988d71b29" dependencies = [ "arrow", "async-trait", @@ -2308,9 +2308,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cd518f85b3a027d913047b4ab0811111a2b734fdf84a51b78234455f10d8fa" +checksum = "7c521ae91b6d899d7190812881f45fe9463c222c11d90149f5332a2e56555904" dependencies = [ "arrow", "arrow-schema", @@ -2340,15 +2340,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c5919b560862dd1cb20a5c75b3e48c3d9e49719c91d78da5c38eaad5495c4" +checksum = "0c2569c2e74a34fcae3a369ab3912d61e41a062592510f11e31b1dae2725ecdc" [[package]] name = "datafusion-execution" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7620ad535af8730186bdec05d720e40d9a4584b40c7eb75f72c43aa7978e5484" +checksum = "39cfdd5ac0769d1e74a885450e49fde8f63ca7ab41eb479a068564aaea80b837" dependencies = [ "arrow", "arrow-buffer", @@ -2373,9 +2373,9 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea229af8b73e566fc6ebb1c8dd42ab250f8855c5304bd173f4b50ead78b64304" +checksum = "86560fd7df4dda924f5474b588c56e7ae8e96dcf158c7327c9cfe9daff448ec5" dependencies = [ "arrow", "arrow-schema", @@ -2389,15 +2389,16 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap 2.14.2", "itertools 0.15.0", + "serde", "serde_json", "sqlparser", ] [[package]] name = "datafusion-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a18baa840f77760dcc9eaad281329bc57536c5c39e65e8649596f2a51216be7" +checksum = "e82e313cf655cf6bc2dce6cbd23bcb0d20236a54a1b09c436615e4bac4d8c49f" dependencies = [ "arrow", "datafusion-common", @@ -2407,9 +2408,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763bbfeec40b55a8a961d7349c0fcfcec4f3f66cc234d6295250b95371b85e34" +checksum = "939ae2771da8a3fdc23ab4de610ae8fef9515476626b49d96d0d5c782b105771" dependencies = [ "arrow", "arrow-buffer", @@ -2439,9 +2440,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c0c2b2e856883ecf6a4fa8b23e7c6041183af8a68d0583dc95d1d21e4eaee0" +checksum = "fd6c3aee9f2bc65716787d651e023743a90b4efbc57a2f679f036632e16fb110" dependencies = [ "arrow", "datafusion-common", @@ -2460,9 +2461,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e35e42c87dfbff06934ea0e99cfed3440ac614c3f3bfb8c33be35ca0587d4140" +checksum = "988b5d3dfb4939d81c2f9fab267063610fe52a843620f2fa6cb2a1598e43369c" dependencies = [ "arrow", "datafusion-common", @@ -2472,9 +2473,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08079b0604cdcc43c7d6cf006cfd5d3ed9c235ade0c1ebe91085a62f875603df" +checksum = "08149d433a97c9ea3cdf80de06480ab5c7f5453485f1e452d5b3583c81d5c0dd" dependencies = [ "arrow", "arrow-ord", @@ -2497,9 +2498,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe2070ecf5638a93c81c8babc7858339ef52165a9f54da1411afd69adf7d85ce" +checksum = "2c767567cce2782d4d0bdfa55fab010aaba0bdbe533509e9e0b9a6a3ee29ed59" dependencies = [ "arrow", "async-trait", @@ -2513,9 +2514,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9a164ebe1998c5c3f9e15c9a0c4e6b599a0c8739b5471bfb089752441963f1" +checksum = "0a311c1f558d1f792f1f631c3670da12ee72a20fbf5d6ca151325d66191c03c1" dependencies = [ "arrow", "datafusion-common", @@ -2530,9 +2531,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98965b183f92201c16798994820ae032649c4821e089c4c80e7d0f4485d41df5" +checksum = "fffd54b8acf6ecef284cc0677e01e1b8c71be64fec50612ba00a0ff0d9b41a52" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2540,9 +2541,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a711bd5a06ec85d9683e142d28a98ad0679a7220669f5c1b8589e57e9b29e46" +checksum = "14d6df9763c8fb968b2417d797aef4c55ac2e4a895aaef38c72c16a287aa4a40" dependencies = [ "datafusion-doc", "quote", @@ -2551,9 +2552,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3256b62ff44616ec83c7f04741d21124970833bb377ed56d9abc5ec6f707ec" +checksum = "acfb8cea6e05e2e75bae183340cf7d53cccb3db0449e2f1fbc07654d980fd19f" dependencies = [ "arrow", "chrono", @@ -2570,9 +2571,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1e37a25b5aa1f3a1db114f0ee3605bf01639b83a831dcde35119ee61f0b11d6" +checksum = "9f002cef849e10c0597a217ec7891c2796c0f8b42d739f5d24abc24496373406" dependencies = [ "arrow", "arrow-schema", @@ -2592,9 +2593,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a4d12787d42e88b26b0d3ec006192fa23ceb832ef2a9e87439ce820a8e038cc" +checksum = "82761522c4536e7059d3b185cd99f10525c2bf48f1bd26a0a57e73e19eadc0fb" dependencies = [ "arrow", "datafusion-common", @@ -2607,9 +2608,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d13bf7f50fd902f5afc46d50759f5185b5620c40b24054f80f7c648897d3e5a6" +checksum = "f8c5c67aa252886f4d6b331c8114797f5b140464e774c4da36889269a1bf1bd8" dependencies = [ "arrow", "chrono", @@ -2624,9 +2625,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11fbe7e6665072957d44ed9587676e17a1f9ab3acfadfc287dac23750585f18a" +checksum = "f1a49dd3a76e4879012682c59dc3735282e60640c97483546a4113564a60bf8f" dependencies = [ "arrow", "datafusion-common", @@ -2643,9 +2644,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" +checksum = "b1451d6817e5ac0dd21819b19949d1c420ab51d748f55dd9a96ecb5418a24c97" dependencies = [ "arrow", "arrow-data", @@ -2672,15 +2673,16 @@ dependencies = [ "num-traits", "parking_lot", "pin-project-lite", + "serde", "serde_json", "tokio", ] [[package]] name = "datafusion-pruning" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +checksum = "f11f0dbe2b5a642f6106e44de8bbd6654fa7220ca14f32313d4efd5514a86968" dependencies = [ "arrow", "datafusion-common", @@ -2694,9 +2696,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7616b17b1d2fee3e9b13c651eac91b768d153265e6ba88e06e0b4af7883f69c" +checksum = "62dc5ce37a5ce1ff5362949776829aded2529a11dbaa13d24b0ffa6a059e106a" dependencies = [ "arrow-schema", "async-trait", @@ -2709,23 +2711,22 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7809ecefb16fcb950c6b22ee3df99cb412dca6fae72c329d64c938ae53254f9" +checksum = "0c3d6f76463984c4ed8587c52ed4f21f142fd96fa55d8f128f673f2db4c46942" dependencies = [ "arrow", "bigdecimal", "chrono", "crc32fast", "datafusion", - "datafusion-catalog", "datafusion-common", - "datafusion-execution", "datafusion-expr", "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-aggregate-common", "datafusion-functions-nested", + "datafusion-session", "log", "num-traits", "percent-encoding", @@ -2739,9 +2740,9 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b4848be66dd008d55adbe1dbc8cbd6a2502588ad0c0b8483bab0f43fd4495d0" +checksum = "816c224a0ef5469a32a562b5d4830544f714f97c5f50142109286d40a288c7db" dependencies = [ "arrow", "bigdecimal", @@ -6587,7 +6588,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", diff --git a/native/Cargo.toml b/native/Cargo.toml index dac17e5a289..defbf41ae27 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -44,11 +44,11 @@ arrow-select = { version = "59.2.0" } async-trait = { version = "0.1" } bytes = { version = "1.11.1" } parquet = { version = "59.2.0", default-features = false, features = ["experimental"] } -datafusion = { version = "55.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } -datafusion-datasource = { version = "55.1.0" } -datafusion-functions = { version = "55.1.0", default-features = false, features = ["datetime_expressions"] } -datafusion-physical-expr-adapter = { version = "55.1.0" } -datafusion-spark = { version = "55.1.0", features = ["core"] } +datafusion = { version = "55.2.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } +datafusion-datasource = { version = "55.2.0" } +datafusion-functions = { version = "55.2.0", default-features = false, features = ["datetime_expressions"] } +datafusion-physical-expr-adapter = { version = "55.2.0" } +datafusion-spark = { version = "55.2.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } datafusion-comet-common = { path = "common" } datafusion-comet-jni-bridge = { path = "jni-bridge" } diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 2de67226df0..b21c5bc75a8 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -111,7 +111,7 @@ jni = { version = "0.22.4", features = ["invocation"] } lazy_static = "1.4" assertables = "10" hex = "0.4.3" -datafusion-functions-nested = { version = "55.1.0" } +datafusion-functions-nested = { version = "55.2.0" } comet-test-udfs = { path = "../comet-test-udfs" } # Test-only helpers for a hand-rolled in-memory HTTP client that exercises the web-identity STS # retry path without a socket, plus reqsign's SigV4 signer for the signing round-trip test. Already diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 15bafdd2d63..dbde337306e 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -52,9 +52,8 @@ //! `ignoreMissing` is set, and it walks the raw `MessageType` to decide. That walk has to run //! over the Parquet schema rather than the Arrow schema the schema adapter is handed later, //! because an id on a repeated `list` or `key_value` group, or on the message root, never -//! reaches an Arrow field. Until apache/datafusion#24790, which is not in DataFusion 55.1.0, -//! the INT96 coercion also rebuilt container fields without their metadata, so a struct id could -//! vanish on the way to Arrow as well. +//! reaches an Arrow field. DataFusion 55.2.0 preserves container metadata through INT96 coercion +//! (apache/datafusion#24790), so struct ids no longer vanish on the way to Arrow. //! //! The second is the Variant footer rewrite, `with_spark_arrow_schema`, which replaces the //! Arrow schema hint in the footer for scans that project Variant. From 78b0eb0df342e631a619c4951fd9b531f0375a3e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 02:38:56 +0000 Subject: [PATCH 54/89] fix: check native UDF types at registration and explain planning-time evaluation (#6777) Registering a native UDF with a type Comet has no native representation for, such as a user-defined type, succeeded, and every call then fell back to Spark, which cannot evaluate it. CometNativeUDF.register now refuses such a signature before it loads the library. Spark also evaluates a call itself while planning a query: over local data, in a filter on partition columns, and to sample the keys of a global sort. The error a native UDF raises there said that Comet had not taken the operator holding the call, which sent users to the explain output for a fallback that never happened. The error, the scaladoc and the Rust UDF guide now list those cases, and the guide says to sort on a column holding a UDF's result rather than on the call. Split out of #6697. --- docs/source/user-guide/latest/rust_udfs.md | 17 +++++++++ .../org/apache/comet/udf/CometNativeUDF.scala | 20 ++++++++--- .../apache/comet/udf/CometUdfExceptions.scala | 11 +++--- .../org/apache/comet/udf/NativeUdfCall.scala | 7 ++-- .../apache/comet/CometNativeUdfSuite.scala | 35 +++++++++++++++++++ 5 files changed, 76 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index ba98cd6acbb..147a2069c94 100644 --- a/docs/source/user-guide/latest/rust_udfs.md +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -158,6 +158,7 @@ CometNativeUDF.register( Registration loads the library on the driver and verifies that a function with that name exists, so a bad path or a missing function fails immediately with a clear error rather than at execution time. +It also checks that Comet can carry the argument and return types natively. The function is then callable from SQL or the DataFrame API like any other: @@ -181,6 +182,19 @@ way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux and paths may be loaded, so it makes no difference to the trust decision described under [Limitations](#limitations). +### When Spark evaluates the call + +A call runs only inside Comet's native execution. Spark cannot run the function, so wherever Spark +has to evaluate the call itself, the query fails with `CometUdfNotEvaluatedException`. That happens: + +- when Comet does not take the operator holding the call, for example because another expression + in it is not supported, or because the call is an argument of an ordinary Scala or Java UDF. The + query's extended explain output gives the reason. +- while Spark plans the query: over local data such as `VALUES`, which the optimizer evaluates + eagerly, in a filter on partition columns, which Spark evaluates to prune partitions, and in the + sort keys of a global sort, which Spark evaluates on a sample of rows to choose range bounds. To + sort on a UDF's result, select it as a column and sort on that column instead. + ## Return types A UDF does not have one fixed return type. Its `return_field` is called with the actual argument @@ -277,6 +291,9 @@ the first panic, and the SDK cannot prevent it. This feature is at an early stage. The current limitations are: - **Scalar functions only.** Aggregate, window, and table functions are not supported. +- **Spark cannot evaluate a call.** A query fails wherever Spark would have to run the function + itself, which includes a few places while it plans the query. See + [When Spark evaluates the call](#when-spark-evaluates-the-call). - **Immutable functions only.** A UDF must return the same output for the same input. Comet plans every Rust UDF with DataFusion's `Volatility::Immutable`, so a function that reads a clock, draws from an RNG, or carries state across batches may be folded at plan time, evaluated once and diff --git a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala index 5c2a9608307..0feb9605922 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.comet.CometUdfErrors import org.apache.spark.sql.types.DataType +import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.shims.ShimSessionFunctionRegistry /** @@ -46,9 +47,10 @@ object CometNativeUDF { /** * Register a single native UDF with an explicit signature. * - * Validates the library on the driver (loads it, confirms a UDF named `name` exists), then - * installs `name` as a temporary function of the session, as `spark.udf.register` would: other - * sessions do not see it, and registering another function under the same name replaces it. + * Checks that Comet can carry the argument and return types natively and validates the library + * on the driver (loads it, confirms a UDF named `name` exists), then installs `name` as a + * temporary function of the session, as `spark.udf.register` would: other sessions do not see + * it, and registering another function under the same name replaces it. * * Executors need no registration: the library path travels with the plan in the * `NativeScalarUdf` proto, and each executor loads the library itself on first use. The path @@ -58,8 +60,9 @@ object CometNativeUDF { * argument types differ, other than in nullability, and inserts no casts, so cast the arguments * in the query instead. * - * Spark cannot evaluate the call itself: if Comet does not take the operator holding it, the - * query fails. + * Spark cannot evaluate the call itself, so a query fails where Spark would have to: in an + * operator Comet does not take, and while planning over local data, in a filter on partition + * columns, or to sample the keys of a global sort. * * `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a * nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently @@ -84,6 +87,13 @@ object CometNativeUDF { "eliminated as a common subexpression. " + "See https://github.com/apache/datafusion-comet/issues/5249") } + // With a type Comet cannot represent natively, every call would fall back to Spark, which + // cannot evaluate it. + (inputTypes :+ returnType).find(QueryPlanSerde.serializeDataType(_).isEmpty).foreach { t => + throw new IllegalArgumentException( + s"native UDF '$name': Comet has no native representation for type " + + s"${t.catalogString}, so no call to it could run.") + } validateLibrary(libraryPath, name) ShimSessionFunctionRegistry .functionRegistry(spark) diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala index eeda2447b1f..44e0554653d 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala @@ -21,12 +21,11 @@ package org.apache.comet.udf import org.apache.comet.CometRuntimeException -/** - * Thrown when Spark itself evaluates a call to a UDF registered with Comet, which means Comet did - * not take the operator holding it. - */ +/** Thrown when Spark itself evaluates a call to a native UDF, which only Comet can evaluate. */ class CometUdfNotEvaluatedException(name: String) extends CometRuntimeException( s"UDF '$name' is registered with Comet and runs only inside Comet's native execution, but " + - "Spark evaluated it, which means Comet did not take the operator holding the call. The " + - "query's extended explain output gives the reason that operator fell back to Spark.") + "Spark evaluated it. Either Comet did not take the operator holding the call, and the " + + "query's extended explain output gives the reason, or Spark evaluated the call itself " + + "while planning: over local data, in a filter on partition columns, or to sample the " + + "keys of a global sort.") diff --git a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala index 22c71039a2d..d1f444cc4ba 100644 --- a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala +++ b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala @@ -29,9 +29,10 @@ import org.apache.spark.sql.types.DataType * in the session builds one for each call it resolves. * * Comet replaces it with a `NativeScalarUdf` that runs the function named `name` in the library - * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws. That happens only when - * Comet does not take the operator holding the call, and failing there makes the fallback visible - * rather than silent. + * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws, which makes a fallback + * visible rather than silent. Spark evaluates a call where Comet does not take the operator + * holding it, and while planning in a few places: over local data, in a filter on partition + * columns, and to sample the keys of a global sort. * * `argumentTypes` is the registered signature, one type per child. Spark's analyzer checks each * call against it, disregarding nullability, and inserts no casts. diff --git a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala index 9f36d18ff78..4bad30e2c96 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala @@ -446,6 +446,41 @@ class CometNativeUdfSuite extends CometTestBase { } } + test("Spark evaluating a native UDF while planning fails the query") { + // Spark's optimizer evaluates a projection over local data itself, so this fails even though + // Comet would take every operator. + CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType) + val e = intercept[Exception] { + sql("SELECT add_one_c(x) FROM VALUES (1L), (2L) AS t(x)").collect() + } + assert( + stackTraceContains(e, "Spark evaluated the call itself while planning"), + s"unhelpful error: $e") + } + + test("a global sort on a native UDF's result runs when the sort is on a column") { + // Spark evaluates a global sort's keys on a sample of rows to choose range bounds, so the + // sort has to be on a column holding the result rather than on the call itself. + CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType) + val df = spark.range(0, 20, 1, 4).selectExpr("add_one_c(19 - id) AS y").orderBy("y") + assert(df.collect().map(_.getLong(0)).toSeq == (1L to 20L)) + } + + test("a type Comet cannot carry natively is refused at registration") { + // Earlier tests register `echo_c`, so drop it first to see that the refusal installs nothing. + sql("DROP TEMPORARY FUNCTION IF EXISTS echo_c") + val e = intercept[IllegalArgumentException] { + CometNativeUDF.register( + spark, + "echo_c", + libPath, + Seq(ObjectType(classOf[String])), + LongType) + } + assert(e.getMessage.contains("Comet has no native representation"), e.getMessage) + assert(!spark.catalog.functionExists("echo_c")) + } + // A native UDF call has to reach its serde wherever Comet converts expressions, not just in a // projection. Filters, join conditions, grouping keys and window partitioning each route // through a different Comet operator, and a regression in any one of them would show up only as From 9e13649146bba11ad146cc65ef575ebad32b503f Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Sun, 11 Oct 2026 06:13:23 +0000 Subject: [PATCH 55/89] test: move bit_count boundary coverage to SQL fixtures (#6655) --- .../expressions/bitwise/bit_count_types.sql | 50 +++++++++++++++++++ .../comet/CometBitwiseExpressionSuite.scala | 23 --------- 2 files changed, 50 insertions(+), 23 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/bitwise/bit_count_types.sql diff --git a/spark/src/test/resources/sql-tests/expressions/bitwise/bit_count_types.sql b/spark/src/test/resources/sql-tests/expressions/bitwise/bit_count_types.sql new file mode 100644 index 00000000000..65329056a1f --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/bitwise/bit_count_types.sql @@ -0,0 +1,50 @@ +-- 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. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- Spark 3.4.3 emits invalid Java for bit_count(boolean). Use interpreted Spark +-- reference execution; plain queries still assert Comet operator replacement. +-- Config: spark.sql.codegen.wholeStage=false +-- Config: spark.sql.codegen.factoryMode=NO_CODEGEN + +statement +CREATE TABLE test_bit_count_types(l bigint, i int, s smallint, b tinyint, flag boolean) USING parquet + +statement +INSERT INTO test_bit_count_types +SELECT /*+ COALESCE(1) */ l, i, s, b, flag +FROM VALUES + (1111, 2222, 17, 7, true), + (9223372036854775807, 2147483647, 32767, 127, true), + (-9223372036854775808, -2147483648, -32768, -128, false), + (-1, -1, -1, -1, false), + (0, 0, 0, 0, false), + (1, 1, 1, 1, true), + (4294967293, 0, 0, 0, true), + (4294967294, 0, 0, 0, false), + (4294967295, 0, 0, 0, true), + (4294967296, 0, 0, 0, false), + (NULL, NULL, NULL, NULL, NULL) +AS input(l, i, s, b, flag) + +query +SELECT bit_count(l), bit_count(i), bit_count(s), bit_count(b), bit_count(flag) +FROM test_bit_count_types + +query +SELECT bit_count(true), bit_count(false) FROM test_bit_count_types diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index db72532ee50..36540bc4997 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -116,29 +116,6 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe } } - test("bitwise_count - min/max values") { - Seq(false, true).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "bitwise_count_test" - withTable(table) { - sql(s"create table $table(col1 long, col2 int, col3 short, col4 byte) using parquet") - sql(s"insert into $table values(1111, 2222, 17, 7)") - sql( - s"insert into $table values(${Long.MaxValue}, ${Int.MaxValue}, ${Short.MaxValue}, ${Byte.MaxValue})") - sql( - s"insert into $table values(${Long.MinValue}, ${Int.MinValue}, ${Short.MinValue}, ${Byte.MinValue})") - - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(col1) FROM $table")) - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(col2) FROM $table")) - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(col3) FROM $table")) - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(col4) FROM $table")) - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(true) FROM $table")) - checkSparkAnswerAndOperator(sql(s"SELECT bit_count(false) FROM $table")) - } - } - } - } - test("bitwise_count - random values (spark gen)") { withTempDir { dir => val path = new Path(dir.toURI.toString, "test.parquet") From d735166ad81846f394f944bbb6c75a6cd91dcd93 Mon Sep 17 00:00:00 2001 From: dustin Date: Sun, 11 Oct 2026 10:28:13 +0000 Subject: [PATCH 56/89] test: cover container field ids next to INT96 timestamps in the native Parquet scan (#6862) DataFusion 55.2.0 keeps field metadata when coercing INT96 timestamps, so a field id on a struct, list, map or Variant column now survives into the native scan's id matching. Add tests for root and nested containers that read the column under a new name, so only the id can match it. --- .../comet/parquet/ParquetReadSuite.scala | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 2a595398367..687edf1450a 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -42,8 +42,10 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn} import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec} +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils +import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -2598,6 +2600,114 @@ abstract class ParquetReadSuite extends CometTestBase { } } + // Spark writes timestamps as INT96 by default, and the native scan's INT96 coercion rebuilds + // struct, list and map fields, which must keep their metadata or an id on the container is + // lost and the column reads as nulls. The read renames the column so only its id matches it. + private val int96Ts = Timestamp.valueOf("2020-01-01 00:00:00") + private val int96Inner = new StructType().add("a", IntegerType).add("ts", TimestampType) + + Seq[(String, DataType, Any)]( + ("struct", int96Inner, Row(1, int96Ts)), + ("list", ArrayType(int96Inner), Seq(Row(1, int96Ts), Row(2, int96Ts))), + ("map", MapType(StringType, int96Inner), Map("k" -> Row(1, int96Ts)))).foreach { + case (label, dataType, value) => + test(s"an id on a $label column holding an INT96 timestamp matches it natively") { + withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") { + withTempPath { dir => + val rows = Seq(Row(value), Row(null)) + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "INT96") { + spark + .createDataFrame( + spark.sparkContext.parallelize(rows), + new StructType().add("c", dataType, true, withId(1))) + .write + .parquet(dir.getCanonicalPath) + } + val readSchema = new StructType().add("renamed", dataType, true, withId(1)) + val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath) + checkSparkAnswerAndOperator(df) + checkAnswer(df, rows) + } + } + } + } + + // Ids on a struct below a struct or a list, and on its INT96 leaf, match renamed fields. + // Spark's vectorized reader before 4.0 rejects the renamed struct below the list, so the + // comparison with Spark runs from 4.0 on and the pinned rows hold everywhere. + test("an id on a struct nested in a struct or list with an INT96 leaf matches it natively") { + def schema(innerName: String, tsName: String): StructType = { + def inner(innerId: Int, tsId: Int): StructType = new StructType().add( + innerName, + new StructType().add("a", IntegerType).add(tsName, TimestampType, true, withId(tsId)), + true, + withId(innerId)) + new StructType().add("s", inner(2, 3)).add("l", ArrayType(inner(4, 5))) + } + val rows = Seq( + Row(Row(Row(1, int96Ts)), Seq(Row(Row(2, int96Ts)), Row(Row(3, int96Ts)))), + Row(Row(null), Seq(Row(Row(4, null)), null))) + + withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") { + withTempPath { dir => + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "INT96") { + spark + .createDataFrame(spark.sparkContext.parallelize(rows), schema("inner", "ts")) + .write + .parquet(dir.getCanonicalPath) + } + val df = spark.read + .schema(schema("renamed_inner", "renamed_ts")) + .parquet(dir.getCanonicalPath) + if (isSpark40Plus) { + checkSparkAnswerAndOperator(df) + } + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert( + collect(plan) { case scan: CometNativeScanExec => scan }.nonEmpty, + s"expected CometNativeScanExec in the plan:\n$plan") + checkAnswer(df, rows) + } + } + } + + // A Variant column is a struct of value and metadata in the Arrow schema, so the same coercion + // rebuilds it when the file holds an INT96 column anywhere, and its id has to survive. + test("an id on a Variant column next to an INT96 timestamp matches it natively") { + assume(Utils.variantType.isDefined, "VariantType requires Spark 4.0+") + withSQLConf( + SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true", + "spark.sql.variant.allowReadingShredded" -> "true", + "spark.sql.variant.pushVariantIntoScan" -> "false") { + withTempPath { dir => + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "INT96") { + spark + .sql("""SELECT parse_json(j) AS v, t AS ts FROM VALUES + | ('{"a":1}', TIMESTAMP'2020-01-01 00:00:00'), (NULL, NULL) AS input(j, t) + |""".stripMargin) + .select(col("v").as("v", withId(1)), col("ts")) + .coalesce(1) + .write + .parquet(dir.getCanonicalPath) + } + val readSchema = new StructType() + .add("renamed", Utils.variantType.get, true, withId(1)) + .add("ts", TimestampType) + val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath) + checkSparkAnswerAndOperator(df) + checkAnswer( + df.selectExpr("to_json(renamed)", "ts"), + Row("""{"a":1}""", int96Ts) :: Row(null, null) :: Nil) + } + } + } + // Spark checks each file on its own. A directory holding one file with ids and one without // raises on the second, and with `ignoreMissing` the file without ids reads as nulls because // no root field of it carries the requested id. Each side is written as one file. Spread over From 52a3a336a79c8736b5824274f26d4d9d1d1666aa Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:33:22 +0000 Subject: [PATCH 57/89] fix: match Spark percentile ordering for negative NaNs (#6829) --- .../src/agg_funcs/quantile_summaries.rs | 124 ++++++++++++++++-- .../aggregate/approx_percentile_nan.sql | 64 +++++++++ 2 files changed, 176 insertions(+), 12 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile_nan.sql diff --git a/native/spark-expr/src/agg_funcs/quantile_summaries.rs b/native/spark-expr/src/agg_funcs/quantile_summaries.rs index e6dd8fb9291..664d974d4be 100644 --- a/native/spark-expr/src/agg_funcs/quantile_summaries.rs +++ b/native/spark-expr/src/agg_funcs/quantile_summaries.rs @@ -21,9 +21,10 @@ //! //! Reference: `org.apache.spark.sql.catalyst.util.QuantileSummaries`. //! -//! Results are bit-identical to Spark for deterministic plans. What is -//! load-bearing for that identity, and must not change: +//! Results match Spark for deterministic plans (NaN bit patterns are not +//! specified). What is load-bearing for that identity, and must not change: //! - the 50000-value head-buffer flush boundary (`DEFAULT_HEAD_SIZE`), +//! - Java's primitive-array ordering of that buffer (all NaNs last, -0.0 before 0.0), //! - the integer arithmetic: `floor(2 * relative_error * count)` for `delta`, //! `2 * relative_error * count` for the compress/merge threshold, `/ 2` for //! `target_error`, and `ceil(percentile * count)` for the query rank, @@ -158,10 +159,24 @@ impl QuantileSummaries { } let mut current_count = self.count; let mut sorted = std::mem::take(&mut self.head_sampled); - // Spark relies on `Array[Double].sorted`. The sort key is the value - // itself, so equal keys are bit-equal elements and stability is - // irrelevant; `total_cmp` gives a deterministic total order. + // Spark's `Array[Double].sorted` puts every NaN last, regardless of + // sign or payload. IEEE total order agrees for every other value, + // including -0.0 before 0.0, but puts negative NaNs at the front. sorted.sort_unstable_by(|a, b| a.total_cmp(b)); + // Move that leading NaN block behind the already sorted values and + // positive NaNs. Java compares all NaNs equal, so their suffix order + // does not affect ranks. This preserves their original bits and adds + // no per-comparison cost. Negative NaNs form a prefix in total order, + // so finding its end needs only a binary search, even for all-NaN input. + // Most buffers have no negative NaN. Avoid the binary search in that + // case, especially for the many small buffers of grouped aggregates. + if sorted + .first() + .is_some_and(|v| v.is_nan() && v.is_sign_negative()) + { + let leading_nans = sorted.partition_point(|v| v.is_nan() && v.is_sign_negative()); + sorted.rotate_left(leading_nans); + } scratch.sampled.clear(); scratch.sampled.reserve(self.sampled.len() + sorted.len()); @@ -597,14 +612,99 @@ mod tests { #[test] fn signed_zero_ordering_is_deterministic() { - // `total_cmp` orders -0.0 before 0.0; inserting both must not panic and - // must produce a value drawn from the input. + // Java's primitive-array ordering puts -0.0 before 0.0, including at + // the percentile endpoints. SQL ordering treats the two zeros as equal. let qs = summary_of(&[-0.0, 0.0, -0.0, 0.0, 1.0]); - let got = qs.query(&[0.5]).unwrap()[0]; - assert!( - got == 0.0 || got == 1.0, - "median of signed zeros produced {got}" - ); + assert_eq!(qs.query(&[0.0]).unwrap()[0].to_bits(), (-0.0f64).to_bits()); + let zeros = summary_of(&[0.0, -0.0, 0.0, -0.0]); + assert_eq!(zeros.query(&[1.0]).unwrap()[0].to_bits(), 0.0f64.to_bits()); + } + + #[test] + fn non_finite_values_use_java_primitive_array_order() { + let qs = summary_of(&[ + -f64::NAN, + f64::INFINITY, + 1.5, + -0.0, + f64::from_bits(0xfff0_0000_0000_0001), + f64::NEG_INFINITY, + 0.0, + -1.5, + f64::NAN, + ]); + let expected = [f64::NEG_INFINITY, -1.5, -0.0, 0.0, 1.5, f64::INFINITY]; + assert_eq!(qs.sampled.len(), 9); + for (sample, value) in qs.sampled.iter().zip(expected) { + assert_eq!(sample.value.to_bits(), value.to_bits()); + } + assert!(qs.sampled[expected.len()..] + .iter() + .all(|s| s.value.is_nan())); + let mut nan_bits: Vec<_> = qs.sampled[expected.len()..] + .iter() + .map(|s| s.value.to_bits()) + .collect(); + nan_bits.sort_unstable(); + let mut expected_nan_bits = [ + (-f64::NAN).to_bits(), + f64::NAN.to_bits(), + 0xfff0_0000_0000_0001, + ]; + expected_nan_bits.sort_unstable(); + assert_eq!(nan_bits, expected_nan_bits); + } + + #[test] + fn nan_sign_and_payload_match_spark_across_head_flushes() { + // Spark's result for 300000 ascending values with the first replaced + // by NaN, at the default accuracy. Six head-buffer flushes exercise + // both the initial sort and its interleave with an existing summary. + let percentages = [0.0, 0.25, 0.5, 0.75, 1.0]; + let expected = [1.0, 75001.0, 150000.0, 225013.0]; + for bits in [ + 0x7ff8_0000_0000_0000, + 0xfff8_0000_0000_0000, + 0x7ff8_0000_0000_0042, + 0xfff8_0000_0000_0042, + 0x7ff0_0000_0000_0001, + 0xfff0_0000_0000_0001, + ] { + let mut values: Vec = (0..300_000).map(|i| i as f64).collect(); + values[0] = f64::from_bits(bits); + let qs = summary_of(&values); + let got = qs.query(&percentages).unwrap(); + assert_eq!(qs.count(), values.len() as i64); + assert_eq!(&got[..4], &expected, "NaN bits: {bits:#x}"); + assert!(got[4].is_nan(), "maximum for NaN bits {bits:#x}: {got:?}"); + } + } + + #[test] + fn nan_sign_does_not_change_serialized_partial_merges() { + let percentages = [0.0, 0.25, 0.5, 0.75, 1.0]; + let mut merged = Vec::new(); + for nan in [f64::NAN, -f64::NAN] { + let mut summary = + QuantileSummaries::new(QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, EPS); + let mut scratch = QuantileSummariesScratch::default(); + for partition in 0..3 { + let start = partition * 100_000; + let mut values: Vec = (start..start + 100_000).map(|i| i as f64).collect(); + values[0] = nan; + let partial = summary_of(&values); + let decoded = QuantileSummaries::from_bytes( + QuantileSummaries::DEFAULT_COMPRESS_THRESHOLD, + &partial.to_bytes(), + ); + summary.merge(&decoded, &mut scratch); + } + assert_eq!(summary.count(), 300_000); + let got = summary.query(&percentages).unwrap(); + assert!(got[4].is_nan(), "maximum after merging: {got:?}"); + merged.push(got); + } + assert_eq!(&merged[0][..4], &merged[1][..4]); } #[test] diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile_nan.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile_nan.sql new file mode 100644 index 00000000000..ab3b005d236 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/approx_percentile_nan.sql @@ -0,0 +1,64 @@ +-- 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. + +-- One input partition, with more than the 50000-value head-buffer threshold +-- for both scalar and grouped aggregates. Negate NaN after the Parquet scan: +-- Spark canonicalizes NaNs on write, so storing -NaN would miss this bug. +-- Plain queries assert native execution as well as exact Spark results. +-- Config: spark.sql.files.maxPartitionBytes=134217728 + +statement +CREATE TABLE test_approx_percentile_nan(d double, f float, g int) USING parquet + +statement +INSERT INTO test_approx_percentile_nan +SELECT IF(id = 0, double('NaN'), cast(id AS double)), + IF(id = 0, float('NaN'), cast(id AS float)), + cast(id % 2 AS int) +FROM range(0, 300000, 1, 1) + +-- Canonical NaN control, with both numeric input types and array output. +query +SELECT percentile_approx(d, array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)), + approx_percentile(f, array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)) +FROM test_approx_percentile_nan + +-- Negative-sign NaN must not change any rank or replace the maximum with a +-- finite value. Six head-buffer flushes reproduce issue #6519. +query +SELECT percentile_approx(IF(isnan(d), -d, d), array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)) +FROM test_approx_percentile_nan + +query +SELECT approx_percentile(IF(isnan(f), -f, f), array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)) +FROM test_approx_percentile_nan + +-- Grouped accumulators and their partial-to-final serialized states. Each +-- group exceeds the head-buffer threshold; only group 0 contains a NaN. +query +SELECT g, + percentile_approx(IF(isnan(d), -d, d), array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)), + approx_percentile(IF(isnan(f), -f, f), array(0.0D, 0.25D, 0.5D, 0.75D, 1.0D)) +FROM test_approx_percentile_nan GROUP BY g ORDER BY g + +-- Scalar output uses the same ordering, including the NaN endpoint. +query +SELECT percentile_approx(IF(isnan(d), -d, d), 0.5D), + isnan(percentile_approx(IF(isnan(d), -d, d), 1.0D)), + approx_percentile(IF(isnan(f), -f, f), 0.5D), + isnan(approx_percentile(IF(isnan(f), -f, f), 1.0D)) +FROM test_approx_percentile_nan From 50d73b37e461f7aa2ced66287939bf3537d160b7 Mon Sep 17 00:00:00 2001 From: nat-openai Date: Sun, 11 Oct 2026 10:58:19 +0000 Subject: [PATCH 58/89] perf: integral sequence optimisation (#6833) * perf: optimise integral sequence generation Reduce the per-row and per-element cost of integral sequence(start, stop, step) while preserving the existing Spark-compatible results and errors. - Borrow typed scalar arguments instead of materialising full input columns. Keep all-scalar results as List scalars and preserve constant-folding and caller expansion behaviour. - Calculate ordinary sequence lengths with checked i64 arithmetic, including unit-step fast paths. Retain exact wide arithmetic for overflow so Spark's boundary, collection-size and internal-error cases report the same errors. - Reserve the child values buffer once and initialise fixed slices using independent index arithmetic. This removes per-element capacity checks and lets LLVM vectorise stores. Wrapping intermediates preserve values near integer bounds; publish the vector only after every value is initialised. - Union and reuse Arrow validity buffers once per batch instead of checking each argument and rebuilding validity for every row. Return all-null lists directly after input type and length validation, preserving null suppression of invalid underlying values and supporting sliced input bitmaps. - Retain per-row lengths between sizing and generation so short sequences do not repeat length arithmetic. This avoids the short-row regressions found during measurement without introducing separate generation kernels. Preserve validation order: every row is checked before the aggregate Arrow List offset ceiling, so a later Spark row error still precedes that ceiling. Keep existing scalar result shape, null placement, allocator-failure handling and i32 offset limits. This change introduces no sequence byte budget, buffer ownership policy, cancellation polling or JNI callbacks. Extend behavioural tests across all scalar/array argument combinations and integral widths, extreme values and steps, sliced null bitmaps, empty and mismatched arrays, length overflow and row-versus-batch error precedence. Extend the Criterion benchmark to 32 shapes covering one-element and long sequences, default and explicit steps, mixed directions, scalar expansion, null densities and errors. Record the performance audit: long column-input sequences improve by approximately 5x on Apple M3 Max, with no reproducible regression in the measured shapes. Repeated all-scalar batch expansion is effectively unchanged. * test: restore integral sequence benchmark coverage Keep the optimisation split's benchmark coverage aligned with the original sequence proposal while retaining the existing execution and memory policy. - Add the JVM benchmark queries and matching input columns for sequences of 2, 1,000 and 50,000 elements. These fill the short and long column-input gaps without changing Spark session setup or introducing byte-budget configuration or cancellation-related settings. - Add 2-element and 10,000-element all-scalar cases to the native benchmark, retaining the one-element case. Exercise both one-row and 8,192-row output expansion through spark_sequence(...).into_array(rows), increasing native coverage from 32 to 36 shapes while measuring equivalent caller work. - Clarify that the recorded approximately 5x speedup compares column-valued inputs against the previous native kernel, rather than against Spark's implementation or all-scalar batch expansion. The production kernel is unchanged. Validate the Scala benchmark with the reactor test-compile and Spotless checks, compile the native release benchmark, and check Rust and Markdown formatting. The restored cases are ready for benchmark execution; this commit makes no new performance claims for them. --- .../expression-audits/array_funcs.md | 1 + native/spark-expr/benches/sequence.rs | 72 ++- native/spark-expr/src/array_funcs/sequence.rs | 498 ++++++++++++++---- .../benchmark/CometSequenceBenchmark.scala | 6 + 4 files changed, 461 insertions(+), 116 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index 50ac45b852d..9c918cbd870 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -180,6 +180,7 @@ ## sequence +- Performance (tuned 2026-10-09): checked i64 length arithmetic, borrowed scalar inputs, bulk validity union, and vectorisable child-buffer stores; approximately 5× faster than the previous native kernel for column-valued inputs producing 365-element and 10,000-element sequences on arm64, with short/null-heavy shapes also improved. Benchmark: `benches/sequence.rs`. - Spark 3.4.3 (audited 2026-08-29): `Sequence(start, stop, stepOpt, timeZoneId)`; `Sequence.impl` selects the implementation from `dataType.elementType`, so the integral/temporal split is knowable at plan time. Codegen for the integral path checks boundaries with a plain `IllegalArgumentException("Illegal sequence boundaries: ...")`, then calls the static `Sequence.sequenceLength`, which raises `SparkRuntimeException(_LEGACY_ERROR_TEMP_2161)` past `MAX_ROUNDED_ARRAY_LENGTH` and `internalError("Unreachable code reached.")` when `stop - start` overflows Long but the exact length is within the limit. Default step is per-row `start <= stop ? 1 : -1`. - Spark 3.5.8 (audited 2026-08-29): internal refactors only (`DataTypeUtils.sameType`, `PhysicalIntegralType.integral`); runtime semantics identical to 3.4.3. - Spark 4.0.1 (audited 2026-08-29): boundary error becomes `SparkIllegalArgumentException(_LEGACY_ERROR_TEMP_3243)` and the length error becomes `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER` (now carrying the function name); adds `throwable` optimizer hint. `sequenceLength` itself is unchanged. diff --git a/native/spark-expr/benches/sequence.rs b/native/spark-expr/benches/sequence.rs index 4a67a5f33e5..200500a0a7e 100644 --- a/native/spark-expr/benches/sequence.rs +++ b/native/spark-expr/benches/sequence.rs @@ -18,6 +18,7 @@ use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field}; use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::common::ScalarValue; use datafusion::physical_plan::ColumnarValue; use datafusion_comet_spark_expr::spark_sequence; use std::hint::black_box; @@ -57,7 +58,7 @@ fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("sequence"); // Short sequences: per-row overhead dominates. - for elems in [2i64, 5] { + for elems in [1i64, 2, 5] { let args = args_with_len(elems, None); group.bench_function(format!("short_{elems}_elems"), |b| { b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) @@ -111,6 +112,33 @@ fn criterion_benchmark(c: &mut Criterion) { }); } + // Explicit steps on one-element rows, and non-unit steps in both directions. + for (label, elems) in [ + ("explicit_step_one_element", 1i64), + ("mixed_non_unit_steps", 17), + ] { + let start = Int64Array::from(vec![0i64; NUM_ROWS]); + let step = Int64Array::from( + (0..NUM_ROWS) + .map(|i| if i % 2 == 0 { 3i64 } else { -3 }) + .collect::>(), + ); + let stop = Int64Array::from( + step.values() + .iter() + .map(|step| step * (elems - 1)) + .collect::>(), + ); + let args = vec![ + ColumnarValue::Array(Arc::new(start)), + ColumnarValue::Array(Arc::new(stop)), + ColumnarValue::Array(Arc::new(step)), + ]; + group.bench_function(label, |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + // Error path: the boundary check rejects the first row. { let start = Int64Array::from(vec![0i64; NUM_ROWS]); @@ -126,6 +154,48 @@ fn criterion_benchmark(c: &mut Criterion) { }); } + // Scalar steps avoid the cost of materialising an input column. + for elems in [1i64, 2, 5, 365, 10_000] { + let mut args = args_with_len(elems, None); + args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))); + group.bench_function(format!("scalar_step_{elems}_elems"), |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Include the caller's expansion so scalar and array results do equivalent work. + for rows in [1, NUM_ROWS] { + for elems in [1i64, 2, 5, 365, 10_000] { + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(elems - 1))), + ]; + group.bench_function(format!("all_scalar_{rows}_rows_{elems}_elems"), |b| { + b.iter(|| { + black_box( + spark_sequence(&args, &return_type) + .unwrap() + .into_array(rows) + .unwrap(), + ) + }) + }); + } + } + + for (label, every) in [ + ("sparse_nulls", 10usize), + ("dense_nulls", 2), + ("all_nulls", 1), + ] { + for elems in [1i64, 2, 5] { + let args = args_with_len(elems, Some(every)); + group.bench_function(format!("{label}_{elems}_elems"), |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + } + group.finish(); } diff --git a/native/spark-expr/src/array_funcs/sequence.rs b/native/spark-expr/src/array_funcs/sequence.rs index 4b1335f20fe..17d519c7ce6 100644 --- a/native/spark-expr/src/array_funcs/sequence.rs +++ b/native/spark-expr/src/array_funcs/sequence.rs @@ -29,19 +29,19 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, ListArray, NullBufferBuilder, PrimitiveArray}; -use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::array::{Array, ArrayRef, ListArray, PrimitiveArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ ArrowPrimitiveType, DataType, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type, }; use datafusion::common::cast::as_primitive_array; -use datafusion::common::{exec_err, DataFusionError, Result, ScalarValue}; +use datafusion::common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; use datafusion::logical_expr::ColumnarValue; use crate::SparkError; /// Spark's ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH (Integer.MAX_VALUE - 15). -const MAX_ROUNDED_ARRAY_LENGTH: i128 = (i32::MAX - 15) as i128; +const MAX_ROUNDED_ARRAY_LENGTH: i64 = (i32::MAX - 15) as i64; pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result { let child_field = match data_type { @@ -58,22 +58,29 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result Some(array.len()), + ColumnarValue::Scalar(_) => None, + }) + .unwrap_or(1); + for arg in args { + if let ColumnarValue::Array(array) = arg { + if array.len() != num_rows { + return internal_err!( + "Arguments has mixed length. Expected length: {num_rows}, found length: {}", + array.len() + ); + } + } + } let result = match child_field.data_type() { - DataType::Int8 => { - sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i8) - } - DataType::Int16 => { - sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i16) - } - DataType::Int32 => { - sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i32) - } - DataType::Int64 => { - sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v) - } + DataType::Int8 => sequence_integral::(args, num_rows, child_field), + DataType::Int16 => sequence_integral::(args, num_rows, child_field), + DataType::Int32 => sequence_integral::(args, num_rows, child_field), + DataType::Int64 => sequence_integral::(args, num_rows, child_field), other => exec_err!("spark_sequence does not support element type {other:?}"), }?; @@ -86,59 +93,114 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result( - start: &ArrayRef, - stop: &ArrayRef, - step: Option<&ArrayRef>, +trait Integral: ArrowPrimitiveType { + fn scalar(value: &ScalarValue) -> Result>; + fn from_i64(value: i64) -> Self::Native; +} +macro_rules! integral { + ($ty:ty, $variant:ident, $native:ty) => { + impl Integral for $ty { + fn scalar(value: &ScalarValue) -> Result> { + match value { + ScalarValue::$variant(value) => Ok(*value), + _ => exec_err!( + "spark_sequence expects {} arguments, got {}", + Self::DATA_TYPE, + value.data_type() + ), + } + } + fn from_i64(value: i64) -> Self::Native { + value as $native + } + } + }; +} +integral!(Int8Type, Int8, i8); +integral!(Int16Type, Int16, i16); +integral!(Int32Type, Int32, i32); +integral!(Int64Type, Int64, i64); + +enum Input<'a, T: Integral> { + Scalar(Option), + Array(&'a PrimitiveArray), +} +impl<'a, T: Integral> Input<'a, T> +where + T::Native: Into, +{ + fn new(value: &'a ColumnarValue) -> Result { + match value { + ColumnarValue::Scalar(value) => Ok(Self::Scalar(T::scalar(value)?)), + ColumnarValue::Array(array) => Ok(Self::Array(as_primitive_array::(array)?)), + } + } + #[inline(always)] + fn value(&self, row: usize) -> i64 { + match self { + Self::Scalar(value) => value.unwrap().into(), + Self::Array(array) => array.values()[row].into(), + } + } + fn nulls(&self, len: usize) -> Option { + match self { + Self::Scalar(None) => Some(NullBuffer::new_null(len)), + Self::Scalar(Some(_)) => None, + Self::Array(array) => array.nulls().cloned(), + } + } +} + +fn sequence_integral( + args: &[ColumnarValue], + num_rows: usize, child_field: FieldRef, - from_i64: impl Fn(i64) -> T::Native, ) -> Result where T::Native: Into, { - let start = as_primitive_array::(start)?; - let stop = as_primitive_array::(stop)?; - let step = step.map(|arr| as_primitive_array::(arr)).transpose()?; - let num_rows = start.len(); - - let row_is_null = |row: usize| { - start.is_null(row) || stop.is_null(row) || step.is_some_and(|arr| arr.is_null(row)) - }; - // With no explicit step, Spark uses `start <= stop ? 1 : -1` per row, so the direction - // always matches the bounds and the boundary check below cannot fail. + let start = Input::::new(&args[0])?; + let stop = Input::::new(&args[1])?; + let step = args.get(2).map(Input::::new).transpose()?; + // Reuse validity once per batch instead of checking three inputs and rebuilding a + // bitmap row by row. Arrow unions sliced bitmaps with bulk word operations. + let nulls = NullBuffer::union( + start.nulls(num_rows).as_ref(), + stop.nulls(num_rows).as_ref(), + ); + let step_nulls = step.as_ref().and_then(|step| step.nulls(num_rows)); + let nulls = NullBuffer::union(nulls.as_ref(), step_nulls.as_ref()) + .filter(|nulls| nulls.null_count() != 0); + if nulls + .as_ref() + .is_some_and(|nulls| nulls.null_count() == num_rows) + { + return Ok(Arc::new(ListArray::new_null(child_field, num_rows))); + } let row_step = |row: usize, start: i64, stop: i64| -> i64 { - match step { - Some(arr) => arr.value(row).into(), - None => { - if start <= stop { - 1 - } else { - -1 - } - } + match &step { + Some(step) => step.value(row), + None if start <= stop => 1, + None => -1, } }; - // First pass: compute per-row lengths so the child buffer can be reserved once for the - // whole batch. Valid rows always produce at least one element, so length 0 marks a null row. - let mut lengths: Vec = Vec::with_capacity(num_rows); - let mut total: usize = 0; + // Validate every row before checking the batch total, preserving the first Spark error. + // Retain lengths so short sequences do not repeat per-row arithmetic during generation. + let mut lengths = Vec::with_capacity(num_rows); + let mut total = 0usize; for row in 0..num_rows { - if row_is_null(row) { + if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row)) { lengths.push(0); continue; } - let s: i64 = start.value(row).into(); - let e: i64 = stop.value(row).into(); + let s = start.value(row); + let e = stop.value(row); let len = sequence_length(s, e, row_step(row, s, e))?; total += len; lengths.push(len); } - // Comet-specific ceiling: the sum of every row's length in one Arrow batch must fit in - // the i32 offset buffer. Spark has no equivalent guard because it stores each row as its - // own `long[]`, so the user may hit this on a query Spark itself would run. Report it via - // a dedicated error that names `spark.comet.batchSize` as the actionable knob rather than - // Spark's per-array size limit. + // Arrow's List offsets must fit i32; Spark's per-row array limit is checked above. if total > i32::MAX as usize { return Err(DataFusionError::External(Box::new( SparkError::SequenceBatchTooLarge { @@ -147,94 +209,99 @@ where ))); } - // Second pass: write elements straight into the child buffer and push offsets. The - // batch-total check above guarantees `values.len() <= i32::MAX` at every iteration, so - // the offset push cannot overflow. `try_reserve_exact` returns a query error on allocator - // failure so an oversized reservation cannot abort the executor. let mut values: Vec = Vec::new(); values.try_reserve_exact(total).map_err(|_| { DataFusionError::External(Box::new(SparkError::SequenceBatchTooLarge { total_elements: total.to_string(), })) })?; - let mut offsets: Vec = Vec::with_capacity(num_rows + 1); + let slots = &mut values.spare_capacity_mut()[..total]; + let mut offsets = Vec::with_capacity(num_rows + 1); offsets.push(0); - let mut nulls = NullBufferBuilder::new(num_rows); + let mut written = 0; for (row, &len) in lengths.iter().enumerate() { - if len == 0 { - nulls.append_null(); - } else { - nulls.append_non_null(); - let s: i64 = start.value(row).into(); - let e: i64 = stop.value(row).into(); + if len != 0 { + let s = start.value(row); + let e = stop.value(row); let step = row_step(row, s, e); - // Every element pushed lies between start and stop inclusive, so the widened - // arithmetic cannot overflow; only the final unused increment may wrap. - let mut v = s; - for _ in 0..len { - values.push(from_i64(v)); - v = v.wrapping_add(step); + // Fixed bounds and independent arithmetic let LLVM vectorise the stores without + // a capacity check for each element. Wrapping intermediates preserve i64 extremes. + for (i, slot) in slots[written..written + len].iter_mut().enumerate() { + slot.write(T::from_i64(s.wrapping_add(step.wrapping_mul(i as i64)))); } + written += len; } - offsets.push(values.len() as i32); + offsets.push(written as i32); } - + // SAFETY: the sizing pass reserved `total` slots and the loop initialised each exactly + // once. No fallible operation exposes the partially initialised vector. + unsafe { values.set_len(total) }; let values = PrimitiveArray::::new(ScalarBuffer::from(values), None); - let list = ListArray::try_new( + Ok(Arc::new(ListArray::try_new( child_field, OffsetBuffer::new(offsets.into()), Arc::new(values), - nulls.finish(), - )?; - Ok(Arc::new(list)) + nulls, + )?)) } -/// Number of elements of `sequence(start, stop, step)`, matching Spark's boundary check and -/// `Sequence.sequenceLength` (byte-identical from Spark 3.4.3 through 4.1.1), including which -/// of the three failure paths fires and the exact length value each of them reports. +/// Match released Spark 3.4--4.1 Sequence.sequenceLength, including its wide-arithmetic +/// fallback error when subtraction overflows even though the resulting length would fit. +#[inline(always)] fn sequence_length(start: i64, stop: i64, step: i64) -> Result { if !((step > 0 && start <= stop) || (step < 0 && start >= stop) || (step == 0 && start == stop)) { - return Err(DataFusionError::External(Box::new( - SparkError::SequenceIllegalBoundaries { - start: start.to_string(), - stop: stop.to_string(), - step: step.to_string(), - }, - ))); + return Err(boundary_error(start, stop, step)); } if stop == start { return Ok(1); } - // Spark computes stop - start with Math.subtractExact and special-cases - // Long.MinValue / -1; both raise ArithmeticException, which reroutes the length through a - // BigInt fallback. i128 arithmetic gives the same exact value on every path. - // - // Max delta magnitude is |i64::MAX - i64::MIN| = 2 * i64::MAX + 1, which fits in i128 but - // overflows i64, so the `> i64::MAX` / `< i64::MIN` predicates below are both live. - let delta = stop as i128 - start as i128; - let overflowed = delta > i64::MAX as i128 - || delta < i64::MIN as i128 - || (delta == i64::MIN as i128 && step == -1); - let len = 1 + delta / step as i128; - if len > MAX_ROUNDED_ARRAY_LENGTH { - return Err(DataFusionError::External(Box::new( - SparkError::CollectionSizeLimitExceeded { - num_elements: len.to_string(), - max_elements: MAX_ROUNDED_ARRAY_LENGTH as i64, - function_name: "sequence".to_string(), - }, - ))); + if let Some(len) = stop + .checked_sub(start) + .and_then(|delta| match step { + 1 => Some(delta), + -1 => delta.checked_neg(), + _ => delta.checked_div(step), + }) + .and_then(|quotient| quotient.checked_add(1)) + { + if len <= MAX_ROUNDED_ARRAY_LENGTH { + return Ok(len as usize); + } + return Err(length_error(len as i128)); } - if overflowed { - // Spark's BigInt fallback lands on `internalError("Unreachable code reached.")` when - // the exact length is within the limit (e.g. sequence spanning more than Long.MaxValue - // with a large step); reproduce it for error parity. - return Err(DataFusionError::External(Box::new(SparkError::Internal( - "Unreachable code reached.".to_string(), - )))); + sequence_length_overflow(start, stop, step) +} + +#[cold] +fn boundary_error(start: i64, stop: i64, step: i64) -> DataFusionError { + DataFusionError::External(Box::new(SparkError::SequenceIllegalBoundaries { + start: start.to_string(), + stop: stop.to_string(), + step: step.to_string(), + })) +} + +#[cold] +fn sequence_length_overflow(start: i64, stop: i64, step: i64) -> Result { + // Only the rare overflow/error path needs wide division. There is no separate add-overflow + // branch: its exact length necessarily exceeds Spark's maximum and is rejected here. + let len = 1 + (stop as i128 - start as i128) / step as i128; + if len > MAX_ROUNDED_ARRAY_LENGTH as i128 { + return Err(length_error(len)); } - Ok(len as usize) + Err(DataFusionError::External(Box::new(SparkError::Internal( + "Unreachable code reached.".to_string(), + )))) +} + +#[cold] +fn length_error(len: i128) -> DataFusionError { + DataFusionError::External(Box::new(SparkError::CollectionSizeLimitExceeded { + num_elements: len.to_string(), + max_elements: MAX_ROUNDED_ARRAY_LENGTH, + function_name: "sequence".to_string(), + })) } #[cfg(test)] @@ -390,4 +457,205 @@ mod tests { }; assert_eq!(row_values(&list, 0), vec![1, 2, 3]); } + #[test] + fn scalar_and_array_combinations_preserve_each_integral_type() { + macro_rules! verify { + ($ty:ty, $variant:ident, $native:ty) => {{ + for (start, stop, step) in [ + (<$native>::MIN, <$native>::MIN + 3, 1), + (<$native>::MAX, <$native>::MAX - 3, -1), + (0, 0, 0), + (-3, 4, 2), + (<$native>::MIN, -1, <$native>::MAX), + (0, <$native>::MIN, <$native>::MIN), + ] { + for mask in 0..8 { + let args = [start, stop, step] + .into_iter() + .enumerate() + .map(|(index, value)| { + let scalar = ScalarValue::$variant(Some(value)); + if mask & (1 << index) == 0 { + ColumnarValue::Scalar(scalar) + } else { + ColumnarValue::Array(scalar.to_array_of_size(3).unwrap()) + } + }) + .collect::>(); + let result = spark_sequence(&args, &list_of(<$ty>::DATA_TYPE)).unwrap(); + assert_eq!(matches!(result, ColumnarValue::Scalar(_)), mask == 0); + let array = result.into_array(3).unwrap(); + let list = as_primitive_list(&array); + let expected_len = if start == stop { + 1 + } else { + (1 + (stop as i128 - start as i128) / step as i128) as usize + }; + for row in 0..3 { + let values = list.value(row); + let values = as_primitive_array::<$ty>(&values).unwrap(); + assert_eq!(values.len(), expected_len); + for (i, actual) in values.values().iter().enumerate() { + assert_eq!( + *actual as i128, + start as i128 + step as i128 * i as i128 + ); + } + } + } + } + }}; + } + verify!(Int8Type, Int8, i8); + verify!(Int16Type, Int16, i16); + verify!(Int32Type, Int32, i32); + verify!(Int64Type, Int64, i64); + } + + #[test] + fn sliced_null_inputs_skip_invalid_underlying_values() { + use arrow::buffer::NullBuffer; + // Slice at a non-byte-aligned offset. The null start hides invalid boundaries; + // the null stop hides a length overflow; the null step hides zero-step boundaries. + let arrays = [ + Int64Array::new( + vec![99, 5, i64::MIN, 1, 7, 99].into(), + Some(NullBuffer::from(vec![true, false, true, true, true, true])), + ), + Int64Array::new( + vec![99, 1, i64::MAX, 5, 7, 99].into(), + Some(NullBuffer::from(vec![true, true, false, true, true, true])), + ), + Int64Array::new( + vec![99, 1, 1, 0, 0, 99].into(), + Some(NullBuffer::from(vec![true, true, true, false, true, true])), + ), + ]; + let args = arrays + .into_iter() + .map(|array| ColumnarValue::Array(Arc::new(array.slice(1, 4)))) + .collect::>(); + let result = spark_sequence(&args, &list_of(DataType::Int64)) + .unwrap() + .into_array(4) + .unwrap(); + let list = as_primitive_list(&result); + assert_eq!( + list.nulls().unwrap(), + &NullBuffer::from(vec![false, false, false, true]) + ); + assert_eq!(list.value_offsets(), &[0, 0, 0, 0, 1]); + assert_eq!(row_values(&list, 3), vec![7]); + assert_eq!(list.values().null_count(), 0); + + for index in 0..3 { + let mut args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(5))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(-1))), + ]; + args[index] = ColumnarValue::Scalar(ScalarValue::Int64(None)); + let result = spark_sequence(&args, &list_of(DataType::Int64)).unwrap(); + assert!( + matches!(result, ColumnarValue::Scalar(ScalarValue::List(ref list)) if list.is_null(0)) + ); + } + } + + #[test] + fn empty_and_mismatched_arrays() { + let empty = ColumnarValue::Array(Arc::new(Int64Array::from(Vec::::new()))); + let scalar = ColumnarValue::Scalar(ScalarValue::Int64(Some(1))); + let result = spark_sequence(&[empty.clone(), scalar.clone()], &list_of(DataType::Int64)) + .unwrap() + .into_array(0) + .unwrap(); + let list = as_primitive_list(&result); + assert!(list.is_empty()); + assert_eq!(list.value_offsets(), &[0]); + assert!(list.values().is_empty()); + let one = ColumnarValue::Array(Arc::new(Int64Array::from(vec![1]))); + assert!(spark_sequence(&[empty, scalar, one], &list_of(DataType::Int64)).is_err()); + } + + #[test] + fn row_errors_precede_aggregate_limit() { + let error = run_i64( + vec![Some(0), Some(0), Some(5)], + vec![Some(1_073_741_823), Some(1_073_741_823), Some(1)], + Some(vec![Some(1), Some(1), Some(1)]), + ) + .unwrap_err() + .to_string(); + assert!( + error.contains("Illegal sequence boundaries: 5 to 1 by 1"), + "{error}" + ); + let error = run_i64( + vec![Some(0), Some(0)], + vec![Some(1_073_741_823), Some(1_073_741_823)], + None, + ) + .unwrap_err() + .to_string(); + assert!( + error.contains("2147483648") && error.contains("spark.comet.batchSize"), + "{error}" + ); + } + + #[test] + fn checked_length_matches_exact_released_spark_arithmetic() { + let values = [ + i64::MIN, + i64::MIN + 1, + -(i32::MAX as i64), + -1, + 0, + 1, + i32::MAX as i64, + i64::MAX - 1, + i64::MAX, + ]; + for start in values { + for stop in values { + for step in values { + let result = sequence_length(start, stop, step); + let legal = (step > 0 && start <= stop) + || (step < 0 && start >= stop) + || (step == 0 && start == stop); + if !legal { + assert!(result + .unwrap_err() + .to_string() + .contains("Illegal sequence boundaries")); + } else if start == stop { + assert_eq!(result.unwrap(), 1); + } else { + let delta = stop as i128 - start as i128; + let len = 1 + delta / step as i128; + if len > MAX_ROUNDED_ARRAY_LENGTH as i128 { + let error = result.unwrap_err().to_string(); + assert!( + error.contains(&len.to_string()), + "{start}, {stop}, {step}: {error}" + ); + } else if delta < i64::MIN as i128 || delta > i64::MAX as i128 { + assert!(result + .unwrap_err() + .to_string() + .contains("Unreachable code reached.")); + } else { + assert_eq!(result.unwrap(), len as usize); + } + } + } + } + } + assert_eq!( + sequence_length(0, MAX_ROUNDED_ARRAY_LENGTH - 1, 1).unwrap(), + MAX_ROUNDED_ARRAY_LENGTH as usize + ); + assert!(sequence_length(0, MAX_ROUNDED_ARRAY_LENGTH, 1).is_err()); + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala index 1b559ee6ddc..0674890f1a9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala @@ -37,9 +37,12 @@ package org.apache.spark.sql.benchmark object CometSequenceBenchmark extends CometBenchmarkBase { private val sequenceQueries = List( + ("seq_short_2_elems", "SELECT sequence(c_start, c_stop_2) FROM parquetV1Table"), ("seq_short_5_elems", "SELECT sequence(c_start, c_stop_5) FROM parquetV1Table"), ("seq_spine_365_elems", "SELECT sequence(c_start, c_stop_365) FROM parquetV1Table"), + ("seq_long_1000_elems", "SELECT sequence(c_start, c_stop_1000) FROM parquetV1Table"), ("seq_long_10000_elems", "SELECT sequence(c_start, c_stop_10000) FROM parquetV1Table"), + ("seq_long_50000_elems", "SELECT sequence(c_start, c_stop_50000) FROM parquetV1Table"), ("seq_descending_default_step", "SELECT sequence(c_stop_365, c_start) FROM parquetV1Table"), ("seq_explicit_step_7", "SELECT sequence(c_start, c_stop_365, 7L) FROM parquetV1Table"), ( @@ -59,9 +62,12 @@ object CometSequenceBenchmark extends CometBenchmarkBase { prepareTable( dir, spark.sql("SELECT CAST(PMOD(value, 100000) AS BIGINT) AS c_start," + + " CAST(PMOD(value, 100000) AS BIGINT) + 1 AS c_stop_2," + " CAST(PMOD(value, 100000) AS BIGINT) + 4 AS c_stop_5," + " CAST(PMOD(value, 100000) AS BIGINT) + 364 AS c_stop_365," + + " CAST(PMOD(value, 100000) AS BIGINT) + 999 AS c_stop_1000," + " CAST(PMOD(value, 100000) AS BIGINT) + 9999 AS c_stop_10000," + + " CAST(PMOD(value, 100000) AS BIGINT) + 49999 AS c_stop_50000," + " CASE WHEN PMOD(value, 10) = 0 THEN CAST(NULL AS BIGINT)" + " ELSE CAST(PMOD(value, 100000) AS BIGINT) END AS c_null_start," + " CASE WHEN PMOD(value, 10) = 0 THEN CAST(NULL AS BIGINT)" + From 19b0dd264cd39c2d34a7bffa98694d44f48c8163 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 11 Oct 2026 11:35:27 +0000 Subject: [PATCH 59/89] feat: filter decoded probe batches before intermediate joins (#6433) * feat: filter decoded probe batches before intermediate joins * fix: address early join runtime filter review feedback --- docs/source/user-guide/latest/metrics.md | 23 +- .../user-guide/latest/tuning/operators.md | 25 ++ .../operators/dynamic_filter/early.rs | 142 +++++++ .../operators/dynamic_filter/join.rs | 36 +- .../operators/dynamic_filter/join/tests.rs | 394 ++++++++++++++++++ .../execution/operators/dynamic_filter/mod.rs | 62 ++- .../dynamic_filter/parquet_reader.rs | 8 + .../operators/dynamic_filter/tests.rs | 120 ++++++ native/core/src/execution/planner.rs | 42 +- native/operators/src/lib.rs | 2 + native/operators/src/projection.rs | 295 +++++++++++++ .../spark/sql/comet/CometMetricNode.scala | 10 +- .../sql-tests/join/dynamic_filter_chain.sql | 72 ++++ .../apache/comet/exec/CometJoinSuite.scala | 104 ++++- 14 files changed, 1295 insertions(+), 40 deletions(-) create mode 100644 native/core/src/execution/operators/dynamic_filter/early.rs create mode 100644 native/operators/src/projection.rs create mode 100644 spark/src/test/resources/sql-tests/join/dynamic_filter_chain.sql diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index 56491281db7..75d9d1c2820 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -93,14 +93,21 @@ With `spark.comet.exec.join.dynamicFilter.enabled=true`, native broadcast and sh report these additional metric keys. See [Join Runtime Filters](tuning/operators.md#join-runtime-filters) for eligibility and reader restrictions. -| Metric | Description | -| -------------------------------------- | ----------------------------------------------------------------- | -| `dynamic_filter_join_rows_evaluated` | Probe rows evaluated by the runtime filter. | -| `dynamic_filter_join_rows_pruned` | Probe rows rejected by that filter before the hash probe. | -| `dynamic_filter_join_rows_bypassed` | Probe rows passed through while the runtime filter is inactive. | -| `dynamic_filter_join_eval_time` | Time evaluating the runtime filter. | -| `dynamic_filter_join_filters_attached` | Executions that attach their runtime filter to a native reader. | -| `dynamic_filter_join_filters_skipped` | Executions whose probe input is ineligible for reader attachment. | +| Metric | Description | +| -------------------------------------- | ------------------------------------------------------------------------------------------------------------- | +| `dynamic_filter_join_rows_evaluated` | Probe rows evaluated by the runtime filter. | +| `dynamic_filter_join_rows_pruned` | Probe rows rejected by that filter before the hash probe. | +| `dynamic_filter_join_rows_bypassed` | Probe rows passed through while the runtime filter is inactive. | +| `dynamic_filter_join_eval_time` | Time evaluating the runtime filter. | +| `dynamic_filter_join_filters_attached` | Executions that attach their runtime filter to a native reader. | +| `dynamic_filter_join_filters_skipped` | Executions whose probe input is ineligible for reader attachment. | +| `dynamic_filter_early_rows_evaluated` | Rows evaluated before an intermediate join on this join's probe path. | +| `dynamic_filter_early_rows_pruned` | Rows rejected there before the intermediate join does its probe work. | +| `dynamic_filter_early_rows_bypassed` | Rows passed through while early filtering is inactive or has stopped after two nonempty all-matching batches. | +| `dynamic_filter_early_eval_time` | Time evaluating that early filter. | + +See [early filtering in join chains](tuning/operators.md#early-filtering-in-join-chains) +for placement and adaptive bypass behavior. The row counters measure residual filtering of decoded probe batches. They exclude rows skipped by the reader. An attached filter does not guarantee that any row groups are pruned: compare the diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index cb538fb63d4..94bb76d5adb 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -67,6 +67,31 @@ boundaries. A shuffled hash join can still filter probe batches after shuffle, b its filter back to an earlier scan stage. Compare the [runtime-filter and scan metrics](../metrics.md#hash-joins) with the setting disabled to distinguish reduced hash-probe work from reader I/O savings. +#### Early filtering in join chains + +With `spark.comet.exec.join.dynamicFilter.enabled=true`, an eligible join can also +use its completed build keys to filter decoded rows before an intermediate inner +join. For example, when a final join accepts only customer 42, it can remove other +customers before an earlier join looks up their orders. This reduces intermediate +join work while the final join still checks every match and preserves duplicates. + +Placement follows direct probe-side columns through ordinary inner joins already +eligible for runtime filtering (a single signed-integer key), +column-only projections, and direct column null checks above the first intermediate join. +After crossing a join, the early consumer stays above inferred null checks so null +keys do not prevent adaptive bypass. It stops at computed +expressions, other predicates, limits, unsupported joins, and execution boundaries. +These early filters operate on decoded batches; they do not introduce additional +Parquet pruning or skip schema conversions. Each intermediate join retains its +existing reader-filter checks. + +Early filtering stops evaluating after two consecutive nonempty evaluated batches +remove no rows, avoiding repeated membership checks on unselective inputs. That +decision is local to one execution. The final join remains responsible for matching +rows, including if later batches become more selective. The `dynamic_filter_early_*` +metrics belong to the join supplying the filter; intermediate joins and projections +retain their own input/output metrics. + ## Adaptive Partial Aggregation Set `spark.comet.exec.aggregate.skipPartial.enabled=true` to let Comet bypass partial hash diff --git a/native/core/src/execution/operators/dynamic_filter/early.rs b/native/core/src/execution/operators/dynamic_filter/early.rs new file mode 100644 index 00000000000..6f23c54d3ca --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter/early.rs @@ -0,0 +1,142 @@ +// 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. + +//! Place a live ancestor filter before an intermediate join's probe work. +//! +//! This is decoded-batch filtering only. It leaves scan/schema conversion and +//! arbitrary expressions in place, and never propagates into an intermediate +//! build side. The downstream join remains the authority for matching rows. + +use std::any::Any; +use std::sync::Arc; + +use datafusion::common::{internal_err, Result}; +use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_comet_operators::CometFilterExec; + +use super::parquet_reader::is_direct_column_null_checks; +use super::{DynamicFilterExec, DynamicFilterJoinExec}; +use datafusion_comet_operators::CometProjectionExec; + +pub(super) fn place_early_filter( + input: &Arc, + predicate: Arc, + metrics: &ExecutionPlanMetricsSet, +) -> Result> { + Ok(place(input, predicate, metrics, false)?.unwrap_or_else(|| Arc::clone(input))) +} + +fn remap( + predicate: Arc, + column: Arc, +) -> Result> { + // Derived expressions share producer updates. Taking current() here would + // capture the initial TRUE placeholder instead of the completed build domain. + let mapped: Arc = predicate.with_new_children(vec![column])?; + mapped.downcast::().map_err(|_| { + datafusion::common::DataFusionError::Internal( + "Dynamic filter remapping changed type".into(), + ) + }) +} + +fn place( + input: &Arc, + predicate: Arc, + metrics: &ExecutionPlanMetricsSet, + crossed_join: bool, +) -> Result>> { + let children = predicate.children(); + let [key] = children.as_slice() else { + return internal_err!("Early join filtering requires one key"); + }; + let Some(key) = key.downcast_ref::() else { + return internal_err!("Early join filtering requires a column key"); + }; + if input.fetch().is_none() { + if let Some(join) = input.downcast_ref::() { + let template = join.template(); + // This wrapper is already limited to ordinary inner, single-key joins. + // Do not skip fallible join residuals, or guess an embedded projection. + if template.filter().is_none() && !template.contains_projection() { + let build_columns = template.left().schema().fields().len(); + if let Some(index) = key.index().checked_sub(build_columns) { + if let Some(field) = template.right().schema().fields().get(index) { + let mapped = remap( + Arc::clone(&predicate), + Arc::new(Column::new(field.name(), index)), + )?; + if let Some(probe) = place(template.right(), mapped, metrics, true)? { + return Ok(Some(join.with_execution_probe(probe)?)); + } + } + } + } + } else if let Some(projection) = input.downcast_ref::() { + let exprs = projection.projection().expr(); + // Even an unrelated computed expression may fail or be stateful. + if exprs.iter().all(|expr| expr.expr.is::()) { + if let Some(expr) = exprs.get(key.index()) { + let mapped = remap(Arc::clone(&predicate), Arc::clone(&expr.expr))?; + if let Some(child) = place(projection.input(), mapped, metrics, crossed_join)? { + return Ok(Some(projection.with_execution_input(child)?)); + } + } + } + } else if let Some(filter) = input.downcast_ref::() { + // Once below a join, leave the terminal consumer above Spark's + // inferred null checks. Rejected nulls otherwise make an unselective + // ancestor look selective and prevent adaptive bypass. Early consumers + // do not prune readers, so descending farther offers no I/O benefit. + if !crossed_join + && !filter.has_projection() + && is_direct_column_null_checks(filter.predicate()) + { + if let Some(child) = place( + filter.input(), + Arc::clone(&predicate), + metrics, + crossed_join, + )? { + return Ok(Some(filter.with_execution_input(child)?)); + } + } + } else if let Some(filter) = input.downcast_ref::() { + if let Some(child) = + place(&filter.input, Arc::clone(&predicate), metrics, crossed_join)? + { + return Ok(Some(filter.with_execution_input(child))); + } + } + } + // A terminal consumer is useful only after crossing a join. Above other + // boundaries the existing final consumer already performs the same work. + Ok(crossed_join.then(|| { + Arc::new( + DynamicFilterExec::new( + Arc::clone(input), + predicate, + metrics.clone(), + "dynamic_filter_early", + ) + .adaptive(), + ) as Arc + })) +} diff --git a/native/core/src/execution/operators/dynamic_filter/join.rs b/native/core/src/execution/operators/dynamic_filter/join.rs index ba2e2759608..c2688781b13 100644 --- a/native/core/src/execution/operators/dynamic_filter/join.rs +++ b/native/core/src/execution/operators/dynamic_filter/join.rs @@ -45,6 +45,7 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; +use super::early::place_early_filter; use super::parquet_reader::try_attach_parquet_reader_filter; use super::DynamicFilterExec; @@ -84,19 +85,38 @@ impl DynamicFilterJoinExec { }) } + pub(super) fn template(&self) -> &HashJoinExec { + &self.template + } + + pub(super) fn with_execution_probe( + &self, + probe: Arc, + ) -> Result> { + Ok(Arc::new(Self { + template: self + .template + .builder() + .reset_state() + .with_new_children(vec![Arc::clone(self.template.left()), probe])? + .build()?, + config: self.config.clone(), + metrics: self.metrics.clone(), + })) + } + fn build_runtime_join(&self) -> Result { let predicate = Arc::new(DynamicFilterPhysicalExpr::new( vec![Arc::clone(&self.template.on()[0].1)], lit(true), )); - let reader = try_attach_parquet_reader_filter( - self.template.right(), - Arc::clone(&predicate), - &self.config, - )?; + let probe = + place_early_filter(self.template.right(), Arc::clone(&predicate), &self.metrics)?; + let reader = + try_attach_parquet_reader_filter(&probe, Arc::clone(&predicate), &self.config)?; let reader_filter_attached = reader.is_some(); let consumer = Arc::new(DynamicFilterExec::new( - reader.unwrap_or_else(|| Arc::clone(self.template.right())), + reader.unwrap_or(probe), Arc::clone(&predicate), self.metrics.clone(), "dynamic_filter_join", @@ -192,6 +212,10 @@ impl ExecutionPlan for DynamicFilterJoinExec { self.template.maintains_input_order() } + fn fetch(&self) -> Option { + self.template.fetch() + } + fn children(&self) -> Vec<&Arc> { self.template.children() } diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index 14242b4e4fd..fd2248a053f 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs @@ -1564,3 +1564,397 @@ fn wrapper_preserves_join_statistics_and_distribution() { ); } } + +fn chain_join( + build: Arc, + probe: Arc, + build_key: usize, + probe_key: usize, +) -> HashJoinExec { + let left = Arc::new(Column::new( + build.schema().field(build_key).name(), + build_key, + )); + let right = Arc::new(Column::new( + probe.schema().field(probe_key).name(), + probe_key, + )); + HashJoinExec::try_new( + build, + probe, + vec![(left, right)], + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap() +} + +#[tokio::test] +async fn ancestor_filter_prunes_before_intermediate_join_and_keeps_metric_owners() { + use datafusion_comet_operators::CometProjectionExec; + let ctx = SessionContext::new(); + let fact = input( + vec![Some(0), Some(1), Some(2), Some(2), Some(3), None], + &DataType::Int32, + 0, + ); + let dimension = input( + vec![Some(0), Some(1), Some(2), Some(2), Some(3)], + &DataType::Int32, + 0, + ); + let ancestor_build = input(vec![Some(2), Some(2)], &DataType::Int32, 0); + let inner = chain_join(dimension, fact, 0, 0); + let inner = Arc::new( + DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ); + let projection = Arc::new(CometProjectionExec::from_datafusion( + ProjectionExec::try_new( + vec![ + ( + Arc::new(Column::new("payload", 3)) as Arc, + "payload".into(), + ), + ( + Arc::new(Column::new("key", 2)) as Arc, + "key".into(), + ), + ], + Arc::clone(&inner) as Arc, + ) + .unwrap(), + )); + let outer = chain_join( + ancestor_build, + Arc::clone(&projection) as Arc, + 0, + 1, + ); + let expected = collect( + Arc::new(outer.builder().reset_state().build().unwrap()), + ctx.task_ctx(), + ) + .await + .unwrap(); + assert_eq!(row_count(&expected), 8); + // Start fresh metric owners after executing the ordinary reference plan. + let before_inner = inner.metrics().unwrap().output_rows().unwrap_or(0); + let before_projection = projection.metrics().unwrap().output_rows().unwrap_or(0); + let outer = Arc::new( + DynamicFilterJoinExec::try_new(&outer, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ) as Arc; + for execution in 1..=2 { + let actual = collect(Arc::clone(&outer), ctx.task_ctx()).await.unwrap(); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + assert_eq!( + metric(&outer, "dynamic_filter_early_rows_evaluated"), + 6 * execution + ); + assert_eq!( + metric(&outer, "dynamic_filter_early_rows_pruned"), + 4 * execution + ); + assert_eq!( + inner.metrics().unwrap().output_rows().unwrap(), + before_inner + 4 * execution + ); + assert_eq!( + projection.metrics().unwrap().output_rows().unwrap(), + before_projection + 4 * execution + ); + } +} + +#[tokio::test] +async fn early_filter_preserves_probe_limit_selection() { + use datafusion::physical_plan::limit::LocalLimitExec; + let ctx = SessionContext::new(); + let fact = Arc::new(LocalLimitExec::new( + input(vec![Some(0), Some(1), Some(2)], &DataType::Int32, 0), + 2, + )); + let dimension = input(vec![Some(0), Some(1), Some(2)], &DataType::Int32, 0); + let inner = chain_join(dimension, fact, 0, 0); + let inner = Arc::new( + DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ); + let outer = chain_join(input(vec![Some(2)], &DataType::Int32, 0), inner, 0, 2); + let expected = collect( + Arc::new(outer.builder().reset_state().build().unwrap()), + ctx.task_ctx(), + ) + .await + .unwrap(); + let outer = Arc::new( + DynamicFilterJoinExec::try_new(&outer, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ) as Arc; + let actual = collect(Arc::clone(&outer), ctx.task_ctx()).await.unwrap(); + assert_eq!(row_count(&expected), 0); + assert_eq!(row_count(&actual), 0); + // The early consumer sees exactly the two already-selected rows. + assert_eq!(metric(&outer, "dynamic_filter_early_rows_evaluated"), 2); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_pruned"), 2); +} + +#[tokio::test] +async fn early_filter_does_not_hide_computed_projection_errors() { + use datafusion_comet_operators::CometProjectionExec; + let ctx = SessionContext::new(); + let inner = chain_join( + input(vec![Some(0), Some(2)], &DataType::Int32, 0), + input(vec![Some(0), Some(2)], &DataType::Int32, 0), + 0, + 0, + ); + let inner = Arc::new( + DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ); + let projection = Arc::new(CometProjectionExec::from_datafusion( + ProjectionExec::try_new( + vec![ + ( + Arc::new(BinaryExpr::new( + lit(10i32), + Operator::Divide, + Arc::new(Column::new("key", 2)), + )) as Arc, + "quotient".into(), + ), + ( + Arc::new(Column::new("key", 2)) as Arc, + "key".into(), + ), + ], + inner, + ) + .unwrap(), + )); + let outer = chain_join(input(vec![Some(2)], &DataType::Int32, 0), projection, 0, 1); + for enabled in [false, true] { + let plan = PhysicalPlanner::apply_join_dynamic_filter( + Arc::new(outer.builder().reset_state().build().unwrap()), + enabled, + ctx.copied_config().options(), + ) + .unwrap(); + let error = collect(plan, ctx.task_ctx()).await.unwrap_err(); + assert!(error.to_string().to_lowercase().contains("zero"), "{error}"); + } +} + +#[test] +fn early_filter_stops_at_build_columns_and_join_residuals() { + use super::super::early::place_early_filter; + use datafusion::physical_plan::joins::utils::JoinFilter; + let config = ConfigOptions::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let inner = plain_join(); + let wrapped = Arc::new( + DynamicFilterJoinExec::try_new(&inner, &config) + .unwrap() + .unwrap(), + ) as Arc; + // Column 1 belongs to this intermediate join's build; do not filter its probe. + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("key", 1))], + lit(true), + )); + assert!(Arc::ptr_eq( + &wrapped, + &place_early_filter(&wrapped, predicate, &metrics).unwrap() + )); + let filtered = inner + .builder() + .with_filter(Some(JoinFilter::new( + lit(true), + vec![], + Arc::new(Schema::empty()), + ))) + .build() + .unwrap(); + let wrapped = Arc::new( + DynamicFilterJoinExec::try_new(&filtered, &config) + .unwrap() + .unwrap(), + ) as Arc; + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("key", 2))], + lit(true), + )); + assert!(Arc::ptr_eq( + &wrapped, + &place_early_filter(&wrapped, predicate, &metrics).unwrap() + )); +} + +#[tokio::test] +async fn early_consumers_release_ancestor_domains_when_cancelled() { + for after_output in [false, true] { + let (session, pool) = limited_session(8 * 1024 * 1024); + let inner = chain_join( + input((0..4096).map(Some).collect(), &DataType::Int32, 0), + input((0..4096).map(Some).collect(), &DataType::Int32, 0), + 0, + 0, + ); + let inner = Arc::new(DynamicFilterJoinExec::new(&inner, ConfigOptions::default()).unwrap()); + let outer = chain_join( + input(vec![Some(3)], &DataType::Int32, 0), + Arc::clone(&inner) as Arc, + 0, + 2, + ); + let outer = DynamicFilterJoinExec::new(&outer, ConfigOptions::default()).unwrap(); + let runtime = outer.build_runtime_join().unwrap(); + let predicate = Arc::downgrade(&produced_join_filter(&runtime.join)); + let final_consumer = runtime + .join + .right() + .downcast_ref::() + .unwrap(); + let intermediate = final_consumer + .input + .downcast_ref::() + .unwrap(); + let early_consumer = intermediate + .template + .right() + .downcast_ref::() + .unwrap(); + let early_predicate = Arc::downgrade(&early_consumer.predicate); + let mut stream = outer + .execute_runtime_join(runtime.join, 0, session.task_ctx()) + .unwrap(); + assert!(predicate.upgrade().is_some()); + if after_output { + loop { + if stream.next().await.unwrap().unwrap().num_rows() > 0 { + break; + } + } + assert!(pool.reserved() > 0); + assert!( + outer + .metrics() + .unwrap() + .sum_by_name("dynamic_filter_early_rows_pruned") + .unwrap() + .as_usize() + > 0 + ); + } + drop(stream); + assert_eq!(pool.reserved(), 0); + assert!(predicate.upgrade().is_none()); + assert!(early_predicate.upgrade().is_none()); + assert!(inner.template.dynamic_expressions_produced().is_empty()); + assert!(outer.template.dynamic_expressions_produced().is_empty()); + } +} + +#[tokio::test] +async fn nullable_unselective_ancestor_bypasses_above_inferred_null_checks() { + let ctx = SessionContext::new(); + // Every decoded batch has a null FK, as in nullable fact-table keys. Null + // rejection must belong to Spark's inferred filter, not the adaptation sample. + let fact = input( + vec![ + Some(0), + None, + Some(1), + None, + Some(2), + None, + Some(3), + None, + Some(4), + None, + ], + &DataType::Int32, + 0, + ); + let checked = Arc::new(CometFilterExec::from_datafusion( + FilterExec::try_new( + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))), + fact, + ) + .unwrap() + .with_batch_size(1) + .unwrap(), + )); + let inner = chain_join( + input((0..5).map(Some).collect(), &DataType::Int32, 0), + checked, + 0, + 0, + ); + let inner = Arc::new( + DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ); + let outer = chain_join( + input((0..5).map(Some).collect(), &DataType::Int32, 0), + inner, + 0, + 2, + ); + let expected = collect( + Arc::new(outer.builder().reset_state().build().unwrap()), + ctx.task_ctx(), + ) + .await + .unwrap(); + let outer = Arc::new( + DynamicFilterJoinExec::try_new(&outer, ctx.copied_config().options()) + .unwrap() + .unwrap(), + ) as Arc; + let actual = collect(Arc::clone(&outer), ctx.task_ctx()).await.unwrap(); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + assert_eq!(row_count(&actual), 5); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_evaluated"), 2); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_pruned"), 0); + assert_eq!(metric(&outer, "dynamic_filter_early_rows_bypassed"), 3); +} + +#[test] +fn early_filter_stops_at_intermediate_join_fetch() { + use super::super::early::place_early_filter; + let config = ConfigOptions::default(); + let limited = plain_join().with_fetch(Some(1)).unwrap(); + let wrapped = Arc::new( + DynamicFilterJoinExec::try_new(limited.downcast_ref::().unwrap(), &config) + .unwrap() + .unwrap(), + ) as Arc; + assert_eq!(wrapped.fetch(), Some(1)); + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("key", 2))], + lit(true), + )); + assert!(Arc::ptr_eq( + &wrapped, + &place_early_filter(&wrapped, predicate, &ExecutionPlanMetricsSet::new()).unwrap() + )); +} diff --git a/native/core/src/execution/operators/dynamic_filter/mod.rs b/native/core/src/execution/operators/dynamic_filter/mod.rs index 3dc33ce15d4..ab6d329c66d 100644 --- a/native/core/src/execution/operators/dynamic_filter/mod.rs +++ b/native/core/src/execution/operators/dynamic_filter/mod.rs @@ -17,6 +17,7 @@ //! Runtime-filter wiring and shared filtering of decoded batches. +mod early; mod join; mod parquet_reader; mod topk; @@ -45,12 +46,13 @@ use datafusion::physical_plan::{ use futures::StreamExt; /// A task-local consumer of a live runtime predicate. -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) struct DynamicFilterExec { input: Arc, predicate: Arc, metrics: ExecutionPlanMetricsSet, metric_prefix: &'static str, + adaptive: bool, } impl DynamicFilterExec { @@ -65,8 +67,21 @@ impl DynamicFilterExec { predicate, metrics, metric_prefix, + adaptive: false, } } + + fn adaptive(mut self) -> Self { + self.adaptive = true; + self + } + + fn with_execution_input(&self, input: Arc) -> Arc { + Arc::new(Self { + input, + ..self.clone() + }) + } } impl DisplayAs for DynamicFilterExec { @@ -122,12 +137,11 @@ impl ExecutionPlan for DynamicFilterExec { if children.len() != 1 { return internal_err!("CometDynamicFilterExec requires one child"); } - Ok(Arc::new(Self::new( - children.remove(0), - Arc::clone(&self.predicate), - ExecutionPlanMetricsSet::new(), - self.metric_prefix, - ))) + Ok(Arc::new(Self { + input: children.remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..self.as_ref().clone() + })) } fn reset_state(self: Arc) -> Result> { @@ -138,12 +152,11 @@ impl ExecutionPlan for DynamicFilterExec { self.predicate.children().into_iter().cloned().collect(), lit(true), )); - Ok(Arc::new(Self::new( - Arc::clone(&self.input), + Ok(Arc::new(Self { predicate, - ExecutionPlanMetricsSet::new(), - self.metric_prefix, - ))) + metrics: ExecutionPlanMetricsSet::new(), + ..self.as_ref().clone() + })) } fn execute( @@ -172,8 +185,23 @@ impl ExecutionPlan for DynamicFilterExec { // add its input/output counts or elapsed time to the join's existing metrics. let eval_time = MetricBuilder::new(&self.metrics) .subset_time(format!("{}_eval_time", self.metric_prefix), partition); + // Early filtering duplicates the final consumer. Stop that extra work after + // two nonempty evaluated batches remove nothing. The downstream join still + // verifies every row, so later selectivity changes only lose an optimization. + // Exact zero avoids a workload-specific break-even ratio: even a small + // reduction may save substantial work in an intermediate join. Permanent + // bypass bounds duplicate evaluation; clustered inputs may lose later + // pruning, but the final consumer preserves correctness. Resampling and + // ratio thresholds would need workload evidence before adding more policy. + // Keep this decision per stream; an inactive TRUE placeholder is not a sample. + let adaptive = self.adaptive; + let mut unselective_batches = 0; let stream = input.map(move |batch| { let batch = batch?; + if adaptive && unselective_batches >= 2 { + bypassed.add(batch.num_rows()); + return Ok(batch); + } let _timer = eval_time.timer(); // AND may prefilter its input before evaluating hash membership. A // zero-copy key projection keeps payload columns out of that temporary @@ -187,12 +215,22 @@ impl ExecutionPlan for DynamicFilterExec { Ok(batch) } ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => { + if adaptive { + unselective_batches = 0; + } evaluated.add(batch.num_rows()); pruned.add(batch.num_rows()); Ok(batch.slice(0, 0)) } ColumnarValue::Array(mask) => { let filtered = filter_record_batch(&batch, as_boolean_array(&mask)?)?; + if adaptive && batch.num_rows() > 0 { + unselective_batches = if filtered.num_rows() == batch.num_rows() { + unselective_batches + 1 + } else { + 0 + }; + } evaluated.add(batch.num_rows()); pruned.add(batch.num_rows() - filtered.num_rows()); Ok(filtered) diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs index f40fbfc6cb5..b5dd28d3f31 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs @@ -62,6 +62,14 @@ pub(super) fn try_attach_parquet_reader_filter( log::debug!("Join dynamic filter reader pushdown skipped: probe has a fetch limit"); return Ok(None); } + // An ancestor join may already filter decoded batches below this join. Keep + // that consumer while attaching this join's own predicate to its reader. + if let Some(filter) = input.downcast_ref::() { + return Ok( + try_attach_parquet_reader_filter(&filter.input, predicate, config)? + .map(|reader| filter.with_execution_input(reader)), + ); + } // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND // chains of inferred null checks. A reader predicate can cross those direct // checks because both operations only discard rows. Keep every other filter diff --git a/native/core/src/execution/operators/dynamic_filter/tests.rs b/native/core/src/execution/operators/dynamic_filter/tests.rs index e610f8b4a62..76e2fb963cb 100644 --- a/native/core/src/execution/operators/dynamic_filter/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/tests.rs @@ -123,3 +123,123 @@ async fn placeholder_updates_and_errors_are_not_hidden() { let error = collect(wrapper, task).await.unwrap_err(); assert!(error.to_string().contains("must evaluate to a Boolean")); } + +#[tokio::test] +async fn early_filter_bypasses_only_after_nonempty_evaluated_batches_and_resets_per_stream() { + use datafusion::physical_expr::expressions::IsNotNullExpr; + let ctx = SessionContext::new(); + let plan = input((0..10).map(Some).collect(), &DataType::Int32, 0); + let key = Arc::new(Column::new("key", 0)) as Arc; + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&key)], + lit(true), + )); + let metrics = ExecutionPlanMetricsSet::new(); + let filter = Arc::new( + DynamicFilterExec::new( + plan, + Arc::clone(&predicate), + metrics, + "dynamic_filter_early", + ) + .adaptive(), + ) as Arc; + // An unpopulated producer does not use up the adaptation samples. + let mut stream = filter.execute(0, ctx.task_ctx()).unwrap(); + assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 2); + assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 2); + predicate + .update(Arc::new(IsNotNullExpr::new(Arc::clone(&key)))) + .unwrap(); + while let Some(batch) = stream.next().await { + assert_eq!(batch.unwrap().num_rows(), 2); + } + assert_eq!(metric(&filter, "dynamic_filter_early_rows_evaluated"), 4); + assert_eq!(metric(&filter, "dynamic_filter_early_rows_bypassed"), 6); + // Adaptation belongs to a stream, and the downstream join remains responsible + // for matching rows if a later producer update would become more selective. + predicate + .update(Arc::new(BinaryExpr::new(key, Operator::Eq, lit(7i32)))) + .unwrap(); + let rows = collect(Arc::clone(&filter), ctx.task_ctx()).await.unwrap(); + assert_eq!(rows.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(metric(&filter, "dynamic_filter_early_rows_evaluated"), 14); + assert_eq!(metric(&filter, "dynamic_filter_early_rows_pruned"), 9); +} + +#[tokio::test] +async fn empty_batches_do_not_disable_early_filtering() { + use datafusion::physical_expr::expressions::IsNotNullExpr; + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let first = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![Some(1)]))], + ) + .unwrap(); + let last = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![None::]))], + ) + .unwrap(); + let plan = memory_exec(vec![first, RecordBatch::new_empty(schema), last]); + let key = Arc::new(Column::new("key", 0)) as Arc; + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&key)], + Arc::new(IsNotNullExpr::new(key)), + )); + let filter = Arc::new( + DynamicFilterExec::new( + plan, + predicate, + ExecutionPlanMetricsSet::new(), + "dynamic_filter_early", + ) + .adaptive(), + ) as Arc; + let output = collect(Arc::clone(&filter), ctx.task_ctx()).await.unwrap(); + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(metric(&filter, "dynamic_filter_early_rows_pruned"), 1); +} + +#[test] +fn adaptive_mode_survives_rebuilding() { + let source = input(vec![Some(1)], &DataType::Int32, 0); + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("key", 0))], + lit(true), + )); + let original = Arc::new( + DynamicFilterExec::new( + Arc::clone(&source), + Arc::clone(&predicate), + ExecutionPlanMetricsSet::new(), + "dynamic_filter_early", + ) + .adaptive(), + ); + let execution = original.with_execution_input(Arc::clone(&source)); + let replaced = Arc::clone(&original) + .replace_children( + vec![source], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + let reset = original.reset_state().unwrap(); + for plan in [&execution, &replaced, &reset] { + let filter = plan.downcast_ref::().unwrap(); + assert!(filter.adaptive); + assert_eq!(filter.metric_prefix, "dynamic_filter_early"); + } + assert!(Arc::ptr_eq( + &execution + .downcast_ref::() + .unwrap() + .predicate, + &predicate + )); + assert!(!Arc::ptr_eq( + &reset.downcast_ref::().unwrap().predicate, + &predicate + )); +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 96802ca805d..cc9157f8468 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -84,8 +84,8 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_operators::{ - range_exec, CometFilterExec, ExpandExec, ExplodeExec, PartitionedRankLimitExec, SampleExec, - WindowFnKind, + range_exec, CometFilterExec, CometProjectionExec, ExpandExec, ExplodeExec, + PartitionedRankLimitExec, SampleExec, WindowFnKind, }; use datafusion_comet_spark_expr::{ create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle, @@ -2616,15 +2616,15 @@ impl PhysicalPlanner { partition_count, )?; - // Reader attachment replaces the probe filter's child per execution. - // Keep its metrics owned by the same Spark filter node. + // Reader attachment and early placement recursively replace probe-side + // filters and projections. Keep metrics owned by their Spark nodes. if join.dynamic_filter_enabled && !join.null_aware_anti_join { let probe = if join.build_side == BuildSide::BuildLeft as i32 { &mut join_params.right } else { &mut join_params.left }; - *probe = Self::prepare_probe_filter_for_runtime_reader(Arc::clone(probe)); + *probe = Self::prepare_probe_for_runtime_filters(Arc::clone(probe))?; } let left = Arc::clone(&join_params.left.native_plan); @@ -2909,14 +2909,32 @@ impl PhysicalPlanner { } } - /// Keep the Spark filter's metric identity when its reader is replaced for an execution. - fn prepare_probe_filter_for_runtime_reader(plan: Arc) -> Arc { - let Some(filter) = plan.native_plan.downcast_ref::() else { - return plan; - }; + /// Keep Spark metric identities for the small set of nodes whose children + /// runtime-filter placement can replace. Stop at native/Spark tree boundaries. + fn prepare_probe_for_runtime_filters( + plan: Arc, + ) -> Result, ExecutionError> { + let mut native: Arc = + if let Some(filter) = plan.native_plan.downcast_ref::() { + Arc::new(CometFilterExec::from_datafusion(filter.clone())) + } else if let Some(projection) = plan.native_plan.downcast_ref::() { + Arc::new(CometProjectionExec::from_datafusion(projection.clone())) + } else { + return Ok(plan); + }; let mut prepared = plan.as_ref().clone(); - prepared.native_plan = Arc::new(CometFilterExec::from_datafusion(filter.clone())); - Arc::new(prepared) + if let ([child], [only]) = (plan.children.as_slice(), native.children().as_slice()) { + if Arc::ptr_eq(only, &child.native_plan) { + let child = Self::prepare_probe_for_runtime_filters(Arc::clone(child))?; + native = native.replace_children( + vec![Arc::clone(&child.native_plan)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + prepared.children = vec![child]; + } + } + prepared.native_plan = native; + Ok(Arc::new(prepared)) } /// Attach after choosing the final build side, including the projection emitted diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs index cfc92a0eb69..b5f762025ad 100644 --- a/native/operators/src/lib.rs +++ b/native/operators/src/lib.rs @@ -28,6 +28,7 @@ mod expand; mod explode; mod filter; +mod projection; mod range; mod rank_limit; mod sample; @@ -35,6 +36,7 @@ mod sample; pub use expand::ExpandExec; pub use explode::ExplodeExec; pub use filter::CometFilterExec; +pub use projection::CometProjectionExec; pub use range::range_exec; pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind}; pub use sample::SampleExec; diff --git a/native/operators/src/projection.rs b/native/operators/src/projection.rs new file mode 100644 index 00000000000..b35afcae31d --- /dev/null +++ b/native/operators/src/projection.rs @@ -0,0 +1,295 @@ +// 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. + +//! A DataFusion projection whose metrics remain owned by its Spark plan node. +//! +//! Comet normally keeps a one-to-one Spark/native plan tree so native metric +//! handles map back to the corresponding Spark operator. Some execution-local +//! rewrites need to replace a projection's child. DataFusion gives that replacement +//! a new private metric set, so this adapter owns the stable metric set and +//! registers the handles from the projection that actually executes. + +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result, Statistics}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions, +}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; + +#[derive(Debug)] +pub struct CometProjectionExec { + projection: ProjectionExec, + metrics: ExecutionPlanMetricsSet, +} + +impl CometProjectionExec { + pub fn from_datafusion(projection: ProjectionExec) -> Self { + Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + pub fn input(&self) -> &Arc { + self.projection.input() + } + + pub fn projection(&self) -> &ProjectionExec { + &self.projection + } + + fn replace_input( + &self, + input: Arc, + options: ReplaceChildrenOptions, + ) -> Result { + let replaced = Arc::new(self.projection.clone()).replace_children(vec![input], options)?; + let Some(projection) = replaced.downcast_ref::() else { + return internal_err!("ProjectionExec child replacement changed its plan type"); + }; + Ok(projection.clone()) + } + + /// Replace the child for one execution while keeping the metric identity + /// owned by the permanent Spark projection node. + pub fn with_execution_input( + &self, + input: Arc, + ) -> Result> { + Ok(Arc::new(Self { + projection: self.replace_input( + input, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + metrics: self.metrics.clone(), + })) + } +} + +impl DisplayAs for CometProjectionExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.projection.fmt_as(t, f) + } +} + +impl ExecutionPlan for CometProjectionExec { + fn name(&self) -> &str { + self.projection.name() + } + + fn properties(&self) -> &Arc { + self.projection.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![self.projection.input()] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.projection.apply_expressions(f) + } + + fn benefits_from_input_partitioning(&self) -> Vec { + self.projection.benefits_from_input_partitioning() + } + + fn maintains_input_order(&self) -> Vec { + self.projection.maintains_input_order() + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + if children.len() != 1 { + return internal_err!("CometProjectionExec requires one child"); + } + Ok(Arc::new(Self { + projection: self.replace_input(children.remove(0), options)?, + metrics: ExecutionPlanMetricsSet::new(), + })) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + // Every execution gets fresh DataFusion projection metrics. Register their + // handles on the stable set before returning the stream, including when + // opening the child fails. + let projection = self.replace_input( + Arc::clone(self.projection.input()), + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let result = projection.execute(partition, context); + for metric in projection.metrics().unwrap_or_default().iter() { + self.metrics.register(Arc::clone(metric)); + } + result + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + self.projection.child_stats_requests(partition) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + self.projection.statistics_from_inputs(input_stats, args) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.projection.cardinality_effect() + } + + fn fetch(&self) -> Option { + self.projection.fetch() + } + + fn with_fetch(&self, fetch: Option) -> Option> { + let projection = self.projection.with_fetch(fetch)?; + let projection = projection.downcast_ref::()?.clone(); + Some(Arc::new(Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + })) + } + + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + let projection = self.projection.with_preserve_order(preserve_order)?; + let projection = projection.downcast_ref::()?.clone(); + Some(Arc::new(Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion::common::stats::Precision; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{binary, col, lit}; + use datafusion::physical_plan::statistics::StatisticsContext; + + fn projection() -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let partitions = [1..11, 11..41] + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(values))], + ) + .map(|batch| vec![batch]) + }) + .collect::, _>>()?; + let input = MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&schema), None)?; + ProjectionExec::try_new( + vec![ + ( + binary(col("key", &schema)?, Operator::Plus, lit(5i32), &schema)?, + "offset".into(), + ), + (col("key", &schema)?, "key".into()), + ], + input, + ) + } + + #[test] + fn expression_visitor_preserves_roots_and_recursion() -> Result<()> { + let wrapper = CometProjectionExec::from_datafusion(projection()?); + assert_eq!(wrapper.projection().expr()[0].expr.children().len(), 2); + for recursion in [TreeNodeRecursion::Continue, TreeNodeRecursion::Stop] { + let mut visits = 0; + let result = wrapper.apply_expressions(&mut |expression| { + assert!(Arc::ptr_eq( + expression, + &wrapper.projection().expr()[visits].expr + )); + visits += 1; + Ok(recursion) + })?; + assert_eq!( + visits, + if recursion == TreeNodeRecursion::Stop { + 1 + } else { + 2 + } + ); + assert_eq!(result, recursion); + } + Ok(()) + } + + #[test] + fn statistics_match_datafusion_for_each_partition() -> Result<()> { + let projection = projection()?; + let wrapper = CometProjectionExec::from_datafusion(projection.clone()); + let context = StatisticsContext::new(); + let mut row_counts = Vec::new(); + for partition in [None, Some(0), Some(1)] { + let args = StatisticsArgs::new().with_partition(partition); + let expected = context.compute(&projection, &args)?; + let actual = context.compute(&wrapper, &args)?; + assert_eq!(actual, expected); + assert_ne!(actual.num_rows, Precision::Absent); + row_counts.push(actual.num_rows); + } + assert_ne!(row_counts[0], row_counts[1]); + assert_ne!(row_counts[1], row_counts[2]); + Ok(()) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 7d86de43c4d..e038eb78d85 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -526,7 +526,15 @@ object CometMetricNode { "dynamic_filter_join_filters_attached" -> SQLMetrics.createMetric(sc, "Join runtime filters attached to native readers"), "dynamic_filter_join_filters_skipped" -> - SQLMetrics.createMetric(sc, "Join runtime filters not eligible for native readers")) + SQLMetrics.createMetric(sc, "Join runtime filters not eligible for native readers"), + "dynamic_filter_early_rows_evaluated" -> + SQLMetrics.createMetric(sc, "Rows evaluated before an intermediate join"), + "dynamic_filter_early_rows_pruned" -> + SQLMetrics.createMetric(sc, "Rows rejected before an intermediate join"), + "dynamic_filter_early_rows_bypassed" -> + SQLMetrics.createMetric(sc, "Rows bypassing inactive or unselective early filtering"), + "dynamic_filter_early_eval_time" -> + SQLMetrics.createNanoTimingMetric(sc, "Time evaluating the early join runtime filter")) } /** diff --git a/spark/src/test/resources/sql-tests/join/dynamic_filter_chain.sql b/spark/src/test/resources/sql-tests/join/dynamic_filter_chain.sql new file mode 100644 index 00000000000..5237e225bf9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/join/dynamic_filter_chain.sql @@ -0,0 +1,72 @@ +-- 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. + +-- Compare complete join results with Spark with runtime filtering enabled and disabled. +-- Keep the join order fixed so both broadcast build sides and stacked ancestor consumers run. +-- Config: spark.sql.adaptive.enabled=false +-- Config: spark.sql.autoBroadcastJoinThreshold=-1 +-- Config: spark.sql.cbo.joinReorder.enabled=false +-- Config: spark.comet.batchSize=16 +-- ConfigMatrix: spark.comet.exec.join.dynamicFilter.enabled=false,true + +statement +CREATE TABLE chain_fact(k INT, payload INT) USING parquet + +statement +INSERT INTO chain_fact VALUES (1, 10), (1, 11), (2, 20), (3, 30), (4, 40), (NULL, 50) + +statement +CREATE TABLE chain_dimension(k INT, detail INT) USING parquet + +statement +INSERT INTO chain_dimension VALUES (1, 100), (1, 101), (2, 200), (3, 300), (NULL, 500) + +statement +CREATE TABLE chain_selection(k INT, selected INT) USING parquet + +statement +INSERT INTO chain_selection VALUES (1, 1000), (1, 1001), (2, 2000), (NULL, 5000) + +statement +CREATE TABLE chain_final(k INT, selected INT) USING parquet + +statement +INSERT INTO chain_final VALUES (1, 10000), (1, 10001), (3, 30000), (NULL, 50000) + +-- Duplicate keys multiply matches, while NULL keys never match each other. +query +SELECT /*+ BROADCAST(d), BROADCAST(s) */ f.*, d.detail, s.selected +FROM chain_fact f JOIN chain_dimension d ON f.k = d.k +JOIN chain_selection s ON f.k = s.k + +query +SELECT /*+ BROADCAST(d), BROADCAST(s) */ f.*, d.detail, s.selected +FROM chain_dimension d JOIN chain_fact f ON f.k = d.k +JOIN chain_selection s ON f.k = s.k + +-- Two ancestors install early consumers below the innermost join. +query +SELECT /*+ BROADCAST(d), BROADCAST(s), BROADCAST(t) */ f.*, d.detail, s.selected, t.selected +FROM chain_fact f JOIN chain_dimension d ON f.k = d.k +JOIN chain_selection s ON f.k = s.k +JOIN chain_final t ON f.k = t.k + +query +SELECT /*+ BROADCAST(d), BROADCAST(s), BROADCAST(t) */ f.*, d.detail, s.selected, t.selected +FROM chain_dimension d JOIN chain_fact f ON f.k = d.k +JOIN chain_selection s ON f.k = s.k +JOIN chain_final t ON f.k = t.k diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 369c2517636..36e7c5c4f4a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Dynam import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException @@ -435,6 +435,108 @@ class CometJoinSuite extends CometTestBase { } } + test("join dynamic filter rejects rows before an intermediate join") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.JOIN_REORDER_ENABLED.key -> "false", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "128") { + withParquetTable((0 until 1000).map(i => (i, i.toLong)), "early_fact") { + withParquetTable((0 until 2000).map(i => (i % 1000, i)), "early_dimension") { + withParquetTable(Seq((42, 1), (42, 2)), "early_selection") { + for (buildLeft <- Seq(false, true); enabled <- Seq(false, true)) { + val from = if (buildLeft) { + "early_dimension d JOIN early_fact f" + } else { + "early_fact f JOIN early_dimension d" + } + val query = "SELECT /*+ BROADCAST(d), BROADCAST(s) */ " + + "f._1 AS selected_key, f._2 AS payload, d._2 AS detail, s._2 AS selection " + + s"FROM $from ON f._1 = d._1 " + + "JOIN early_selection s ON f._1 = s._1" + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> enabled.toString) { + val (_, plan) = checkSparkAnswerAndOperator( + sql(query), + Seq(classOf[CometBroadcastHashJoinExec])) + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 2, s"Expected two native broadcast joins:\n$plan") + val outputs = joins.map(_.metrics("output_rows").value).sorted + assert(outputs == (if (enabled) Seq(2L, 4L) else Seq(4L, 2000L))) + if (enabled) { + assert( + joins.map(_.metrics("dynamic_filter_early_rows_pruned").value).sum == 999L) + assert( + joins + .map(_.metrics("dynamic_filter_early_rows_evaluated").value) + .sum == 1000L) + // The ancestor consumer must not hide the scan from the intermediate join's + // own reader filter. collect visits the outer join before the intermediate one. + assert(joins.last.metrics("dynamic_filter_join_filters_attached").value > 0L) + val projections = collect(plan) { case project: CometProjectExec => project } + assert(projections.nonEmpty) + assert(projections.forall(_.metrics("output_rows").value > 0L)) + } + } + } + } + } + } + } + } + + test("join dynamic filters stack across three joins and preserve reader attachment") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.JOIN_REORDER_ENABLED.key -> "false", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "128") { + withParquetTable((0 until 1000).map(i => (i, i.toLong)), "stacked_fact") { + withParquetTable((0 until 2000).map(i => (i % 1000, i)), "stacked_dimension") { + withParquetTable(Seq((42, 1), (43, 2)), "stacked_selection") { + withParquetTable(Seq((42, 1), (42, 2)), "stacked_final") { + for (buildLeft <- Seq(false, true); enabled <- Seq(false, true)) { + val from = if (buildLeft) { + "stacked_dimension d JOIN stacked_fact f" + } else { + "stacked_fact f JOIN stacked_dimension d" + } + val query = "SELECT /*+ BROADCAST(d), BROADCAST(s), BROADCAST(t) */ " + + "f._1 AS selected_key, f._2 AS payload, d._2 AS detail, " + + "s._2 AS selection, t._2 AS final_selection " + + s"FROM $from ON f._1 = d._1 " + + "JOIN stacked_selection s ON f._1 = s._1 " + + "JOIN stacked_final t ON f._1 = t._1" + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> enabled.toString) { + val (_, plan) = checkSparkAnswerAndOperator( + sql(query), + Seq(classOf[CometBroadcastHashJoinExec])) + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 3, s"Expected three native broadcast joins:\n$plan") + val outputs = joins.map(_.metrics("output_rows").value).sorted + assert(outputs == (if (enabled) Seq(2L, 2L, 4L) else Seq(4L, 4L, 2000L))) + if (enabled) { + // Each ancestor installs a consumer beneath the first join, so both must + // evaluate rows even when one has already filtered most of the batch. + assert( + joins.count(_.metrics("dynamic_filter_early_rows_evaluated").value > 0L) + == 2) + assert( + joins.map(_.metrics("dynamic_filter_early_rows_pruned").value).sum == 999L) + assert(joins.last.metrics("dynamic_filter_join_filters_attached").value > 0L) + } + } + } + } + } + } + } + } + } + test("join dynamic filter prunes a projected Parquet reader through null-check conjunctions") { withTempPath { probePath => withSQLConf( From 9a7b1c1b99ad9a1ddd8657b659723f6fdacf35dd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 12:56:06 +0000 Subject: [PATCH 60/89] docs: list the S3A settings the Iceberg native-write gate ignores (#6828) The user guide and the contributor guide both name the Hadoop S3A settings that the native Iceberg write gate ignores, and both list only the vectored-read and downgrade.syncable.exceptions settings. #6808 adds the S3A committer, bulk.delete.page.size and experimental.input.fadvise settings to that list. Name them in both guides, point the S3A row of the eligibility table at the list, state when a key belongs on it, and note that the list matches the plain key, so a per-bucket spelling still falls back. --- docs/source/contributor-guide/iceberg-writes.md | 11 ++++++++--- docs/source/user-guide/latest/iceberg-writes.md | 9 +++++++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 1ef7539329a..2cdfb9038b5 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -312,9 +312,14 @@ Per-bucket `fs.s3a.bucket..*` keys are split into a complete bucket name suffix. `fs.s3a.bucket.target.other.endpoint` belongs to bucket `target.other`, so it does not make a write to `target` ineligible. The Hadoop check excludes keys whose only source is Hadoop's built-in `core-default.xml`, while -retaining programmatic, site-XML, and custom `*-default.xml` settings. It also ignores Spark's -session-wide S3A vectored-read and `downgrade.syncable.exceptions` settings, which cannot affect a -data-file write request. Keep those allowlists aligned when adding storage support. Reaching +retaining programmatic, site-XML, and custom `*-default.xml` settings. It also ignores the keys in +`IgnoredHadoopS3Keys`: Spark's session-wide S3A vectored-read and `downgrade.syncable.exceptions` +settings, plus the committer, `bulk.delete.page.size` and `experimental.input.fadvise` settings that +deployments set cluster-wide (Spark 4.1 sets two of the committer keys itself). A key belongs on +that list only if S3A would not apply it to a data-file write request. Settings that S3A does apply, +such as encryption, ACLs, storage class, signing and the credentials provider, must keep falling +back. The list matches the exact key, so a per-bucket spelling is not ignored. Keep those +allowlists aligned when adding storage support. Reaching `FileIOBuilder.with_prop` is not proof of support: the builder accepts unknown properties and the backend ignores them. A fallback reason must contain property names only, never values. diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index a04f6be61cf..2d84caeaca0 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -201,7 +201,7 @@ A write is eligible only when ALL of the following hold: | `write.target-file-size-bytes` | any value (the two writers can choose different roll points; see accepted divergences) | | data location URI scheme | `file`, `memory`, `s3`, `s3a`, `gs`, matched case-sensitively (`S3://` falls back). `s3`, `s3a` and `gs` need a bucket in the authority (`s3://bucket/...`), so a hostless form such as `s3:/bucket/key` falls back. `gs` only when the `FileIO` opening the data location is a `GCSFileIO`; see below | | resolved `table.locationProvider()` | Iceberg's built-in `DefaultLocationProvider` | -| Hadoop S3A settings for an `s3` / `s3a` data location | only `fs.s3a.access.key`, `secret.key`, `session.token`, `endpoint`, `endpoint.region`, and `path.style.access`, including their `fs.s3a.bucket..*` forms; any other effective `fs.s3a.*` setting falls back | +| Hadoop S3A settings for an `s3` / `s3a` data location | only `fs.s3a.access.key`, `secret.key`, `session.token`, `endpoint`, `endpoint.region`, and `path.style.access`, including their `fs.s3a.bucket..*` forms; any other effective `fs.s3a.*` setting falls back, apart from the ignored settings listed below | | Iceberg `FileIO` S3 settings for an `s3` / `s3a` data location | the S3 endpoint, region, static/session credentials, path-style, SSE (`none`, `s3`, `kms`, or `custom`; not `dsse-kms`), assume-role, anonymous/config-chain settings parsed by the pinned iceberg-rust version, plus Comet's credential-provider class and built-in web-identity properties. When a custom provider is configured, its vendor-owned `s3.*` / `client.*` properties are also forwarded; unsupported Iceberg-defined S3 settings still fall back | | partition spec | any, except an identity partition on a `float` or `double` column, or a `void` field whose source column was dropped beside a live field (see below) | | column types | any except `uuid` (Spark plans it as a string; no Arrow cast reaches `fixed(16)`) and the v3 types `variant`, `unknown`, `timestamp_ns`, `geometry` and `geography` | @@ -242,7 +242,12 @@ properties; session and catalog Hadoop options are neither checked nor forwarded `core-default.xml` values are not treated as explicit settings, but programmatic settings and values from site or custom `*-default.xml` resources are. Spark's session-wide S3A vectored-read and `downgrade.syncable.exceptions` compatibility settings are also -ignored because they cannot alter an Iceberg data-file write request. Unknown explicit +ignored, as are the S3A committer (`committer.name`, `committer.magic.enabled`, +`committer.threads`), `bulk.delete.page.size` and `experimental.input.fadvise` settings, because +none of them can alter an Iceberg data-file write request. Spark 4.1 sets `committer.name` and +`committer.magic.enabled` itself when `spark-hadoop-cloud` is on the classpath. Ignored settings are +matched by their plain `fs.s3a.` keys, so a per-bucket `fs.s3a.bucket..` +form still falls back. Unknown explicit `fs.s3a.*`, `s3.*`, or `client.*` settings therefore fall back at planning time instead of being silently ignored by the native storage backend. The exception is a vendor-owned `s3.*` / `client.*` property when From 8081c7c0fb13f60962ca9e069db75abd3e0ac26b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 12:56:29 +0000 Subject: [PATCH 61/89] test: count the partitions and batches the native cache scan reads (#6836) * test: count the partitions and batches the native cache scan reads Spark's PartitionBatchPruningSuite checks cache pruning through the readPartitions and readBatches accumulators of InMemoryTableScanExec, which Spark's scan fills under spark.sql.inMemoryTableScanStatistics.enable. CometInMemoryTableScanExec had none, so the Spark SQL diffs made the suite assert only the answer when a Comet scan was present. That verifies nothing about pruning: the answer is the same whether a batch was pruned or decoded and filtered. With the cache on by default since #5634, Spark's forty-odd pruning cases ran against Comet's format without checking that any batch was skipped. Give the Comet scan the same two accumulators, counted after the statistics filter as Spark counts them, and restore the suite's original assertions in the 3.5.9, 4.0.4, 4.1.3 and 4.2.0 diffs with a case for the Comet scan. The 3.4.3 diff never touched the suite because the cache is opt-in there. The new CometInMemoryCacheSuite test checks the counts with pruning on and off, and with AQE on and off, since on Spark 3.5+ the scan then sits in a table cache stage and the instance that ran has to be the one the final plan holds. * test: return the read counts through a local, since withSQLConf returns Unit on Spark 3.x --- dev/diffs/3.5.9.diff | 72 +++++-------------- dev/diffs/4.0.4.diff | 72 +++++-------------- dev/diffs/4.1.3.diff | 72 +++++-------------- dev/diffs/4.2.0.diff | 72 +++++-------------- .../comet/CometInMemoryTableScanExec.scala | 33 ++++++++- .../comet/exec/CometInMemoryCacheSuite.scala | 51 +++++++++++++ 6 files changed, 151 insertions(+), 221 deletions(-) diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index 4ee26ebe93c..7549996767a 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -2461,7 +2461,7 @@ index 2c73622739a..5d0efeb263a 100644 val d1 = spark.range(1).withColumn("key", expr("id % 100")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala -index 885286843a1..f0f805a6cd9 100644 +index 885286843a1..6e9d6690374 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala @@ -17,6 +17,7 @@ @@ -2472,60 +2472,22 @@ index 885286843a1..f0f805a6cd9 100644 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - val result = df.collect().map(_(0)).toArray - assert(result.length === 1) - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- assert(readPartitions === 5) -- assert(readBatches === 10) -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ assert(scans.head.readPartitions.value === 5) -+ assert(scans.head.readBatches.value === 10) -+ } - } - - def checkBatchPruning( -@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - df.collect().map(_(0)).toArray - } - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- -- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") -- assert( -- readPartitions === expectedReadPartitions, -- s"Wrong number of read partitions: $queryExecution") -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan prunes on the same statistics, but has none of these test-only -+ // accumulators to read, so only the answer above is checked for it. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ val readPartitions = scans.head.readPartitions.value -+ val readBatches = scans.head.readBatches.value -+ assert( -+ readBatches === expectedReadBatches, -+ s"Wrong number of read batches: $queryExecution") -+ assert( -+ readPartitions === expectedReadPartitions, -+ s"Wrong number of read partitions: $queryExecution") -+ } - } - } - } +@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + assert(readPartitions === 5) + assert(readBatches === 10) +@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + + assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd52d038ca6..154c800be67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 895f6bd3db8..2133589aa92 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -3228,7 +3228,7 @@ index 2c73622739a..5d0efeb263a 100644 val d1 = spark.range(1).withColumn("key", expr("id % 100")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala -index 88ff51d0ff4..7d00fca3b0a 100644 +index 88ff51d0ff4..3763fad99c8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala @@ -17,6 +17,7 @@ @@ -3239,60 +3239,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - val result = df.collect().map(_(0)).toArray - assert(result.length === 1) - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- assert(readPartitions === 5) -- assert(readBatches === 10) -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ assert(scans.head.readPartitions.value === 5) -+ assert(scans.head.readBatches.value === 10) -+ } - } - - def checkBatchPruning( -@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - df.collect().map(_(0)).toArray - } - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- -- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") -- assert( -- readPartitions === expectedReadPartitions, -- s"Wrong number of read partitions: $queryExecution") -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan prunes on the same statistics, but has none of these test-only -+ // accumulators to read, so only the answer above is checked for it. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ val readPartitions = scans.head.readPartitions.value -+ val readBatches = scans.head.readBatches.value -+ assert( -+ readBatches === expectedReadBatches, -+ s"Wrong number of read batches: $queryExecution") -+ assert( -+ readPartitions === expectedReadPartitions, -+ s"Wrong number of read partitions: $queryExecution") -+ } - } - } - } +@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + assert(readPartitions === 5) + assert(readBatches === 10) +@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + + assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index 0b814aa1b10..929fbfd09da 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -3447,7 +3447,7 @@ index 2c73622739a..5d0efeb263a 100644 val d1 = spark.range(1).withColumn("key", expr("id % 100")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala -index 88ff51d0ff4..7d00fca3b0a 100644 +index 88ff51d0ff4..3763fad99c8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala @@ -17,6 +17,7 @@ @@ -3458,60 +3458,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - val result = df.collect().map(_(0)).toArray - assert(result.length === 1) - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- assert(readPartitions === 5) -- assert(readBatches === 10) -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ assert(scans.head.readPartitions.value === 5) -+ assert(scans.head.readBatches.value === 10) -+ } - } - - def checkBatchPruning( -@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - df.collect().map(_(0)).toArray - } - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- -- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") -- assert( -- readPartitions === expectedReadPartitions, -- s"Wrong number of read partitions: $queryExecution") -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan prunes on the same statistics, but has none of these test-only -+ // accumulators to read, so only the answer above is checked for it. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ val readPartitions = scans.head.readPartitions.value -+ val readBatches = scans.head.readBatches.value -+ assert( -+ readBatches === expectedReadBatches, -+ s"Wrong number of read batches: $queryExecution") -+ assert( -+ readPartitions === expectedReadPartitions, -+ s"Wrong number of read partitions: $queryExecution") -+ } - } - } - } +@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + assert(readPartitions === 5) + assert(readBatches === 10) +@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + + assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index f58f210fbef..549799d941a 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -3518,7 +3518,7 @@ index 2c73622739a..5d0efeb263a 100644 val d1 = spark.range(1).withColumn("key", expr("id % 100")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala -index 88ff51d0ff4..7d00fca3b0a 100644 +index 88ff51d0ff4..3763fad99c8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala @@ -17,6 +17,7 @@ @@ -3529,60 +3529,22 @@ index 88ff51d0ff4..7d00fca3b0a 100644 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - val result = df.collect().map(_(0)).toArray - assert(result.length === 1) - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- assert(readPartitions === 5) -- assert(readBatches === 10) -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ assert(scans.head.readPartitions.value === 5) -+ assert(scans.head.readBatches.value === 10) -+ } - } - - def checkBatchPruning( -@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl - df.collect().map(_(0)).toArray - } - -- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { -- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) -- }.head -- -- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") -- assert( -- readPartitions === expectedReadPartitions, -- s"Wrong number of read partitions: $queryExecution") -+ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } -+ // Comet's cache scan prunes on the same statistics, but has none of these test-only -+ // accumulators to read, so only the answer above is checked for it. -+ if (scans.isEmpty) { -+ assert(collect(df.queryExecution.executedPlan) { -+ case c: CometInMemoryTableScanExec => c -+ }.nonEmpty) -+ } else { -+ val readPartitions = scans.head.readPartitions.value -+ val readBatches = scans.head.readBatches.value -+ assert( -+ readBatches === expectedReadBatches, -+ s"Wrong number of read batches: $queryExecution") -+ assert( -+ readPartitions === expectedReadPartitions, -+ s"Wrong number of read partitions: $queryExecution") -+ } - } - } - } +@@ -183,6 +184,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + assert(readPartitions === 5) + assert(readBatches === 10) +@@ -204,6 +206,7 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + + val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { + case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) ++ case in: CometInMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) + }.head + + assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd8d1308e99..0e1f80045a3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index 207513f0a77..97f041033d0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.LongAccumulator import org.apache.comet.CometConf import org.apache.comet.serde.CometOperatorSerde @@ -68,6 +69,16 @@ case class CometInMemoryTableScanExec( override lazy val metrics: Map[String, SQLMetric] = Map( "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + // The partitions and batches a scan read after pruning, as Spark's InMemoryTableScanExec counts + // them under spark.sql.inMemoryTableScanStatistics.enable. Spark's PartitionBatchPruningSuite + // reads these to check that pruning skipped what it should have, and only they can show that: + // the answer is the same whether a batch was pruned or decoded and filtered. + lazy val enableAccumulatorsForTest: Boolean = conf.inMemoryTableScanStatisticsEnabled + + lazy val readPartitions: LongAccumulator = sparkContext.longAccumulator + + lazy val readBatches: LongAccumulator = sparkContext.longAccumulator + // `scanOutput` always equals this, including when it is empty. An empty-output scan // (`SELECT count(*)`) emits genuinely zero-column batches carrying only a row count: widening it // to a placeholder column, or to the whole cache schema, makes the emitted batches disagree with @@ -125,6 +136,11 @@ case class CometInMemoryTableScanExec( // stats bug, for instance -- and silently ignoring it would make Comet diverge from Spark on a // knob a user reaching for it is specifically trying to control. override def doExecuteColumnar(): RDD[ColumnarBatch] = { + if (enableAccumulatorsForTest) { + readPartitions.setValue(0) + readBatches.setValue(0) + } + val numOutputRows = longMetric("numOutputRows") // Resolved here rather than at planning time. CachedRDDBuilder.cachedColumnBuffers is not a @@ -142,8 +158,23 @@ case class CometInMemoryTableScanExec( cachedBuffers } + // Counted after pruning, as Spark counts them, so a pruned batch is not a read one. + val countedBuffers = if (enableAccumulatorsForTest) { + filteredBuffers.mapPartitionsInternal { iter => + if (iter.hasNext) { + readPartitions.add(1) + } + iter.map { batch => + readBatches.add(1) + batch + } + } + } else { + filteredBuffers + } + serializer - .convertCachedBatchToColumnarBatch(filteredBuffers, relationOutput, scanOutput, conf) + .convertCachedBatchToColumnarBatch(countedBuffers, relationOutput, scanOutput, conf) .map { cb => numOutputRows += cb.numRows().toLong cb diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index aba97d80219..be0eb109db5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1214,6 +1214,57 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache counts the partitions and batches a scan reads") { + // Spark's PartitionBatchPruningSuite reads these accumulators to check that pruning skipped + // the batches it should have; the answer alone cannot tell a pruned batch from a decoded and + // filtered one. + // With AQE the scan sits inside a table cache stage on Spark 3.5+, and the instance that ran + // has to be the one the final plan holds. + Seq(false, true).foreach { adaptive => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.IN_MEMORY_TABLE_SCAN_STATISTICS_ENABLED.key -> "true", + SQLConf.COLUMN_BATCH_SIZE.key -> "50", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + try { + // 10 partitions of 2 batches each. + spark + .range(0, 1000, 1, 10) + .selectExpr("id as key", "id % 7 as value") + .createOrReplaceTempView("counted_cache") + spark.catalog.cacheTable("counted_cache") + spark.table("counted_cache").count() + + // withSQLConf returns Unit on Spark 3.x, so the counts leave it through a local. + def readCounts(query: String, pruning: Boolean): (Long, Long) = { + var counts: (Long, Long) = null + withSQLConf(SQLConf.IN_MEMORY_PARTITION_PRUNING.key -> pruning.toString) { + val df = spark.sql(query) + df.collect() + val scans = collect(df.queryExecution.executedPlan) { + case s: CometInMemoryTableScanExec => s + } + assert(scans.size == 1, df.queryExecution.executedPlan) + counts = (scans.head.readPartitions.value, scans.head.readBatches.value) + } + counts + } + + val narrow = "SELECT key FROM counted_cache WHERE key >= 900 AND key < 905" + assert(readCounts(narrow, pruning = true) == ((1L, 1L))) + assert(readCounts(narrow, pruning = false) == ((10L, 20L))) + assert( + readCounts("SELECT key FROM counted_cache WHERE key < 0", pruning = true) == (( + 0L, + 0L))) + assert(readCounts("SELECT key FROM counted_cache", pruning = true) == ((10L, 20L))) + } finally { + spark.catalog.clearCache() + } + } + } + } + test("Comet in-memory cache honors inMemoryColumnarStorage.partitionPruning=false") { // CometInMemoryTableScanExec applies the serializer's stats filter before decoding, the same // way Spark's InMemoryTableScanExec.filteredCachedBatches does. Spark gates that on From 6a5a607d11d563926278086aed4a017ea5f2b6c6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 12:56:50 +0000 Subject: [PATCH 62/89] perf: generate the columnar-to-row projection once per executor instead of per partition (#6858) * perf: generate the columnar-to-row projection once per executor CometColumnarToRowExec.doExecute built its UnsafeProjection with UnsafeProjection.create in every partition, which regenerates the projection's Java source each time. The plan runs that path whenever it is outside whole-stage codegen, which a schema with more nested fields than spark.sql.codegen.maxFields always is. For 100 columns nested three levels deep, generating the source took about 19 ms per partition, half the time of a 201-partition JVM shuffle of 1,000 rows. Cache the compiled class per executor by column layout and create a new instance for each partition. * ci: run CometUnsafeProjectionSuite on Linux and macOS --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../sql/comet/CometColumnarToRowExec.scala | 3 +- .../sql/comet/CometUnsafeProjection.scala | 146 +++++++++++++++ .../comet/CometUnsafeProjectionSuite.scala | 177 ++++++++++++++++++ 5 files changed, 327 insertions(+), 1 deletion(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometUnsafeProjection.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometUnsafeProjectionSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 49eafc6e3a1..e672c214689 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -560,6 +560,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.spark.sql.comet.CometUnsafeProjectionSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 22bd6251dcf..3e0e6a8cdbc 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -264,6 +264,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.spark.sql.comet.CometUnsafeProjectionSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala index da1d9ba296f..089ad6fbf6d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala @@ -77,7 +77,8 @@ case class CometColumnarToRowExec(child: SparkPlan) // plan (this) in the closure. val localOutput = this.output child.executeColumnar().mapPartitionsInternal { batches => - val toUnsafe = UnsafeProjection.create(localOutput, localOutput) + // Outside whole-stage codegen this runs once per partition, so reuse the generated class. + val toUnsafe = CometUnsafeProjection.create(localOutput) batches.flatMap { batch => numInputBatches += 1 numOutputRows += batch.numRows().toLong diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometUnsafeProjection.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnsafeProjection.scala new file mode 100644 index 00000000000..d787460a8b6 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnsafeProjection.scala @@ -0,0 +1,146 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.util.{LinkedHashMap => JLinkedHashMap, Map => JMap} +import java.util.concurrent.atomic.AtomicLong + +import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodeAndComment, CodeFormatter, CodegenContext, CodeGenerator, GeneratedClass, GenerateUnsafeProjection} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.DataType + +/** + * Creates the projection that copies each row of a batch into an UnsafeRow, as + * `UnsafeProjection.create(output, output)` does, but generates and compiles its class only once + * per executor for each column layout. + * + * `UnsafeProjection.create` generates the projection's Java source on every call, and only the + * compiled class is cached. `CometColumnarToRowExec` calls it once per partition whenever it runs + * outside whole-stage codegen, which Spark skips for a schema with more fields than + * `spark.sql.codegen.maxFields`, nested fields included. Spark's own scans fall back to rows for + * such schemas, so its `ColumnarToRowExec` rarely meets one, but Comet's operators always produce + * batches. The source grows with the number of nested fields, and each level of nesting repeats + * the work of splitting its children into methods: for 100 columns nested three levels deep, + * generating it took about 19 ms in every partition. + * + * Every call returns a new instance of the shared class, so callers own their projection as they + * do one from `UnsafeProjection.create`. + */ +private[comet] object CometUnsafeProjection + extends CodeGeneratorWithInterpretedFallback[Seq[BoundReference], UnsafeProjection] { + + /** + * What the generated source depends on: the type and nullability of each column, and the method + * size at which `CodegenContext.splitExpressions` splits the field writes. + */ + private case class Layout(columns: Seq[(DataType, Boolean)], methodSplitThreshold: Int) + + /** The same bound as Spark's compiled class cache, `spark.sql.codegen.cache.maxEntries`. */ + private[comet] val MaxCachedClasses = 100 + + /** Least recently used first. Guarded by `classes.synchronized`. */ + private val classes = new JLinkedHashMap[Layout, GeneratedClass](16, 0.75f, true) { + override def removeEldestEntry(eldest: JMap.Entry[Layout, GeneratedClass]): Boolean = + size() > MaxCachedClasses + } + + private val classesGenerated = new AtomicLong(0) + + /** How many projection classes this executor has generated, for tests. */ + private[comet] def generatedClassCount: Long = classesGenerated.get() + + /** A projection of rows with the columns of `output` to UnsafeRows. */ + def create(output: Seq[Attribute]): UnsafeProjection = + createObject(output.zipWithIndex.map { case (attr, ordinal) => + BoundReference(ordinal, attr.dataType, attr.nullable) + }) + + override protected def createCodeGeneratedObject( + columns: Seq[BoundReference]): UnsafeProjection = { + val layout = + Layout(columns.map(c => (c.dataType, c.nullable)), SQLConf.get.methodSplitThreshold) + val cached = classes.synchronized(classes.get(layout)) + if (cached != null) { + cached.generate(Array.empty[Any]).asInstanceOf[UnsafeProjection] + } else { + // Generated outside the lock: tasks that miss on the same layout at once each generate + // it, as every task does with UnsafeProjection.create. + val (generated, references) = generate(columns) + // A stored class is instantiated without references, so store only one that needs none. + // GenerateUnsafeProjection references no objects for bound columns. + if (references.isEmpty) { + val _ = classes.synchronized(classes.putIfAbsent(layout, generated)) + } + generated.generate(references).asInstanceOf[UnsafeProjection] + } + } + + override protected def createInterpretedObject(columns: Seq[BoundReference]): UnsafeProjection = + InterpretedUnsafeProjection.createProjection(columns) + + /** + * Generates and compiles the class that `GenerateUnsafeProjection.create` does, from the same + * template, which is repeated here because that method returns only an instance. Returns the + * class with the objects its instances reference. + */ + private def generate(columns: Seq[BoundReference]): (GeneratedClass, Array[Any]) = { + val ctx = new CodegenContext + val eval = GenerateUnsafeProjection.createCode(ctx, columns) + val body = + s""" + |public java.lang.Object generate(Object[] references) { + | return new SpecificUnsafeProjection(references); + |} + | + |class SpecificUnsafeProjection extends ${classOf[UnsafeProjection].getName} { + | + | private Object[] references; + | ${ctx.declareMutableStates()} + | + | public SpecificUnsafeProjection(Object[] references) { + | this.references = references; + | ${ctx.initMutableStates()} + | } + | + | public void initialize(int partitionIndex) { + | ${ctx.initPartition()} + | } + | + | // Scala.Function1 need this + | public java.lang.Object apply(java.lang.Object row) { + | return apply((InternalRow) row); + | } + | + | public UnsafeRow apply(InternalRow ${ctx.INPUT_ROW}) { + | ${eval.code} + | return ${eval.value}; + | } + | + | ${ctx.declareAddedFunctions()} + |} + """.stripMargin + val code = CodeFormatter.stripOverlappingComments( + new CodeAndComment(body, ctx.getPlaceHolderToComments())) + val (generatedClass, _) = CodeGenerator.compile(code) + classesGenerated.incrementAndGet() + (generatedClass, ctx.references.toArray) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometUnsafeProjectionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometUnsafeProjectionSuite.scala new file mode 100644 index 00000000000..07c6371f3e8 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometUnsafeProjectionSuite.scala @@ -0,0 +1,177 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.sql.Timestamp +import java.util.UUID + +import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, InterpretedUnsafeProjection, UnsafeProjection} +import org.apache.spark.sql.execution.WholeStageCodegenExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +class CometUnsafeProjectionSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + /** A schema with every writer GenerateUnsafeProjection nests: structs, arrays and maps. */ + private def schema(nestedName: String = "a"): StructType = StructType( + Seq( + StructField("i", IntegerType), + StructField("s", StringType), + StructField("d", DecimalType(38, 10)), + StructField( + "st", + StructType( + Seq(StructField(nestedName, LongType), StructField("b", ArrayType(StringType))))), + StructField( + "m", + MapType(StringType, ArrayType(StructType(Seq(StructField("x", DoubleType)))))), + StructField("aa", ArrayType(ArrayType(IntegerType))), + StructField("bin", BinaryType), + StructField("ts", TimestampType))) + + private def attributes(nestedName: String = "a"): Seq[AttributeReference] = + schema(nestedName).fields.toSeq.map(f => AttributeReference(f.name, f.dataType, f.nullable)()) + + /** Attributes whose layout no earlier call has cached, since a nested name is part of it. */ + private def newLayout(): Seq[AttributeReference] = + attributes(s"a_${UUID.randomUUID().toString.replace("-", "")}") + + private lazy val rows: Seq[InternalRow] = { + val toCatalyst = CatalystTypeConverters.createToCatalystConverter(schema()) + Seq( + Row( + 1, + "a", + BigDecimal("1.5"), + Row(2L, Seq("x", null)), + Map("k" -> Seq(Row(1.0), null)), + Seq(Seq(1, null), null), + Array[Byte](1, 2), + new Timestamp(0)), + Row(null, null, null, null, null, null, null, null), + Row( + 3, + "字" * 100, + BigDecimal("-1234567890123456789012345678.1234567890"), + Row(null, Seq.empty), + Map.empty, + Seq.empty, + Array.emptyByteArray, + null)).map(toCatalyst(_).asInstanceOf[InternalRow]) + } + + private def withCodegenOnly(confs: (String, String)*)(f: => Unit): Unit = + withSQLConf((SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY") +: confs: _*)(f) + + test("rows match UnsafeProjection.create") { + // A small threshold splits the field writes into many methods. + Seq("1024", "64").foreach { threshold => + withCodegenOnly(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> threshold) { + val attrs = newLayout() + val expected = UnsafeProjection.create(attrs, attrs) + val actual = CometUnsafeProjection.create(attrs) + rows.foreach(row => assert(actual(row) == expected(row))) + } + } + } + + test("NO_CODEGEN creates an interpreted projection") { + withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN") { + val attrs = attributes() + val projection = CometUnsafeProjection.create(attrs) + assert(projection.isInstanceOf[InterpretedUnsafeProjection]) + val expected = UnsafeProjection.create(attrs, attrs) + rows.foreach(row => assert(projection(row) == expected(row))) + } + } + + test("generates the class once per layout and a new projection per call") { + withCodegenOnly() { + val attrs = newLayout() + val before = CometUnsafeProjection.generatedClassCount + val first = CometUnsafeProjection.create(attrs) + val second = CometUnsafeProjection.create(attrs) + assert(CometUnsafeProjection.generatedClassCount == before + 1) + assert(first ne second) + // Each projection writes into its own row, so the first row survives the second call. + val firstRow = first(rows.head) + val secondRow = second(rows(1)) + val expected = UnsafeProjection.create(attrs, attrs) + assert(firstRow == expected(rows.head)) + assert(secondRow == expected(rows(1))) + } + } + + test("a different method split threshold generates its own class") { + val attrs = newLayout() + val before = CometUnsafeProjection.generatedClassCount + Seq("1024", "64", "1024").foreach { threshold => + withCodegenOnly(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> threshold) { + CometUnsafeProjection.create(attrs) + } + } + assert(CometUnsafeProjection.generatedClassCount == before + 2) + } + + test("keeps the most recently used classes") { + withCodegenOnly() { + val layouts = Seq.fill(CometUnsafeProjection.MaxCachedClasses + 1) { + val name = s"f_${UUID.randomUUID().toString.replace("-", "")}" + Seq(AttributeReference("c", StructType(Seq(StructField(name, IntegerType))))()) + } + layouts.foreach(CometUnsafeProjection.create) + val before = CometUnsafeProjection.generatedClassCount + CometUnsafeProjection.create(layouts.last) + assert(CometUnsafeProjection.generatedClassCount == before) + CometUnsafeProjection.create(layouts.head) + assert(CometUnsafeProjection.generatedClassCount == before + 1) + } + } + + test("columnar to row outside whole-stage codegen reuses the generated class") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(200) + .selectExpr( + "id", + "named_struct('a', id, 'b', array(cast(id as string), null)) as st", + "map(cast(id as string), array(id)) as m") + .repartition(4) + .write + .parquet(path) + // More fields than this keep a plan out of whole-stage codegen, as a wide schema would. + withSQLConf(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2") { + val df = spark.read.parquet(path) + checkSparkAnswerAndOperator(df) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case c: CometColumnarToRowExec => c }.nonEmpty, plan) + assert(collect(plan) { case w: WholeStageCodegenExec => w }.isEmpty, plan) + // The first run generated the class, so another run over every partition generates none. + val before = CometUnsafeProjection.generatedClassCount + assert(df.collect().length == 200) + assert(CometUnsafeProjection.generatedClassCount == before) + } + } + } +} From afa477f2265ad272988853d668e1d2a01fb95cb5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 12:57:29 +0000 Subject: [PATCH 63/89] feat: add lz4 as an opt-in codec for Comet's in-memory cache (#6859) spark.comet.exec.inMemoryCache.compression.codec accepts lz4 alongside zstd and none. Buffers are written as Arrow's standard LZ4_FRAME with lz4-java, the JNI library behind spark.io.compression.codec, rather than Arrow's own pure-Java LZ4 codec, and LZ4_FRAME batches are read with the same codec. lz4 decompresses several times faster than zstd, which closes most of the gap between Spark operators reading Comet's cache format and reading Spark's own. It stays opt-in because it takes around three times zstd's footprint, and on sequential longs more than Spark's own format. --- .../user-guide/latest/in-memory-cache.md | 62 ++++-- .../scala/org/apache/comet/CometConf.scala | 26 +-- .../execution/arrow/CachedBatchIpc.scala | 20 +- .../arrow/Lz4FrameCompressionCodec.scala | 209 ++++++++++++++++++ .../comet/exec/CometInMemoryCacheSuite.scala | 114 +++++++++- .../CometInMemoryCacheBenchmark.scala | 111 +++++++--- .../arrow/CometCachedBatchHelper.scala | 8 + 7 files changed, 469 insertions(+), 81 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index d6867defeab..f0c2cd6a31e 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -91,36 +91,47 @@ column out of six does roughly a sixth of the decompression work, and a `SELECT selects no columns at all, answers from the row count stored beside the payload without touching it. -Compression defaults to `zstd`, for footprint rather than for speed. Over the same 5M-row, -six-column relation the tables under [Performance](#performance) use — and measured by the same -benchmark — it holds the data in a sixth of the memory and pays for that on both sides: about 40% -longer to materialize, and, on a read wide enough to inflate everything, close to five times -longer. A narrow projection pays far less, because it only inflates the columns it asked for. +Three codecs are offered: `zstd`, the default, `lz4` and `none`. Measured by the same benchmark +over the same 5M-row, six-column relation the tables under [Performance](#performance) use (Apple +M3 Ultra, JDK 17, Spark 4.1, release build): | Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | | ------ | ----------: | --------: | ----------: | ----------: | -| `zstd` | 1507 ms | 55 MiB | 45 ms | 295 ms | -| `none` | 1081 ms | 315 MiB | 35 ms | 64 ms | +| `zstd` | 1232 ms | 51 MiB | 36 ms | 248 ms | +| `lz4` | 1115 ms | 145 MiB | 33 ms | 110 ms | +| `none` | 857 ms | 315 MiB | 31 ms | 54 ms | + +`zstd` is the default, for footprint rather than for speed. It holds the data in a sixth of the +memory `none` takes and pays for that on both sides: about 45% longer to materialize, and, on a read +wide enough to inflate everything, more than four times longer. A narrow projection pays far less, +because it only inflates the columns it asked for. + +`lz4` decompresses several times faster than `zstd`, so a read of every column takes less than half +as long, at nearly three times the footprint. It is written with lz4-java, the JNI-accelerated +library behind `spark.io.compression.codec`, in Arrow's standard `LZ4_FRAME` format. Arrow's own +LZ4 codec is not used: it is commons-compress's pure-Java implementation, and is orders of +magnitude slower to write than `zstd` while also producing larger output. + +What `lz4` saves depends more on the data than what `zstd` saves does. It holds this relation in +less memory than Spark's own cache format does (217 MiB), but not a relation of sequential `bigint` +columns, which Spark's format delta-encodes: each of the wide relations under +[Limitations](#limitations) takes about 74 MiB with `lz4`, against 26 to 27 MiB in Spark's format +and with `zstd`. `none` is the better setting for a relation that fits in memory uncompressed and is read at close -to full width. The default is the other way round because a cache that does not fit costs more than -one that is slower to read, and Spark's own cache format compresses by default too. - -Arrow's other IPC codec, LZ4, is deliberately not offered and the config rejects it. It is -commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind -`spark.io.compression.codec`, and is orders of magnitude slower to write than `zstd` while also -producing larger output. +to full width. The default is `zstd` because a cache that does not fit costs more than one that is +slower to read, and Spark's own cache format compresses by default too. Dictionary-encoded columns are decoded before they are stored. A payload with no schema message has nowhere to record either that a column is dictionary encoded or the dictionary itself. ## Configuration -| Config | Default | Description | -| ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | -| `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. | -| `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | +| Config | Default | Description | +| ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | +| `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd`, `lz4` or `none`. Affects newly cached data only: a batch records the codec it was written with. | +| `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | ## Performance @@ -205,6 +216,12 @@ all of them longs, is about 10% slower under either kind of operator. That cost decompression. With the `none` codec, the same read is 2.7x faster than Spark's format with Comet operators above the scan, and 1.6x faster with a Spark operator above it. +Measured on an Apple M3 Ultra, Comet's format with `lz4` is faster than Spark's in every shape +under both kinds of operator. The three-column read is 1.4x faster than Spark's format with either +kind of operator above the scan, where `zstd` is about 15% slower on the same machine, and the full +projection is 2.5x faster with Comet operators above the scan and 1.7x faster with a Spark operator +above it. + Building the cache is measured the same way, from the same source, with its rows produced either by Comet operators or by Spark operators. Comet's format is written straight from the Arrow batches of Comet operators, while Spark's is always built from rows, so above Comet operators it first converts @@ -292,6 +309,13 @@ the others. The row reader takes up to 1.6 times as long, and it is the only rea wider than `spark.sql.codegen.maxFields`: reading every column of relations of 100, 200 and 1500 nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. A Spark operator above Comet's native cache scan does not pay this; see [Performance](#performance). + +Most of what the row reader pays is decompression, so `lz4` narrows these gaps. Measured on an +Apple M3 Ultra, the row reader takes up to 1.25 times as long as Spark's format with `lz4`, against +up to 1.7 times with `zstd`, and the fused reader is faster than Spark's format for every read. +Reading every column of the wide relations takes 1.3 to 2.0 times as long as from Spark's format +with `lz4`, against 2.1 to 2.6 times with `zstd`. `lz4` takes more memory than Spark's format for +those relations, though; see [Storage format](#storage-format). [#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks these gaps. Comet's serializer exists because Spark's own Arrow cache format diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 4aab7c72a11..2beaf8f6a3c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -388,7 +388,7 @@ object CometConf extends ShimCometConf { "without the fused reader. Relations whose schema Comet's Arrow writer does not " + "support are always cached in Spark's default " + "format. Each cached batch is stored as one Arrow IPC record batch with per-buffer " + - "zstd compression, and a scan copies out only the buffers of the columns it " + + "compression, and a scan copies out only the buffers of the columns it " + "projected, so the unselected ones are never decompressed. Eligible Spark " + "whole-stage codegen consumers read cached vectors directly when vectorized cache " + "reading is enabled; other Spark row consumers use a reusable row buffer. Decoding " + @@ -409,20 +409,20 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "The Arrow IPC compression codec used when Comet's cache serializer writes cached " + - "data. Unlike spark.io.compression.codec, this compresses each Arrow buffer " + - "separately rather than the batch as a whole, which is what lets a projected scan " + - "decompress only the columns it selected. zstd is the default for footprint rather " + - "than speed: it stores cached data in a fraction of the memory, but makes " + - "materializing slower, and a read slower in proportion to how many columns it " + - "selects. Set to none when a relation fits in memory uncompressed and is read at " + - "close to full width. " + + "data: zstd, lz4 or none. Unlike spark.io.compression.codec, this compresses each " + + "Arrow buffer separately rather than the batch as a whole, which is what lets a " + + "projected scan decompress only the columns it selected. zstd is the default for " + + "footprint rather than speed: it stores cached data in a fraction of the memory, but " + + "makes materializing slower, and a read slower in proportion to how many columns it " + + "selects. lz4 decompresses several times faster than zstd, which makes wide reads " + + "much faster, at around three times zstd's footprint. It uses lz4-java, the " + + "JNI-accelerated library behind spark.io.compression.codec, rather than Arrow's own " + + "pure-Java lz4. Set to none when a relation fits in memory uncompressed and is read " + + "at close to full width. " + "Only affects newly cached data; the codec a batch was written with is recorded in " + - "the batch itself and is what the read path uses. Arrow's lz4 is deliberately not " + - "offered: it is a pure-Java implementation, unrelated to the JNI-accelerated lz4 " + - "behind spark.io.compression.codec, and is orders of magnitude slower to write than " + - "zstd while also producing larger output.") + "the batch itself and is what the read path uses.") .stringConf - .checkValues(Set("none", "zstd")) + .checkValues(Set("none", "zstd", "lz4")) .createWithDefault("zstd") val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL: ConfigEntry[Int] = diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 06c3e527592..16dc6813899 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -73,24 +73,28 @@ private[comet] object CachedBatchIpc { // Constructed directly rather than through CompressionCodec.Factory, which ignores the level // and always builds a codec at zstd's default. case "zstd" => new ZstdCompressionCodec(zstdLevel) - // Arrow's other codec, LZ4_FRAME, is not offered. It is commons-compress's pure-Java LZ4 -- - // no relation to the JNI-accelerated lz4-java behind spark.io.compression.codec -- and - // measures three orders of magnitude slower to write than zstd while also producing larger - // output, so nothing prefers it. Reads still accept it, since the factory the read path uses - // handles whatever codec a batch records. + // Arrow's LZ4_FRAME format, written with lz4-java rather than Arrow's own pure-Java codec; + // see Lz4FrameCompressionCodec. + case "lz4" => Lz4FrameCompressionCodec case other => throw new SparkException( s"Unsupported Arrow compression codec for Comet's cache: $other. " + - "Supported values: none, zstd") + "Supported values: none, zstd, lz4") } // Decompressors are stateless and shared. Resolving one per cached batch would allocate a codec - // per batch on every scan, and the enum lookup walks the CodecType values each time. + // per batch on every scan, and the enum lookup walks the CodecType values each time. LZ4_FRAME + // is read with the codec that writes it rather than with Arrow's pure-Java one, which inflates + // through heap copies and streams. private val readCodecs: Map[CompressionUtil.CodecType, CompressionCodec] = CompressionUtil.CodecType .values() .filter(_ != CompressionUtil.CodecType.NO_COMPRESSION) - .map(t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t)) + .map { + case CompressionUtil.CodecType.LZ4_FRAME => + CompressionUtil.CodecType.LZ4_FRAME -> Lz4FrameCompressionCodec + case t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t) + } .toMap /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala new file mode 100644 index 00000000000..b0cc91501b4 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala @@ -0,0 +1,209 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} +import org.apache.arrow.vector.compression.{AbstractCompressionCodec, CompressionUtil} +import org.apache.spark.SparkException +import org.apache.spark.internal.Logging + +import net.jpountz.lz4.{LZ4Compressor, LZ4Exception, LZ4Factory, LZ4SafeDecompressor} + +/** + * Arrow's LZ4_FRAME compression codec, on lz4-java. + * + * Arrow ships an LZ4_FRAME codec of its own, but it is commons-compress's pure-Java LZ4, which + * writes orders of magnitude more slowly than zstd and produces larger output. This writes the + * same format with lz4-java, the JNI-accelerated library behind `spark.io.compression.codec`: + * each buffer is one standard LZ4 frame behind Arrow's 8-byte uncompressed-length prefix, which + * `AbstractCompressionCodec` reads and writes, so any Arrow reader can decode it. lz4-java comes + * from Spark's own classpath, as zstd-jni does for Arrow's zstd codec. + * + * Every frame has the same header: independent blocks of up to 4 MiB, and no checksums, content + * size or dictionary. The reader accepts no other, since it only ever reads frames this codec + * wrote. It checks every length the frame records before reading through it, because Comet turns + * off Arrow's own bounds checks (`arrow.enable_unsafe_memory_access`), so a corrupt length would + * otherwise read or write past the buffers. + * + * Lengths are written in the platform's byte order, as the rest of the payload is, which is the + * little-endian order the format requires on every platform Comet runs on. + */ +private[comet] object Lz4FrameCompressionCodec extends AbstractCompressionCodec with Logging { + + private val MaxBlockSize = 4 * 1024 * 1024 + + // The magic number, then the frame descriptor: FLG (version 01, independent blocks), BD (a + // 4 MiB maximum block size) and HC, the second byte of the xxHash32 of FLG and BD with a zero + // seed. HC is a constant because FLG and BD are. + private val Magic = 0x184d2204 + private val Flg: Byte = 0x60 + private val Bd: Byte = 0x70 + private val Hc: Byte = 0x73 + private val HeaderSize = 7 + + // Each block is led by its length, whose high bit marks a block stored as it was because + // compressing it would not have made it smaller. A zero length ends the frame. + private val LengthSize = 4 + private val StoredBlock = 0x80000000 + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + // fastestInstance() takes lz4-java's JNI implementation only when lz4-java was loaded by the + // system class loader, and silently settles for its Java one otherwise, which is the case under + // `mvn exec:java`. Ask for the JNI one directly, and use Java only where it cannot load. + private lazy val factory: LZ4Factory = + try LZ4Factory.nativeInstance() + catch { + case e @ (NonFatal(_) | _: LinkageError) => + logWarning( + "Could not load lz4-java's native library, so Comet's in-memory cache uses its " + + "slower Java implementation of lz4", + e) + LZ4Factory.fastestJavaInstance() + } + private lazy val compressor: LZ4Compressor = factory.fastCompressor() + private lazy val decompressor: LZ4SafeDecompressor = factory.safeDecompressor() + + override def getCodecType: CompressionUtil.CodecType = CompressionUtil.CodecType.LZ4_FRAME + + override protected def doCompress( + allocator: BufferAllocator, + uncompressed: ArrowBuf): ArrowBuf = { + val length = uncompressed.writerIndex() + val fullBlocks = length / MaxBlockSize + val lastBlock = (length % MaxBlockSize).toInt + val bound = Prefix + HeaderSize + + fullBlocks * (LengthSize + compressor.maxCompressedLength(MaxBlockSize)) + + (if (lastBlock > 0) LengthSize + compressor.maxCompressedLength(lastBlock) else 0) + + LengthSize + val frame = allocator.buffer(bound) + try { + frame.setInt(Prefix, Magic) + frame.setByte(Prefix + 4, Flg) + frame.setByte(Prefix + 5, Bd) + frame.setByte(Prefix + 6, Hc) + var at = Prefix + HeaderSize + var from = 0L + while (from < length) { + val n = math.min(MaxBlockSize.toLong, length - from).toInt + val room = compressor.maxCompressedLength(n) + val packed = compressor.compress( + uncompressed.nioBuffer(from, n), + 0, + n, + frame.nioBuffer(at + LengthSize, room), + 0, + room) + if (packed < n) { + frame.setInt(at, packed) + at += LengthSize + packed + } else { + frame.setInt(at, n | StoredBlock) + frame.setBytes(at + LengthSize, uncompressed, from, n.toLong) + at += LengthSize + n + } + from += n + } + frame.setInt(at, 0) + frame.writerIndex(at + LengthSize) + frame + } catch { + case NonFatal(e) => + frame.close() + throw e + } + } + + override protected def doDecompress( + allocator: BufferAllocator, + compressed: ArrowBuf): ArrowBuf = { + val length = readUncompressedLength(compressed) + val end = compressed.writerIndex() + if (end < Prefix + HeaderSize || + compressed.getInt(Prefix) != Magic || + compressed.getByte(Prefix + 4) != Flg || + compressed.getByte(Prefix + 5) != Bd || + compressed.getByte(Prefix + 6) != Hc) { + throw corrupt("the buffer does not begin with the frame header Comet writes") + } + + val plain = allocator.buffer(length) + try { + var at = Prefix + HeaderSize + var to = 0L + var ended = false + while (!ended) { + if (end - at < LengthSize) { + throw corrupt("the frame ends before its end mark") + } + val word = compressed.getInt(at) + at += LengthSize + if (word == 0) { + ended = true + } else { + val n = word & ~StoredBlock + if (n > end - at) { + throw corrupt(s"a block of $n bytes runs past the end of the frame") + } + if ((word & StoredBlock) != 0) { + if (n > length - to) { + throw corrupt(s"the frame holds more than its recorded $length bytes") + } + plain.setBytes(to, compressed, at, n.toLong) + to += n + } else { + // A block never decodes to more than the maximum block size, and must not decode + // past the length the frame records. + val room = math.min(length - to, MaxBlockSize.toLong).toInt + val decoded = + try { + decompressor.decompress( + compressed.nioBuffer(at, n), + 0, + n, + plain.nioBuffer(to, room), + 0, + room) + } catch { + case e: LZ4Exception => + throw corrupt(s"a block does not decode within $room bytes", e) + } + to += decoded + } + at += n + } + } + if (to != length) { + throw corrupt(s"the frame decodes to $to bytes rather than its recorded $length") + } + plain.writerIndex(length) + plain + } catch { + case NonFatal(e) => + plain.close() + throw e + } + } + + private def corrupt(reason: String, cause: Throwable = null): SparkException = + new SparkException(s"Comet cached batch holds a corrupt LZ4 frame: $reason", cause) +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index be0eb109db5..b657cf1526f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -20,6 +20,7 @@ package org.apache.comet.exec import java.{util => ju} +import java.io.ByteArrayInputStream import java.nio.charset.StandardCharsets import scala.concurrent.{Await, Future} @@ -32,8 +33,8 @@ import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, VarBinaryVector, VarCharVector} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} import org.apache.arrow.vector.types.pojo.ArrowType +import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin -import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} @@ -57,6 +58,8 @@ import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plu import org.apache.comet.rules.CometCacheColumnarRule import org.apache.comet.vector.{CometPlainVector, CometVector} +import net.jpountz.lz4.LZ4FrameInputStream + class CometInMemoryCacheSuite extends CometTestBase { import testImplicits._ @@ -2180,7 +2183,7 @@ class CometInMemoryCacheSuite extends CometTestBase { // Every codec the config accepts, not just the default. `none` takes a different path on read // -- the payload records no codec, so nothing is decompressed -- and shipped broken for a // while because the only tests that ran were on the default codec. - Seq("none", "zstd").foreach { codec => + Seq("none", "zstd", "lz4").foreach { codec => withConversions( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", @@ -2227,6 +2230,103 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache writes lz4 buffers as standard LZ4 frames") { + // Arrow's LZ4_FRAME codec is the standard LZ4 frame format, so check what the codec writes + // against lz4-java's own frame reader rather than against the codec that wrote it. The inputs + // span the codec's 4 MiB block: a buffer of one block, and one of three whose first block is + // incompressible, which a frame stores verbatim. + val random = new scala.util.Random(7) + def compressible(n: Int): Array[Byte] = Array.tabulate[Byte](n)(i => (i / 16).toByte) + val inputs = Seq( + "one block" -> compressible(100 * 1000), + "three blocks" -> + (Array.fill[Byte](4 << 20)(random.nextInt().toByte) ++ compressible(5 << 20))) + val codec = CometCachedBatchHelper.compressionCodec("lz4") + assert(codec.getCodecType == CompressionUtil.CodecType.LZ4_FRAME) + + val allocator = CometArrowAllocator.newChildAllocator("lz4-frames", 0, Long.MaxValue) + try { + inputs.foreach { case (label, data) => + val plain = allocator.buffer(data.length.toLong) + plain.setBytes(0, data) + plain.writerIndex(data.length.toLong) + // Takes ownership of `plain`. + val packed = codec.compress(allocator, plain) + try { + // Arrow prefixes a compressed buffer with its uncompressed length, and a buffer it + // stored verbatim with -1. + assert(packed.getLong(0) == data.length, s"$label was stored verbatim") + val frame = new Array[Byte]((packed.writerIndex() - 8).toInt) + packed.getBytes(8, frame) + val in = new LZ4FrameInputStream(new ByteArrayInputStream(frame)) + val decoded = + try in.readAllBytes() + finally in.close() + assert(decoded.sameElements(data), s"lz4-java read $label back differently") + } finally { + packed.close() + } + } + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } + } + + test("Comet in-memory cache rejects a corrupt lz4 frame without leaking") { + // A corrupt frame has to fail as itself, naming the codec, and release the buffer it was + // decoding into. Offsets are from the LZ4 frame format behind Arrow's 8-byte + // uncompressed-length prefix: the frame's 4-byte magic number at 8, then a 3-byte frame + // descriptor, then the first block's 4-byte length at 15. A frame ends with a 4-byte end + // mark. Truncated frames are decoded from copies of exactly their length, so reading past + // one is an error of its own rather than a read of the bytes that follow. + val data = Array.tabulate[Byte](100 * 1000)(i => (i / 16).toByte) + val codec = CometCachedBatchHelper.compressionCodec("lz4") + + val allocator = CometArrowAllocator.newChildAllocator("lz4-corrupt", 0, Long.MaxValue) + try { + def truncated(frame: ArrowBuf, length: Long): ArrowBuf = { + val copy = allocator.buffer(length) + copy.setBytes(0, frame, 0, length) + copy.writerIndex(length) + copy + } + // Each corruption returns the buffer to decode: the frame itself, or a truncated copy. + val corruptions: Seq[(String, ArrowBuf => ArrowBuf)] = Seq( + "a buffer that is not an LZ4 frame" -> { b => b.setInt(8, 0); b }, + "a frame cut off inside its header" -> (b => truncated(b, 12)), + "a frame cut off before its end mark" -> (b => truncated(b, b.writerIndex() - 4)), + "a block that runs past the end of the frame" -> { b => b.setInt(15, 0x7fffff00); b }, + "a frame that decodes to less than its recorded length" -> { b => + b.setLong(0, data.length + 1L); b + }, + "a frame that decodes to more than its recorded length" -> { b => + b.setLong(0, data.length - 1L); b + }) + + corruptions.foreach { case (label, corrupt) => + val plain = allocator.buffer(data.length.toLong) + plain.setBytes(0, data) + plain.writerIndex(data.length.toLong) + val packed = codec.compress(allocator, plain) + assert(packed.getLong(0) == data.length, "the test needs a compressed buffer") + val input = corrupt(packed) + try { + val thrown = intercept[SparkException](codec.decompress(allocator, input).close()) + assert(thrown.getMessage.contains("LZ4"), s"$label failed with: $thrown") + } finally { + // A failed decompress leaves its input to the caller, as Arrow's own codecs do. + Seq(packed, input).distinct.foreach { b => + if (b.getReferenceManager.getRefCount > 0) b.close() + } + } + assert(allocator.getAllocatedMemory == 0, s"$label leaked") + } + } finally { + allocator.close() + } + } + test("Comet in-memory cache rejects a payload that records an unknown compression codec") { // `CodecType.fromCompressionType` answers NO_COMPRESSION for any byte outside its enum, so a // reader that took its word for it would read a compressed body as plain bytes and hand back @@ -2822,8 +2922,12 @@ class CometInMemoryCacheSuite extends CometTestBase { // reachable from nothing the failure path can see. The second is the one that catches a leak // in `VectorLoader`; the first is the one that catches a cleanup path releasing the shared // body twice. Both run over a multi-chunk payload too, where the copy into the body walks - // several chunks before the failure. - chunkSizes.foreach { case (_, chunkSize) => + // several chunks before the failure, and under each codec that compresses, since each fails + // in a decoder of its own. + for { + codec <- Seq("zstd", "lz4") + (_, chunkSize) <- chunkSizes + } withSQLConf(CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> codec) { withProjectionCache(chunkSize) { (relation, batches) => val cacheSchema = Utils.fromAttributes(relation.output) val pristine = CometCachedBatchHelper.snapshotPayloads(batches) @@ -2852,7 +2956,7 @@ class CometInMemoryCacheSuite extends CometTestBase { } assert( CometArrowAllocator.getAllocatedMemory == before, - s"everything allocated before a failure in $where must be released") + s"everything allocated before a failure in $where under $codec must be released") } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index 230126a7025..a7b8412eb72 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -19,8 +19,11 @@ package org.apache.spark.sql.benchmark +import java.io.PrintStream import java.nio.charset.StandardCharsets +import scala.collection.mutable + import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.SparkSession @@ -73,11 +76,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { Seq("id", "sc", "deep", "wide", "tail", "d"), nestedNumRows) - // Every value spark.comet.exec.inMemoryCache.compression.codec accepts, default first. Arrow's - // other IPC codec, LZ4_FRAME, is not one of them: it is commons-compress's pure-Java LZ4 rather - // than the JNI-accelerated lz4-java behind spark.io.compression.codec, and the write path - // rejects it. - private val codecs = Seq("zstd", "none") + // Every value spark.comet.exec.inMemoryCache.compression.codec accepts, default first. + private val codecs = Seq("zstd", "lz4", "none") // A sink for the benchmarked call's result: written but never read, so that neither the // compiler nor the JIT can treat `gatherColumnStats` as dead code. Not `private`, because @@ -444,19 +444,29 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * CometInMemoryTableScan's batches through CometColumnarToRow instead, which * runAdaptiveBenchmark measures. * + * Comet's format is cached under zstd, the default, and under lz4, since most of what a Spark + * operator pays to read it is decompression. + * * Both formats are cached from the same relation, one copy at a time as in runCodecBenchmark, * and each case checks which serializer cached the relation it reads and which reader it uses. */ private def runSparkOperatorBenchmark(relation: CachedRelation): Unit = { val view = s"${relation.table}_spark_operators" - val reads = Seq( - SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), - SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf), - SparkOperatorRead( - "Comet's cache format, fused reader", - cometSerializer, - fusedReaderConf, - fused = true)) + val reads = SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf) +: + Seq("zstd", "lz4").flatMap { codec => + Seq( + SparkOperatorRead( + s"Comet's cache format, $codec, row reader", + cometSerializer, + sparkOperatorConf, + codec = codec), + SparkOperatorRead( + s"Comet's cache format, $codec, fused reader", + cometSerializer, + fusedReaderConf, + fused = true, + codec = codec)) + } spark.catalog.clearCache() withTempTable(view) { @@ -465,14 +475,16 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .createOrReplaceTempView(view) val cache = new OneCachedCopy(view) - readShapes(view).foreach { case (label, query, scanned) => + val benchmarks = readShapes(view).map { case (label, query, scanned) => val benchmark = new Benchmark( s"in-memory cache read by Spark operators, $label", relation.rows.toLong, output = output) reads.foreach(addSparkOperatorCase(benchmark, cache, query, scanned, _)) benchmark.run() + benchmark } + cache.printFootprints(benchmarks.last.out) spark.catalog.uncacheTable(view) } @@ -488,12 +500,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * spark.sql.codegen.maxFields (100 by default) that is how any Spark operator reads a cached * relation, because the scan stops offering columnar output and nothing can fuse with it. Every * width holds about the same number of values, so the cases differ in width rather than in data - * volume. + * volume. The columns are sequential longs, which Spark's format delta-encodes, so the + * footprint of each format is printed too. */ private def runWideSparkOperatorBenchmark(): Unit = { - val reads = Seq( - SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), - SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf)) + val reads = SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf) +: + Seq("zstd", "lz4").map { codec => + SparkOperatorRead( + s"Comet's cache format, $codec, row reader", + cometSerializer, + sparkOperatorConf, + codec = codec) + } Seq(100, 200, 1500).foreach { width => val rows = 20L * 1000 * 1000 / width @@ -513,37 +531,57 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { output = output) reads.foreach(addSparkOperatorCase(benchmark, cache, s"SELECT * FROM $view", width, _)) benchmark.run() + cache.printFootprints(benchmark.out) spark.catalog.uncacheTable(view) } } } - /** How Spark operators read a cached relation: its format and the settings of the read. */ + /** + * How Spark operators read a cached relation: its format, the codec Comet's format is written + * with, and the settings of the read. + */ private case class SparkOperatorRead( name: String, serializer: String, conf: Seq[(String, String)], - fused: Boolean = false) + fused: Boolean = false, + codec: String = "zstd") /** * Holds one cached copy of `view` at a time, cached under `conf`. Two copies could not coexist * anyway: the cache manager keys on the plan rather than the name, so a second one would find * the first. + * + * Records what each format occupies as it is cached, since a format that reads faster is only + * worth having at a footprint the application can afford. */ private class OneCachedCopy(view: String, conf: Seq[(String, String)] = sparkOperatorConf) { - private var cachedBy: String = _ + private var cachedBy: (String, String) = _ + private val footprints = mutable.LinkedHashMap.empty[String, Long] - def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { - spark.catalog.uncacheTable(view) - cachedBy = null - withCacheSerializer(serializer) { - withSQLConf(conf: _*) { - spark.catalog.cacheTable(view) - spark.table(view).count() + // The codec only matters to Comet's serializer, and Spark's ignores its config. + def cacheBy(serializer: String, codec: String = "zstd"): Unit = + if (cachedBy != ((serializer, codec))) { + spark.catalog.uncacheTable(view) + cachedBy = null + withCacheSerializer(serializer) { + withSQLConf(conf ++ codecConf(codec): _*) { + spark.catalog.cacheTable(view) + spark.table(view).count() + } } + cachedBy = (serializer, codec) + val format = + if (serializer == sparkSerializer) "Spark's cache format" + else s"Comet's cache format, $codec" + footprints(format) = cachedBytes(view) } - cachedBy = serializer + + // As in runCodecBenchmark, footprint has no column in a Benchmark table. + def printFootprints(out: PrintStream): Unit = footprints.foreach { case (format, bytes) => + out.println(f"Cached footprint ($format): ${bytes / (1024.0 * 1024.0)}%.1f MiB") } } @@ -557,7 +595,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // Re-caching in this case's format is setup, so it is outside the timer, and it only happens // on the case's first call, which is a warmup iteration. benchmark.addTimerCase(read.name) { timer => - cache.cacheBy(read.serializer) + cache.cacheBy(read.serializer, read.codec) withSQLConf(read.conf: _*) { if (!verified) { verifySparkOperatorRead(query, scanned, read) @@ -573,8 +611,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { /** * What the feature changes for a query that runs with Comet, against Spark's own cache format, * with AQE on and Comet's other settings at their defaults. This is the comparison an - * application gets from turning the feature on or off. Both formats are cached from the same - * relation, one copy at a time as in runSparkOperatorBenchmark. + * application gets from turning the feature on or off, and, through lz4, from changing the + * codec. Both formats are cached from the same relation, one copy at a time as in + * runSparkOperatorBenchmark. * * Two shapes of plan read the cache. With Comet operators above the cache scan, Comet's format * runs the whole query natively, while Spark's leaves the operators directly above its scan on @@ -586,8 +625,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { */ private def runAdaptiveBenchmark(relation: CachedRelation): Unit = { val view = s"${relation.table}_adaptive" - val formats = - Seq("Spark's cache format" -> sparkSerializer, "Comet's cache format" -> cometSerializer) + val formats = ("Spark's cache format", sparkSerializer, "zstd") +: + Seq("zstd", "lz4").map(codec => (s"Comet's cache format, $codec", cometSerializer, codec)) val operatorsAbove = Seq( "Comet operators" -> Seq.empty[(String, String)], "a Spark operator" -> Seq(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false")) @@ -607,12 +646,12 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { s"in-memory cache with AQE, $operators above the scan, $label", relation.rows.toLong, output = output) - formats.foreach { case (name, serializer) => + formats.foreach { case (name, serializer, codec) => var verified = false // Re-caching in this case's format is setup, so it is outside the timer, and it only // happens on the case's first call, which is a warmup iteration. benchmark.addTimerCase(name) { timer => - cache.cacheBy(serializer) + cache.cacheBy(serializer, codec) withSQLConf(adaptiveConf ++ operatorConf: _*) { if (!verified) { verifyAdaptiveRead(query, scanned, serializer, operatorConf.nonEmpty) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala index 47e876b5f86..7789e13ed7f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -125,6 +125,14 @@ object CometCachedBatchHelper { chunkSize: Int = 1024 * 1024): ChunkedByteBuffer = CachedBatchIpc.serialize(batch, codec, allocator, chunkSize)._1 + /** + * The codec the cache writer uses for a `spark.comet.exec.inMemoryCache.compression.codec` + * value. A thin shim, for tests about a codec's own encoding rather than about a cached + * relation. + */ + def compressionCodec(name: String): CompressionCodec = + CachedBatchIpc.compressionCodec(name, zstdLevel = 1) + /** A payload [[serialize]] wrote, as the cached batch the writer would have stored it in. */ def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch = CometCachedBatch(numRows, InternalRow.empty, payload) From 7321b42f3b4aceb1e86c9bae6f6e3b2276471664 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Sun, 11 Oct 2026 13:37:48 +0000 Subject: [PATCH 64/89] chore: re-enable unused-parameter warnings in strict profile (#6846) --- pom.xml | 84 +++++- .../apache/comet/expressions/CometCast.scala | 7 +- .../CometSpark34AqeDppFallbackRule.scala | 9 +- .../apache/comet/serde/QueryPlanSerde.scala | 3 - .../org/apache/comet/serde/aggregates.scala | 34 +-- .../org/apache/comet/serde/arithmetic.scala | 10 - .../scala/org/apache/comet/serde/arrays.scala | 3 - .../org/apache/comet/serde/bitwise.scala | 5 - .../org/apache/comet/serde/predicates.scala | 13 - .../org/apache/comet/serde/strings.scala | 2 - .../spark/sql/comet/CometScanExec.scala | 10 +- .../comet/CometIcebergWriteActionSuite.scala | 259 ++++++++---------- .../CometIcebergWriteDetectionSuite.scala | 197 +++++-------- .../exec/CometNativeColumnarToRowSuite.scala | 2 +- .../CometConditionalExpressionBenchmark.scala | 18 +- 15 files changed, 282 insertions(+), 374 deletions(-) diff --git a/pom.xml b/pom.xml index a020f0719f9..a2970a2c6ff 100644 --- a/pom.xml +++ b/pom.xml @@ -884,18 +884,16 @@ under the License. execution's `args` replaces the plugin-level list instead of appending to it, so each list below is self-contained. - Two lints are deliberately absent from both lists: + `-Ywarn-unused:params` is enabled for both main and test sources. Unused + helper parameters are removed. Category-scoped `-Wconf` filters retain required + signatures for JNI declarations, cross-version shims, extension hooks, scan + factories and reflective test fixtures. - `-Ywarn-unused:params` reports parameters that are unused because a signature - requires them rather than because they are dead: the `@native` declarations in - `Native.scala`, which have no body to use them in; cross-version shims under - `src/main/spark-*` that match the signature of the Spark version they shim; and - overridable defaults and serde helpers whose parameter is part of their public - shape. Suppressing these one by one with `@nowarn` does not work across profiles: - some shared sources warn under Scala 2.12 but not 2.13 (`CometScanContrib.scala`, - for example), so an annotation that silences one profile is unused on the other, - which `-Xlint:_` reports through `-Ywarn-unused:nowarn` and `-Xfatal-warnings` - turns into a build failure. + These filters avoid adding profile-dependent `@nowarn` annotations: an annotation + needed under one Scala version can be unused under another. Existing shuffle + annotations for Scala 2.13 private-type warnings have a scoped unused-nowarn + exception on Scala 2.12. `-Xlint:unused` enables + `-Ywarn-unused:nowarn`, which `-Xfatal-warnings` promotes to an error. `-Ywarn-value-discard` stays on for main sources, where a discarded result is usually a dropped builder or a swallowed return, but is off for test sources, @@ -922,7 +920,37 @@ under the License. -Ywarn-dead-code -Ywarn-numeric-widen -Ywarn-value-discard - -Ywarn-unused:imports,patvars,privates,locals,-implicits + -Ywarn-unused:imports,patvars,privates,locals,params,-implicits + + -Wconf:cat=unused-params&site=org\.apache\.comet\.Native\..*:s + -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometExpressionSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometOperatorSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsSparkPartialToNativeFinal\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsNativePartialToSparkFinal\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometScanWithPlanData\.withDynamicPruningFilters\..*:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.CometExecIterator\.InputFailure\.unapply\.nativeFailure:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.udf\.CometUDF\.evaluate\.(partitionIndex|planId):s + -Wconf:cat=unused-params&site=org\.apache\.comet\.udf\.CometUDF\.releasePlan\.planId:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometBaseAggregate\.operatorSupportLevel\.op:s + + -Wconf:cat=unused-nowarn&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.Comet(Celeborn)?ShuffleManager\.get(Reader|Writer):s + + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometIcebergNativeScanExec\.apply\.session:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometNativeScanExec\.apply\.session:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometScanExec\.apply\.session:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.NoNoArgCtorScanContrib\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.RecordingCryptoAwareCelebornClientFactory\.get\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.RecordingLegacyCelebornClientFactory\.get\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.LegacyMapperEndCelebornPushClient\.cleanup\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.WrongReturnTypeCelebornPushClient\.pushOrMergeData\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.CometCelebornFallbackTestClient\.reportShuffleFetchFailure\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.CometCelebornFallbackTestClient\.cleanupShuffle\..*:s -Xfatal-warnings @@ -937,7 +965,37 @@ under the License. -Xlint:_ -Ywarn-dead-code -Ywarn-numeric-widen - -Ywarn-unused:imports,patvars,privates,locals,-implicits + -Ywarn-unused:imports,patvars,privates,locals,params,-implicits + + -Wconf:cat=unused-params&site=org\.apache\.comet\.Native\..*:s + -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometExpressionSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometOperatorSerde\.getSupportLevel\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsSparkPartialToNativeFinal\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsNativePartialToSparkFinal\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometScanWithPlanData\.withDynamicPruningFilters\..*:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.CometExecIterator\.InputFailure\.unapply\.nativeFailure:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.udf\.CometUDF\.evaluate\.(partitionIndex|planId):s + -Wconf:cat=unused-params&site=org\.apache\.comet\.udf\.CometUDF\.releasePlan\.planId:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometBaseAggregate\.operatorSupportLevel\.op:s + + -Wconf:cat=unused-nowarn&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.Comet(Celeborn)?ShuffleManager\.get(Reader|Writer):s + + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometIcebergNativeScanExec\.apply\.session:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometNativeScanExec\.apply\.session:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometScanExec\.apply\.session:s + + -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.NoNoArgCtorScanContrib\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.RecordingCryptoAwareCelebornClientFactory\.get\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.RecordingLegacyCelebornClientFactory\.get\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.LegacyMapperEndCelebornPushClient\.cleanup\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.shuffle\.WrongReturnTypeCelebornPushClient\.pushOrMergeData\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.CometCelebornFallbackTestClient\.reportShuffleFetchFailure\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.execution\.shuffle\.CometCelebornFallbackTestClient\.cleanupShuffle\..*:s -Xfatal-warnings diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index b6fcdf32806..d37c621f18e 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -243,7 +243,7 @@ object CometCast case (_: DecimalType, _: DecimalType) => Compatible() case (DataTypes.StringType, _) => - canCastFromString(toType, timeZoneId, evalMode) + canCastFromString(toType) case (_, DataTypes.StringType) => canCastToString(fromType, timeZoneId, evalMode) case (DataTypes.TimestampType, _) => @@ -301,10 +301,7 @@ object CometCast } } - private def canCastFromString( - toType: DataType, - timeZoneId: Option[String], - evalMode: CometEvalMode.Value): SupportLevel = { + private def canCastFromString(toType: DataType): SupportLevel = { toType match { case DataTypes.BooleanType => Compatible() diff --git a/spark/src/main/scala/org/apache/comet/rules/CometSpark34AqeDppFallbackRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometSpark34AqeDppFallbackRule.scala index afedc1601c9..ea8653071dc 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometSpark34AqeDppFallbackRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometSpark34AqeDppFallbackRule.scala @@ -126,8 +126,8 @@ case object CometSpark34AqeDppFallbackRule // AQE re-optimization path: see case 4 in the class-level docstring. On subsequent // re-optimize cycles Spark's PlanAdaptiveSubqueries hands us a SubqueryBroadcastExec (not // the original SAB), and the freshly-planned main-BHJ build BE has no tag carried over. - sbScans.foreach { case (scan, sb) => - tagForSubqueryBroadcast(plan, scan, sb) + sbScans.foreach { case (_, sb) => + tagForSubqueryBroadcast(plan, sb) } // This rule only tags; it never rewrites the plan structurally. @@ -227,10 +227,7 @@ case object CometSpark34AqeDppFallbackRule * plan. No peer-scan tagging path is needed because the SMJ self-join case (SPARK-32509) never * produces a `SubqueryBroadcastExec`. */ - private def tagForSubqueryBroadcast( - plan: SparkPlan, - scan: SparkPlan, - sb: SubqueryBroadcastExec): Unit = { + private def tagForSubqueryBroadcast(plan: SparkPlan, sb: SubqueryBroadcastExec): Unit = { val keyIds: Set[Any] = sb.buildKeys.flatMap(_.references.map(_.exprId)).toSet if (keyIds.isEmpty) { logWarning(s"SubqueryBroadcast '${sb.name}' has empty buildKeys; skipping") diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index c0bfb3489bd..43960ad5604 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1140,7 +1140,6 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * Some(Expr) or None if not supported */ def createUnaryExpr( - expr: Expression, child: Expression, inputs: Seq[Attribute], binding: Boolean, @@ -1166,7 +1165,6 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } def createBinaryExpr( - expr: Expression, left: Expression, right: Expression, inputs: Seq[Attribute], @@ -1208,7 +1206,6 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * each combined `BinaryExpr` as `And` or `Or`. */ def createBalancedBinaryExpr( - expr: Expression, operands: Seq[Expression], inputs: Seq[Attribute], binding: Boolean, 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 1ed8af82d73..1222bdd9476 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -542,8 +542,7 @@ trait CometCovBase { nullOnDivideByZero: Boolean, statsType: Int, inputs: Seq[Attribute], - binding: Boolean, - conf: SQLConf): Option[ExprOuterClass.AggExpr] = { + binding: Boolean): Option[ExprOuterClass.AggExpr] = { val child1Expr = exprToProto(cov.left, inputs, binding) val child2Expr = exprToProto(cov.right, inputs, binding) val dataType = serializeDataType(cov.dataType) @@ -575,14 +574,7 @@ object CometCovSample extends CometAggregateExpressionSerde[CovSample] with Come inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertCov( - aggExpr, - covSample, - covSample.nullOnDivideByZero, - 0, - inputs, - binding, - conf: SQLConf) + convertCov(aggExpr, covSample, covSample.nullOnDivideByZero, 0, inputs, binding) } } @@ -593,20 +585,12 @@ object CometCovPopulation extends CometAggregateExpressionSerde[CovPopulation] w inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertCov( - aggExpr, - covPopulation, - covPopulation.nullOnDivideByZero, - 1, - inputs, - binding, - conf: SQLConf) + convertCov(aggExpr, covPopulation, covPopulation.nullOnDivideByZero, 1, inputs, binding) } } trait CometVariance { def convertVariance( - aggExpr: AggregateExpression, expr: CentralMomentAgg, nullOnDivideByZero: Boolean, statsType: Int, @@ -641,7 +625,7 @@ object CometVarianceSamp extends CometAggregateExpressionSerde[VarianceSamp] wit inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertVariance(aggExpr, variance, variance.nullOnDivideByZero, 0, inputs, binding) + convertVariance(variance, variance.nullOnDivideByZero, 0, inputs, binding) } } @@ -652,19 +636,17 @@ object CometVariancePop extends CometAggregateExpressionSerde[VariancePop] with inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertVariance(aggExpr, variance, variance.nullOnDivideByZero, 1, inputs, binding) + convertVariance(variance, variance.nullOnDivideByZero, 1, inputs, binding) } } trait CometStddev { def convertStddev( - aggExpr: AggregateExpression, stddev: CentralMomentAgg, nullOnDivideByZero: Boolean, statsType: Int, inputs: Seq[Attribute], - binding: Boolean, - conf: SQLConf): Option[ExprOuterClass.AggExpr] = { + binding: Boolean): Option[ExprOuterClass.AggExpr] = { val child = stddev.child val childExpr = exprToProto(child, inputs, binding) val dataType = serializeDataType(stddev.dataType) @@ -694,7 +676,7 @@ object CometStddevSamp extends CometAggregateExpressionSerde[StddevSamp] with Co inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertStddev(aggExpr, stddev, stddev.nullOnDivideByZero, 0, inputs, binding, conf: SQLConf) + convertStddev(stddev, stddev.nullOnDivideByZero, 0, inputs, binding) } } @@ -705,7 +687,7 @@ object CometStddevPop extends CometAggregateExpressionSerde[StddevPop] with Come inputs: Seq[Attribute], binding: Boolean, conf: SQLConf): Option[ExprOuterClass.AggExpr] = { - convertStddev(aggExpr, stddev, stddev.nullOnDivideByZero, 1, inputs, binding, conf: SQLConf) + convertStddev(stddev, stddev.nullOnDivideByZero, 1, inputs, binding) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala index 23bca9fa15f..751495be2c0 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arithmetic.scala @@ -30,7 +30,6 @@ import org.apache.comet.shims.CometEvalModeUtil trait MathBase { def createMathExpression( - expr: Expression, left: Expression, right: Expression, inputs: Seq[Attribute], @@ -115,7 +114,6 @@ trait MathBase { * single `returnType` / `evalMode` is correct for every inner node. */ def createBalancedMathExpr( - expr: Expression, operands: Seq[Expression], inputs: Seq[Attribute], binding: Boolean, @@ -166,7 +164,6 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { { case _: Add => true; case _ => false }, { case a: Add => (a.left, a.right) }) createBalancedMathExpr( - expr, operands, inputs, binding, @@ -175,7 +172,6 @@ object CometAdd extends CometExpressionSerde[Add] with MathBase { (builder, mathExpr) => builder.setAdd(mathExpr)) } else { createMathExpression( - expr, expr.left, expr.right, inputs, @@ -197,7 +193,6 @@ object CometSubtract extends CometExpressionSerde[Subtract] with MathBase { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createMathExpression( - expr, expr.left, expr.right, inputs, @@ -225,7 +220,6 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { { case _: Multiply => true; case _ => false }, { case m: Multiply => (m.left, m.right) }) createBalancedMathExpr( - expr, operands, inputs, binding, @@ -234,7 +228,6 @@ object CometMultiply extends CometExpressionSerde[Multiply] with MathBase { (builder, mathExpr) => builder.setMultiply(mathExpr)) } else { createMathExpression( - expr, expr.left, expr.right, inputs, @@ -269,7 +262,6 @@ object CometDivide extends CometExpressionSerde[Divide] with MathBase { val rightExpr = if (expr.evalMode != EvalMode.ANSI) nullIfWhenPrimitive(expr.right) else expr.right createMathExpression( - expr, expr.left, rightExpr, inputs, @@ -310,7 +302,6 @@ object CometIntegralDivide extends CometExpressionSerde[IntegralDivide] with Mat } val divideExpr = createMathExpression( - expr, left, rightExpr, inputs, @@ -356,7 +347,6 @@ object CometRemainder extends CometExpressionSerde[Remainder] with MathBase { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createMathExpression( - expr, expr.left, expr.right, inputs, diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 67ab955296d..a16d42c94da 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -86,7 +86,6 @@ object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBas // IS NOT NULL does not care about element nullability, so use the un-widened source and skip // serializing a redundant cast. val isNotNullExpr = createUnaryExpr( - expr, expr.children.head, inputs, binding, @@ -861,7 +860,6 @@ object CometElementAt extends CometExpressionSerde[ElementAt] { // CometArraysZip. if (needsNullGuard(expr)) { val isNotNullExpr = createUnaryExpr( - expr, expr.left, inputs, binding, @@ -974,7 +972,6 @@ object CometSize extends CometExpressionSerde[Size] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createUnaryExpr( - expr, expr.child, inputs, binding, diff --git a/spark/src/main/scala/org/apache/comet/serde/bitwise.scala b/spark/src/main/scala/org/apache/comet/serde/bitwise.scala index ec06c712535..7fdad7c49f5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/bitwise.scala +++ b/spark/src/main/scala/org/apache/comet/serde/bitwise.scala @@ -37,7 +37,6 @@ object CometBitwiseAnd extends CometExpressionSerde[BitwiseAnd] { { case _: BitwiseAnd => true; case _ => false }, { case b: BitwiseAnd => (b.left, b.right) }) createBalancedBinaryExpr( - expr, operands, inputs, binding, @@ -67,7 +66,6 @@ object CometBitwiseOr extends CometExpressionSerde[BitwiseOr] { { case _: BitwiseOr => true; case _ => false }, { case b: BitwiseOr => (b.left, b.right) }) createBalancedBinaryExpr( - expr, operands, inputs, binding, @@ -85,7 +83,6 @@ object CometBitwiseXor extends CometExpressionSerde[BitwiseXor] { { case _: BitwiseXor => true; case _ => false }, { case b: BitwiseXor => (b.left, b.right) }) createBalancedBinaryExpr( - expr, operands, inputs, binding, @@ -107,7 +104,6 @@ object CometShiftRight extends CometExpressionSerde[ShiftRight] { } createBinaryExpr( - expr, expr.left, rightExpression, inputs, @@ -130,7 +126,6 @@ object CometShiftLeft extends CometExpressionSerde[ShiftLeft] { } createBinaryExpr( - expr, expr.left, rightExpression, inputs, diff --git a/spark/src/main/scala/org/apache/comet/serde/predicates.scala b/spark/src/main/scala/org/apache/comet/serde/predicates.scala index 129979cff71..446568d44f0 100644 --- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/predicates.scala @@ -47,7 +47,6 @@ object CometNot extends CometExpressionSerde[Not] { expr.child match { case inner: EqualTo if hasNativePath(CometEqualTo, inner) => createBinaryExpr( - inner, inner.left, inner.right, inputs, @@ -55,7 +54,6 @@ object CometNot extends CometExpressionSerde[Not] { (builder, binaryExpr) => builder.setNeq(binaryExpr)) case inner: EqualNullSafe if hasNativePath(CometEqualNullSafe, inner) => createBinaryExpr( - inner, inner.left, inner.right, inputs, @@ -71,7 +69,6 @@ object CometNot extends CometExpressionSerde[Not] { // native, or returns None, which cascades this Not to None and falls the enclosing // operator back to Spark. createUnaryExpr( - expr, expr.child, inputs, binding, @@ -93,7 +90,6 @@ object CometAnd extends CometExpressionSerde[And] { { case _: And => true; case _ => false }, { case a: And => (a.left, a.right) }) createBalancedBinaryExpr( - expr, operands, inputs, binding, @@ -111,7 +107,6 @@ object CometOr extends CometExpressionSerde[Or] { { case _: Or => true; case _ => false }, { case o: Or => (o.left, o.right) }) createBalancedBinaryExpr( - expr, operands, inputs, binding, @@ -125,7 +120,6 @@ object CometEqualTo extends CollationAwareBinaryPredicate[EqualTo] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -140,7 +134,6 @@ object CometEqualNullSafe extends CollationAwareBinaryPredicate[EqualNullSafe] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -155,7 +148,6 @@ object CometGreaterThan extends CollationAwareBinaryPredicate[GreaterThan] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -170,7 +162,6 @@ object CometGreaterThanOrEqual extends CollationAwareBinaryPredicate[GreaterThan inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -185,7 +176,6 @@ object CometLessThan extends CollationAwareBinaryPredicate[LessThan] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -200,7 +190,6 @@ object CometLessThanOrEqual extends CollationAwareBinaryPredicate[LessThanOrEqua inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -215,7 +204,6 @@ object CometIsNull extends CometExpressionSerde[IsNull] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createUnaryExpr( - expr, expr.child, inputs, binding, @@ -229,7 +217,6 @@ object CometIsNotNull extends CometExpressionSerde[IsNotNull] { inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { createUnaryExpr( - expr, expr.child, inputs, binding, diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 68d0e252df8..26ebad05a96 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -418,7 +418,6 @@ object CometLike extends CometExpressionSerde[Like] with CodegenDispatchFallback override def convert(expr: Like, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { createBinaryExpr( - expr, expr.left, expr.right, inputs, @@ -507,7 +506,6 @@ object CometRLike val allowIncompat = CometConf.isExprAllowIncompat(getExprConfigName(expr)) if (provablyCompatible(expr) || (allowIncompat && nativeApplicable(expr))) { return createBinaryExpr( - expr, expr.left, expr.right, inputs, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala index 3e1beec716d..26464c2ab70 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala @@ -253,10 +253,7 @@ case class CometScanExec( */ def getFilePartitions(): Seq[FilePartition] = { val filePartitions = if (bucketedScan) { - createFilePartitionsForBucketedScan( - relation.bucketSpec.get, - dynamicallySelectedPartitions, - relation) + createFilePartitionsForBucketedScan(relation.bucketSpec.get, dynamicallySelectedPartitions) } else { createFilePartitionsForNonBucketedScan(dynamicallySelectedPartitions, relation) } @@ -271,13 +268,10 @@ case class CometScanExec( * the bucketing spec. * @param selectedPartitions * Hive-style partition that are part of the read. - * @param fsRelation - * [[HadoopFsRelation]] associated with the read. */ private def createFilePartitionsForBucketedScan( bucketSpec: BucketSpec, - selectedPartitions: Array[PartitionDirectory], - fsRelation: HadoopFsRelation): Seq[FilePartition] = { + selectedPartitions: Array[PartitionDirectory]): Seq[FilePartition] = { logInfo(s"Planning with ${bucketSpec.numBuckets} buckets") val filesGroupedToBuckets = selectedPartitions diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 36a7a235246..5f19ac20430 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -105,8 +105,8 @@ class CometIcebergWriteActionSuite test("AppendData unpartitioned INSERT INTO routes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "append_unpart", partitionSpec = "") + withIcebergCatalog { _ => + createTable("append_unpart", partitionSpec = "") val snapshot = captureWrite("append_unpart") { spark.sql( "INSERT INTO cat.db.append_unpart VALUES " + @@ -119,8 +119,8 @@ class CometIcebergWriteActionSuite test("spark.comet.enabled=false keeps Spark's own write plan with the split flag on") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "comet_disabled", partitionSpec = "") + withIcebergCatalog { _ => + createTable("comet_disabled", partitionSpec = "") // withSQLConf returns Unit before Spark 4.0, so the assertions run inside it. withSQLConf(CometConf.COMET_ENABLED.key -> "false") { val snapshot = captureWrite("comet_disabled") { @@ -141,8 +141,8 @@ class CometIcebergWriteActionSuite test("AppendData partitioned INSERT INTO routes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "append_part", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("append_part", partitionSpec = "PARTITIONED BY (region)") val snapshot = captureWrite("append_part") { spark.sql( "INSERT INTO cat.db.append_part VALUES " + @@ -155,9 +155,9 @@ class CometIcebergWriteActionSuite test("AppendData INSERT FROM SELECT survives the intervening exchange/sort") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "src", partitionSpec = "") - createTable(warehouseDir, "append_from_select", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("src", partitionSpec = "") + createTable("append_from_select", partitionSpec = "PARTITIONED BY (region)") spark.sql( "INSERT INTO cat.db.src VALUES " + "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)") @@ -174,8 +174,8 @@ class CometIcebergWriteActionSuite test("AppendData on an empty source still emits a single commit") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "empty_target", partitionSpec = "") + withIcebergCatalog { _ => + createTable("empty_target", partitionSpec = "") val snapshot = captureWrite("empty_target") { spark.sql( "INSERT INTO cat.db.empty_target SELECT id, region, amount " + @@ -188,8 +188,8 @@ class CometIcebergWriteActionSuite test("AppendData from a zero-partition RDD still runs one write task and commits") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "zero_part", partitionSpec = "") + withIcebergCatalog { _ => + createTable("zero_part", partitionSpec = "") val schema = StructType( Seq( StructField("id", IntegerType), @@ -214,8 +214,8 @@ class CometIcebergWriteActionSuite test("AQE re-plan of the writer subtree writes and commits exactly once") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "aqe_replan", partitionSpec = "") + withIcebergCatalog { _ => + createTable("aqe_replan", partitionSpec = "") val session = spark import session.implicits._ (1 to 100) @@ -257,8 +257,8 @@ class CometIcebergWriteActionSuite test("partitioned write under AQE keeps the clustered writer working across the shuffle") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "aqe_part", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("aqe_part", partitionSpec = "PARTITIONED BY (region)") val session = spark import session.implicits._ (1 to 500) @@ -287,8 +287,8 @@ class CometIcebergWriteActionSuite test("multi-partition partitioned write collects one commit message per task") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "multi_task", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("multi_task", partitionSpec = "PARTITIONED BY (region)") val session = spark import session.implicits._ (1 to 100) @@ -321,8 +321,8 @@ class CometIcebergWriteActionSuite test("cached table reflects an append that follows a schema change on Spark 4.1+") { assume(icebergAvailable, "Iceberg not available in classpath") assume(isSpark41Plus, "name-based cache refresh needs Spark 4.1+") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "cache_refresh", partitionSpec = "") + withIcebergCatalog { _ => + createTable("cache_refresh", partitionSpec = "") spark.sql(s"INSERT INTO $catalog.$ns.cache_refresh VALUES (1, 'us-east', 10.0)") spark.sql(s"CACHE TABLE $catalog.$ns.cache_refresh") assert(spark.sql(s"SELECT * FROM $catalog.$ns.cache_refresh").count() == 1) @@ -347,8 +347,8 @@ class CometIcebergWriteActionSuite test("OverwriteByExpression replaces existing rows via two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "overwrite_static", partitionSpec = "") + withIcebergCatalog { _ => + createTable("overwrite_static", partitionSpec = "") spark.sql( "INSERT INTO cat.db.overwrite_static VALUES " + "(1, 'old', 1.0), (2, 'old', 2.0), (3, 'old', 3.0)") @@ -367,8 +367,8 @@ class CometIcebergWriteActionSuite test("OverwritePartitionsDynamic replaces only touched partitions") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "overwrite_dynamic", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("overwrite_dynamic", partitionSpec = "PARTITIONED BY (region)") spark.sql( "INSERT INTO cat.db.overwrite_dynamic VALUES " + "(1, 'us-east', 1.0), (2, 'us-west', 2.0), (3, 'eu', 3.0)") @@ -390,9 +390,8 @@ class CometIcebergWriteActionSuite test("ReplaceData (CoW DELETE) on a row predicate goes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "cow_delete", partitionSpec = "", properties = Some("'write.delete.mode'='copy-on-write'")) @@ -412,9 +411,8 @@ class CometIcebergWriteActionSuite test("ReplaceData (CoW UPDATE) routes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "cow_update", partitionSpec = "", properties = Some("'write.update.mode'='copy-on-write'")) @@ -435,9 +433,8 @@ class CometIcebergWriteActionSuite test("ReplaceData (CoW MERGE) with matched and unmatched legs routes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "cow_merge", partitionSpec = "", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -462,7 +459,7 @@ class CometIcebergWriteActionSuite assumeNativeAcceleration() assume(isSpark41Plus, "MERGE write summaries need Spark 4.1+") assume(icebergVersionAtLeast(1, 11), "Iceberg snapshot MERGE summaries need Iceberg 1.11+") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => def merge(table: String): Unit = { spark.sql(s""" |MERGE INTO $catalog.$ns.$table t @@ -503,7 +500,6 @@ class CometIcebergWriteActionSuite val sparkTable = s"merge_summary_spark_$adaptiveEnabled" Seq(nativeTable, sparkTable).foreach { table => createTable( - warehouseDir, table, partitionSpec = "", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -548,8 +544,8 @@ class CometIcebergWriteActionSuite test("failed write job aborts and leaves the table unchanged") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "task_fail", partitionSpec = "") + withIcebergCatalog { _ => + createTable("task_fail", partitionSpec = "") coalesceInsert("task_fail", Seq((1, "us-east", 10.0))) val session = spark import session.implicits._ @@ -581,9 +577,8 @@ class CometIcebergWriteActionSuite test("commit-time validation still sees a conflicting concurrent append") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "conflict", partitionSpec = "", properties = Some( @@ -635,11 +630,11 @@ class CometIcebergWriteActionSuite test("a commit-time failure surfaces as the same exception as Spark's own write") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => // A serializable overwrite validated from a snapshot older than a matching append fails // Iceberg's commit-time validation deterministically, without any concurrency. def failCommit(table: String): Throwable = { - createTable(warehouseDir, table, partitionSpec = "") + createTable(table, partitionSpec = "") coalesceInsert(table, Seq((1, "us-east", 10.0))) val validateFrom = spark .sql(s"SELECT snapshot_id FROM $catalog.$ns.$table.snapshots") @@ -689,7 +684,7 @@ class CometIcebergWriteActionSuite spark.read.parquet(srcDir.getAbsolutePath).createOrReplaceTempView("job_fail_src") def failJob(table: String): Throwable = { - createTable(warehouseDir, table, partitionSpec = "") + createTable(table, partitionSpec = "") coalesceInsert(table, Seq((1, "us-east", 10.0))) val e = intercept[Exception] { spark @@ -748,7 +743,7 @@ class CometIcebergWriteActionSuite test(s"WriteDelta $command is recognized and stays on Iceberg JVM DeltaWriter") { assume(icebergAvailable, "Iceberg not available in classpath") assume(isSpark35Plus, "WriteDelta interception starts with Spark 3.5") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val suffix = command.toLowerCase(java.util.Locale.ROOT) val cometTable = s"write_delta_${suffix}_comet" val jvmTable = s"write_delta_${suffix}_jvm" @@ -758,7 +753,6 @@ class CometIcebergWriteActionSuite Seq(cometTable, jvmTable).foreach { table => createTable( - warehouseDir, table, partitionSpec = "PARTITIONED BY (region)", properties = Some(properties)) @@ -874,10 +868,9 @@ class CometIcebergWriteActionSuite test("sanity check: Spark's default DELETE path works against a Hadoop catalog") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { createTable( - warehouseDir, "spark_cow_delete", partitionSpec = "", properties = Some("'write.delete.mode'='copy-on-write'")) @@ -892,8 +885,8 @@ class CometIcebergWriteActionSuite test("disabled config falls through to Spark's V2ExistingTableWriteExec") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "disabled_conf", partitionSpec = "") + withIcebergCatalog { _ => + createTable("disabled_conf", partitionSpec = "") val snapshot = captureWrite("disabled_conf") { withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { @@ -913,9 +906,9 @@ class CometIcebergWriteActionSuite test("native acceleration: AppendData INSERT FROM SELECT") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_source", partitionSpec = "") - createTable(warehouseDir, "native_target", partitionSpec = "") + withIcebergCatalog { _ => + createTable("native_source", partitionSpec = "") + createTable("native_target", partitionSpec = "") spark.sql( "INSERT INTO cat.db.native_source VALUES " + "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)") @@ -928,8 +921,8 @@ class CometIcebergWriteActionSuite test("native acceleration: AppendData unpartitioned VALUES") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_append_values", partitionSpec = "") + withIcebergCatalog { _ => + createTable("native_append_values", partitionSpec = "") assertNativeWriteEngages("native_append_values", Seq(1, 2, 3)) { spark.sql( "INSERT INTO cat.db.native_append_values VALUES " + @@ -941,12 +934,12 @@ class CometIcebergWriteActionSuite for (adaptive <- Seq(false, true)) { test(s"transition-heavy fallback preserves Iceberg writes with AQE=$adaptive") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val suffix = if (adaptive) "aqe" else "no_aqe" val nativeTable = s"transition_native_$suffix" val fallbackTable = s"transition_fallback_$suffix" - createTable(warehouseDir, nativeTable, partitionSpec = "") - createTable(warehouseDir, fallbackTable, partitionSpec = "") + createTable(nativeTable, partitionSpec = "") + createTable(fallbackTable, partitionSpec = "") val values = "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)" withSQLConf( @@ -985,8 +978,8 @@ class CometIcebergWriteActionSuite val fallbackTable = s"transition_cow_fallback_$suffix" val props = Some("'write.delete.mode'='copy-on-write'") val spec = "PARTITIONED BY (region)" - createTable(warehouseDir, nativeTable, spec, props) - createTable(warehouseDir, fallbackTable, spec, props) + createTable(nativeTable, spec, props) + createTable(fallbackTable, spec, props) val seed = Seq( (1, "us-east", 10.0), @@ -1053,11 +1046,11 @@ class CometIcebergWriteActionSuite // otherwise, so pin parity against the JVM writer on the same runtime instead of a literal. test("native acceleration: appended files carry the same sort_order_id as the JVM writer") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => // WRITE ORDERED BY (provided by IcebergSparkSessionExtensions, enabled in this suite) // bumps the table's sort order id to a non-default value (1). Seq("sorted_native", "sorted_jvm").foreach { t => - createTable(warehouseDir, t, partitionSpec = "") + createTable(t, partitionSpec = "") spark.sql(s"ALTER TABLE cat.db.$t WRITE ORDERED BY id") } val insert = (t: String) => @@ -1081,8 +1074,8 @@ class CometIcebergWriteActionSuite test("native acceleration: AppendData partitioned by identity") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_append_part", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("native_append_part", partitionSpec = "PARTITIONED BY (region)") assertNativeWriteEngages("native_append_part", Seq(1, 2, 3)) { spark.sql( "INSERT INTO cat.db.native_append_part VALUES " + @@ -1093,8 +1086,8 @@ class CometIcebergWriteActionSuite test("native acceleration: OverwriteByExpression (INSERT OVERWRITE STATIC)") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_overwrite_static", partitionSpec = "") + withIcebergCatalog { _ => + createTable("native_overwrite_static", partitionSpec = "") spark.sql( "INSERT INTO cat.db.native_overwrite_static VALUES " + "(1, 'old', 1.0), (2, 'old', 2.0), (3, 'old', 3.0)") @@ -1110,8 +1103,8 @@ class CometIcebergWriteActionSuite test("native acceleration: OverwritePartitionsDynamic") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_overwrite_dyn", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("native_overwrite_dyn", partitionSpec = "PARTITIONED BY (region)") spark.sql( "INSERT INTO cat.db.native_overwrite_dyn VALUES " + "(1, 'us-east', 1.0), (2, 'us-west', 2.0), (3, 'eu', 3.0)") @@ -1125,9 +1118,8 @@ class CometIcebergWriteActionSuite test("native acceleration: ReplaceData (CoW DELETE)") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_delete", partitionSpec = "", properties = Some("'write.delete.mode'='copy-on-write'")) @@ -1153,7 +1145,6 @@ class CometIcebergWriteActionSuite withIcebergCatalog { warehouseDir => Seq("v3_native", "v3_jvm").foreach { t => createTable( - warehouseDir, t, partitionSpec = "PARTITIONED BY (region)", properties = Some("'format-version'='3'")) @@ -1224,10 +1215,9 @@ class CometIcebergWriteActionSuite test("native acceleration: format-version=3 row lineage matches iceberg-java") { assumeNativeAcceleration() assume(icebergVersionAtLeast(1, 10), "Spark reads Iceberg row lineage from Iceberg 1.10 on") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => Seq("lineage_native", "lineage_jvm").foreach { t => createTable( - warehouseDir, t, partitionSpec = "", properties = Some("'format-version'='3', 'write.update.mode'='copy-on-write'")) @@ -1290,13 +1280,9 @@ class CometIcebergWriteActionSuite test("native acceleration: a format-version=3 column with a default matches iceberg-java") { assumeNativeAcceleration() assume(icebergVersionAtLeast(1, 10), "Column defaults require Iceberg 1.10+") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => Seq("default_native", "default_jvm").foreach { t => - createTable( - warehouseDir, - t, - partitionSpec = "", - properties = Some("'format-version'='3'")) + createTable(t, partitionSpec = "", properties = Some("'format-version'='3'")) addIcebergIntColumnWithDefault(loadIcebergTable(spark, catalog, ns, t), "c", 7) spark.sql(s"REFRESH TABLE $catalog.$ns.$t") } @@ -1320,9 +1306,9 @@ class CometIcebergWriteActionSuite test("native acceleration: a driver commit failure aborts all completed task files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val table = "native_commit_failure" - createTable(warehouseDir, table, partitionSpec = "") + createTable(table, partitionSpec = "") // Native writes stay off for the seed. The second append is a committed change after // `validateFrom`, so the overwrite below asks Iceberg for serializable validation from // that older snapshot. The check runs inside BatchWrite.commit, after the native tasks @@ -1403,9 +1389,8 @@ class CometIcebergWriteActionSuite // `ColumnarBatch cannot be cast to InternalRow`. test("native acceleration: ReplaceData (CoW DELETE) with AQE disabled") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_delete_no_aqe", partitionSpec = "PARTITIONED BY (region)", properties = Some("'write.delete.mode'='copy-on-write'")) @@ -1447,9 +1432,8 @@ class CometIcebergWriteActionSuite test("native acceleration: ReplaceData (CoW UPDATE)") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_update", partitionSpec = "", properties = Some("'write.update.mode'='copy-on-write'")) @@ -1472,9 +1456,8 @@ class CometIcebergWriteActionSuite test("native acceleration: ReplaceData (CoW MERGE) honors the versioned native contract") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_merge", partitionSpec = "", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -1531,7 +1514,7 @@ class CometIcebergWriteActionSuite test("MERGE matches Spark on partitioned Iceberg copy-on-write and merge-on-read") { assumeNativeAcceleration() assume(isSpark35Plus, "native MergeRows requires Spark 3.5+") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => spark .range(0, 20000, 1, 8) .selectExpr( @@ -1575,7 +1558,6 @@ class CometIcebergWriteActionSuite val sparkTable = s"merge_parity_${suffix}_aqe_${adaptiveEnabled}_spark" Seq(nativeTable, sparkTable).foreach { table => createTable( - warehouseDir, table, partitionSpec = "PARTITIONED BY (region)", properties = Some(s"'format-version'='2', 'write.merge.mode'='$mode'")) @@ -1689,12 +1671,11 @@ class CometIcebergWriteActionSuite test("native MergeRows Iceberg cardinality violation matches Spark") { assumeNativeAcceleration() assume(isSpark35Plus, "native MergeRows requires Spark 3.5+") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val nativeTable = "merge_cardinality_native" val sparkTable = "merge_cardinality_spark" Seq(nativeTable, sparkTable).foreach { table => createTable( - warehouseDir, table, partitionSpec = "PARTITIONED BY (region)", properties = Some("'format-version'='2', 'write.merge.mode'='copy-on-write'")) @@ -1762,9 +1743,8 @@ class CometIcebergWriteActionSuite // shapes, which put different operators between the columnar CoW scan and the write. test("native acceleration: ReplaceData (CoW UPDATE) with AQE disabled") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_update_no_aqe", partitionSpec = "PARTITIONED BY (region)", properties = Some("'write.update.mode'='copy-on-write'")) @@ -1798,9 +1778,8 @@ class CometIcebergWriteActionSuite // The subtree below the write mixes a Spark-columnar `BatchScan`, a row-based MergeRows // path, and Comet operators, so it needs transitions inserted in three different places. assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_cow_merge_no_aqe", partitionSpec = "PARTITIONED BY (region)", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -2081,8 +2060,8 @@ class CometIcebergWriteActionSuite test("native acceleration: RTAS replaces table contents through the native writer") { assumeNativeAcceleration() assume(isSpark35Plus, "RTAS re-plans its inner append only on Spark 3.5+") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "rtas_native", partitionSpec = "") + withIcebergCatalog { _ => + createTable("rtas_native", partitionSpec = "") coalesceInsert("rtas_native", Seq((1, "old", 1.0))) val snapshot = withNativeEnabled { captureWrite("rtas_native") { @@ -2415,9 +2394,8 @@ class CometIcebergWriteActionSuite test("native acceleration: fanout writer handles unsorted partitioned input") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_fanout", partitionSpec = "PARTITIONED BY (region)", properties = Some("'write.spark.fanout.enabled'='true'")) @@ -2450,10 +2428,10 @@ class CometIcebergWriteActionSuite // literal, so an iceberg-java wording change surfaces here as a parity failure. test("native acceleration: clustered writer rejects unclustered input like the JVM writer") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val (nativeTable, jvmTable) = ("unclustered_native", "unclustered_jvm") val tables = Seq(nativeTable, jvmTable) - tables.foreach(createTable(warehouseDir, _, partitionSpec = "PARTITIONED BY (region)")) + tables.foreach(createTable(_, partitionSpec = "PARTITIONED BY (region)")) val session = spark import session.implicits._ // "us-east" is revisited after the writer has moved on to "eu" and closed it. @@ -2574,9 +2552,8 @@ class CometIcebergWriteActionSuite test("native acceleration: target-file-size rolls on iceberg-java's 1000-row cadence") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "native_roll", partitionSpec = "", properties = Some("'write.target-file-size-bytes'='1'")) @@ -2765,8 +2742,8 @@ class CometIcebergWriteActionSuite test("native acceleration: empty append commits exactly once with zero data files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "native_empty", partitionSpec = "") + withIcebergCatalog { _ => + createTable("native_empty", partitionSpec = "") val snapshot = withNativeEnabled { captureWrite("native_empty") { spark.sql( @@ -2793,7 +2770,7 @@ class CometIcebergWriteActionSuite test("native acceleration: empty append to a partitioned table commits with zero data files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => // The unpartitioned test above ends in UnpartitionedWriter::close() with nothing written; // these end in ClusteredWriter::close() (the partitioned default) and FanoutWriter::close() // respectively, with the partition splitter never invoked at all. @@ -2801,7 +2778,7 @@ class CometIcebergWriteActionSuite ("native_empty_clustered", None), ("native_empty_fanout", Some("'write.spark.fanout.enabled'='true'"))).foreach { case (table, props) => - createTable(warehouseDir, table, partitionSpec = "PARTITIONED BY (region)", props) + createTable(table, partitionSpec = "PARTITIONED BY (region)", props) val snapshot = withNativeEnabled { captureWrite(table) { spark.sql(s"INSERT INTO $catalog.$ns.$table SELECT id, region, amount " + @@ -2971,8 +2948,8 @@ class CometIcebergWriteActionSuite test("native acceleration: partition paths are URL-escaped like iceberg-java") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "escaped_native", partitionSpec = "PARTITIONED BY (region)") - createTable(warehouseDir, "escaped_jvm", partitionSpec = "PARTITIONED BY (region)") + createTable("escaped_native", partitionSpec = "PARTITIONED BY (region)") + createTable("escaped_jvm", partitionSpec = "PARTITIONED BY (region)") val regions = Seq("a/b", "c#d", "e?f", "g h", "i=j", "k%l", "m+n", "*-._", "日本", "") val values = regions.zipWithIndex .map { case (region, i) => s"($i, '$region', $i.5)" } @@ -3014,7 +2991,7 @@ class CometIcebergWriteActionSuite // native path must too. test("native acceleration: a failed task deletes the data files it already finalized") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val session = spark import session.implicits._ (1 to 10000) @@ -3033,11 +3010,7 @@ class CometIcebergWriteActionSuite withNativeEnabled(withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "1000") { // Control: the same source and settings without the failure roll into several files, so // the failing run below really does have finalized files to clean up. - createTable( - warehouseDir, - "cleanup_control", - partitionSpec = "", - properties = rollingProps) + createTable("cleanup_control", partitionSpec = "", properties = rollingProps) val controlPlans = capturePlans(spark) { spark.sql( s"INSERT INTO $catalog.$ns.cleanup_control SELECT id, region, amount FROM cleanup_src") @@ -3049,7 +3022,7 @@ class CometIcebergWriteActionSuite val controlFiles = parquetFiles(dataDir("cleanup_control")) assert(controlFiles.size >= 3, s"expected the writer to roll files, got $controlFiles") - createTable(warehouseDir, "cleanup_target", partitionSpec = "", properties = rollingProps) + createTable("cleanup_target", partitionSpec = "", properties = rollingProps) coalesceInsert("cleanup_target", Seq((0, "seed", 0.0))) val committed = parquetFiles(dataDir("cleanup_target")) assert(committed.size == 1) @@ -3148,8 +3121,8 @@ class CometIcebergWriteActionSuite test("native acceleration: a pre-native handoff failure cleans up task files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "pre_handoff_target", partitionSpec = "") + withIcebergCatalog { _ => + createTable("pre_handoff_target", partitionSpec = "") coalesceInsert("pre_handoff_target", Seq((0, "seed", 0.0))) val before = countSnapshots("pre_handoff_target") val root = dataDir("pre_handoff_target").toPath.toAbsolutePath @@ -3221,8 +3194,8 @@ class CometIcebergWriteActionSuite test("native acceleration: a post-native handoff failure cleans up task files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "handoff_target", partitionSpec = "") + withIcebergCatalog { _ => + createTable("handoff_target", partitionSpec = "") coalesceInsert("handoff_target", Seq((0, "seed", 0.0))) val before = countSnapshots("handoff_target") val root = dataDir("handoff_target").toPath.toAbsolutePath @@ -3301,7 +3274,7 @@ class CometIcebergWriteActionSuite test("native acceleration: a mid-write failure retries without orphan files") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val session = spark import session.implicits._ (1 to 10000) @@ -3313,7 +3286,6 @@ class CometIcebergWriteActionSuite // choice and writes nothing before then. A 1000-row page lets it finalize files well before // the failure at row 7000, which the default 20000-row page would not. createTable( - warehouseDir, "retry_target", partitionSpec = "", properties = @@ -3379,9 +3351,9 @@ class CometIcebergWriteActionSuite Seq(true, false).foreach { native => test(s"a failed write job deletes the data files of tasks that completed (native=$native)") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val table = s"job_abort_$native" - createTable(warehouseDir, table, partitionSpec = "") + createTable(table, partitionSpec = "") coalesceInsert(table, Seq((0, "seed", 0.0))) val committed = parquetFiles(dataDir(table)) val before = countSnapshots(table) @@ -3448,8 +3420,8 @@ class CometIcebergWriteActionSuite test("deleteFilesQuietly removes data files through the table FileIO and never throws") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "delete_quietly", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("delete_quietly", partitionSpec = "PARTITIONED BY (region)") spark.sql(s"INSERT INTO $catalog.$ns.delete_quietly VALUES (1, 'us', 1.0), (2, 'eu', 2.0)") val locations = spark .sql(s"SELECT file_path FROM $catalog.$ns.delete_quietly.files") @@ -3477,8 +3449,8 @@ class CometIcebergWriteActionSuite // -- the step the locations would otherwise have to be recovered from -- is still covered. test("the write cleanup listener deletes the locations the native writer reported") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "listener_cleanup", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("listener_cleanup", partitionSpec = "PARTITIONED BY (region)") spark.sql( s"INSERT INTO $catalog.$ns.listener_cleanup VALUES (1, 'us', 1.0), (2, 'eu', 2.0)") val locations = spark @@ -3644,9 +3616,9 @@ class CometIcebergWriteActionSuite test("Comet-written rows round-trip through Spark's reader unchanged") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "parity_comet", partitionSpec = "PARTITIONED BY (region)") - createTable(warehouseDir, "parity_spark", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("parity_comet", partitionSpec = "PARTITIONED BY (region)") + createTable("parity_spark", partitionSpec = "PARTITIONED BY (region)") spark.sql( "INSERT INTO cat.db.parity_comet VALUES " + @@ -3670,8 +3642,8 @@ class CometIcebergWriteActionSuite test("write custom metrics are registered on the committer only") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "metrics_once", partitionSpec = "") + withIcebergCatalog { _ => + createTable("metrics_once", partitionSpec = "") val snapshot = captureWrite("metrics_once") { spark.sql(s"INSERT INTO $catalog.$ns.metrics_once VALUES (1, 'us-east', 10.5)") } @@ -3718,8 +3690,8 @@ class CometIcebergWriteActionSuite test("AppendData writes NULLs in every column") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "nulls_all", partitionSpec = "") + withIcebergCatalog { _ => + createTable("nulls_all", partitionSpec = "") val snapshot = captureWrite("nulls_all") { spark.sql(s"""INSERT INTO $catalog.$ns.nulls_all VALUES (CAST(NULL AS INT), CAST(NULL AS STRING), CAST(NULL AS DOUBLE)), @@ -3760,9 +3732,8 @@ class CometIcebergWriteActionSuite test("CoW MERGE with only a matched DELETE leg routes through two-op") { assume(icebergAvailable, "Iceberg not available in classpath") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "merge_del", partitionSpec = "", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -3785,9 +3756,8 @@ class CometIcebergWriteActionSuite assume(icebergAvailable, "Iceberg not available in classpath") // Iceberg 1.5.x (the Spark 3.4 pairing) rejects the clause in its extensions parser. assume(isSpark35Plus, "NOT MATCHED BY SOURCE needs Iceberg 1.8+") - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => createTable( - warehouseDir, "merge_nmbs", partitionSpec = "", properties = Some("'write.merge.mode'='copy-on-write'")) @@ -3916,7 +3886,6 @@ class CometIcebergWriteActionSuite } private def createTable( - warehouseDir: File, tableName: String, partitionSpec: String, properties: Option[String] = None): Unit = { @@ -4172,13 +4141,9 @@ class CometIcebergWriteActionSuite // luck 1 time in 8!, where with two it would pass half the time. test("native acceleration: a fanout write lists its data files in a stable order") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => + withIcebergCatalog { _ => val fanout = Some("'write.spark.fanout.enabled'='true'") - createTable( - warehouseDir, - "fanout_order", - partitionSpec = "PARTITIONED BY (region)", - properties = fanout) + createTable("fanout_order", partitionSpec = "PARTITIONED BY (region)", properties = fanout) val values = (0 until 8).map(i => s"($i, 'r$i', $i.5)").mkString(", ") assertNativeWriteEngages("fanout_order", 0 until 8) { @@ -4201,16 +4166,14 @@ class CometIcebergWriteActionSuite test("write report records which writer ran each Iceberg write") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "report_parquet", partitionSpec = "PARTITIONED BY (region)") + withIcebergCatalog { _ => + createTable("report_parquet", partitionSpec = "PARTITIONED BY (region)") createTable( - warehouseDir, "report_orc", partitionSpec = "", properties = Some("'write.format.default'='orc'")) if (isSpark35Plus) { createTable( - warehouseDir, "report_delta", partitionSpec = "PARTITIONED BY (region)", properties = Some("'format-version'='2', 'write.delete.mode'='merge-on-read'")) @@ -4272,8 +4235,8 @@ class CometIcebergWriteActionSuite test("write report records a streaming micro-batch write") { assumeNativeAcceleration() - withIcebergCatalog { warehouseDir => - createTable(warehouseDir, "report_stream", partitionSpec = "") + withIcebergCatalog { _ => + createTable("report_stream", partitionSpec = "") withTempIcebergDir { dir => val source = new File(dir, "source").getAbsolutePath val session = spark diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index 7caffd0e7ee..430a92f97a2 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -66,15 +66,15 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("clean parquet V2 table planned as AppendData yields Compatible") { - withDetectionCatalog { dir => - createTable(dir, "ok", partitionSpec = "") + withDetectionCatalog { _ => + createTable("ok", partitionSpec = "") assertSupportLevelIs[Compatible]("ok") } } test("registration tags a fall-back reason on the write exec") { - withDetectionCatalog { dir => - createTable(dir, "tagged", partitionSpec = "") + withDetectionCatalog { _ => + createTable("tagged", partitionSpec = "") val writeExec = insertWriteExec("tagged") val reasons = writeExec.getTagValue(CometExplainInfo.FALLBACK_REASONS) assert( @@ -84,8 +84,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("plan-only mode leaves the write with Spark") { - withDetectionCatalog { dir => - createTable(dir, "plan_only", partitionSpec = "") + withDetectionCatalog { _ => + createTable("plan_only", partitionSpec = "") // IcebergWriteStrategy runs before CometRule, so it needs its own plan-only guard. // withSQLConf returns Unit on Spark 3.4/3.5, hence the var. var plan: SparkPlan = null @@ -104,8 +104,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("SparkWrite reflection helpers all resolve on the current Iceberg runtime") { - withDetectionCatalog { dir => - createTable(dir, "refl_probe", partitionSpec = "") + withDetectionCatalog { _ => + createTable("refl_probe", partitionSpec = "") val sparkWrite = IcebergReflection .getOuterSparkWrite(insertWriteExec("refl_probe").batchWrite) .getOrElse(fail("could not unwrap outer SparkWrite from BatchWrite")) @@ -141,8 +141,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when a session conf overrides the compression codec") { - withDetectionCatalog { dir => - createTable(dir, "session_codec", partitionSpec = "") + withDetectionCatalog { _ => + createTable("session_codec", partitionSpec = "") withSQLConf("spark.sql.iceberg.compression-codec" -> "gzip") { assertSupportLevelIs[Compatible]("session_codec") } @@ -150,9 +150,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.format.default=orc") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "fmt_orc", partitionSpec = "", properties = Some("'write.format.default'='orc'")) @@ -161,8 +160,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: per-write write-format option overrides parquet default") { - withDetectionCatalog { dir => - createTable(dir, "fmt_orc_opt", partitionSpec = "") + withDetectionCatalog { _ => + createTable("fmt_orc_opt", partitionSpec = "") assertUnsupportedContains( dfWriteExec("fmt_orc_opt", "write-format" -> "orc"), "fmt_orc_opt", @@ -172,9 +171,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.object-storage.enabled=true") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "obj_store", partitionSpec = "", properties = Some("'write.object-storage.enabled'='true'")) @@ -183,9 +181,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.location-provider.impl set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "loc_provider", partitionSpec = "", properties = Some("'write.location-provider.impl'='com.example.MyProvider'")) @@ -223,16 +220,16 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("Compatible for a format-version=3 append") { assume(isSpark35Plus, "V3 tables require Iceberg 1.8.1+ (Spark 3.5 profile)") - withDetectionCatalog { dir => - createTable(dir, "v3", partitionSpec = "", properties = Some("'format-version'='3'")) + withDetectionCatalog { _ => + createTable("v3", partitionSpec = "", properties = Some("'format-version'='3'")) assertSupportLevelIs[Compatible]("v3") } } test("fall-back: format-version=4") { assume(icebergVersionAtLeast(1, 10), "V4 tables require Iceberg 1.10+") - withDetectionCatalog { dir => - createTable(dir, "v4", partitionSpec = "", properties = Some("'format-version'='4'")) + withDetectionCatalog { _ => + createTable("v4", partitionSpec = "", properties = Some("'format-version'='4'")) assertUnsupportedContains("v4", "format-version=4") } } @@ -254,14 +251,10 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("fall-back: unknown column in the write schema") { assume(icebergVersionAtLeast(1, 10), "The unknown type requires Iceberg 1.10+") - withDetectionCatalog { dir => + withDetectionCatalog { _ => // Spark DDL cannot declare `unknown` (Iceberg plans it as Spark's NullType), so evolve the // schema through the Iceberg API. - createTable( - dir, - "unknown_col", - partitionSpec = "", - properties = Some("'format-version'='3'")) + createTable("unknown_col", partitionSpec = "", properties = Some("'format-version'='3'")) addIcebergColumn( loadIcebergTable(spark, catalog, ns, "unknown_col"), "u", @@ -275,9 +268,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: encryption.kms-client-impl set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "enc", partitionSpec = "", properties = Some("'encryption.kms-client-impl'='com.example.MyKms'")) @@ -288,10 +280,9 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // Metrics modes are not gated: manifest metrics are re-derived on the JVM with Iceberg's // own MetricsConfig logic before commit, so every mode behaves as it does on the java path. test("Compatible for every write.metadata.metrics mode") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => Seq("counts", "none", "full", "truncate(32)").zipWithIndex.foreach { case (mode, i) => createTable( - dir, s"metrics_mode_$i", partitionSpec = "", properties = Some(s"'write.metadata.metrics.default'='$mode'")) @@ -301,9 +292,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible for per-column metrics modes") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "metrics_col_modes", partitionSpec = "", properties = Some( @@ -316,9 +306,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // The JVM path fails such a write inside parquet-mr's codec setup, so allow the write failure // and pin only the fall-back reason. test("fall-back: non-integer write.parquet.compression-level") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "bad_level", partitionSpec = "", properties = Some("'write.parquet.compression-level'='fast'")) @@ -334,13 +323,12 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // codec-specific parquet-mr writer properties -- so these values write fine on the stock // path but parquet-rs rejects them at writer construction (zstd 1..=22, gzip 0..=9, // brotli 0..=11). They must decline up front rather than fail the task. - withDetectionCatalog { dir => + withDetectionCatalog { _ => val cases = Seq("zstd" -> "0", "zstd" -> "-3", "zstd" -> "23", "gzip" -> "-1", "brotli" -> "12") cases.zipWithIndex.foreach { case ((codec, level), i) => val table = s"bad_level_range_$i" createTable( - dir, table, partitionSpec = "", properties = Some( @@ -353,7 +341,6 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } // A boundary level parquet-rs accepts stays Compatible. createTable( - dir, "good_level", partitionSpec = "", properties = Some( @@ -363,9 +350,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.parquet.bloom-filter-max-bytes set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "bloom_max", partitionSpec = "", properties = Some("'write.parquet.bloom-filter-max-bytes'='524288'")) @@ -374,9 +360,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: per-column bloom filter enabled") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "bloom_col", partitionSpec = "", properties = Some("'write.parquet.bloom-filter-enabled.column.id'='true'")) @@ -388,9 +373,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when the schema exceeds max-inferred-column-defaults") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "too_many_cols", partitionSpec = "", properties = Some("'write.metadata.metrics.max-inferred-column-defaults'='2'")) @@ -399,9 +383,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: row-group-check-min-record-count non-default") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "rg_min", partitionSpec = "", properties = Some("'write.parquet.row-group-check-min-record-count'='500'")) @@ -410,9 +393,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when row-group-check-min-record-count is at default (100)") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "rg_min_default", partitionSpec = "", properties = Some("'write.parquet.row-group-check-min-record-count'='100'")) @@ -421,9 +403,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: row-group-check-max-record-count non-default") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "rg_max", partitionSpec = "", properties = Some("'write.parquet.row-group-check-max-record-count'='50000'")) @@ -432,9 +413,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.parquet.page-version=v2") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "page_v2", partitionSpec = "", properties = Some("'write.parquet.page-version'='v2'")) @@ -443,9 +423,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: parquet.enable.dictionary set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "enable_dict", partitionSpec = "", properties = Some("'parquet.enable.dictionary'='false'")) @@ -454,9 +433,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: per-column write.parquet.stats-enabled. set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "col_stats", partitionSpec = "", properties = Some("'write.parquet.stats-enabled.column.region'='false'")) @@ -465,9 +443,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: unvetted write.parquet.* property") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "unvetted", partitionSpec = "", properties = Some("'write.parquet.bloom-filter-adaptive-enabled'='true'")) @@ -479,9 +456,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: parquet.* table property other than enable.dictionary") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "pq_mr_prop", partitionSpec = "", properties = Some("'parquet.columnindex.truncate.length'='32'")) @@ -490,9 +466,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when a codec level side-channel property is set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "codec_level", partitionSpec = "", properties = Some("'zlib.compress.level'='9'")) @@ -501,8 +476,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: parquet.* key in the session Hadoop configuration") { - withDetectionCatalog { dir => - createTable(dir, "hadoop_conf", partitionSpec = "") + withDetectionCatalog { _ => + createTable("hadoop_conf", partitionSpec = "") withSQLConf("parquet.block.size" -> "1048576") { assertUnsupportedContains("hadoop_conf", "parquet.block.size", "Hadoop configuration") } @@ -515,8 +490,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // disable native Iceberg writes when it happens to be present in the session // Hadoop configuration. test("Compatible when only parquet.hadoop.vectored.io.enabled is set in Hadoop configuration") { - withDetectionCatalog { dir => - createTable(dir, "vectored_io_only", partitionSpec = "") + withDetectionCatalog { _ => + createTable("vectored_io_only", partitionSpec = "") withSQLConf("parquet.hadoop.vectored.io.enabled" -> "true") { assertSupportLevelIs[Compatible]("vectored_io_only") } @@ -524,9 +499,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: io-impl set") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "io_impl", partitionSpec = "", properties = Some("'io-impl'='com.example.MyFileIO'")) @@ -535,9 +509,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: data location URI scheme not supported by the native writer") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "bad_scheme", partitionSpec = "", properties = Some("'write.data.path'='hdfs://nonexistent.invalid/iceberg/db/bad_scheme'")) @@ -548,9 +521,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("fall-back: mixed-case data location scheme (native opens the location verbatim)") { // OpenDAL strips the scheme prefix from a path case-sensitively, so `S3://` cannot be opened // natively even though `s3://` can. The gate must match the scheme verbatim and decline it. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "mixed_case_scheme", partitionSpec = "", properties = @@ -561,9 +533,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("fall-back: hostless hdfs:/ data location is read as hdfs, not file") { // Hadoop normalises `hdfs:///p` to `hdfs:/p`; with no `://` the gate used to call it `file`. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "hostless_hdfs", partitionSpec = "", properties = Some("'write.data.path'='hdfs:/iceberg/db/hostless_hdfs'")) @@ -575,9 +546,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: s3 data location without a bucket in its authority") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "hostless_s3", partitionSpec = "", properties = Some("'write.data.path'='s3:/nonexistent-bucket/iceberg/db/hostless_s3'")) @@ -617,9 +587,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when the data location scheme is s3") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "s3_scheme", partitionSpec = "", properties = Some("'write.data.path'='s3://nonexistent-bucket/iceberg/db/s3_scheme'")) @@ -1029,9 +998,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("fall-back: unsupported Hadoop S3A setting on an S3 data location") { val secret = "SECRET_VALUE_MUST_NOT_APPEAR" withSQLConf("fs.s3a.encryption.algorithm" -> secret) { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "s3a_hadoop_unsupported", partitionSpec = "", properties = @@ -1054,8 +1022,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("S3-only setting gates do not affect a local data location") { - withDetectionCatalog { dir => - createTable(dir, "local_with_s3_conf", partitionSpec = "") + withDetectionCatalog { _ => + createTable("local_with_s3_conf", partitionSpec = "") withSQLConf("fs.s3a.encryption.algorithm" -> "SSE-KMS") { assertSupportLevelIs[Compatible]("local_with_s3_conf") } @@ -1144,9 +1112,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when the data location scheme is memory") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "memory_scheme", partitionSpec = "", properties = Some("'write.data.path'='memory://nonexistent/iceberg/db/memory_scheme'")) @@ -1157,9 +1124,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("fall-back: gs data location under HadoopFileIO (fs.gs.* is not forwarded)") { // The hadoop catalog's table.io() is a HadoopFileIO. Planned only, never executed: running // the write would have the Hadoop GCS connector look for credentials over the network. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "gs_hadoop_io", partitionSpec = "", properties = Some("'write.data.path'='gs://nonexistent/iceberg/db/gs_hadoop_io'")) @@ -1296,9 +1262,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: oss data location scheme (oss.* properties are not forwarded)") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "oss_scheme", partitionSpec = "", properties = Some("'write.data.path'='oss://nonexistent/iceberg/db/oss_scheme'")) @@ -1311,7 +1276,7 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // no values past Int.MaxValue, and parquet-mr rejects non-positive results at write time. // Each such value must fall back so the failure happens on the stock path, never be // silently normalised by the native translation. - withDetectionCatalog { dir => + withDetectionCatalog { _ => val keys = Seq( "write.parquet.row-group-size-bytes", "write.parquet.page-size-bytes", @@ -1321,13 +1286,12 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes keys.zipWithIndex.foreach { case (key, ki) => badValues.zipWithIndex.foreach { case (value, vi) => val table = s"bad_int_${ki}_$vi" - createTable(dir, table, partitionSpec = "", properties = Some(s"'$key'='$value'")) + createTable(table, partitionSpec = "", properties = Some(s"'$key'='$value'")) assertUnsupportedContainsAllowingWriteFailure(table, key) } } // A positive Java int is Compatible, pinning that the gate is not over-broad. createTable( - dir, "good_int", partitionSpec = "", properties = Some("'write.parquet.page-size-bytes'='1048576'")) @@ -1343,7 +1307,7 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // sessionState.newHadoopConf(), so plain fs.s3a.* keys set here are what the gate and // proto assembly see. LocalTableScan conversion is enabled the same way the write action // suite does, so the VALUES insert converts and the built proto is inspectable. - withDetectionCatalog { dir => + withDetectionCatalog { _ => val conf = spark.sessionState.conf conf.setConfString(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key, "true") conf.setConfString("fs.s3a.endpoint", "http://localhost:9000") @@ -1351,7 +1315,6 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes conf.setConfString("fs.s3a.path.style.access", "true") try { createTable( - dir, "s3a_props", partitionSpec = "", properties = Some("'write.data.path'='s3a://probe-bucket/iceberg/db/s3a_props'")) @@ -1416,7 +1379,6 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("Compatible when the data location is an explicit file:// path") { withDetectionCatalog { dir => createTable( - dir, "file_scheme", partitionSpec = "", properties = Some(s"'write.data.path'='file://${dir.getAbsolutePath}/file_scheme_data'")) @@ -1425,9 +1387,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: write.parquet.shred-variants=true") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "shred", partitionSpec = "", properties = Some("'write.parquet.shred-variants'='true'")) @@ -1439,9 +1400,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // Iceberg-Java's MetricsConfig is lenient: it warns and falls back to the default mode on // both paths, so the gate has nothing to protect; the JVM-side metrics assembly goes // through the same lenient parse. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "metrics_typo", partitionSpec = "", properties = Some("'write.metadata.metrics.default'='truncat(16)'")) @@ -1450,9 +1410,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when write.spark.fanout.enabled=true") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "fanout", partitionSpec = "PARTITIONED BY (bucket(4, id))", properties = Some("'write.spark.fanout.enabled'='true'")) @@ -1461,9 +1420,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when write.target-file-size-bytes is non-default") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "target_size", partitionSpec = "", properties = Some("'write.target-file-size-bytes'='1048576'")) @@ -1472,8 +1430,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("no fall-back reason is recorded when the iceberg write feature is disabled") { - withDetectionCatalog { dir => - createTable(dir, "flag_off", partitionSpec = "") + withDetectionCatalog { _ => + createTable("flag_off", partitionSpec = "") withSQLConf(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> "false") { val writeExec = insertWriteExec("flag_off") assert( @@ -1484,8 +1442,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: BatchWrite that is not an Iceberg SparkWrite") { - withDetectionCatalog { dir => - createTable(dir, "plain_write", partitionSpec = "") + withDetectionCatalog { _ => + createTable("plain_write", partitionSpec = "") val stub = new org.apache.spark.sql.connector.write.BatchWrite { override def createBatchWriterFactory( info: org.apache.spark.sql.connector.write.PhysicalWriteInfo) @@ -1504,15 +1462,15 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("Compatible when partitioned by a bucket transform") { - withDetectionCatalog { dir => - createTable(dir, "part_bucket", partitionSpec = "PARTITIONED BY (bucket(4, id))") + withDetectionCatalog { _ => + createTable("part_bucket", partitionSpec = "PARTITIONED BY (bucket(4, id))") assertSupportLevelIs[Compatible]("part_bucket") } } test("Compatible when partitioned by identity on a string column") { - withDetectionCatalog { dir => - createTable(dir, "part_string", partitionSpec = "PARTITIONED BY (region)") + withDetectionCatalog { _ => + createTable("part_string", partitionSpec = "PARTITIONED BY (region)") assertSupportLevelIs[Compatible]("part_string") } } @@ -1554,9 +1512,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // A format-version-1 spec keeps a dropped partition field as a `void` transform, and that // field's source column can be dropped afterwards. The surviving double field must still be // found, whatever the dropped one does to the spec's partition type. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "part_dropped", partitionSpec = "PARTITIONED BY (region, amount)", properties = Some("'format-version'='1'")) @@ -1590,9 +1547,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes test("Compatible when a dropped double partition field remains as void") { // The `void` field only ever holds null, so there are no signed zeros to keep apart. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "part_void", partitionSpec = "PARTITIONED BY (amount)", properties = Some("'format-version'='1'")) @@ -1614,9 +1570,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // an insert that is planned but not run. // https://github.com/apache/datafusion-comet/issues/6141 test("fall-back: void partition field whose source column was dropped, beside a live field") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "void_dropped", partitionSpec = "PARTITIONED BY (region, id)", properties = Some("'format-version'='1'")) @@ -1637,9 +1592,8 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes // Iceberg 1.8 neither writer can write to this table, since iceberg-java cannot bind the // original spec on the executors once its source column is gone, so only the gate's decision // is checked, on an insert that is planned but not run. - withDetectionCatalog { dir => + withDetectionCatalog { _ => createTable( - dir, "all_void_dropped", partitionSpec = "PARTITIONED BY (region)", properties = Some("'format-version'='1'")) @@ -1656,11 +1610,11 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } test("fall-back: uuid column in the write schema") { - withDetectionCatalog { dir => + withDetectionCatalog { _ => // Spark DDL cannot declare `uuid`, so evolve the schema through the Iceberg API. Spark // plans the column as StringType, but the native writer's target Arrow schema demands // FixedSizeBinary(16) with no cast from Utf8 -- detection must decline before execution. - createTable(dir, "uuid_col", partitionSpec = "") + createTable("uuid_col", partitionSpec = "") addIcebergColumn(loadIcebergTable(spark, catalog, ns, "uuid_col"), "u", icebergUuidType()) spark.sql(s"REFRESH TABLE $catalog.$ns.uuid_col") val writeExec = captureWriteExec("uuid_col", allowWriteFailure = false) { @@ -1688,7 +1642,6 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } private def createTable( - warehouseDir: File, tableName: String, partitionSpec: String, properties: Option[String] = None): Unit = { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala index 818be0b1ecd..2dc6fcfd325 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala @@ -451,7 +451,7 @@ class CometNativeColumnarToRowSuite extends CometTestBase with AdaptiveSparkPlan DataGenOptions(generateNegativeZero = false, generateNaN = false, generateInfinity = false) // Test with multiple random deeply nested schemas - for (iteration <- 1 to 3) { + for (_ <- 1 to 3) { val schema = FuzzDataGenerator.generateNestedSchema( random, numCols = 5, diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConditionalExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConditionalExpressionBenchmark.scala index 877318e96b4..b5fd9f277a1 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConditionalExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometConditionalExpressionBenchmark.scala @@ -30,7 +30,7 @@ package org.apache.spark.sql.benchmark */ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { - private def prepareTestTable(values: Int)(f: => Unit): Unit = { + private def prepareTestTable(f: => Unit): Unit = { withTempPath { dir => withTempTable("parquetV1Table") { // Create table with multiple columns for richer test scenarios: @@ -54,7 +54,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def caseWhenLiteralBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { val query = "SELECT CASE WHEN c1 < 0 THEN '<0' WHEN c1 = 0 THEN '=0' ELSE '>0' END FROM parquetV1Table" runExpressionBenchmark("Case When Literal (3 branches)", values.toLong, query) @@ -62,7 +62,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def caseWhenManyBranchesLiteralBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { // 10 branches using c2 (values 0-99) val query = """ SELECT CASE @@ -83,7 +83,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def caseWhenColumnResultBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { // Result expressions are column references, not literals val query = "SELECT CASE WHEN c1 < 0 THEN c3 WHEN c1 = 0 THEN c1 ELSE c3 + c1 END FROM parquetV1Table" @@ -92,7 +92,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def caseWhenManyBranchesColumnResultBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { // 10 branches with column expressions as results val query = """ SELECT CASE @@ -113,14 +113,14 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def ifLiteralBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { val query = "SELECT IF(c1 < 0, '<0', '>=0') FROM parquetV1Table" runExpressionBenchmark("If Literal", values.toLong, query) } } def ifColumnResultBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { // Result expressions are column references val query = "SELECT IF(c1 < 0, c3, c1 + c3) FROM parquetV1Table" runExpressionBenchmark("If Column Result", values.toLong, query) @@ -128,7 +128,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def nestedIfBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { // Nested IF expressions (equivalent to CASE WHEN with multiple branches) val query = """ SELECT IF(c2 < 25, 'a', @@ -141,7 +141,7 @@ object CometConditionalExpressionBenchmark extends CometBenchmarkBase { } def nestedIfColumnResultBenchmark(values: Int): Unit = { - prepareTestTable(values) { + prepareTestTable { val query = """ SELECT IF(c2 < 25, c1, IF(c2 < 50, c3, From 0691bd39f75da61e02aa23ce188de141bd623252 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Sun, 11 Oct 2026 14:11:39 +0000 Subject: [PATCH 65/89] perf: read shuffle blocks in 64 KiB pieces instead of through `Channels.newChannel` (#6805) * perf: read shuffle blocks in 64 KiB pieces instead of through `Channels.newChannel` Both shuffle readers, `NativeBatchDecoderIterator` and `CometShuffleBlockIterator`, read blocks through `Channels.newChannel(in)`. For a stream that is not a plain `FileInputStream`, the JDK channel copies at most 8 KiB per read and calls `in.available()` before every read after the first, which costs an fstat and an lseek on a local shuffle file. Read `in` directly in pieces of up to 64 KiB instead. * Make the shuffle read buffer size configurable and share the read loop Add `spark.comet.shuffle.readBufferSize`, 64 KiB by default. Both shuffle readers now call one `CometShuffleBlockIterator.readFully`, which owns a single buffer per thread, and one test covers both readers. * Pass the read buffer size default as a Long The strict-warnings build rejects the implicit Int to Long widening in `createWithDefault`, so convert the default explicitly. * feedback * Move the local benchmarks out of this PR `CometHashJoinTaskTimeBenchmark` and `CometShuffleReadKillBenchmark` were local tools for measuring this change and its review, and came in with the previous commit by accident. They stay local, and the skew-join benchmark will get its own PR against #6528. --- .../comet/CometShuffleBlockIterator.java | 63 ++++++++++++------ .../scala/org/apache/comet/CometConf.scala | 12 ++++ .../CometBlockStoreShuffleReader.scala | 21 +++++- .../shuffle/CometCelebornShuffleReader.scala | 3 +- .../shuffle/CometShuffledRowRDD.scala | 5 +- .../shuffle/NativeBatchDecoderIterator.scala | 20 ++---- .../CometCelebornShuffleReaderSuite.scala | 33 ++++++++++ ...eBatchDecoderIteratorLifecycleChecks.scala | 65 ++++++++++++++++++- 8 files changed, 182 insertions(+), 40 deletions(-) diff --git a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java index 86220f90093..c19ab285e8c 100644 --- a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java +++ b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java @@ -25,8 +25,6 @@ import java.io.InputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; -import java.nio.channels.Channels; -import java.nio.channels.ReadableByteChannel; import java.util.function.LongConsumer; /** @@ -44,7 +42,12 @@ public class CometShuffleBlockIterator implements Closeable { private static final int INITIAL_BUFFER_SIZE = 128 * 1024; - private final ReadableByteChannel channel; + /** Default of {@code spark.comet.shuffle.readBufferSize}. */ + public static final int DEFAULT_READ_BUFFER_SIZE = 64 * 1024; + + private static final ThreadLocal READ_BUFFER = new ThreadLocal<>(); + + private final int readBufferSize; private final InputStream inputStream; private final LongConsumer recordsReadUpdater; private final ByteBuffer headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN); @@ -53,13 +56,14 @@ public class CometShuffleBlockIterator implements Closeable { private int currentBlockLength = 0; public CometShuffleBlockIterator(InputStream in) { - this(in, records -> {}); + this(in, records -> {}, DEFAULT_READ_BUFFER_SIZE); } - public CometShuffleBlockIterator(InputStream in, LongConsumer recordsReadUpdater) { + public CometShuffleBlockIterator( + InputStream in, LongConsumer recordsReadUpdater, int readBufferSize) { this.inputStream = in; - this.channel = Channels.newChannel(in); this.recordsReadUpdater = recordsReadUpdater; + this.readBufferSize = readBufferSize; } /** @@ -77,17 +81,15 @@ public int hasNext() throws IOException { } // Read 16-byte header: clear() resets position=0, limit=capacity, - // preparing the buffer for channel.read() to fill it + // preparing the buffer for readFully() to fill it headerBuf.clear(); - while (headerBuf.hasRemaining()) { - int bytesRead = channel.read(headerBuf); - if (bytesRead < 0) { - if (headerBuf.position() == 0) { - close(); - return -1; - } - throw new EOFException("Data corrupt: unexpected EOF while reading batch header"); + readFully(inputStream, headerBuf, readBufferSize); + if (headerBuf.hasRemaining()) { + if (headerBuf.position() == 0) { + close(); + return -1; } + throw new EOFException("Data corrupt: unexpected EOF while reading batch header"); } headerBuf.flip(); long compressedLength = headerBuf.getLong(); @@ -121,11 +123,9 @@ public int hasNext() throws IOException { dataBuf.clear(); dataBuf.limit(currentBlockLength); - while (dataBuf.hasRemaining()) { - int bytesRead = channel.read(dataBuf); - if (bytesRead < 0) { - throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch"); - } + readFully(inputStream, dataBuf, readBufferSize); + if (dataBuf.hasRemaining()) { + throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch"); } // Note: native side uses get_direct_buffer_address (base pointer) + currentBlockLength, // not the buffer's position/limit. No flip needed. @@ -133,6 +133,29 @@ public int hasNext() throws IOException { return currentBlockLength; } + /** + * Fills {@code dst} from {@code in}, stopping short only at the end of the stream. Both shuffle + * readers use it. The stream is read directly, in pieces of up to {@code readBufferSize}, through + * a buffer reused per thread. {@code Channels.newChannel} would read 8 KiB at a time and call + * {@code available()} before every piece but the first, which costs an fstat and an lseek on a + * local shuffle file. + */ + public static void readFully(InputStream in, ByteBuffer dst, int readBufferSize) + throws IOException { + byte[] chunk = READ_BUFFER.get(); + if (chunk == null || chunk.length != readBufferSize) { + chunk = new byte[readBufferSize]; + READ_BUFFER.set(chunk); + } + while (dst.hasRemaining()) { + int read = in.read(chunk, 0, Math.min(dst.remaining(), chunk.length)); + if (read < 0) { + return; + } + dst.put(chunk, 0, read); + } + } + /** * Reports a native IPC/codec failure to the remote stream that supplied the current block. Called * by ShuffleScan via JNI; local streams preserve the original native exception. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 2beaf8f6a3c..c6ac50606f3 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -894,6 +894,18 @@ object CometConf extends ShimCometConf { s"Write buffer size must be between 1 and ${Int.MaxValue} bytes") .createWithDefault(1024 * 1024) + val COMET_SHUFFLE_READ_BUFFER_SIZE: ConfigEntry[Long] = + conf("spark.comet.shuffle.readBufferSize") + .category(CATEGORY_SHUFFLE) + .doc( + "Size in bytes of each read from a shuffle block stream. Comet reads a shuffle block " + + "through a reusable buffer of this size per task thread before decoding it.") + .bytesConf(ByteUnit.BYTE) + .checkValue( + v => v > 0 && v <= Int.MaxValue, + s"Read buffer size must be between 1 and ${Int.MaxValue} bytes") + .createWithDefault(CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE.toLong) + val COMET_SHUFFLE_JVM_PREFER_DICTIONARY_RATIO: ConfigEntry[Double] = conf( "spark.comet.shuffle.jvm.preferDictionary.ratio") .withAlternative("spark.comet.shuffle.preferDictionary.ratio") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index 3592095a37c..156b50a9c00 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.io.InputStream +import java.io.{FilterInputStream, InputStream} import org.apache.spark.{InterruptibleIterator, MapOutputTracker, SparkEnv, TaskContext} import org.apache.spark.internal.{config, Logging} @@ -91,6 +91,7 @@ class CometBlockStoreShuffleReader[K, C]( val nativeLib = new Native() val nativeUtil = new NativeUtil() val tracingEnabled = CometConf.COMET_TRACING_ENABLED.get() + val readBufferSize = CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt // Closes last read iterator and shared resources after the task is finished. // We need to close read iterator during iterating input streams, @@ -113,7 +114,8 @@ class CometBlockStoreShuffleReader[K, C]( dep.decodeTime, nativeLib, nativeUtil, - tracingEnabled) + tracingEnabled, + readBufferSize = readBufferSize) currentReadIterator }) .map(b => (0, b)) @@ -158,10 +160,23 @@ class CometBlockStoreShuffleReader[K, C]( */ override def readAsRawStream(): InputStream = { val streams = fetchIterator.map(_._2) - new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] { + val blocks = new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] { override def hasMoreElements: Boolean = streams.hasNext override def nextElement(): InputStream = streams.next() }) + // Native code reads this stream on the task thread, outside any InterruptibleIterator, so + // every read checks whether the task was killed. + new FilterInputStream(blocks) { + override def read(): Int = { + context.killTaskIfInterrupted() + in.read() + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + context.killTaskIfInterrupted() + in.read(buffer, offset, length) + } + } } private def fetchContinuousBlocksInBatch: Boolean = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala index b7deac55b49..945dbe2c094 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReader.scala @@ -122,7 +122,8 @@ private[shuffle] final class CometCelebornShuffleReader[K, C]( new Native(), nativeUtil, CometConf.COMET_TRACING_ENABLED.get(), - Some(serializedSchema)) + Some(serializedSchema), + CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt) } catch { case NonFatal(failure) => Utils.tryWithSafeFinally(throw failure) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index edd0332ac56..dddb9767761 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLShuffleReadMetricsRe import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.CometShuffleBlockIterator +import org.apache.comet.{CometConf, CometShuffleBlockIterator} /** * Different from [[org.apache.spark.sql.execution.ShuffledRowRDD]], this RDD is specialized for @@ -176,7 +176,8 @@ class CometShuffledBatchRDD( } new CometShuffleBlockIterator( readerAndMetrics._1.readAsRawStream(), - readerAndMetrics._2.incRecordsRead) + readerAndMetrics._2.incRecordsRead, + CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt) } override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index 6227da4bf4f..b74732a8c11 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala @@ -21,14 +21,13 @@ package org.apache.spark.sql.comet.execution.shuffle import java.io.{EOFException, InputStream} import java.nio.{ByteBuffer, ByteOrder} -import java.nio.channels.{Channels, ReadableByteChannel} import scala.util.control.NonFatal import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.{CometShuffleReadFailureHandler, Native} +import org.apache.comet.{CometShuffleBlockIterator, CometShuffleReadFailureHandler, Native} import org.apache.comet.vector.NativeUtil /** @@ -42,7 +41,8 @@ case class NativeBatchDecoderIterator( nativeLib: Native, nativeUtil: NativeUtil, tracingEnabled: Boolean, - expectedSchema: Option[Array[Byte]] = None) + expectedSchema: Option[Array[Byte]] = None, + readBufferSize: Int = CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE) extends Iterator[ColumnarBatch] { // One consumer reads this iterator, while task completion may close it from another thread. @@ -60,15 +60,9 @@ case class NativeBatchDecoderIterator( import NativeBatchDecoderIterator._ - private val channel: ReadableByteChannel = if (in != null) { - Channels.newChannel(in) - } else { - null - } - override def hasNext: Boolean = { synchronized { - if (channel == null || isClosed) { + if (in == null || isClosed) { return false } if (batch.isDefined) { @@ -172,7 +166,7 @@ case class NativeBatchDecoderIterator( private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = { // read compressed batch size from header longBuf.clear() - while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {} + CometShuffleBlockIterator.readFully(in, longBuf, readBufferSize) // If we reach the end of the stream, we are done, or if we read partial length // then the stream is corrupted. @@ -189,7 +183,7 @@ case class NativeBatchDecoderIterator( // read field count from header longBuf.clear() - while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {} + CometShuffleBlockIterator.readFully(in, longBuf, readBufferSize) if (longBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading field count") } @@ -214,7 +208,7 @@ case class NativeBatchDecoderIterator( } dataBuf.clear() dataBuf.limit(bytesToRead.toInt) - while (dataBuf.hasRemaining && channel.read(dataBuf) >= 0) {} + CometShuffleBlockIterator.readFully(in, dataBuf, readBufferSize) if (dataBuf.hasRemaining) { throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch") } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index 54448ec2039..0e618b96088 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -931,6 +931,35 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { context.markTaskCompleted(None) } + test("a killed task stops at the next read of a local shuffle's direct-read stream") { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + "spark.sql.adaptive.enabled" -> "false") { + withParquetTable((0 until 12).map(i => (i, s"row-$i")), "local_direct_read_rows") { + val shuffled = sql("SELECT * FROM local_direct_read_rows").repartition(1, $"_1") + // Running the query writes the shuffle that the stream below reads again. + assert(shuffled.collect().length == 12) + val rdd = shuffled.queryExecution.executedPlan + .collectFirst { case exchange: CometShuffleExchangeExec => + exchange.executeColumnar().asInstanceOf[CometShuffledBatchRDD] + } + .getOrElse(fail("Expected a native Comet shuffle exchange")) + val context = TaskContext.empty() + val blocks = rdd.computeAsShuffleBlockIterator(rdd.partitions.head, context) + try { + assert(blocks.hasNext() > 0) + context.markInterrupted("killed after the first batch") + intercept[TaskKilledException](blocks.hasNext()) + } finally { + blocks.close() + context.markTaskCompleted(None) + } + } + } + } + test("malformed native frames invalidate once before allocating or opening another reducer") { def header(length: Long, fields: Long): Array[Byte] = ByteBuffer .allocate(16) @@ -1631,6 +1660,10 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { NativeBatchDecoderIteratorLifecycleChecks.closesPrefetchedBatch() } + test("shuffle readers read in pieces of the read buffer size without asking for available") { + NativeBatchDecoderIteratorLifecycleChecks.readersReadInPiecesOfTheReadBufferSize() + } + test("decoder cleanup releases its delivered batch exactly once") { NativeBatchDecoderIteratorLifecycleChecks.closesDeliveredBatch() } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala index 52d11d94e30..01e93397f37 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.{CometShuffleReadFailureHandler, Native} +import org.apache.comet.{CometShuffleBlockIterator, CometShuffleReadFailureHandler, Native} import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} import org.apache.comet.vector.NativeUtil @@ -117,6 +117,69 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } + /** Both shuffle readers read in pieces of their read buffer size, without `available()`. */ + def readersReadInPiecesOfTheReadBufferSize(): Unit = { + val bodyBytes = 1024 * 1024 + val readBufferSize = 16 * 1024 + val largeFrame = ByteBuffer + .allocate(16 + bodyBytes) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(8L + bodyBytes) + .putLong(0L) + .array() + class CountingInput extends ByteArrayInputStream(largeFrame) { + var availableCalls = 0 + var largestRead = 0 + + override def available(): Int = { + availableCalls += 1 + super.available() + } + + override def read(bytes: Array[Byte], offset: Int, length: Int): Int = { + largestRead = math.max(largestRead, length) + super.read(bytes, offset, length) + } + + def check(): Unit = { + assert(availableCalls == 0, "The reader asked the stream for available bytes") + assert(largestRead == readBufferSize, s"The largest read was $largestRead bytes") + } + } + + val decoderInput = new CountingInput + val batch = new TrackingBatch() + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = Some(batch) + } + try { + val decoder = NativeBatchDecoderIterator( + decoderInput, + new SQLMetric("nsTiming", 0L), + null, + util, + tracingEnabled = false, + readBufferSize = readBufferSize) + assert(decoder.next() eq batch) + assert(!decoder.hasNext) + } finally { + util.close() + } + decoderInput.check() + + val blockInput = new CountingInput + val blocks = new CometShuffleBlockIterator(blockInput, (_: Long) => (), readBufferSize) + try { + assert(blocks.hasNext() == bodyBytes) + assert(blocks.hasNext() == -1) + } finally { + blocks.close() + } + blockInput.check() + } + def closesPrefetchedBatch(): Unit = { val batch = new TrackingBatch() withDecoder(batch) { (decoder, inputCloseCalls) => From 6296625b9081b1ec0186d23c9c7f4222cd905f30 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Sun, 11 Oct 2026 15:08:00 +0000 Subject: [PATCH 66/89] perf: presize collections and reuse buffers in native loops (#6861) Give collections a capacity where their final size is known, and stop allocating a fresh structure in every iteration of per-row and per-batch loops where one can be reused, borrowed or moved instead. --- native/common/src/bin/analyze_trace.rs | 10 ++-- .../execution/operators/iceberg_dictionary.rs | 2 +- .../src/execution/operators/merge_rows.rs | 3 +- .../src/execution/operators/parquet_writer.rs | 3 +- native/core/src/execution/planner.rs | 4 +- native/core/src/execution/spark_plan.rs | 6 +-- native/core/src/parquet/parquet_support.rs | 5 +- native/core/src/parquet/schema_adapter.rs | 6 +-- native/operators/src/expand.rs | 2 +- native/operators/src/explode.rs | 53 +++++++++---------- .../src/partitioners/multi_partition.rs | 10 ++-- native/shuffle/src/schema_align.rs | 4 +- native/shuffle/src/spark_unsafe/row.rs | 46 ++++++++++++---- native/spark-expr/src/agg_funcs/percentile.rs | 6 ++- .../src/bloom_filter/bloom_filter_agg.rs | 8 +++ .../src/bloom_filter/spark_bit_array.rs | 4 +- .../src/bloom_filter/spark_bloom_filter.rs | 17 +++--- .../src/json_funcs/json_array_length.rs | 4 +- native/spark-expr/src/jvm_udf/mod.rs | 3 +- native/spark-expr/src/string_funcs/split.rs | 26 +++++---- 20 files changed, 129 insertions(+), 93 deletions(-) diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs index b2cdb9eba93..644fd1d4746 100644 --- a/native/common/src/bin/analyze_trace.rs +++ b/native/common/src/bin/analyze_trace.rs @@ -121,7 +121,7 @@ fn main() { } let file = File::open(&args[1]).expect("Failed to open trace file"); - let reader = BufReader::new(file); + let mut reader = BufReader::new(file); // Index into ALLOCATED_COUNTERS of the counter being analyzed, once one has been seen let mut source: Option = None; @@ -147,8 +147,12 @@ fn main() { // Each line is one JSON event, possibly with a trailing comma. // The file starts with "[ " on the first event line or as a prefix. - for line in reader.lines() { - let line = line.expect("Failed to read line"); + let mut line = String::new(); + loop { + line.clear(); + if reader.read_line(&mut line).expect("Failed to read line") == 0 { + break; + } let trimmed = line.trim(); // Skip empty lines or bare array brackets diff --git a/native/core/src/execution/operators/iceberg_dictionary.rs b/native/core/src/execution/operators/iceberg_dictionary.rs index 0d6dc8165e1..1540e3fd747 100644 --- a/native/core/src/execution/operators/iceberg_dictionary.rs +++ b/native/core/src/execution/operators/iceberg_dictionary.rs @@ -122,7 +122,7 @@ impl DictionaryChooser { .with_coerce_types(base.coerce_types()) .convert(schema) .map_err(DataFusionError::from)?; - let mut leaves = Vec::new(); + let mut leaves = Vec::with_capacity(parquet_schema.num_columns()); for (column, field) in schema.fields().iter().enumerate() { let mut leaf = Leaf { column, diff --git a/native/core/src/execution/operators/merge_rows.rs b/native/core/src/execution/operators/merge_rows.rs index 9166127168f..869f6141b94 100644 --- a/native/core/src/execution/operators/merge_rows.rs +++ b/native/core/src/execution/operators/merge_rows.rs @@ -16,6 +16,7 @@ // under the License. use arrow::array::{Array, ArrayRef, BooleanArray, Int64Array, RecordBatch}; +use arrow::buffer::BooleanBuffer; use arrow::compute::kernels::boolean::{and, and_not, not}; use arrow::compute::{filter_record_batch, prep_null_mask_filter}; use arrow::datatypes::{DataType, SchemaRef}; @@ -503,7 +504,7 @@ fn run_group( // predicates that Spark would never evaluate after an earlier clause matched. let fire = match instr.condition.evaluate(¤t)? { ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => { - BooleanArray::from(vec![true; current.num_rows()]) + BooleanArray::new(BooleanBuffer::new_set(current.num_rows()), None) } ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => continue, value => value diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index cc2da616e54..81fed5e5994 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -547,9 +547,10 @@ impl ExecutionPlan for ParquetWriterExec { // the placeholder Scan. Both schemas use the same Catalyst data types, and // disabling field-name matching still recursively validates nested nullability; // only nested field names and metadata are ignored. + let (_, columns, _) = batch.into_parts(); RecordBatch::try_new_with_options( Arc::clone(&schema_for_write), - batch.columns().to_vec(), + columns, &RecordBatchOptions::new().with_match_field_names(false), ) .map_err(|e| { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index cc9157f8468..ae5916092b1 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -504,6 +504,8 @@ impl PhysicalPlanner { object_store_options: &HashMap, ) -> Result, ExecutionError> { let mut files = Vec::with_capacity(partition.partitioned_file.len()); + // Create an empty input schema for partition values because they are all literals. + let empty_schema = Arc::new(Schema::empty()); partition.partitioned_file.iter().try_for_each(|file| { assert!(file.start + file.length <= file.file_size); @@ -523,8 +525,6 @@ impl PhysicalPlanner { partitioned_file.object_meta.location = path; // Process partition values - // Create an empty input schema for partition values because they are all literals. - let empty_schema = Arc::new(Schema::empty()); let partition_values: Result, _> = file .partition_values .iter() diff --git a/native/core/src/execution/spark_plan.rs b/native/core/src/execution/spark_plan.rs index 2487c76a140..7f5c645a19f 100644 --- a/native/core/src/execution/spark_plan.rs +++ b/native/core/src/execution/spark_plan.rs @@ -57,15 +57,11 @@ impl SparkPlan { children: Vec>, additional_native_plans: Vec>, ) -> Self { - let mut accum: Vec> = vec![]; - for plan in &additional_native_plans { - accum.push(Arc::clone(plan)); - } Self { plan_id, native_plan, children, - additional_native_plans: accum, + additional_native_plans, } } diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index cc17429a72b..2a2c1737473 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -443,7 +443,7 @@ pub(crate) fn match_struct_fields( // `None` marks an id that more than one file field carries. let from_id_to_index: HashMap> = if should_match_by_id { - let mut map = HashMap::new(); + let mut map = HashMap::with_capacity(from_fields.len()); for (i, field) in from_fields.iter().enumerate() { if let Some(id) = field_id(field) { map.entry(id).and_modify(|m| *m = None).or_insert(Some(i)); @@ -465,7 +465,8 @@ pub(crate) fn match_struct_fields( // Group file field indices by folded name so a case-insensitive collision is detected // (Spark's `caseInsensitiveParquetFieldMap`) rather than silently overwritten. - let mut folded_to_indices: HashMap<&str, Vec> = HashMap::new(); + let mut folded_to_indices: HashMap<&str, Vec> = + HashMap::with_capacity(from_folded.len()); for (i, folded) in from_folded.iter().enumerate() { folded_to_indices .entry(folded.as_str()) diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index 495eeb545ec..fb37111ec5f 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -229,7 +229,7 @@ fn remap_physical_schema( // Pre-build id -> first matching logical field for the per-physical rename pass below. let id_to_logical: HashMap = if should_match_by_id { - let mut map = HashMap::new(); + let mut map = HashMap::with_capacity(logical_schema.fields().len()); for lf in logical_schema.fields() { if let Some(id) = field_id(lf) { map.entry(id).or_insert(lf); @@ -944,7 +944,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { let original_folded = original_folded.as_ref().unwrap_or(&physical_folded); // Only allocate per-column index vectors when a folded name actually collides. - let mut seen = HashSet::new(); + let mut seen = HashSet::with_capacity(original_folded.len()); let mut collisions = HashSet::new(); for name in original_folded { if !seen.insert(name.as_str()) { @@ -982,7 +982,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { }; let id_duplicate_roots = if should_match_by_id { - let mut exact_names = HashSet::new(); + let mut exact_names = HashSet::with_capacity(physical_file_schema.fields().len()); let mut duplicate_names = HashSet::new(); for field in physical_file_schema.fields() { if !exact_names.insert(field.name()) { diff --git a/native/operators/src/expand.rs b/native/operators/src/expand.rs index 22bdc6d2196..13d08518978 100644 --- a/native/operators/src/expand.rs +++ b/native/operators/src/expand.rs @@ -220,7 +220,7 @@ impl ExpandStream { batch: &RecordBatch, projection: &[Arc], ) -> Result { - let mut columns = vec![]; + let mut columns = Vec::with_capacity(projection.len()); projection.iter().try_for_each(|expr| { let column = expr.evaluate(batch)?; diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 5af1e052d7b..6aae0e2d349 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -567,25 +567,23 @@ fn flatten_struct_cols( struct_column_indices: &HashSet, ) -> Result { // horizontal expansion because of struct unnest - let columns_expanded = input_batch - .iter() - .enumerate() - .map(|(idx, column_data)| match struct_column_indices.get(&idx) { + let mut columns_expanded = Vec::with_capacity(schema.fields().len()); + for (idx, column_data) in input_batch.iter().enumerate() { + match struct_column_indices.get(&idx) { Some(_) => match column_data.data_type() { DataType::Struct(_) => { let struct_arr = column_data.as_any().downcast_ref::().unwrap(); - Ok(struct_arr.columns().to_vec()) + columns_expanded.extend_from_slice(struct_arr.columns()); + } + data_type => { + return internal_err!( + "expecting column {idx} from input plan to be a struct, got {data_type}" + ) } - data_type => internal_err!( - "expecting column {idx} from input plan to be a struct, got {data_type}" - ), }, - None => Ok(vec![Arc::clone(column_data)]), - }) - .collect::>>()? - .into_iter() - .flatten() - .collect(); + None => columns_expanded.push(Arc::clone(column_data)), + } + } Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } @@ -791,7 +789,7 @@ fn build_batch( let transformed = match list_type_columns.len() { 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), _ => { - let mut temp_unnested_result = HashMap::new(); + let mut temp_unnested_result = HashMap::with_capacity(list_type_columns.len()); let max_recursion = list_type_columns .iter() .fold(0, |highest_depth, ListUnnest { depth, .. }| { @@ -831,7 +829,7 @@ fn build_batch( } let unnested_array_map: HashMap> = temp_unnested_result.into_iter().fold( - HashMap::new(), + HashMap::with_capacity(list_type_columns.len()), |mut acc, ( ListUnnest { @@ -894,19 +892,16 @@ fn build_batch( ) .collect::>(); - let ret = flatten_arrs - .into_iter() - .enumerate() - .flat_map(|(col_idx, arr)| { - // Convert original column into its unnested version(s) - // Plural because one column can be unnested with different recursion level - // and into separate output columns - match multi_unnested_per_original_index.remove(&col_idx) { - Some(unnested_arrays) => unnested_arrays, - None => vec![arr], - } - }) - .collect::>(); + let mut ret = Vec::with_capacity(flatten_arrs.len() + list_type_columns.len()); + for (col_idx, arr) in flatten_arrs.into_iter().enumerate() { + // Convert original column into its unnested version(s) + // Plural because one column can be unnested with different recursion level + // and into separate output columns + match multi_unnested_per_original_index.remove(&col_idx) { + Some(unnested_arrays) => ret.extend(unnested_arrays), + None => ret.push(arr), + } + } flatten_struct_cols(&ret, schema, struct_column_indices) } diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 05b14b53243..f217165d38f 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -23,7 +23,7 @@ use crate::partitioners::partitioned_batch_iterator::{ use crate::partitioners::ShufflePartitioner; use crate::writers::PartitionWriter; use crate::{comet_partitioning, CometPartitioning, RoundRobinStrategy}; -use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch}; +use arrow::array::{Array, ArrayData, RecordBatch}; use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::{DataFusionError, HashSet}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -424,16 +424,14 @@ impl MultiPartitionShuffleRepartitioner { } else { (*max_hash_columns).min(input.num_columns()) }; - let columns_to_hash: Vec = (0..num_columns_to_hash) - .map(|i| Arc::clone(input.column(i))) - .collect(); + let columns_to_hash = &input.columns()[..num_columns_to_hash]; // Use identical seed as Spark hash partitioning. let hashes_buf = &mut scratch.hashes_buf[..num_rows]; hashes_buf.fill(42_u32); // Compute hash for selected columns - create_murmur3_hashes(&columns_to_hash, hashes_buf)?; + create_murmur3_hashes(columns_to_hash, hashes_buf)?; // Assign partition IDs based on hash (same as hash partitioning) let partition_ids = &mut scratch.partition_ids[..num_rows]; @@ -736,7 +734,7 @@ impl Debug for MultiPartitionShuffleRepartitioner { #[cfg(test)] mod tests { use super::*; - use arrow::array::Int64Array; + use arrow::array::{ArrayRef, Int64Array}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; #[derive(Default)] diff --git a/native/shuffle/src/schema_align.rs b/native/shuffle/src/schema_align.rs index e07e92e5820..d9089fd1f17 100644 --- a/native/shuffle/src/schema_align.rs +++ b/native/shuffle/src/schema_align.rs @@ -232,11 +232,11 @@ struct SchemaAlignStream { impl SchemaAlignStream { fn align(&self, batch: RecordBatch) -> Result { - let num_rows = batch.num_rows(); + let (_, columns, num_rows) = batch.into_parts(); cast_and_stamp_schema( "CometSchemaAlignExec", &self.target_schema, - batch.columns().to_vec(), + columns, num_rows, ) } diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index ce18a398df5..9246665f454 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -37,7 +37,7 @@ use arrow::array::{ Array, ArrayRef, RecordBatch, RecordBatchOptions, }; use arrow::compute::cast; -use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow::error::ArrowError; use datafusion::physical_plan::metrics::Time; use datafusion_comet_jni_bridge::errors::CometError; @@ -380,6 +380,10 @@ fn append_nested_struct_fields_field_major( ) -> Result<(), CometError> { let num_rows = row_addresses.len(); let mut row = SparkUnsafeRow::new_with_num_fields(fields.len()); + // Reused by every nested struct field. The recursive call only borrows them. + let mut nested_addresses: Vec = Vec::new(); + let mut nested_sizes: Vec = Vec::new(); + let mut nested_is_null: Vec = Vec::new(); // Helper macro for processing primitive fields macro_rules! process_field { @@ -517,9 +521,12 @@ fn append_nested_struct_fields_field_major( let nested_builder = get_field_builder!(struct_builder, StructBuilder, field_idx); // Collect nested struct addresses and sizes in one pass, building validity - let mut nested_addresses: Vec = Vec::with_capacity(num_rows); - let mut nested_sizes: Vec = Vec::with_capacity(num_rows); - let mut nested_is_null: Vec = Vec::with_capacity(num_rows); + nested_addresses.clear(); + nested_sizes.clear(); + nested_is_null.clear(); + nested_addresses.reserve(num_rows); + nested_sizes.reserve(num_rows); + nested_is_null.reserve(num_rows); for row_idx in 0..num_rows { if struct_is_null[row_idx] { @@ -837,6 +844,11 @@ fn append_struct_fields_field_major( } } + // Reused by every nested struct field. The recursive call only borrows them. + let mut nested_addresses: Vec = Vec::new(); + let mut nested_sizes: Vec = Vec::new(); + let mut nested_is_null: Vec = Vec::new(); + // Helper macro for processing primitive fields macro_rules! process_field { ($builder_type:ty, $field_idx:expr, $get_value:expr) => {{ @@ -970,9 +982,12 @@ fn append_struct_fields_field_major( let nested_builder = get_field_builder!(struct_builder, StructBuilder, field_idx); // Collect nested struct addresses and sizes in one pass, building validity - let mut nested_addresses: Vec = Vec::with_capacity(num_rows); - let mut nested_sizes: Vec = Vec::with_capacity(num_rows); - let mut nested_is_null: Vec = Vec::with_capacity(num_rows); + nested_addresses.clear(); + nested_sizes.clear(); + nested_is_null.clear(); + nested_addresses.reserve(num_rows); + nested_sizes.reserve(num_rows); + nested_is_null.reserve(num_rows); for (row_idx, i) in (row_start..row_end).enumerate() { if struct_is_null[row_idx] { @@ -1374,7 +1389,7 @@ pub fn process_sorted_row_partition( // Create builders once and reuse them across batches. // After finish() is called, builders are reset and can be reused. - let mut data_builders: Vec> = vec![]; + let mut data_builders: Vec> = Vec::with_capacity(schema.len()); schema.iter().try_for_each(|dt| { make_builders(dt, batch_size, prefer_dictionary_ratio) .map(|builder| data_builders.push(builder))?; @@ -1395,6 +1410,8 @@ pub fn process_sorted_row_partition( // One context for every batch this call encodes; the JVM calls in once per sorted // partition, so there is no wider native scope to hoist it to. let mut codec_context = ShuffleCodecContext::default(); + // The block writer pre-encodes the batch schema, so it is kept until the schema changes. + let mut block_writer: Option<(SchemaRef, ShuffleBlockWriter)> = None; while current_row < row_num { let n = std::cmp::min(batch_size, row_num - current_row); @@ -1427,8 +1444,17 @@ pub fn process_sorted_row_partition( frozen.clear(); let mut cursor = Cursor::new(&mut frozen); - let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec.clone())?; - written += block_writer.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time)?; + // The schema changes between batches when a dictionary-encoded column falls back to a + // plain array (see `builder_to_array`). + if block_writer + .as_ref() + .is_none_or(|(cached, _)| cached != batch.schema_ref()) + { + let writer = ShuffleBlockWriter::try_new(batch.schema_ref(), codec.clone())?; + block_writer = Some((batch.schema(), writer)); + } + let (_, writer) = block_writer.as_ref().expect("block writer is set above"); + written += writer.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time)?; if let Some(checksum) = &mut current_checksum { checksum.update(&mut cursor)?; diff --git a/native/spark-expr/src/agg_funcs/percentile.rs b/native/spark-expr/src/agg_funcs/percentile.rs index 87219db2534..7fd66a4a691 100644 --- a/native/spark-expr/src/agg_funcs/percentile.rs +++ b/native/spark-expr/src/agg_funcs/percentile.rs @@ -247,7 +247,11 @@ impl GroupsAccumulator for SparkPercentileGroupsAccumulator { offsets.push(len); } - let values = emit_group_values.into_iter().flatten().collect::>(); + // `len` is the total number of values, which `flatten().collect()` has no size hint for. + let mut values = Vec::with_capacity(len as usize); + for group in emit_group_values { + values.extend_from_slice(&group); + } let values = Float64Array::new(ScalarBuffer::from(values), None); let list = ListArray::new( Arc::new(Field::new_list_field(DataType::Float64, true)), diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs b/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs index 8920b30c9de..adce71233a2 100644 --- a/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs +++ b/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs @@ -24,6 +24,7 @@ use crate::bloom_filter::spark_bloom_filter::{SparkBloomFilter, SparkBloomFilter use arrow::array::ArrayRef; use arrow::array::BinaryArray; +use arrow::array::StringArray; use datafusion::common::{downcast_value, ScalarValue}; use datafusion::error::{DataFusionError, Result}; use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; @@ -118,6 +119,13 @@ impl Accumulator for SparkBloomFilter { return Ok(()); } let arr = &values[0]; + // Hash strings in place rather than materializing a `ScalarValue::Utf8` per row. + if let Some(strings) = arr.as_any().downcast_ref::() { + for value in strings.iter().flatten() { + self.put_binary(value.as_bytes()); + } + return Ok(()); + } (0..arr.len()).try_for_each(|index| { let v = ScalarValue::try_from_array(arr, index)?; diff --git a/native/spark-expr/src/bloom_filter/spark_bit_array.rs b/native/spark-expr/src/bloom_filter/spark_bit_array.rs index b55567402be..214a832c064 100644 --- a/native/spark-expr/src/bloom_filter/spark_bit_array.rs +++ b/native/spark-expr/src/bloom_filter/spark_bit_array.rs @@ -62,8 +62,8 @@ impl SparkBitArray { self.data.len() } - pub fn data(&self) -> Vec { - self.data.clone() + pub fn data(&self) -> &[u64] { + &self.data } /// Number of set bits in the array. Mirrors Spark's `BitArray.cardinality()`. diff --git a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs index 0c5dbea6c8a..114b6be4045 100644 --- a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs +++ b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs @@ -16,7 +16,6 @@ // under the License. use arrow::array::{ArrowNativeTypeOp, BooleanArray, Int64Array}; -use arrow::datatypes::ToByteSlice; use datafusion::common::{DataFusionError, Result as DFResult}; use std::cmp; @@ -157,19 +156,19 @@ impl SparkBloomFilter { /// Serializes a SparkBloomFilter to a byte array conforming to Spark's BloomFilter /// binary format. The output format follows the filter's `version`. pub fn spark_serialization(&self) -> Vec { - let mut out: Vec = (self.version.to_int() as u32).to_be_bytes().to_vec(); - out.append(&mut self.num_hash_functions.to_be_bytes().to_vec()); + // At most a 16-byte header (version, numHashFunctions, seed, numWords), then the words. + let mut out: Vec = Vec::with_capacity(16 + self.bits.word_size() * 8); + out.extend_from_slice(&(self.version.to_int() as u32).to_be_bytes()); + out.extend_from_slice(&self.num_hash_functions.to_be_bytes()); if let SparkBloomFilterVersion::V2 = self.version { // Spark's BloomFilterImplV2.writeTo writes the seed between // numHashFunctions and the bit array. - out.append(&mut (self.seed as u32).to_be_bytes().to_vec()); + out.extend_from_slice(&(self.seed as u32).to_be_bytes()); } - out.append(&mut (self.bits.word_size() as u32).to_be_bytes().to_vec()); - let mut filter_state: Vec = self.bits.data(); - for i in filter_state.iter_mut() { - *i = i.to_be(); + out.extend_from_slice(&(self.bits.word_size() as u32).to_be_bytes()); + for word in self.bits.data() { + out.extend_from_slice(&word.to_be_bytes()); } - out.append(&mut Vec::from(filter_state.to_byte_slice())); out } diff --git a/native/spark-expr/src/json_funcs/json_array_length.rs b/native/spark-expr/src/json_funcs/json_array_length.rs index 272b0b72d5a..468fff426df 100644 --- a/native/spark-expr/src/json_funcs/json_array_length.rs +++ b/native/spark-expr/src/json_funcs/json_array_length.rs @@ -105,9 +105,7 @@ fn spark_json_array_length_scalar(scalar: &ScalarValue) -> Result { } fn spark_json_array_length_scalar_inner(json_str: &Option) -> Result { - let array_length = json_str - .clone() - .and_then(|json_str| get_json_array_length(&json_str)); + let array_length = json_str.as_deref().and_then(get_json_array_length); Ok(ScalarValue::Int32(array_length)) } diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 6355b1a9da6..e0de5fe3693 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -208,7 +208,6 @@ impl PhysicalExpr for JvmScalarUdfExpr { let out_arr_ptr = out_array.as_mut() as *mut FFI_ArrowArray as i64; let out_sch_ptr = out_schema.as_mut() as *mut FFI_ArrowSchema as i64; - let class_name = self.class_name.clone(); let n_args = arrays.len(); JVMClasses::with_env(|env| { @@ -221,7 +220,7 @@ impl PhysicalExpr for JvmScalarUdfExpr { })?; let jclass_name = env - .new_string(&class_name) + .new_string(&self.class_name) .map_err(|e| CometError::JNI { source: e })?; let in_arr_java = env diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index 64033440caf..685e0a4356b 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -445,18 +445,24 @@ fn push_split_parts( builder: &mut GenericStringBuilder, ) { if limit == 0 { - // limit = 0: split all, drop trailing empties. Need to know the end - // before pushing, so collect borrowed slices first (no string copies). - let mut parts: Vec<&str> = regex.split(string).collect(); - while parts.last().is_some_and(|s| s.is_empty()) { - parts.pop(); + // limit = 0: split all, drop trailing empties. Empty parts are held back until a + // non-empty part follows, so the trailing ones are never appended. + let mut pending_empty = 0; + let mut appended = false; + for part in regex.split(string) { + if part.is_empty() { + pending_empty += 1; + } else { + for _ in 0..pending_empty { + builder.append_value(""); + } + pending_empty = 0; + builder.append_value(part); + appended = true; + } } - if parts.is_empty() { + if !appended { builder.append_value(""); - } else { - for p in parts { - builder.append_value(p); - } } } else if limit > 0 { // limit > 0: at most limit-1 splits. From d8d37fc527e9498dcab5c1a52a470edb8fa11718 Mon Sep 17 00:00:00 2001 From: dustin Date: Sun, 11 Oct 2026 16:08:37 +0000 Subject: [PATCH 67/89] fix: resolve Azure ABFS authentication through Hadoop and keep ambient credentials out (#6059) * fix: let explicit Hadoop Azure auth outrank ambient AZURE_* environment variables The native Azure store seeded its builder from the environment and object_store's fixed credential order let an ambient bearer token, account key or webhook token file outrank or hybridise the identity the job configured through fs.azure keys. With a Hadoop auth mechanism configured, or a provider type naming MsiTokenProvider, no environment variable is consulted at all, so the configured identity can neither be replaced nor redirected. Without one, the environment applies as before, keeping AKS workload identity working. * fix: validate every Hadoop Azure auth mechanism and borrow only the token file for a named identity Reject Hadoop configuration the native scan cannot honour instead of letting object_store fall through to the Azure CLI or managed identity: a partial OAuth principal, an explicit auth.type the translated keys do not satisfy, a provider class other than the MSI, workload identity or client credentials providers (matched by exact simple name), the untranslated SAS provider, key provider, refresh token and user name or password keys, and blank values. Translate fs.azure.sas.fixed.token below the container-scoped SAS key, take the account from the host lowercased, and when Hadoop names only an identity borrow AZURE_FEDERATED_TOKEN_FILE alone, erroring when no token file exists. Every error carries one prefix and names the key, never a credential value. * fix: read only the Hadoop keys of the auth mechanism the account selects Hadoop's `AbfsConfiguration.getAuthType` resolves the account-scoped `fs.azure.account.auth.type` over the global one, and `AzureBlobFileSystemStore.initializeClient` then reads only that mechanism's keys. The native scan translated and validated every credential key instead, so an account that selects `SharedKey` while a global OAuth secret remains in the configuration was rejected for missing the OAuth client id and tenant. An explicit auth type now decides which keys are read: only the selected mechanism's keys are translated into the builder and validated, and the other mechanisms' keys, provider class and unsupported-mechanism keys are ignored as Hadoop ignores them, blank or not. Incomplete active credentials are still rejected, and no auth type at all still reads every key. * fix: apply Hadoop's default SharedKey when an account key is set without an auth type `AbfsConfiguration.getAuthType` falls back to `SharedKey` when no `fs.azure.account.auth.type` is set, and `initializeClient` then reads only the account key. The native scan read every mechanism's keys in that case, so a valid account key beside an unused OAuth client secret was rejected for the missing client id and tenant. An account key entry without an auth type now selects `SharedKey`: the OAuth and SAS keys beside it are neither translated nor validated, a blank key is still reported, `fs.azure.account.keyprovider` is still rejected, and the environment contributes nothing. With neither an auth type nor a key, every mechanism is read as before, so OAuth-only and environment-only configurations keep working. * fix: keep the Azure account as written and accept blank MSI client id and tenant Hadoop reads the account from the raw URL authority and appends it to each account-scoped key as written, so a lowercased account missed the key Hadoop finds and fell through to the environment. extract_account now keeps the host label as is. Under MsiTokenProvider Hadoop accepts an empty client id and tenant, which is how a system-assigned identity is configured. Those two keys are treated as absent for that provider only, so the builder proceeds to the managed identity endpoint with no client id. A blank client id under any other provider is still an error naming the key. * fix: read only the OAuth keys the selected Azure provider reads A shared Hadoop configuration can carry OAuth settings for several providers. Hadoop reads only the keys of the provider class it selects, but the native scan validated and translated every fs.azure.account.oauth2 key. A complete client credentials setup was rejected over an unused password or a blank MSI endpoint, and a leftover tenant, secret or token file could change the identity the store used. The provider's key list is now resolved once and both validation and translation read only those keys. With no provider class the rules are as before. An explicit SimpleKeyProvider, Hadoop's own default, is accepted with an account key and rejected without one. Other key providers are still rejected, and the class name must match exactly because Hadoop loads it untrimmed. * fix: follow Hadoop's defaults for blank Azure keys and prefer the fixed SAS token Under MsiTokenProvider, Hadoop reads fs.azure.account.oauth2.msi.endpoint and msi.authority through getTrimmedPasswordString(key, default), so a blank value means the IMDS endpoint and https://login.microsoftonline.com/. The native scan rejected the blank key while the driver read fine. WorkloadIdentityTokenProvider defaults msi.authority and token.file the same way. Blank values of those keys now count as unset under the class that defaults them, and stay errors without a provider class. The token file falls to AZURE_FEDERATED_TOKEN_FILE as an unset one does, and no ambient endpoint or authority host stands in for the blank key. ABFS reads fs.azure.sas.fixed.token and never the container-scoped fs.azure.sas.. key, which belongs to WASB, so a job that set both listed files with one token and read them natively with the other. The fixed token now comes first and the container-scoped key is the fallback. Whichever key is selected must be non-blank. The tenant id section of the Azure docs now says the provider class decides which of msi.tenant and client.endpoint is read. * fix: resolve Azure SAS settings the way Hadoop ABFS does A global fs.azure.sas.token.provider.type now counts only when the global auth type (SharedKey when unset) matches the account's, as in Hadoop's getTokenProviderClass. A configuration with a global OAuth auth type, an account-scoped SAS auth type and an unused global provider class was rejected even though Hadoop accepts it. The SAS token lookup now probes the container-scoped fixed token fs.azure.sas.fixed.token.. first, under the URL's own host with any explicit port, as Hadoop 3.4.2 and later do. Then come the account-scoped and global fixed tokens and the WASB container key. A blank container-scoped token is an error. The user guide notes that older Hadoop does not read that key and that the native scan reads container-scoped forms only for the SAS fixed token. * fix: trim the Azure SAS token before handing it to object_store Hadoop reads the fixed token with getTrimmedPasswordString and drops a leading '?'. The native scan checked the trimmed value for blankness but passed the original string on, so a padded token such as ' ?sig=...' reached object_store with the '?' after the whitespace and the parsed credential had a '?sig' key. Trim whichever SAS key was selected before translating it. * fix: read the container-scoped SAS token only when the driver's Hadoop does Hadoop 3.4.2 added a container-scoped lookup to AbfsConfiguration (containerConf), so fs.azure.sas.fixed.token.. is read before the account-level token. Hadoop 3.4.1 and older never read it, and Spark 3.4, 3.5 and 4.0 ship those versions. Reading it unconditionally broke or changed scans that work there: a blank container key next to a valid global token failed the scan, and a different value replaced the token Hadoop uses. NativeConfig now checks whether the hadoop-azure on the classpath has containerConf and passes the answer to native code as comet.azure.containerScopedSasToken in the object store options for abfs and abfss. The native scan reads and validates the container-scoped key only when it is true; otherwise it follows Hadoop 3.4.1 (account-scoped fixed token, then global) with the WASB key as its own fallback. When the key is present but not read, ambient AZURE_* credentials stay out, and if it is the only credential the scan fails with an error naming the key instead of building a store that would fall back to managed identity. Hadoop fails on that configuration too. * fix: send Azure client-credential token requests to the configured endpoint host Hadoop's ClientCredsTokenProvider posts to the full fs.azure.account.oauth2.client.endpoint. The native store took only the tenant from it, so object_store built the token URL on login.microsoftonline.com and a configured authentication proxy or sovereign authority was bypassed. When the tenant comes from the endpoint, a client secret is set and no token file is, the endpoint's origin is now translated to object_store's authority host as well. Without a secret, or with a token file (which object_store prefers for workload identity), the endpoint is not a token URL and its host receives no other credential. * fix: keep a client endpoint's path prefix in the Azure token URL An endpoint mounted under a path, such as https://auth-proxy.example/aad//oauth2/v2.0/token, gave its first segment as the tenant and only its origin as the authority host, so object_store requested /aad/oauth2/v2.0/token and dropped the tenant. The tenant is now the segment just before oauth2, and the authority host is the origin plus every segment ahead of it, so object_store requests the configured URL. Endpoints without a later oauth2 segment keep the first segment as the tenant, and standard and v1 endpoints are unchanged. * fix: take the Azure token endpoint's tenant from before its last oauth2 segment An endpoint whose path prefix itself contains oauth2, such as /gateway/oauth2//oauth2/v2.0/token, gave gateway as the tenant because the first oauth2 segment decided the split, so object_store dropped the real tenant from the token URL. Use the last oauth2 segment, which is the one in the token suffix; endpoints without a later oauth2 segment still take the first segment as the tenant. * fix: read Azure account-scoped keys under the URL host alone, as Hadoop does AbfsConfiguration.accountConf builds one scoped key, ., from the URL authority as written, port included. The native scan also probed the other endpoint suffix and the bare account name, and once the lookup chose the auth mechanism those extra forms changed the answer: a SharedKey auth type scoped to the blob host overrode a global OAuth for a dfs URL, and an override scoped to host:443 was dropped. Read only Hadoop's form, then the global key, for every account-scoped setting, the container SAS keys and the WASB fallback included. * fix: read Azure credential keys in Hadoop's container-scoped form Hadoop 3.4.2 reads the OAuth keys and the SAS fixed token through getPasswordString, which probes .. before the account-scoped and global forms. The native scan read that form only for the SAS fixed token, so a client endpoint set per container was missed once the provider class decided which keys are read, and a configuration that works on Hadoop failed. Every key Hadoop reads that way now takes the container form first when the driver's hadoop-azure has containerConf, under the probe the SAS token already used. The account key and the mechanism keys stay account-then-global, as in Hadoop. On older Hadoop a container-scoped credential key that is the only credential is an error naming the key rather than a fall-through to the environment, unless a workload identity is named. * feat: resolve Azure ABFS authentication through Hadoop on the driver The native scan reimplemented the ABFS driver's fs.azure.* lookup rules, and every review round found another place where the two disagreed: scoping by container, account or host, the auth-type guard for provider classes, trimming, blank values, and what each hadoop-azure release changed. NativeConfig now asks the classpath's AbfsConfiguration the questions AzureBlobFileSystemStore.initializeClient asks, through its public getters, and forwards the answers: comet.azure.* markers for the outcome, auth type, provider class and any error, plus the resolved values under the global key names. Mechanisms the native scan cannot run are declined by name, never instantiated. Any reflection failure or Hadoop exception becomes an error marker; nothing falls through. Only the selected account's resolved values travel in the plan, instead of every fs.azure.* key. hadoop-azure is a test dependency pinned to each profile's Hadoop line so the resolver runs against 3.3.4, 3.4.1, 3.4.2 and 3.5.0 in the suite. * refactor: build the native Azure store from the driver's resolved authentication The native scan no longer reads fs.azure.* settings by Hadoop's lookup rules. The driver resolves the mechanism and its values through the ABFS driver and forwards them with comet.azure.* markers, so this side only maps them: none reads the environment, resolved applies exactly the keys the mechanism owns, declined and error fail with the driver's verdict, and a missing or unknown marker fails as well. The scoping, capability probe, unread-key guard, provider-class filtering, blank checks and environment policy that reimplemented Hadoop are gone. * docs: describe Azure authentication as resolved by Hadoop on the driver The Azure section now explains that the driver asks the ABFS driver which mechanism and values apply, lists which mechanisms the native scan can build and what happens for the rest, states when the environment is read, and records the changes from 1.1. The account scoped lookup order, the endpoint suffix limitation and the WASB container key, which no longer apply, are gone. * test: pin the Azure resolver against hadoop-azure and the real FileSystem AbfsAuthParitySuite runs 137 configurations through the resolver and compares each with an expectation recorded for the hadoop-azure version on the classpath (3.3.4, 3.4.1, 3.4.2 or 3.5.0, chosen by capability probes), with hand-written expectations for the cases where versions differ or Hadoop's rules surprise. Every case also goes through FileSystem.newInstance offline and the outcome class and root cause must agree. Record mode writes the expectation file for the running version and refuses a partial run. Writing the cross-check found one gap: SimpleKeyProvider reads the account key through a configuration without a container, so a key set as fs.azure.account.key.null. initializes the FileSystem while the resolver reported no configuration and would have sent the native scan to the environment. The resolver now probes the key the same way. * fix: let unused unreadable ABFS keys pass the configuration check Deciding whether an account has ABFS authentication configured read every auth key through Hadoop. A key the selected mechanism never uses, such as an OAuth secret with a cyclic ${...} reference next to a valid SharedKey, threw there and failed the scan although hadoop-azure initializes fine. A key whose read throws inside Hadoop now counts as configured. The chosen mechanism reads only its own keys, so a broken key it needs still fails in Hadoop's words. * fix: fall back for ABFS scans that span containers or accounts The native scan resolved ABFS authentication for the first file's container and account and applied it to every file, so a scan over two accounts with different keys sent the first account's key to the second. Matching credentials are not enough either: the native scan builds one object store per file partition from that partition's first file and keeps only each file's path, while Spark packs files from different roots into one partition. A scan whose ABFS files span more than one container or account, or that mixes abfs/abfss with another scheme, now falls back to Spark for both the Parquet and the CSV native scans. The reason is also tagged on the Spark scan that stays in the plan, and a declined CSV scan now falls back to that scan instead of keeping the Comet wrapper. AbfsScanFallbackSuite reads two fake accounts and two containers end to end through a local abfss filesystem. AbfsAuthParitySuite now creates java.io.tmpdir before using it, since a fresh checkout may not have it. * perf: stop reading ABFS credentials once the account is known to be configured Deciding whether an account has ABFS authentication read every auth key through Hadoop's getPassword, only so a later failure message could be redacted against all of them. With a credential provider configured, each read builds and loads the keystore, so a plain SharedKey scan made 48 lookups where Hadoop makes 3, about half a second of driver planning per scan. The check now stops at the first key that configures the account, reading the plain keys first. The remaining credentials are read only when resolution fails, before the error message and the driver log are redacted. Also make the multi-account fallback tests pass the strict Scala warnings job. --- .github/workflows/pr_build_linux.yml | 3 + .github/workflows/pr_build_macos.yml | 3 + dev/release/rat_exclude_files.txt | 1 + docs/source/contributor-guide/development.md | 32 + docs/source/user-guide/latest/datasources.md | 105 +- native/core/src/parquet/objectstore/azure.rs | 1167 ++++++++++++----- native/core/src/parquet/parquet_support.rs | 6 +- pom.xml | 13 + spark/pom.xml | 33 + .../comet/objectstore/AbfsAuthResolver.scala | 517 ++++++++ .../comet/objectstore/AbfsReflection.scala | 299 +++++ .../comet/objectstore/NativeConfig.scala | 54 +- .../apache/comet/rules/CometExecRule.scala | 2 +- .../apache/comet/rules/CometScanRule.scala | 9 +- .../serde/operator/CometNativeScan.scala | 37 +- .../sql/comet/CometCsvNativeScanExec.scala | 22 +- .../comet/hadoop/fs/FakeAbfsFileSystem.java | 59 + .../abfs-auth-parity/expected-3.3.4.json | 876 +++++++++++++ .../abfs-auth-parity/expected-3.4.1.json | 936 +++++++++++++ .../abfs-auth-parity/expected-3.4.2.json | 947 +++++++++++++ .../abfs-auth-parity/expected-3.5.0.json | 943 +++++++++++++ .../objectstore/AbfsAuthParityCases.scala | 1014 ++++++++++++++ .../objectstore/AbfsAuthParitySuite.scala | 422 ++++++ .../objectstore/AbfsAuthResolverSuite.scala | 703 ++++++++++ .../objectstore/AbfsScanFallbackSuite.scala | 149 +++ .../comet/objectstore/NativeConfigSuite.scala | 179 ++- 26 files changed, 8126 insertions(+), 405 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/objectstore/AbfsAuthResolver.scala create mode 100644 spark/src/main/scala/org/apache/comet/objectstore/AbfsReflection.scala create mode 100644 spark/src/test/java/org/apache/comet/hadoop/fs/FakeAbfsFileSystem.java create mode 100644 spark/src/test/resources/abfs-auth-parity/expected-3.3.4.json create mode 100644 spark/src/test/resources/abfs-auth-parity/expected-3.4.1.json create mode 100644 spark/src/test/resources/abfs-auth-parity/expected-3.4.2.json create mode 100644 spark/src/test/resources/abfs-auth-parity/expected-3.5.0.json create mode 100644 spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParityCases.scala create mode 100644 spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParitySuite.scala create mode 100644 spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthResolverSuite.scala create mode 100644 spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index e672c214689..192dc8ffb32 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -574,6 +574,9 @@ jobs: org.apache.comet.vector.CometVectorUtilsSuite org.apache.comet.udf.CometUdfBridgeSuite org.apache.comet.objectstore.NativeConfigSuite + org.apache.comet.objectstore.AbfsAuthResolverSuite + org.apache.comet.objectstore.AbfsAuthParitySuite + org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 3e0e6a8cdbc..0729cb26f67 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -278,6 +278,9 @@ jobs: org.apache.comet.vector.CometVectorUtilsSuite org.apache.comet.udf.CometUdfBridgeSuite org.apache.comet.objectstore.NativeConfigSuite + org.apache.comet.objectstore.AbfsAuthResolverSuite + org.apache.comet.objectstore.AbfsAuthParitySuite + org.apache.comet.objectstore.AbfsScanFallbackSuite org.apache.comet.serde.operator.CometIcebergNativeScanSuite org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index 597a7894cd7..30e822a6e4e 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -30,5 +30,6 @@ spark/src/test/resources/tpch-query-results/*.out spark/src/test/resources/tpch-extended/q*.sql spark/src/test/resources/test-data/*.csv spark/src/test/resources/test-data/*.ndjson +spark/src/test/resources/abfs-auth-parity/*.json spark/inspections/CometTPC*results.txt benchmarks/tpc/queries/**/*.sql diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index 4bcd1834e8e..d1190124016 100644 --- a/docs/source/contributor-guide/development.md +++ b/docs/source/contributor-guide/development.md @@ -574,6 +574,38 @@ SPARK_GENERATE_GOLDEN_FILES=1 ./mvnw -Dsuites="org.apache.spark.sql.comet.CometT SPARK_GENERATE_GOLDEN_FILES=1 ./mvnw -Dsuites="org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite" -Pspark-4.0 -nsu test ``` +### ABFS Authentication Parity Files + +`AbfsAuthParitySuite` pins how `AbfsAuthResolver` resolves Azure ABFS authentication on each +hadoop-azure release. `spark/src/test/resources/abfs-auth-parity/expected-.json` holds, +for every case in `AbfsAuthParityCases`, the outcome kind (`resolved`, `unconfigured`, `declined`, +`failed`), the auth type, the provider class, the error class and the forwarded values, with +secrets stored as SHA-256 digests. The suite picks the file from the hadoop-azure on the test +classpath, so each Spark profile checks one version: + +| Profile | hadoop-azure | +| ------------- | ------------ | +| `-Pspark-3.5` | 3.3.4 | +| `-Pspark-4.0` | 3.4.1 | +| `-Pspark-4.1` | 3.4.2 | +| `-Pspark-4.2` | 3.5.0 | + +A new case or a hadoop-azure bump requires re-recording all four files. Run the suite once per +profile with `-Dcomet.abfs.parity.record=` and copy the written files into the resources +directory: + +```sh +./mvnw -Pspark-3.5 test -Dtest=none -Dsuites=org.apache.comet.objectstore.AbfsAuthParitySuite -Dsurefire.failIfNoSpecifiedTests=false -Dcomet.abfs.parity.record= +./mvnw -Pspark-4.0 test -Dtest=none -Dsuites=org.apache.comet.objectstore.AbfsAuthParitySuite -Dsurefire.failIfNoSpecifiedTests=false -Dcomet.abfs.parity.record= +./mvnw -Pspark-4.1 test -Dtest=none -Dsuites=org.apache.comet.objectstore.AbfsAuthParitySuite -Dsurefire.failIfNoSpecifiedTests=false -Dcomet.abfs.parity.record= +./mvnw -Pspark-4.2 test -Dtest=none -Dsuites=org.apache.comet.objectstore.AbfsAuthParitySuite -Dsurefire.failIfNoSpecifiedTests=false -Dcomet.abfs.parity.record= +``` + +The per-version pins in the cases and the cross-check against a real `AzureBlobFileSystem` +still assert while recording, so a recorded file cannot contradict what hadoop-azure does. The +file is written only when every case ran. Review the diff of the regenerated files before +committing them: each line that changes is a behaviour change on that hadoop-azure release. + ## Benchmark There's a `make` command to run micro benchmarks in the repo. For diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index f351aff5b6f..0a5fe4aded4 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -394,65 +394,50 @@ Comet's S3 support has the following limitations: Comet's Parquet scan reads Azure Data Lake Storage Gen2 (ADLS Gen2) through the [`object_store` crate](https://crates.io/crates/object_store), using the `abfs` and `abfss` URL schemes. -As with S3, standard Hadoop ABFS configurations (the `fs.azure.*` keys you already set in `core-site.xml` or via `spark.hadoop.*`) are translated into the `object_store` crate's format, so existing configurations continue to work. +The Hadoop ABFS configuration you already have (the `fs.azure.*` keys in `core-site.xml` or under `spark.hadoop.*`) authenticates the native scan as well, so an existing setup keeps working. -URLs use the same shape Spark and Hadoop emit: `abfss://@.dfs.core.windows.net/`. The account is taken from the host and the container from the URL user-info. The `wasb`, `wasbs`, `az`, `azure`, and `adl` schemes are not supported by the native scan. +URLs use the shape Spark and Hadoop emit: `abfss://@.dfs.core.windows.net/`. The `wasb`, `wasbs`, `az`, `azure`, and `adl` schemes are not supported by the native scan. ### Root CA Certificates Azure scans discover Root CA Certificates the same way S3 scans do. See [Root CA Certificates](#root-ca-certificates) above. The Rust-based scan uses system Root CA Certificates rather than the Java Trust Store. -### Supported Authentication +### Authentication is resolved by Hadoop -Comet first calls `MicrosoftAzureBuilder::from_env()`, so any `AZURE_*` environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_FEDERATED_TOKEN_FILE`, `AZURE_AUTHORITY_HOST`, `AZURE_STORAGE_*`) are honored out of the box. This is what makes AKS Workload Identity work in a stock pod with no extra configuration. Any Hadoop `fs.azure.*` keys below are then applied on top, overriding the environment. +For each `abfs` or `abfss` path, the Spark driver asks the hadoop-azure library on its classpath which authentication mechanism applies to that container and account, and which values it needs. The question goes to the same `AbfsConfiguration` the ABFS `FileSystem` uses, so every Hadoop rule holds: account-scoped keys (`.`), container-scoped keys (`..`, hadoop-azure 3.4.2 and later), credential providers behind `hadoop.security.credential.provider.path` (a JCEKS keystore, for example) and `${...}` substitution behave exactly as they do for Spark's own reads. This holds for the hadoop-azure matching each Spark line's Hadoop: 3.3.4 with Spark 3.4 and 3.5, 3.4.1 with 4.0, 3.4.2 with 4.1, 3.5.0 with 4.2. -| Authentication method | Hadoop keys | -| ------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Shared account key | `fs.azure.account.key` | -| OAuth2 client credentials | `fs.azure.account.oauth2.client.id`, `fs.azure.account.oauth2.client.secret`, `fs.azure.account.oauth2.client.endpoint` (tenant id is extracted from the endpoint URL) | -| Managed identity (MSI) | `fs.azure.account.oauth2.msi.tenant`, `fs.azure.account.oauth2.msi.endpoint`, `fs.azure.account.oauth2.msi.authority` | -| Workload Identity | `fs.azure.account.oauth2.client.id`, `fs.azure.account.oauth2.msi.tenant`, `fs.azure.account.oauth2.token.file` | -| SAS token | `fs.azure.sas..` | +The native scan receives the resolved values and builds the `object_store` client from them. It reads no `fs.azure.*` key of its own. A scan whose files span more than one ABFS container or account, or that mixes `abfs`/`abfss` paths with other schemes, falls back to Spark's own scan. -### Hadoop-to-object_store key mapping +hadoop-azure must be on the driver classpath. It already is whenever Spark can list the path. When it is missing, the scan fails with an error that says so. -Internally, each Hadoop key is translated into a specific `AzureConfigKey` on the underlying `object_store` `MicrosoftAzureBuilder`: +#### Mechanisms -| Hadoop key (account-scoped suffix omitted) | `AzureConfigKey` | -| ------------------------------------------ | ------------------------ | -| `fs.azure.account.key` | `AccessKey` | -| `fs.azure.account.oauth2.client.id` | `ClientId` | -| `fs.azure.account.oauth2.client.secret` | `ClientSecret` | -| `fs.azure.account.oauth2.client.endpoint` | `AuthorityId` (from URL) | -| `fs.azure.account.oauth2.msi.tenant` | `AuthorityId` | -| `fs.azure.account.oauth2.msi.endpoint` | `MsiEndpoint` | -| `fs.azure.account.oauth2.msi.authority` | `AuthorityHost` | -| `fs.azure.account.oauth2.token.file` | `FederatedTokenFile` | -| `fs.azure.sas..` | `SasKey` | +| Hadoop mechanism | Native scan | Keys the native scan receives | +| ------------------------------------------------------------------------------------------------------------------------------ | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `SharedKey` (Hadoop's default when `fs.azure.account.auth.type` is unset), with any configured `fs.azure.account.keyprovider` | Supported | `fs.azure.account.key`. Hadoop's key provider runs on the driver and the key it returns is forwarded | +| `OAuth` with `ClientCredsTokenProvider` | Supported | `fs.azure.account.oauth2.client.id`, `fs.azure.account.oauth2.client.secret`, `fs.azure.account.oauth2.client.endpoint` (the tenant and authority host come from the endpoint URL) | +| `OAuth` with `MsiTokenProvider` | Supported | `fs.azure.account.oauth2.msi.endpoint`, plus `fs.azure.account.oauth2.client.id` when set. `fs.azure.account.oauth2.msi.tenant` and `fs.azure.account.oauth2.msi.authority` are accepted but unused: `object_store`'s IMDS provider sends only the client id to the endpoint | +| `OAuth` with `WorkloadIdentityTokenProvider` (hadoop-azure 3.4.1 and later) | Supported | `fs.azure.account.oauth2.client.id`, `fs.azure.account.oauth2.msi.tenant`, `fs.azure.account.oauth2.token.file`, `fs.azure.account.oauth2.msi.authority` (the last two take Hadoop's defaults when unset) | +| `SAS` with the fixed token (hadoop-azure 3.4.1 and later) | Supported | `fs.azure.sas.fixed.token` | +| `OAuth` with `RefreshTokenBasedTokenProvider` or `UserPasswordTokenProvider` | Not supported | | +| `SAS` with a `fs.azure.sas.token.provider.type` class | Not supported | | +| `fs.azure.account.auth.type=Custom` | Not supported | | +| `WorkloadIdentityTokenProvider` with a `fs.azure.account.oauth2.client.assertion.provider.type` (hadoop-azure 3.5.0 and later) | Not supported | | +| `UserboundSASWithOAuth` (hadoop-azure 3.5.0 and later) | Not supported | | -Anything beyond these keys is not translated and falls through to whatever `from_env()` or the URL itself provided. +Custom provider classes hand each token request to Java code the native scan cannot call, and the refresh token and user password flows have no `object_store` counterpart. When Hadoop selects one of these, the scan fails with an error naming the auth type and, where one is configured, the class. Configure a supported mechanism for that account, or keep the path on Spark's own scan. Values arrive as Hadoop resolved them, trimmed and defaulted where the ABFS driver trims and defaults them, so the store sees what the driver would have used. -### Tenant id resolution +Client credentials and Workload Identity token requests go over HTTPS only. The managed identity request goes to the IMDS endpoint over HTTP, as Hadoop's does. For client credentials, Hadoop posts to `fs.azure.account.oauth2.client.endpoint` as given, while `object_store` posts to `//oauth2/v2.0/token`. Comet takes the tenant from the path segment before the last `oauth2` and everything before that segment as the authority host, so a v1 `/oauth2/token` endpoint and a proxy with a path prefix both resolve. The token request path is always the v2.0 one. An `http://` endpoint fails. -`AuthorityId` (the AAD tenant id) can be supplied in two ways: +### Environment variables -- **Directly**, via `fs.azure.account.oauth2.msi.tenant`. -- **Indirectly**, via `fs.azure.account.oauth2.client.endpoint`, which is a full token URL such as `https://login.microsoftonline.com//oauth2/token`. Comet extracts the tenant id from the first path segment of that URL. +The `AZURE_*` variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_FEDERATED_TOKEN_FILE`, `AZURE_STORAGE_ACCOUNT_KEY`, `AZURE_USE_AZURE_CLI`, and the others `object_store` reads) and `IDENTITY_ENDPOINT` are read only when Hadoop configures no authentication for the account at all: none of `fs.azure.account.auth.type`, `fs.azure.account.key`, `fs.azure.account.keyprovider`, `fs.azure.shellkeyprovider.script`, `fs.azure.account.oauth.provider.type`, `fs.azure.sas.token.provider.type`, `fs.azure.sas.fixed.token` or any `fs.azure.account.oauth2.*` key resolves for it. That case is what lets AKS Workload Identity work with no Hadoop configuration for tables whose files Spark does not list through Hadoop, such as Iceberg tables. A plain Parquet path still needs Hadoop configured for Spark's own listing. -If both are set, the value from `msi.tenant` wins. +Once Hadoop names anything, only the resolved values reach the store. A resolution error (a missing mandatory key, an invalid account key, an unknown provider class) fails the scan with Hadoop's exception. Messages that could quote a credential, such as the key provider's, are kept in the Spark driver log and the scan error names the exception class. The environment is never a fallback. One read happens inside `object_store` rather than Comet: under a resolved managed identity, its IMDS provider picks up `IDENTITY_HEADER` when it fetches a token. -### Account-scoped key lookup order +### Credentials in transit -Account-scoped keys take precedence over global ones, mirroring Hadoop ABFS's own precedence. For each Hadoop key above, Comet probes the following names in order and uses the first match: - -1. `..dfs.core.windows.net` -2. `..blob.core.windows.net` -3. `.` -4. `` (unscoped / global) - -The SAS namespace follows the same shape with the container inlined into the key name: `fs.azure.sas...dfs.core.windows.net`, then `fs.azure.sas...blob.core.windows.net`, then `fs.azure.sas..`. - -The Hadoop values, when present, override anything already picked up from the `AZURE_*` environment. +The resolved values travel in the native plan from the driver to the executors, as the static `fs.s3a.access.key` and `fs.s3a.secret.key` values of an S3 configuration do. Spark's RPC channel is plaintext by default. `spark.network.crypto.enabled` or `spark.ssl.rpc.enabled` protects it. See [Wire encryption](s3-credential-providers.md#wire-encryption). ### Examples @@ -467,45 +452,65 @@ $SPARK_HOME/bin/spark-shell \ **Example 2: Workload Identity (AKS)** -In an AKS pod with Workload Identity enabled, the `AZURE_*` environment variables injected by the webhook are picked up automatically, so no Comet-specific configuration is required. To set the values explicitly instead: +In an AKS pod with Workload Identity enabled and no `fs.azure.*` authentication configured, the `AZURE_*` variables the webhook injects are picked up, so nothing Comet-specific is needed. To configure it through Hadoop instead (hadoop-azure 3.4.1 and later): ```shell $SPARK_HOME/bin/spark-shell \ ... +--conf spark.hadoop.fs.azure.account.auth.type.myaccount.dfs.core.windows.net=OAuth \ +--conf spark.hadoop.fs.azure.account.oauth.provider.type.myaccount.dfs.core.windows.net=org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider \ --conf spark.hadoop.fs.azure.account.oauth2.client.id.myaccount.dfs.core.windows.net= \ ---conf spark.hadoop.fs.azure.account.oauth2.msi.tenant.myaccount.dfs.core.windows.net= \ ---conf spark.hadoop.fs.azure.account.oauth2.token.file.myaccount.dfs.core.windows.net=/var/run/secrets/azure/tokens/azure-identity-token +--conf spark.hadoop.fs.azure.account.oauth2.msi.tenant.myaccount.dfs.core.windows.net= ... ``` -**Example 3: OAuth2 client credentials with tenant embedded in the endpoint URL** +The token file defaults to `/var/run/secrets/azure/tokens/azure-identity-token`, the path AKS mounts. Set `fs.azure.account.oauth2.token.file` for another location. -If only `fs.azure.account.oauth2.client.endpoint` is set, the tenant id is parsed from the endpoint path automatically — there is no need to set `msi.tenant` separately. +**Example 3: OAuth2 client credentials** + +The tenant comes from the endpoint URL, so `fs.azure.account.oauth2.msi.tenant` is not needed: ```shell $SPARK_HOME/bin/spark-shell \ ... +--conf spark.hadoop.fs.azure.account.auth.type.myaccount.dfs.core.windows.net=OAuth \ +--conf spark.hadoop.fs.azure.account.oauth.provider.type.myaccount.dfs.core.windows.net=org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider \ --conf spark.hadoop.fs.azure.account.oauth2.client.id.myaccount.dfs.core.windows.net= \ --conf spark.hadoop.fs.azure.account.oauth2.client.secret.myaccount.dfs.core.windows.net= \ --conf spark.hadoop.fs.azure.account.oauth2.client.endpoint.myaccount.dfs.core.windows.net=https://login.microsoftonline.com//oauth2/token ... ``` -**Example 4: SAS token scoped to a container** +**Example 4: SAS token** + +Hadoop's ABFS driver (3.4.1 and later) reads the fixed token under `fs.azure.account.auth.type=SAS`, and the native scan receives the same value, so one configuration serves the driver and the executors: ```shell $SPARK_HOME/bin/spark-shell \ ... ---conf spark.hadoop.fs.azure.sas.mycontainer.myaccount.dfs.core.windows.net='sv=2020-08-04&sig=...' +--conf spark.hadoop.fs.azure.account.auth.type.myaccount.dfs.core.windows.net=SAS \ +--conf spark.hadoop.fs.azure.sas.fixed.token.myaccount.dfs.core.windows.net='sv=2020-08-04&sig=...' ... ``` +To give one container its own token on hadoop-azure 3.4.2 and later, set `fs.azure.sas.fixed.token.mycontainer.myaccount.dfs.core.windows.net`. Hadoop reads it before the account-level key, and the native scan receives whichever one Hadoop picked. The OAuth keys take the same container-scoped form on those versions. + +### Changes from Comet 1.1 + +Comet 1.1 started from the environment and overlaid the Hadoop keys it found by probing the account name under two endpoint suffixes and bare. Now Hadoop decides. What this changes: + +- Container-scoped keys work on hadoop-azure 3.4.2 and later, as they do for Spark. +- Transport variables such as `AZURE_STORAGE_ENDPOINT` and `AZURE_ALLOW_HTTP` no longer apply beside a Hadoop mechanism. They are read only when Hadoop configures none. +- The WASB key `fs.azure.sas..`, which the ABFS driver never read, is no longer read by the native scan either. Use `fs.azure.sas.fixed.token`, which the native scan now reads. +- A lower-case `oauth` auth type, or OAuth keys without `fs.azure.account.auth.type`, fail the way they fail for Spark. In the second case Hadoop defaults to `SharedKey` and the key provider fails. On hadoop-azure 3.4.1 and later the driver log names the missing key. 3.3.4 reports only `Failure to initialize configuration`. +- A configured mechanism the native scan cannot build fails with an error instead of silently reading the environment. + ### Limitations -1. **Partial Hadoop ABFS configuration support**: Only the `fs.azure.*` keys listed above are translated and applied to the underlying `object_store` crate. +1. **Authentication keys only**: Only the keys in the table above reach the native store. Other `fs.azure.*` settings (read-ahead, retries, HTTP tuning) do not apply to the native scan. 2. **Supported schemes**: Only `abfs` and `abfss` are routed to the native Azure store. `wasb[s]`, `az`, `azure`, and `adl` are not supported. `wasb[s]` is not recognised by `object_store` at all; `az`, `azure`, and `adl` are recognised by `object_store` but treat the URL host as the _container_ rather than the _account_, which is incompatible with Hadoop's account-scoped configuration keys. 3. **URL shape**: URLs must include the account in the host, i.e. `abfss://@.dfs.core.windows.net/`. Bare `abfs:///` (fsspec-style, no account in the URL) is not supported because Comet cannot resolve the storage account name. -4. **Endpoint suffixes for account-scoped lookup**: Comet probes for account-scoped keys under `dfs.core.windows.net` and `blob.core.windows.net`. Custom or sovereign-cloud endpoint suffixes (e.g. `dfs.core.chinacloudapi.cn`, Fabric endpoints) are not probed; use the unscoped `.` form for those. +4. **Endpoint hosts**: `object_store` builds an Azure store only for hosts under `dfs.core.windows.net`, `blob.core.windows.net`, `dfs.fabric.microsoft.com` and `blob.fabric.microsoft.com`. Sovereign-cloud endpoints such as `dfs.core.chinacloudapi.cn` are not supported, whatever keys are configured for them. diff --git a/native/core/src/parquet/objectstore/azure.rs b/native/core/src/parquet/objectstore/azure.rs index 6fed99f1642..4b49cc07eb3 100644 --- a/native/core/src/parquet/objectstore/azure.rs +++ b/native/core/src/parquet/objectstore/azure.rs @@ -15,54 +15,63 @@ // specific language governing permissions and limitations // under the License. -//! Construct a `MicrosoftAzure` object store from a Hadoop ABFS configuration. +//! Construct a `MicrosoftAzure` object store from the ABFS authentication the driver resolved. //! -//! Comet's native scans run outside the JVM, so they bypass Hadoop's -//! `AzureBlobFileSystem` driver entirely. This module bridges the gap by translating the -//! Hadoop `fs.azure.*` configuration namespace (the same keys users already put in -//! `core-site.xml` or `spark.hadoop.*`) into the `object_store` crate's `AzureConfigKey` -//! options, then layering them on top of any AZURE_* environment variables that AKS's -//! Workload Identity webhook injects. +//! Comet's native scans run outside the JVM, so they cannot use Hadoop's `AzureBlobFileSystem`. +//! Which mechanism applies to an `abfs[s]://@/` URL, and which values +//! it needs, is decided on the Spark driver by `AbfsAuthResolver`, which asks the classpath's +//! hadoop-azure the same questions `AzureBlobFileSystemStore.initializeClient` asks. Account-, +//! container- and globally-scoped keys, credential providers and `${...}` references are all +//! Hadoop's business there. This module trusts that answer and only maps it onto +//! `MicrosoftAzureBuilder`. The answer arrives in the object store options as `comet.azure.*` +//! markers plus, for a resolved mechanism, the values under their global Hadoop key names. //! -//! Only `abfs[s]://@./` URLs are supported — -//! the same URL shape Spark and Hadoop emit. `wasb[s]://` is not routed here because -//! `object_store::ObjectStoreScheme::parse` does not recognise it (it only accepts -//! `az | adl | azure | abfs | abfss`). The `az` / `azure` / `adl` schemes are also out of -//! scope: `object_store`'s `MicrosoftAzureBuilder::parse_url` treats the URL host as the -//! *container* for those schemes rather than the *account*, which does not match the -//! account-scoped Hadoop key layout that `NativeConfig` forwards. +//! `comet.azure.resolution` decides everything: //! -//! Supported authentication, in the priority order applied to the builder: +//! - `resolved`: `MicrosoftAzureBuilder::new()` plus exactly the keys the mechanism owns. +//! - `none`: `MicrosoftAzureBuilder::from_env()`, the only path in Comet's own code that reads +//! `AZURE_*` or `IDENTITY_ENDPOINT`. (object_store's IMDS provider still reads +//! `IDENTITY_HEADER` at token-fetch time under a resolved managed identity; that is its +//! behaviour, not a Comet lookup.) +//! - `declined`: an error naming the auth type and the provider class Hadoop configured. +//! - `error`: an error carrying the driver's exception class and message. +//! - missing or anything else: an error, so a plan that did not run the resolver cannot pass. //! -//! 1. `MicrosoftAzureBuilder::from_env()` — picks up `AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, -//! `AZURE_FEDERATED_TOKEN_FILE`, `AZURE_AUTHORITY_HOST`, `AZURE_STORAGE_*`, etc. -//! This is what makes Workload Identity work out of the box in AKS pods. -//! 2. Account-scoped Hadoop keys (`fs.azure.account.X..dfs.core.windows.net`). -//! 3. Global Hadoop keys (`fs.azure.account.X`). +//! `comet.azure.auth.type` is Hadoop's `AuthType` name, `comet.azure.oauth.provider.class` the +//! OAuth provider class, `comet.azure.sas.provider.class` a custom SAS provider and +//! `comet.azure.oauth.assertion.provider.class` a Workload Identity client assertion provider. +//! Under `resolved` each mechanism maps the keys its Hadoop provider constructor reads: //! -//! Items 2 and 3 are forwarded via `MicrosoftAzureBuilder::with_config`, which overrides -//! whatever `from_env()` produced. The account-scoped variant wins over the global one, -//! mirroring Hadoop ABFS's own `AbfsConfiguration` precedence. +//! | Mechanism | Hadoop key (global form) | `AzureConfigKey` | +//! | --- | --- | --- | +//! | `SharedKey` | `fs.azure.account.key` | `AccessKey` | +//! | `ClientCredsTokenProvider` | `fs.azure.account.oauth2.client.id` | `ClientId` | +//! | | `fs.azure.account.oauth2.client.secret` | `ClientSecret` | +//! | | `fs.azure.account.oauth2.client.endpoint` | `AuthorityId`, `AuthorityHost` | +//! | `MsiTokenProvider` | `fs.azure.account.oauth2.msi.endpoint` | `MsiEndpoint` | +//! | | `fs.azure.account.oauth2.client.id` | `ClientId` (optional) | +//! | | `fs.azure.account.oauth2.msi.tenant` | `AuthorityId` (optional) | +//! | | `fs.azure.account.oauth2.msi.authority` | `AuthorityHost` (optional) | +//! | `WorkloadIdentityTokenProvider` | `fs.azure.account.oauth2.client.id` | `ClientId` | +//! | | `fs.azure.account.oauth2.msi.tenant` | `AuthorityId` | +//! | | `fs.azure.account.oauth2.token.file` | `FederatedTokenFile` | +//! | | `fs.azure.account.oauth2.msi.authority` | `AuthorityHost` | +//! | `SAS` (fixed token) | `fs.azure.sas.fixed.token` | `SasKey` | //! -//! The translated keys cover the auth schemes that ABFS users actually configure: -//! -//! | Hadoop key (account-scoped suffix omitted) | `AzureConfigKey` | -//! | ------------------------------------------------------- | ---------------------- | -//! | `fs.azure.account.key` | `AccessKey` | -//! | `fs.azure.account.oauth2.client.id` | `ClientId` | -//! | `fs.azure.account.oauth2.client.secret` | `ClientSecret` | -//! | `fs.azure.account.oauth2.client.endpoint` | `AuthorityId` (from URL) | -//! | `fs.azure.account.oauth2.msi.tenant` | `AuthorityId` | -//! | `fs.azure.account.oauth2.msi.endpoint` | `MsiEndpoint` | -//! | `fs.azure.account.oauth2.msi.authority` | `AuthorityHost` | -//! | `fs.azure.account.oauth2.token.file` | `FederatedTokenFile` | -//! | `fs.azure.sas..` | `SasKey` | -//! -//! Anything beyond these falls through to whatever `from_env()` or the URL itself provided. +//! An optional key that is absent or blank counts as unset, as Hadoop's `MsiTokenProvider` adds +//! the client id and tenant to its request only when they are non-empty. The MSI tenant and +//! authority are accepted but unused: object_store's IMDS provider reads only the client id and +//! endpoint. `RefreshTokenBasedTokenProvider` and `UserPasswordTokenProvider` have no +//! `object_store` counterpart and fail as unsupported. A value a mechanism needs but that is +//! absent (or, for the SAS token and the client credentials endpoint, unusable) is an error +//! naming the key; other `fs.azure.*` keys are ignored. Values are applied as forwarded (Hadoop +//! already trimmed and defaulted them), except that the trailing slash Hadoop appends to +//! `msi.authority` is removed because `object_store` joins `//...` +//! itself. No error message or log line carries a value, only key names. use log::debug; use std::collections::HashMap; -use url::Url; +use url::{Position, Url}; use object_store::{ azure::{AzureConfigKey, MicrosoftAzureBuilder}, @@ -70,6 +79,27 @@ use object_store::{ ObjectStore, ObjectStoreScheme, }; +const MARKER_RESOLUTION: &str = "comet.azure.resolution"; +const MARKER_AUTH_TYPE: &str = "comet.azure.auth.type"; +const MARKER_OAUTH_PROVIDER_CLASS: &str = "comet.azure.oauth.provider.class"; +const MARKER_OAUTH_ASSERTION_PROVIDER_CLASS: &str = "comet.azure.oauth.assertion.provider.class"; +const MARKER_SAS_PROVIDER_CLASS: &str = "comet.azure.sas.provider.class"; +const MARKER_ERROR_CLASS: &str = "comet.azure.error.class"; +const MARKER_ERROR_MESSAGE: &str = "comet.azure.error.message"; + +const AUTH_TYPE_SHARED_KEY: &str = "SharedKey"; +const AUTH_TYPE_OAUTH: &str = "OAuth"; +const AUTH_TYPE_SAS: &str = "SAS"; + +const PROVIDER_CLIENT_CREDS: &str = "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider"; +const PROVIDER_MSI: &str = "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider"; +const PROVIDER_WORKLOAD_IDENTITY: &str = + "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider"; +const PROVIDER_REFRESH_TOKEN: &str = + "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider"; +const PROVIDER_USER_PASSWORD: &str = + "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider"; + const HADOOP_KEY: &str = "fs.azure.account.key"; const HADOOP_OAUTH_CLIENT_ID: &str = "fs.azure.account.oauth2.client.id"; const HADOOP_OAUTH_CLIENT_SECRET: &str = "fs.azure.account.oauth2.client.secret"; @@ -78,11 +108,11 @@ const HADOOP_MSI_TENANT: &str = "fs.azure.account.oauth2.msi.tenant"; const HADOOP_MSI_ENDPOINT: &str = "fs.azure.account.oauth2.msi.endpoint"; const HADOOP_MSI_AUTHORITY: &str = "fs.azure.account.oauth2.msi.authority"; const HADOOP_WI_TOKEN_FILE: &str = "fs.azure.account.oauth2.token.file"; -const HADOOP_SAS_PREFIX: &str = "fs.azure.sas."; +const HADOOP_SAS_FIXED_TOKEN: &str = "fs.azure.sas.fixed.token"; -const ENDPOINT_SUFFIXES: &[&str] = &["dfs.core.windows.net", "blob.core.windows.net"]; +type Settings = Vec<(AzureConfigKey, String)>; -/// Build a `MicrosoftAzure` `ObjectStore` for `url` using `configs`. +/// Build a `MicrosoftAzure` `ObjectStore` for `url` from the driver's resolved authentication. /// /// `url` must use the `abfs[s]://` scheme; the returned `Path` is the URL's resource path /// (container-relative), suitable for direct use with `ObjectStore::get`. @@ -92,338 +122,879 @@ pub fn create_store( ) -> Result<(Box, Path), object_store::Error> { let (scheme, path) = ObjectStoreScheme::parse(url)?; if scheme != ObjectStoreScheme::MicrosoftAzure { - return Err(object_store::Error::Generic { - store: "MicrosoftAzure", - source: format!("Scheme of URL is not Azure: {url}").into(), - }); + return Err(config_error(format!("Scheme of URL is not Azure: {url}"))); } let path = Path::parse(path)?; - - let account = extract_account(url); - let container = extract_container(url); - - // Start from the environment so AKS Workload Identity (AZURE_CLIENT_ID, - // AZURE_TENANT_ID, AZURE_FEDERATED_TOKEN_FILE, AZURE_AUTHORITY_HOST) and any - // explicit AZURE_STORAGE_* variables are honoured without further configuration. - // `with_url` then fills in account/container from the URL itself. - let mut builder = MicrosoftAzureBuilder::from_env().with_url(url.to_string()); - - let translated = translate_hadoop_configs(configs, account.as_deref(), container.as_deref()); - debug!( - "Azure configs for account={:?}, container={:?}: keys={:?}", - account, - container, - translated - .iter() - .map(|(k, _)| k.as_ref()) - .collect::>() - ); - for (key, value) in translated { - builder = builder.with_config(key, value); - } - - let store = builder.build()?; + let store = builder_for(url, configs)?.build()?; Ok((Box::new(store), path)) } -/// Translate a Hadoop ABFS configuration map into `(AzureConfigKey, value)` pairs. -/// -/// `account` and `container` are extracted from the URL and used to resolve account-scoped -/// keys (`fs.azure.X..`) and the SAS namespace -/// (`fs.azure.sas..`). Account-scoped keys win over global ones. -fn translate_hadoop_configs( +/// The builder the markers in `configs` call for. Only `none` touches the environment. +fn builder_for( + url: &Url, configs: &HashMap, - account: Option<&str>, - container: Option<&str>, -) -> Vec<(AzureConfigKey, String)> { - let mut out: Vec<(AzureConfigKey, String)> = Vec::new(); - - let mappings: &[(&str, AzureConfigKey)] = &[ - (HADOOP_KEY, AzureConfigKey::AccessKey), - (HADOOP_OAUTH_CLIENT_ID, AzureConfigKey::ClientId), - (HADOOP_OAUTH_CLIENT_SECRET, AzureConfigKey::ClientSecret), - (HADOOP_MSI_TENANT, AzureConfigKey::AuthorityId), - (HADOOP_MSI_ENDPOINT, AzureConfigKey::MsiEndpoint), - (HADOOP_MSI_AUTHORITY, AzureConfigKey::AuthorityHost), - (HADOOP_WI_TOKEN_FILE, AzureConfigKey::FederatedTokenFile), - ]; - - for (hadoop_base, azure_key) in mappings { - if let Some(value) = account_scoped_value(configs, hadoop_base, account) { - out.push((*azure_key, value)); +) -> Result { + let marker = |key: &str| configs.get(key).map(String::as_str); + match marker(MARKER_RESOLUTION) { + Some("none") => { + debug!("Azure authentication for {url}: none configured in Hadoop, reading AZURE_*"); + Ok(MicrosoftAzureBuilder::from_env().with_url(url.to_string())) + } + Some("resolved") => { + let settings = resolved_settings(configs)?; + debug!( + "Azure authentication for {url}: {} with keys {:?}", + marker(MARKER_AUTH_TYPE).unwrap_or_default(), + settings.iter().map(|(k, _)| k.as_ref()).collect::>() + ); + let builder = MicrosoftAzureBuilder::new().with_url(url.to_string()); + Ok(settings + .into_iter() + .fold(builder, |b, (key, value)| b.with_config(key, value))) } + Some("declined") => { + let class = marker(MARKER_OAUTH_ASSERTION_PROVIDER_CLASS) + .or_else(|| marker(MARKER_SAS_PROVIDER_CLASS)) + .or_else(|| marker(MARKER_OAUTH_PROVIDER_CLASS)); + Err(unsupported( + marker(MARKER_AUTH_TYPE).unwrap_or(""), + class, + )) + } + Some("error") => Err(config_error(format!( + "Hadoop ABFS authentication failed on the driver for {}://{}: {}: {}", + url.scheme(), + &url[Position::BeforeUsername..Position::AfterPort], + marker(MARKER_ERROR_CLASS).unwrap_or(""), + marker(MARKER_ERROR_MESSAGE).unwrap_or(""), + ))), + Some(_) => Err(config_error(format!( + "unknown `{MARKER_RESOLUTION}` value; the native ABFS scan needs the driver's Hadoop \ + resolution" + ))), + None => Err(config_error(format!( + "`{MARKER_RESOLUTION}` is missing; the native ABFS scan needs the driver's Hadoop \ + resolution" + ))), } +} - // `fs.azure.account.oauth2.client.endpoint` is a full token URL of the form - // `https://login.microsoftonline.com//oauth2/token`. object_store wants the - // tenant id directly (`AuthorityId`), so extract it if AuthorityId hasn't already - // been set from `fs.azure.account.oauth2.msi.tenant`. - let has_authority_id = out - .iter() - .any(|(k, _)| matches!(k, AzureConfigKey::AuthorityId)); - if !has_authority_id { - if let Some(endpoint) = account_scoped_value(configs, HADOOP_OAUTH_CLIENT_ENDPOINT, account) - { - if let Some(tenant) = tenant_from_oauth_endpoint(&endpoint) { - out.push((AzureConfigKey::AuthorityId, tenant)); +/// The builder settings for the mechanism the auth type and provider class markers name: the +/// keys its Hadoop provider constructor reads, and nothing else. +fn resolved_settings(configs: &HashMap) -> Result { + let marker = |key: &str| configs.get(key).map(String::as_str); + let auth_type = marker(MARKER_AUTH_TYPE) + .ok_or_else(|| config_error(format!("`{MARKER_AUTH_TYPE}` is missing")))?; + if let Some(class) = + marker(MARKER_OAUTH_ASSERTION_PROVIDER_CLASS).or_else(|| marker(MARKER_SAS_PROVIDER_CLASS)) + { + return Err(unsupported(auth_type, Some(class))); + } + match (auth_type, marker(MARKER_OAUTH_PROVIDER_CLASS)) { + (AUTH_TYPE_SHARED_KEY, _) => Ok(vec![( + AzureConfigKey::AccessKey, + required(configs, HADOOP_KEY)?, + )]), + (AUTH_TYPE_SAS, _) => { + let token = required(configs, HADOOP_SAS_FIXED_TOKEN)?; + if token.is_empty() { + return Err(config_error(format!("`{HADOOP_SAS_FIXED_TOKEN}` is empty"))); } + Ok(vec![(AzureConfigKey::SasKey, token)]) } - } - - // SAS tokens are scoped to `fs.azure.sas..[.]`. - if let (Some(container), Some(account)) = (container, account) { - if let Some(sas) = sas_value(configs, container, account) { - out.push((AzureConfigKey::SasKey, sas)); + (AUTH_TYPE_OAUTH, Some(PROVIDER_CLIENT_CREDS)) => client_creds_settings(configs), + (AUTH_TYPE_OAUTH, Some(PROVIDER_MSI)) => msi_settings(configs), + (AUTH_TYPE_OAUTH, Some(PROVIDER_WORKLOAD_IDENTITY)) => workload_identity_settings(configs), + (AUTH_TYPE_OAUTH, Some(class @ (PROVIDER_REFRESH_TOKEN | PROVIDER_USER_PASSWORD))) => { + Err(unsupported(auth_type, Some(class))) } + (AUTH_TYPE_OAUTH, Some(_)) => Err(config_error(format!( + "unknown `{MARKER_OAUTH_PROVIDER_CLASS}` for the native ABFS scan" + ))), + (AUTH_TYPE_OAUTH, None) => Err(config_error(format!( + "`{MARKER_OAUTH_PROVIDER_CLASS}` is missing" + ))), + _ => Err(config_error(format!( + "unknown `{MARKER_AUTH_TYPE}` for the native ABFS scan" + ))), } +} + +/// `ClientCredsTokenProvider` posts to the endpoint itself, so the endpoint supplies both the +/// tenant and the authority host; an endpoint without either is an error, not a default. +fn client_creds_settings( + configs: &HashMap, +) -> Result { + let endpoint = required(configs, HADOOP_OAUTH_CLIENT_ENDPOINT)?; + let (tenant, host) = oauth_endpoint_parts(&endpoint).ok_or_else(|| { + config_error(format!( + "`{HADOOP_OAUTH_CLIENT_ENDPOINT}` has no `scheme://host` origin or no tenant segment \ + before `oauth2`" + )) + })?; + Ok(vec![ + ( + AzureConfigKey::ClientId, + required(configs, HADOOP_OAUTH_CLIENT_ID)?, + ), + ( + AzureConfigKey::ClientSecret, + required(configs, HADOOP_OAUTH_CLIENT_SECRET)?, + ), + (AzureConfigKey::AuthorityId, tenant), + (AzureConfigKey::AuthorityHost, host), + ]) +} - out +/// `MsiTokenProvider`: the endpoint is required (Hadoop defaults it); the identity is optional. +fn msi_settings(configs: &HashMap) -> Result { + let mut out = vec![( + AzureConfigKey::MsiEndpoint, + required(configs, HADOOP_MSI_ENDPOINT)?, + )]; + out.extend(optional( + configs, + HADOOP_OAUTH_CLIENT_ID, + AzureConfigKey::ClientId, + )); + out.extend(optional( + configs, + HADOOP_MSI_TENANT, + AzureConfigKey::AuthorityId, + )); + out.extend(optional( + configs, + HADOOP_MSI_AUTHORITY, + AzureConfigKey::AuthorityHost, + )); + Ok(out) } -/// Look up `base_key`, preferring account-scoped variants. -/// -/// Probes (in order): `..`, `.`, -/// then the unscoped ``. Returns the first hit. -fn account_scoped_value( +/// `WorkloadIdentityTokenProvider`: every key is mandatory or defaulted in Hadoop. +fn workload_identity_settings( configs: &HashMap, - base_key: &str, - account: Option<&str>, -) -> Option { - if let Some(acc) = account { - for suffix in ENDPOINT_SUFFIXES { - let scoped = format!("{base_key}.{acc}.{suffix}"); - if let Some(v) = configs.get(&scoped) { - return Some(v.clone()); - } - } - let bare = format!("{base_key}.{acc}"); - if let Some(v) = configs.get(&bare) { - return Some(v.clone()); - } - } - configs.get(base_key).cloned() +) -> Result { + Ok(vec![ + ( + AzureConfigKey::ClientId, + required(configs, HADOOP_OAUTH_CLIENT_ID)?, + ), + ( + AzureConfigKey::AuthorityId, + required(configs, HADOOP_MSI_TENANT)?, + ), + ( + AzureConfigKey::FederatedTokenFile, + required(configs, HADOOP_WI_TOKEN_FILE)?, + ), + ( + AzureConfigKey::AuthorityHost, + authority_host(&required(configs, HADOOP_MSI_AUTHORITY)?), + ), + ]) } -/// Resolve the SAS token for `(container, account)`, accepting any of the -/// `fs.azure.sas..[.]` variants. -fn sas_value(configs: &HashMap, container: &str, account: &str) -> Option { - for suffix in ENDPOINT_SUFFIXES { - let key = format!("{HADOOP_SAS_PREFIX}{container}.{account}.{suffix}"); - if let Some(v) = configs.get(&key) { - return Some(v.clone()); - } - } - let bare = format!("{HADOOP_SAS_PREFIX}{container}.{account}"); - configs.get(&bare).cloned() +fn required(configs: &HashMap, key: &str) -> Result { + configs.get(key).cloned().ok_or_else(|| { + config_error(format!( + "`{key}` is missing from the resolved ABFS authentication" + )) + }) } -/// Extract the storage account name from an `abfs[s]://` URL. -/// -/// ABFS hostnames are `.` (e.g. `myacct.dfs.core.windows.net`), -/// so the account is the first label of the host. -fn extract_account(url: &Url) -> Option { - let host = url.host_str()?; - host.split('.').next().map(str::to_string) +/// An optional key, where blank counts as unset (`MsiTokenProvider` sends a client id or tenant +/// only when it is non-empty, untrimmed). The authority loses Hadoop's trailing slash. +fn optional( + configs: &HashMap, + key: &str, + azure_key: AzureConfigKey, +) -> Option<(AzureConfigKey, String)> { + let value = configs.get(key).filter(|v| !v.is_empty())?; + let value = match azure_key { + AzureConfigKey::AuthorityHost => authority_host(value), + _ => value.clone(), + }; + Some((azure_key, value)) } -/// Extract the container name from an `abfs[s]://` URL. -/// -/// ABFS encodes the container as the URL user-info (`container@account.dfs...`). -fn extract_container(url: &Url) -> Option { - let user = url.username(); - if user.is_empty() { - return None; +/// Hadoop appends a slash; object_store joins `//oauth2/v2.0/token`. +fn authority_host(authority: &str) -> String { + authority.trim_end_matches('/').to_string() +} + +fn config_error(message: String) -> object_store::Error { + object_store::Error::Generic { + store: "MicrosoftAzure", + source: message.into(), } - Some(user.to_string()) } -/// Pull the tenant id out of an OAuth token endpoint like -/// `https://login.microsoftonline.com//oauth2/token`. -fn tenant_from_oauth_endpoint(endpoint: &str) -> Option { +fn unsupported(auth_type: &str, class: Option<&str>) -> object_store::Error { + let provider = class + .map(|c| format!(" with provider {c}")) + .unwrap_or_default(); + config_error(format!( + "unsupported authentication mechanism for the native ABFS scan: auth type \ + {auth_type}{provider}" + )) +} + +/// The tenant id and authority host of an OAuth token endpoint like +/// `https://login.microsoftonline.com//oauth2/token`. The tenant is the path segment +/// before the last `oauth2`, or the first segment when no later segment is `oauth2`, so a +/// prefix that itself contains `oauth2` is kept. The authority host is the origin plus every +/// segment before the tenant, so object_store's `//oauth2/v2.0/token` +/// keeps a proxy's path prefix. `None` when the endpoint has no tenant or no `scheme://host` +/// origin. +fn oauth_endpoint_parts(endpoint: &str) -> Option<(String, String)> { let parsed = Url::parse(endpoint).ok()?; - let mut segments = parsed.path_segments()?; - let tenant = segments.next()?; - if tenant.is_empty() { + let segments: Vec<&str> = parsed.path_segments()?.collect(); + let tenant_at = segments + .iter() + .rposition(|segment| *segment == "oauth2") + .filter(|&at| at > 0) + .map_or(0, |at| at - 1); + let tenant = segments.get(tenant_at).filter(|t| !t.is_empty())?; + let origin = parsed.origin(); + if !origin.is_tuple() { return None; } - Some(tenant.to_string()) + let mut host = origin.ascii_serialization(); + for segment in &segments[..tenant_at] { + host.push('/'); + host.push_str(segment); + } + Some((tenant.to_string(), host)) } #[cfg(test)] mod tests { use super::*; + use object_store::client::{ + HttpClient, HttpConnector, HttpError, HttpRequest, HttpResponse, HttpResponseBody, + HttpService, + }; + use object_store::ClientOptions; + use std::sync::Mutex; + + const URL: &str = "abfss://data@myacct.dfs.core.windows.net/path/file.parquet"; + /// Marks every credential value; no error may echo it. + const SECRET: &str = "SECRET-VALUE-MARKER"; + /// A valid base64 account key, since `build()` decodes it. + const KEY: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA="; + const DEFAULT_ENDPOINT: &str = "https://login.microsoftonline.com/hadoop-tenant/oauth2/token"; + const DEFAULT_MSI_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token"; + const DEFAULT_AUTHORITY: &str = "https://login.microsoftonline.com/"; + const PUBLIC_CLOUD: &str = "https://login.microsoftonline.com"; + const TOKEN_FILE: &str = "/var/run/secrets/azure/tokens/azure-identity-token"; + const ENV_KEY: &str = "AZURE_STORAGE_ACCOUNT_KEY"; + const ENV_CLIENT_ID: &str = "AZURE_CLIENT_ID"; + const UNSUPPORTED: &str = "unsupported authentication mechanism"; + + /// Serializes the tests that mutate the process environment. + static ENV_LOCK: Mutex<()> = Mutex::new(()); + + fn lock_env() -> std::sync::MutexGuard<'static, ()> { + ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()) + } - fn url(s: &str) -> Url { - Url::parse(s).unwrap() + fn url() -> Url { + Url::parse(URL).unwrap() } - #[test] - fn extracts_account_and_container_from_abfss_url() { - let u = url("abfss://data@myacct.dfs.core.windows.net/path/file.parquet"); - assert_eq!(extract_account(&u).as_deref(), Some("myacct")); - assert_eq!(extract_container(&u).as_deref(), Some("data")); + type Configs = HashMap; + + fn configs(pairs: &[(&str, &str)]) -> Configs { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + fn with(mut configs: Configs, key: &str, value: &str) -> Configs { + configs.insert(key.into(), value.into()); + configs + } + + fn without(mut configs: Configs, key: &str) -> Configs { + configs.remove(key); + configs + } + + fn resolved(auth_type: &str, pairs: &[(&str, &str)]) -> Configs { + let out = with(configs(pairs), MARKER_RESOLUTION, "resolved"); + with(out, MARKER_AUTH_TYPE, auth_type) + } + + fn shared_key() -> Configs { + resolved(AUTH_TYPE_SHARED_KEY, &[(HADOOP_KEY, SECRET)]) + } + + fn oauth(provider: &str, pairs: &[(&str, &str)]) -> Configs { + with( + resolved(AUTH_TYPE_OAUTH, pairs), + MARKER_OAUTH_PROVIDER_CLASS, + provider, + ) + } + + fn client_creds() -> Configs { + oauth( + PROVIDER_CLIENT_CREDS, + &[ + (HADOOP_OAUTH_CLIENT_ID, "hadoop-client"), + (HADOOP_OAUTH_CLIENT_SECRET, SECRET), + (HADOOP_OAUTH_CLIENT_ENDPOINT, DEFAULT_ENDPOINT), + ], + ) + } + + fn msi() -> Configs { + oauth( + PROVIDER_MSI, + &[ + (HADOOP_MSI_ENDPOINT, DEFAULT_MSI_ENDPOINT), + (HADOOP_MSI_AUTHORITY, DEFAULT_AUTHORITY), + ], + ) + } + + fn workload_identity() -> Configs { + oauth( + PROVIDER_WORKLOAD_IDENTITY, + &[ + (HADOOP_OAUTH_CLIENT_ID, "hadoop-client"), + (HADOOP_MSI_TENANT, "hadoop-tenant"), + (HADOOP_WI_TOKEN_FILE, TOKEN_FILE), + (HADOOP_MSI_AUTHORITY, DEFAULT_AUTHORITY), + ], + ) + } + + fn builder(configs: &Configs) -> MicrosoftAzureBuilder { + builder_for(&url(), configs).expect("builder") + } + + fn value(builder: &MicrosoftAzureBuilder, key: AzureConfigKey) -> Option { + builder.get_config_value(&key) + } + + /// Every credential-bearing key the builder could receive, for "nothing else was set". + const CREDENTIAL_KEYS: &[AzureConfigKey] = &[ + AzureConfigKey::AccessKey, + AzureConfigKey::SasKey, + AzureConfigKey::ClientId, + AzureConfigKey::ClientSecret, + AzureConfigKey::AuthorityId, + AzureConfigKey::AuthorityHost, + AzureConfigKey::MsiEndpoint, + AzureConfigKey::FederatedTokenFile, + AzureConfigKey::Token, + ]; + + fn assert_exactly(configs: &Configs, expected: &[(AzureConfigKey, &str)]) { + let builder = builder(configs); + for key in CREDENTIAL_KEYS { + let want = expected.iter().find(|(k, _)| k == key).map(|(_, v)| *v); + assert_eq!(value(&builder, *key).as_deref(), want, "{key:?}"); + } + } + + /// The error `configs` produce, which must not echo the secret marker. + fn err(configs: &Configs) -> String { + let err = builder_for(&url(), configs) + .err() + .map(|e| e.to_string()) + .expect("an error"); + assert!(!err.contains(SECRET), "value echoed: {err}"); + err } #[test] - fn account_scoped_key_takes_precedence_over_global() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.oauth2.client.id".into(), - "global-client".into(), - ); - configs.insert( - "fs.azure.account.oauth2.client.id.myacct.dfs.core.windows.net".into(), - "scoped-client".into(), - ); - assert_eq!( - account_scoped_value(&configs, HADOOP_OAUTH_CLIENT_ID, Some("myacct")).as_deref(), - Some("scoped-client"), - ); + fn shared_key_sets_only_the_access_key_and_ignores_other_hadoop_keys() { + assert_exactly(&shared_key(), &[(AzureConfigKey::AccessKey, SECRET)]); + let stray = with(shared_key(), HADOOP_SAS_FIXED_TOKEN, "sv=1&sig=x"); + let stray = with(stray, HADOOP_OAUTH_CLIENT_ID, "stray"); + let stray = with(stray, "fs.azure.account.hns.enabled", "true"); + assert_exactly(&stray, &[(AzureConfigKey::AccessKey, SECRET)]); } #[test] - fn account_scoped_lookup_falls_back_to_global() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.oauth2.client.id".into(), - "global-client".into(), - ); - assert_eq!( - account_scoped_value(&configs, HADOOP_OAUTH_CLIENT_ID, Some("myacct")).as_deref(), - Some("global-client"), - ); + fn create_store_checks_the_scheme_and_returns_the_container_path() { + let configs = with(shared_key(), HADOOP_KEY, KEY); + let (_store, path) = create_store(&url(), &configs).expect("store builds"); + assert_eq!(path.as_ref(), "path/file.parquet"); + let s3 = Url::parse("s3://bucket/file.parquet").unwrap(); + let err = create_store(&s3, &configs).unwrap_err().to_string(); + assert!(err.contains("Scheme of URL is not Azure"), "{err}"); } #[test] - fn translates_workload_identity_keys() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.oauth2.client.id.myacct.dfs.core.windows.net".into(), - "client-123".into(), + fn client_creds_take_tenant_and_authority_host_from_the_endpoint_alone() { + assert_exactly( + &client_creds(), + &[ + (AzureConfigKey::ClientId, "hadoop-client"), + (AzureConfigKey::ClientSecret, SECRET), + (AzureConfigKey::AuthorityId, "hadoop-tenant"), + (AzureConfigKey::AuthorityHost, PUBLIC_CLOUD), + ], ); - configs.insert( - "fs.azure.account.oauth2.msi.tenant.myacct.dfs.core.windows.net".into(), - "tenant-abc".into(), + // `msi.*` keys from a shared configuration are not read under this provider. + let proxied = with(client_creds(), HADOOP_MSI_TENANT, "other-tenant"); + let proxied = with( + proxied, + HADOOP_MSI_AUTHORITY, + "https://login.chinacloudapi.cn/", ); - configs.insert( - "fs.azure.account.oauth2.token.file.myacct.dfs.core.windows.net".into(), - "/var/run/secrets/azure/tokens/azure-identity-token".into(), - ); - let translated = translate_hadoop_configs(&configs, Some("myacct"), Some("data")); - - let by_key: HashMap<_, _> = translated.into_iter().collect(); - assert_eq!( - by_key.get(&AzureConfigKey::ClientId).map(String::as_str), - Some("client-123") + let proxied = with( + proxied, + HADOOP_OAUTH_CLIENT_ENDPOINT, + "https://auth-proxy.example:8443/aad/hadoop-tenant/oauth2/v2.0/token", ); + let builder = builder(&proxied); assert_eq!( - by_key.get(&AzureConfigKey::AuthorityId).map(String::as_str), - Some("tenant-abc") + value(&builder, AzureConfigKey::AuthorityId).as_deref(), + Some("hadoop-tenant") ); assert_eq!( - by_key - .get(&AzureConfigKey::FederatedTokenFile) - .map(String::as_str), - Some("/var/run/secrets/azure/tokens/azure-identity-token") + value(&builder, AzureConfigKey::AuthorityHost).as_deref(), + Some("https://auth-proxy.example:8443/aad") ); + // An endpoint without a tenant segment or without a `scheme://host` origin cannot + // fall through to another credential or to the default authority host. + for endpoint in [ + "https://auth-proxy.example/", + "file:///hadoop-tenant/oauth2/token", + ] { + let err = err(&with( + client_creds(), + HADOOP_OAUTH_CLIENT_ENDPOINT, + endpoint, + )); + assert!( + err.contains(HADOOP_OAUTH_CLIENT_ENDPOINT), + "{endpoint}: {err}" + ); + } + } + + /// Records the request URLs a store sends so the token request is observed offline. + #[derive(Debug, Default, Clone)] + struct RecordingConnector(std::sync::Arc>>); + + #[async_trait::async_trait] + impl HttpService for RecordingConnector { + async fn call(&self, req: HttpRequest) -> Result { + self.0.lock().unwrap().push(req.uri().to_string()); + Ok(http::Response::builder() + .status(400) + .body(HttpResponseBody::from(String::new())) + .unwrap()) + } + } + + impl HttpConnector for RecordingConnector { + fn connect(&self, _options: &ClientOptions) -> object_store::Result { + Ok(HttpClient::new(self.clone())) + } + } + + #[tokio::test] + async fn client_creds_token_requests_go_to_the_endpoint_host() { + // Hadoop posts to `client.endpoint` itself, so a proxy mounted under a path keeps that + // path in the authority host and the tenant is the segment before the last `oauth2`. + for (endpoint, token_url) in [ + ( + DEFAULT_ENDPOINT, + "https://login.microsoftonline.com/hadoop-tenant/oauth2/v2.0/token", + ), + ( + "https://auth-proxy.example/gateway/oauth2/hadoop-tenant/oauth2/v2.0/token", + "https://auth-proxy.example/gateway/oauth2/hadoop-tenant/oauth2/v2.0/token", + ), + ] { + let configs = with(client_creds(), HADOOP_OAUTH_CLIENT_ENDPOINT, endpoint); + let recorder = RecordingConnector::default(); + let store = builder(&configs) + .with_http_connector(recorder.clone()) + .build() + .expect("store builds"); + assert!(store.credentials().get_credential().await.is_err()); + let urls = recorder.0.lock().unwrap(); + assert_eq!( + urls.first().map(String::as_str), + Some(token_url), + "{endpoint}" + ); + } } #[test] - fn translates_account_key() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.key.myacct.blob.core.windows.net".into(), - "secret==".into(), + fn msi_sets_the_endpoint_with_an_optional_identity_and_authority() { + assert_exactly( + &msi(), + &[ + (AzureConfigKey::MsiEndpoint, DEFAULT_MSI_ENDPOINT), + (AzureConfigKey::AuthorityHost, PUBLIC_CLOUD), + ], ); - let translated = translate_hadoop_configs(&configs, Some("myacct"), None); - let by_key: HashMap<_, _> = translated.into_iter().collect(); - assert_eq!( - by_key.get(&AzureConfigKey::AccessKey).map(String::as_str), - Some("secret==") + let user_assigned = with(msi(), HADOOP_OAUTH_CLIENT_ID, "user-assigned"); + let user_assigned = with(user_assigned, HADOOP_MSI_TENANT, "hadoop-tenant"); + assert_exactly( + &user_assigned, + &[ + (AzureConfigKey::MsiEndpoint, DEFAULT_MSI_ENDPOINT), + (AzureConfigKey::AuthorityHost, PUBLIC_CLOUD), + (AzureConfigKey::ClientId, "user-assigned"), + (AzureConfigKey::AuthorityId, "hadoop-tenant"), + ], + ); + // Hadoop forwards a blank client id or tenant as "" and sends neither; the authority is + // defaulted in Hadoop but unused by object_store's IMDS provider, so it may be absent. + let blank = with(msi(), HADOOP_OAUTH_CLIENT_ID, ""); + let blank = with(blank, HADOOP_MSI_TENANT, ""); + let blank = without(blank, HADOOP_MSI_AUTHORITY); + assert_exactly( + &blank, + &[(AzureConfigKey::MsiEndpoint, DEFAULT_MSI_ENDPOINT)], ); } #[test] - fn translates_sas_token_for_container() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.sas.data.myacct.dfs.core.windows.net".into(), - "sv=2020-08-04&sig=xyz".into(), - ); - let translated = translate_hadoop_configs(&configs, Some("myacct"), Some("data")); - let by_key: HashMap<_, _> = translated.into_iter().collect(); - assert_eq!( - by_key.get(&AzureConfigKey::SasKey).map(String::as_str), - Some("sv=2020-08-04&sig=xyz") + fn workload_identity_sets_client_tenant_token_file_and_authority() { + assert_exactly( + &workload_identity(), + &[ + (AzureConfigKey::ClientId, "hadoop-client"), + (AzureConfigKey::AuthorityId, "hadoop-tenant"), + (AzureConfigKey::FederatedTokenFile, TOKEN_FILE), + (AzureConfigKey::AuthorityHost, PUBLIC_CLOUD), + ], ); + // Hadoop appends a slash to the authority; object_store joins `//...`. + for (authority, host) in [ + ( + "https://login.chinacloudapi.cn/", + "https://login.chinacloudapi.cn", + ), + ( + "https://login.chinacloudapi.cn", + "https://login.chinacloudapi.cn", + ), + ( + "https://auth-proxy.example/aad/", + "https://auth-proxy.example/aad", + ), + ] { + let configs = with(workload_identity(), HADOOP_MSI_AUTHORITY, authority); + assert_eq!( + value(&builder(&configs), AzureConfigKey::AuthorityHost).as_deref(), + Some(host), + "{authority}" + ); + } } #[test] - fn derives_tenant_from_oauth_endpoint_when_msi_tenant_is_absent() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.oauth2.client.endpoint".into(), - "https://login.microsoftonline.com/00000000-1111-2222-3333-444444444444/oauth2/token" - .into(), - ); - let translated = translate_hadoop_configs(&configs, Some("myacct"), None); - let by_key: HashMap<_, _> = translated.into_iter().collect(); - assert_eq!( - by_key.get(&AzureConfigKey::AuthorityId).map(String::as_str), - Some("00000000-1111-2222-3333-444444444444") - ); + fn fixed_sas_token_is_forwarded_verbatim() { + for token in [ + "sv=2020-08-04&sig=xyz", + "?sv=2020-08-04&sig=xyz", + " sv=1&sig=x ", + ] { + let configs = resolved(AUTH_TYPE_SAS, &[(HADOOP_SAS_FIXED_TOKEN, token)]); + assert_exactly(&configs, &[(AzureConfigKey::SasKey, token)]); + // `split_sas` trims the leading `?`, so the store builds from either spelling. + create_store(&url(), &configs).unwrap_or_else(|e| panic!("{token:?}: {e}")); + } + // An empty token would make object_store send unsigned requests. + let err = err(&resolved(AUTH_TYPE_SAS, &[(HADOOP_SAS_FIXED_TOKEN, "")])); + assert!(err.contains(HADOOP_SAS_FIXED_TOKEN), "{err}"); } #[test] - fn msi_tenant_wins_over_oauth_endpoint_tenant() { - let mut configs = HashMap::new(); - configs.insert( - "fs.azure.account.oauth2.msi.tenant".into(), - "from-msi".into(), - ); - configs.insert( - "fs.azure.account.oauth2.client.endpoint".into(), - "https://login.microsoftonline.com/from-endpoint/oauth2/token".into(), - ); - let translated = translate_hadoop_configs(&configs, Some("myacct"), None); - let by_key: HashMap<_, _> = translated.into_iter().collect(); + fn unsupported_mechanisms_fail_naming_the_auth_type_and_class() { + let custom = "com.example.CustomProvider"; + // (resolution, auth type, marker carrying the class, class) + let cases = [ + // Declined on the driver, where the class was never instantiated either. + ("declined", "OAuth", MARKER_OAUTH_PROVIDER_CLASS, custom), + ( + "declined", + "OAuth", + MARKER_OAUTH_ASSERTION_PROVIDER_CLASS, + custom, + ), + ("declined", "SAS", MARKER_SAS_PROVIDER_CLASS, custom), + ("declined", "Custom", "", ""), + ("declined", "UserboundSASWithOAuth", "", ""), + // Resolved by Hadoop to a provider without an object_store counterpart. + ( + "resolved", + "OAuth", + MARKER_OAUTH_PROVIDER_CLASS, + PROVIDER_REFRESH_TOKEN, + ), + ( + "resolved", + "OAuth", + MARKER_OAUTH_PROVIDER_CLASS, + PROVIDER_USER_PASSWORD, + ), + ]; + for (resolution, auth_type, marker, class) in cases { + let mut configs = configs(&[ + (MARKER_RESOLUTION, resolution), + (MARKER_AUTH_TYPE, auth_type), + (HADOOP_OAUTH_CLIENT_ID, SECRET), + ]); + if marker == MARKER_OAUTH_ASSERTION_PROVIDER_CLASS { + // The assertion provider rides beside the Workload Identity class. + configs = with( + configs, + MARKER_OAUTH_PROVIDER_CLASS, + PROVIDER_WORKLOAD_IDENTITY, + ); + } + if !marker.is_empty() { + configs = with(configs, marker, class); + } + let err = err(&configs); + assert!(err.contains(UNSUPPORTED), "{err}"); + assert!(err.contains(auth_type) && err.contains(class), "{err}"); + } + } + + #[test] + fn bad_or_missing_markers_fail_closed_naming_the_marker() { + let cases = [ + (configs(&[(HADOOP_KEY, SECRET)]), MARKER_RESOLUTION), + ( + with(shared_key(), MARKER_RESOLUTION, "Resolved"), + MARKER_RESOLUTION, + ), + (without(shared_key(), MARKER_AUTH_TYPE), MARKER_AUTH_TYPE), + ( + resolved("Kerberos", &[(HADOOP_KEY, SECRET)]), + MARKER_AUTH_TYPE, + ), + (resolved(AUTH_TYPE_OAUTH, &[]), MARKER_OAUTH_PROVIDER_CLASS), + ( + oauth("com.example.MyTokenProvider", &[]), + MARKER_OAUTH_PROVIDER_CLASS, + ), + ]; + for (configs, marker) in cases { + let err = err(&configs); + assert!(err.contains(marker), "{marker}: {err}"); + } + } + + #[test] + fn a_required_key_missing_under_resolved_fails_naming_it() { + let cases: Vec<(Configs, &[&str])> = vec![ + (shared_key(), &[HADOOP_KEY]), + ( + resolved(AUTH_TYPE_SAS, &[(HADOOP_SAS_FIXED_TOKEN, SECRET)]), + &[HADOOP_SAS_FIXED_TOKEN], + ), + ( + client_creds(), + &[ + HADOOP_OAUTH_CLIENT_ID, + HADOOP_OAUTH_CLIENT_SECRET, + HADOOP_OAUTH_CLIENT_ENDPOINT, + ], + ), + (msi(), &[HADOOP_MSI_ENDPOINT]), + ( + workload_identity(), + &[ + HADOOP_OAUTH_CLIENT_ID, + HADOOP_MSI_TENANT, + HADOOP_WI_TOKEN_FILE, + HADOOP_MSI_AUTHORITY, + ], + ), + ]; + for (complete, required) in cases { + builder(&complete); + for key in required { + let err = err(&without(complete.clone(), key)); + assert!(err.contains(key), "{key}: {err}"); + } + } + } + + #[test] + fn error_marker_surfaces_the_driver_failure() { + let configs = configs(&[ + (MARKER_RESOLUTION, "error"), + ( + MARKER_ERROR_CLASS, + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + ), + (MARKER_ERROR_MESSAGE, "Failure to initialize configuration"), + ]); assert_eq!( - by_key.get(&AzureConfigKey::AuthorityId).map(String::as_str), - Some("from-msi") + err(&configs), + "Generic MicrosoftAzure error: Hadoop ABFS authentication failed on the driver for \ + abfss://data@myacct.dfs.core.windows.net: \ + org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException: \ + Failure to initialize configuration" ); } #[test] - fn create_store_succeeds_with_workload_identity_configs() { - let u = url("abfss://data@myacct.dfs.core.windows.net/path/file.parquet"); - let mut configs = HashMap::new(); - // A typical Workload Identity setup: client id, tenant id, and federated token - // file are all present, and `from_env()` will additionally pick them up from the - // AZURE_* env vars in production. - configs.insert( - "fs.azure.account.oauth2.client.id".into(), - "client-123".into(), - ); - configs.insert( - "fs.azure.account.oauth2.msi.tenant".into(), - "tenant-abc".into(), + fn only_none_reads_the_environment() { + let _guard = lock_env(); + std::env::set_var(ENV_KEY, KEY); + std::env::set_var(ENV_CLIENT_ID, "ambient-client"); + let from_env = builder_for(&url(), &configs(&[(MARKER_RESOLUTION, "none")])); + // Stray Hadoop values beside `none` are not read either. + let from_env_with_stray = builder_for( + &url(), + &configs(&[(MARKER_RESOLUTION, "none"), (HADOOP_KEY, SECRET)]), ); - configs.insert( - "fs.azure.account.oauth2.token.file".into(), - "/var/run/secrets/azure/tokens/azure-identity-token".into(), - ); - let (_store, path) = create_store(&u, &configs).expect("store builds"); - assert_eq!(path.as_ref(), "path/file.parquet"); + let sas = resolved(AUTH_TYPE_SAS, &[(HADOOP_SAS_FIXED_TOKEN, "sv=1&sig=x")]); + let mechanisms = [ + ("shared_key", shared_key(), Some(SECRET), None), + ("client_creds", client_creds(), None, Some("hadoop-client")), + ("msi", msi(), None, None), + ( + "workload_identity", + workload_identity(), + None, + Some("hadoop-client"), + ), + ("sas", sas, None, None), + ]; + let resolved: Vec<_> = mechanisms + .iter() + .map(|(name, configs, key, client)| { + (*name, builder_for(&url(), configs), *key, *client) + }) + .collect(); + let failing = [ + configs(&[ + (MARKER_RESOLUTION, "declined"), + (MARKER_AUTH_TYPE, "Custom"), + ]), + configs(&[(MARKER_RESOLUTION, "error")]), + configs(&[(HADOOP_KEY, SECRET)]), + ]; + let results: Vec> = failing + .iter() + .map(|configs| create_store(&url(), configs).map(|_| ())) + .collect(); + std::env::remove_var(ENV_KEY); + std::env::remove_var(ENV_CLIENT_ID); + + for builder in [from_env, from_env_with_stray] { + let builder = builder.expect("builder"); + assert_eq!( + value(&builder, AzureConfigKey::AccessKey).as_deref(), + Some(KEY) + ); + assert_eq!( + value(&builder, AzureConfigKey::ClientId).as_deref(), + Some("ambient-client") + ); + } + for (name, builder, key, client) in resolved { + let builder = builder.expect(name); + assert_eq!( + value(&builder, AzureConfigKey::AccessKey).as_deref(), + key, + "{name}" + ); + assert_eq!( + value(&builder, AzureConfigKey::ClientId).as_deref(), + client, + "{name}" + ); + } + for (configs, result) in failing.iter().zip(results) { + let err = result.err().map(|e| e.to_string()).expect("an error"); + assert!(!err.contains(KEY), "{configs:?}: {err}"); + } } #[test] - fn create_store_rejects_non_azure_scheme() { - let u = url("s3://bucket/file.parquet"); - let configs = HashMap::new(); - let err = create_store(&u, &configs).expect_err("must fail"); - assert!( - format!("{err}").contains("Scheme of URL is not Azure"), - "unexpected error: {err}" - ); + fn oauth_endpoint_parts_split_tenant_and_authority_host() { + fn parts(tenant: &'static str, host: &'static str) -> Option<(&'static str, &'static str)> { + Some((tenant, host)) + } + for (endpoint, expected) in [ + ( + "https://user:pass@auth-proxy.example/t/oauth2/token", + parts("t", "https://auth-proxy.example"), + ), + ( + "https://auth-proxy.example:443/t/oauth2/token", + parts("t", "https://auth-proxy.example"), + ), + ( + "https://auth-proxy.example:8443/t/oauth2/token", + parts("t", "https://auth-proxy.example:8443"), + ), + ( + "https://[::1]:8443/t/oauth2/token", + parts("t", "https://[::1]:8443"), + ), + ( + "https://auth-proxy.example/aad/t/oauth2/v2.0/token", + parts("t", "https://auth-proxy.example/aad"), + ), + ( + "https://auth-proxy.example:8443/a/b/t/oauth2/v2.0/token", + parts("t", "https://auth-proxy.example:8443/a/b"), + ), + // No `oauth2` after the first segment: the first segment is the tenant. + ( + "https://auth-proxy.example/aad/t/token", + parts("aad", "https://auth-proxy.example"), + ), + ( + "https://auth-proxy.example/oauth2/token", + parts("oauth2", "https://auth-proxy.example"), + ), + ("https://h//t/oauth2/token", parts("t", "https://h/")), + ("https://h/a//oauth2/token", None), + ( + "https://h/aad/t/oauth2/v2.0/token?q=oauth2#f", + parts("t", "https://h/aad"), + ), + ( + "https://h/a/oauth2/b/oauth2/token", + parts("b", "https://h/a/oauth2"), + ), + ( + "https://h/u@evil.example/t/oauth2/token", + parts("t", "https://h/u@evil.example"), + ), + // The URL parser resolves `%2e%2e` and reads `\` as `/`; the host stays `h`. + ("https://h/a/%2e%2e/t/oauth2/token", parts("t", "https://h")), + ("https://h/a\\b/t/oauth2/token", parts("t", "https://h/a/b")), + ( + "https://h/oauth2/t/oauth2/token", + parts("t", "https://h/oauth2"), + ), + // An opaque origin has no host to post to. + ("file:///t/oauth2/token", None), + ("https://auth-proxy.example/", None), + ("urn:t:oauth2", None), + ("not a url", None), + ] { + let actual = oauth_endpoint_parts(endpoint); + assert_eq!( + actual + .as_ref() + .map(|(tenant, host)| (tenant.as_str(), host.as_str())), + expected, + "{endpoint}" + ); + } } } diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 2a2c1737473..8fc807a14cb 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -1340,7 +1340,11 @@ mod tests { fn isolates_azure_containers_in_cache_and_shared_runtime() { use datafusion::execution::runtime_env::RuntimeEnvBuilder; - let options = HashMap::from([("fs.azure.account.key".into(), "c2VjcmV0".into())]); + let options = HashMap::from([ + ("comet.azure.resolution".into(), "resolved".into()), + ("comet.azure.auth.type".into(), "SharedKey".into()), + ("fs.azure.account.key".into(), "c2VjcmV0".into()), + ]); let runtime = RuntimeEnvBuilder::new() .with_object_store_registry(Arc::new(super::CometObjectStoreRegistry::default())) .build_arc() diff --git a/pom.xml b/pom.xml index a2970a2c6ff..0704b7a86c9 100644 --- a/pom.xml +++ b/pom.xml @@ -105,6 +105,11 @@ under the License. Hadoop 3.4.x's CredentialProviderListFactory); the spark-3.4 / spark-3.5 profiles override it down to 3.3.x (SDK v1). hadoop.version itself (hadoop-client) is not changed. --> 3.4.1 + + 3.4.2 ${project.basedir}/../native/target/debug darwin x86_64 @@ -705,6 +710,7 @@ under the License. 3.4.3 3.4 3.3.4 + 3.3.4 1.13.1 4.8.8 2.0.6 @@ -725,6 +731,7 @@ under the License. 3.5.9 3.5 3.3.4 + 3.3.4 1.13.1 4.8.8 2.0.7 @@ -744,6 +751,7 @@ under the License. 2.13 4.0.4 4.0 + 3.4.1 1.15.2 4.13.10 2.0.16 @@ -766,6 +774,7 @@ under the License. 2.13 4.1.3 4.1 + 3.4.2 1.16.0 4.13.10 2.0.17 @@ -786,6 +795,7 @@ under the License. 2.13 4.2.0 4.2 + 3.5.0 1.17.0 4.13.10 2.0.17 @@ -1081,6 +1091,9 @@ under the License. ${comet.test.udfs.lib} + + ${comet.abfs.parity.record} diff --git a/spark/pom.xml b/spark/pom.xml index 18c0ac60914..0dfd7bb259d 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -224,6 +224,39 @@ under the License. + + + org.apache.hadoop + hadoop-azure + ${hadoop-azure.version} + test + + + org.apache.hadoop.thirdparty + hadoop-shaded-guava + + + org.apache.hadoop + hadoop-annotations + + + org.eclipse.jetty + jetty-util-ajax + + + commons-logging + commons-logging + + + software.amazon.awssdk diff --git a/spark/src/main/scala/org/apache/comet/objectstore/AbfsAuthResolver.scala b/spark/src/main/scala/org/apache/comet/objectstore/AbfsAuthResolver.scala new file mode 100644 index 00000000000..a7938ce79a8 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/objectstore/AbfsAuthResolver.scala @@ -0,0 +1,517 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.net.URI + +import scala.collection.mutable.ArrayBuffer +import scala.util.{Failure, Success, Try} +import scala.util.control.NonFatal + +import org.apache.commons.lang3.StringUtils +import org.apache.hadoop.conf.Configuration +import org.apache.spark.internal.Logging + +import org.apache.comet.objectstore.AbfsReflection.Handles + +/** + * Resolves the ABFS authentication for one `abfs[s]://container@account/...` URI on the driver, + * by asking the classpath's hadoop-azure the questions + * `AzureBlobFileSystemStore.initializeClient` asks, and encodes the answer as the options the + * native Azure store reads. + * + * Hadoop's `AbfsConfiguration` decides which mechanism applies and where each value comes from + * (container-, account- or globally-scoped keys, credential providers, `${...}` references). The + * resolver forwards the resolved values under the global key names plus `comet.azure.*` markers, + * never a key for another account, and never the environment. Anything Hadoop throws becomes an + * error marker: the native scan fails with it instead of trying a lookup of its own. + * + * `AbfsAuthParitySuite` pins the outcome of every configuration shape per hadoop-azure version + * (`spark/src/test/resources/abfs-auth-parity/expected-.json`) and cross-checks it + * against the real `AzureBlobFileSystem`; its scaladoc says how to regenerate those files after a + * deliberate change here. + */ +private[comet] object AbfsAuthResolver extends Logging { + + /** Options the native Azure store reads. None of them starts with `fs.`. */ + object Markers { + val RESOLUTION = "comet.azure.resolution" + val AUTH_TYPE = "comet.azure.auth.type" + val OAUTH_PROVIDER_CLASS = "comet.azure.oauth.provider.class" + val SAS_PROVIDER_CLASS = "comet.azure.sas.provider.class" + val OAUTH_ASSERTION_PROVIDER_CLASS = "comet.azure.oauth.assertion.provider.class" + val ERROR_CLASS = "comet.azure.error.class" + val ERROR_MESSAGE = "comet.azure.error.message" + + val RESOLVED = "resolved" + val NONE = "none" + val DECLINED = "declined" + val ERROR = "error" + } + + /** Hadoop ABFS configuration keys, global form (`ConfigurationKeys`). */ + object Keys { + val AUTH_TYPE = "fs.azure.account.auth.type" + val ACCOUNT_KEY = "fs.azure.account.key" + val KEY_PROVIDER = "fs.azure.account.keyprovider" + val SHELL_KEY_PROVIDER_SCRIPT = "fs.azure.shellkeyprovider.script" + val OAUTH_PROVIDER_TYPE = "fs.azure.account.oauth.provider.type" + val CLIENT_ID = "fs.azure.account.oauth2.client.id" + val CLIENT_SECRET = "fs.azure.account.oauth2.client.secret" + val CLIENT_ENDPOINT = "fs.azure.account.oauth2.client.endpoint" + val MSI_TENANT = "fs.azure.account.oauth2.msi.tenant" + val MSI_ENDPOINT = "fs.azure.account.oauth2.msi.endpoint" + val MSI_AUTHORITY = "fs.azure.account.oauth2.msi.authority" + val USER_NAME = "fs.azure.account.oauth2.user.name" + val USER_PASSWORD = "fs.azure.account.oauth2.user.password" + val REFRESH_TOKEN = "fs.azure.account.oauth2.refresh.token" + val REFRESH_TOKEN_ENDPOINT = "fs.azure.account.oauth2.refresh.token.endpoint" + val TOKEN_FILE = "fs.azure.account.oauth2.token.file" + val CLIENT_ASSERTION_PROVIDER_TYPE = "fs.azure.account.oauth2.client.assertion.provider.type" + val SAS_TOKEN_PROVIDER_TYPE = "fs.azure.sas.token.provider.type" + val SAS_FIXED_TOKEN = "fs.azure.sas.fixed.token" + } + + /** The OAuth providers Hadoop accepts under `AuthType.OAuth`, matched by exact class name. */ + object Providers { + val CLIENT_CREDS = "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider" + val USER_PASSWORD = "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider" + val MSI = "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider" + val REFRESH_TOKEN = "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider" + val WORKLOAD_IDENTITY = "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider" + val BUILT_IN: Set[String] = + Set(CLIENT_CREDS, USER_PASSWORD, MSI, REFRESH_TOKEN, WORKLOAD_IDENTITY) + } + + private object AuthTypes { + val SHARED_KEY = "SharedKey" + val OAUTH = "OAuth" + val SAS = "SAS" + } + + private object Exceptions { + private val pkg = "org.apache.hadoop.fs.azurebfs.contracts.exceptions." + val CONFIGURATION_PROPERTY_NOT_FOUND: String = pkg + "ConfigurationPropertyNotFoundException" + val INVALID_CONFIGURATION_VALUE: String = pkg + "InvalidConfigurationValueException" + val INVALID_URI: String = pkg + "InvalidUriException" + val INVALID_URI_AUTHORITY: String = pkg + "InvalidUriAuthorityException" + } + + /** + * Exceptions whose messages name keys, classes or URIs and nothing else, so they are forwarded + * verbatim. Every other message (a `KeyProviderException` quotes the raw account key, custom + * providers say anything) is replaced by a pointer to the driver log. + */ + private val verbatimMessageClasses: Set[String] = Set( + Exceptions.CONFIGURATION_PROPERTY_NOT_FOUND, + Exceptions.INVALID_CONFIGURATION_VALUE, + Exceptions.INVALID_URI, + Exceptions.INVALID_URI_AUTHORITY, + classOf[ClassNotFoundException].getName, + classOf[NoSuchMethodException].getName, + classOf[NoSuchFieldException].getName) + + // An IllegalArgumentException is verbatim only when Hadoop or the JDK enum parser raised it + // ("No enum constant", "Failed to initialize", "Invalid account key."). One thrown + // by a user KeyProvider could carry anything, so it is reported by class only. + private val HADOOP_AZURE_PACKAGE = "org.apache.hadoop.fs.azurebfs." + private val JAVA_ENUM_CLASS = "java.lang.Enum" + + private def isVerbatim(t: Throwable): Boolean = t match { + case _: IllegalArgumentException => + t.getStackTrace.headOption.map(_.getClassName).exists { origin => + origin == JAVA_ENUM_CLASS || origin.startsWith(HADOOP_AZURE_PACKAGE) + } + case _ => verbatimMessageClasses.contains(t.getClass.getName) + } + + private val REDACTED = "" + // Shorter values (`true`, a port) would mangle messages without protecting anything. + private val MIN_REDACTED_LENGTH = 8 + // `SimpleKeyProvider` (3.4.1+) puts the raw account key in its message as `key ="..."`. + private val QUOTED_KEY_PATTERN = "key =\"[^\"]*\"".r + private val MAX_CAUSE_DEPTH = 8 + + // The keys that name a mechanism or carry a credential in hadoop-azure. A configuration where + // none of them resolves for the account is `Unconfigured`; any one of them hands the decision to + // Hadoop, which defaults to SharedKey and fails in its own words when that does not fit. + private[objectstore] val plainAuthKeys: Seq[String] = Seq( + Keys.AUTH_TYPE, + Keys.KEY_PROVIDER, + Keys.SHELL_KEY_PROVIDER_SCRIPT, + Keys.OAUTH_PROVIDER_TYPE, + Keys.SAS_TOKEN_PROVIDER_TYPE) + private[objectstore] val passwordAuthKeys: Seq[String] = Seq( + Keys.ACCOUNT_KEY, + Keys.CLIENT_ID, + Keys.CLIENT_SECRET, + Keys.CLIENT_ENDPOINT, + Keys.MSI_TENANT, + Keys.MSI_ENDPOINT, + Keys.MSI_AUTHORITY, + Keys.USER_NAME, + Keys.USER_PASSWORD, + Keys.REFRESH_TOKEN, + Keys.REFRESH_TOKEN_ENDPOINT, + Keys.TOKEN_FILE, + Keys.CLIENT_ASSERTION_PROVIDER_TYPE, + Keys.SAS_FIXED_TOKEN) + + sealed trait Outcome + + /** Hadoop named a mechanism and every value it needs; `values` use the global key names. */ + final case class Resolved( + authType: String, + providerClass: Option[String], + values: Map[String, String]) + extends Outcome + + /** No auth-bearing key resolves for the account. */ + case object Unconfigured extends Outcome + + /** + * Hadoop named a mechanism the native scan cannot run; `providerClass` was never instantiated. + * `assertionProviderClass` is the trimmed value of + * `fs.azure.account.oauth2.client.assertion.provider.type` that a Workload Identity + * configuration names on hadoop-azure 3.5.0 and later. + */ + final case class Declined( + authType: String, + providerClass: Option[String], + assertionProviderClass: Option[String] = None) + extends Outcome + + /** Hadoop (or the reflection onto it) threw; `message` is already redacted. */ + final case class Failed(exceptionClass: String, message: String) extends Outcome + + def resolve(hadoopConf: Configuration, uri: URI): Outcome = + resolve(hadoopConf, uri, AbfsReflection.handles(AbfsReflection.loaderFor(hadoopConf))) + + private[objectstore] def resolve( + hadoopConf: Configuration, + uri: URI, + handles: Try[Handles]): Outcome = { + val secrets = ArrayBuffer[String]() + // Set once the AbfsConfiguration exists, so a failure can read the credentials it skipped. + var openSession: Option[Session] = None + try { + handles match { + case Failure(e) => failed(uri, e, secrets) + case Success(h) => + authorityParts(uri) match { + case Left(failure) => failure + case Right((container, account)) => + val session = new Session( + h, + h.newAbfsConfiguration(hadoopConf, account, container, uri), + hadoopConf, + account, + secrets) + openSession = Some(session) + if (!session.isConfigured) { + Unconfigured + } else { + val authType: Enum[_] = h.getAuthType(session.abfsConf, account) + authType.name match { + case AuthTypes.SHARED_KEY => sharedKey(session, uri, account) + case AuthTypes.SAS => sas(session, authType) + case AuthTypes.OAUTH => oauth(session, authType) + // Custom and UserboundSASWithOAuth run Java providers per request; never + // instantiated here. + case other => Declined(other, None) + } + } + } + } + } catch { + // NonFatal alone misses NoSuchMethodError and friends, the reflection failures this must + // report as an error marker rather than let escape into planning. + case e @ (_: LinkageError | NonFatal(_)) => + openSession.foreach(_.recordSkippedSecrets()) + failed(uri, e, secrets) + } + } + + /** + * `(container, account)` exactly as `AzureBlobFileSystemStore.authorityParts`: the raw + * authority split once at the first `@`, so port, case and percent-encoding are kept. Failures + * carry Hadoop's exception class and message. + */ + private[objectstore] def authorityParts(uri: URI): Either[Failed, (String, String)] = { + Option(uri.getRawAuthority) match { + case None => + Left(Failed(Exceptions.INVALID_URI_AUTHORITY, s"$uri has invalid authority.")) + case Some(authority) if !authority.contains("@") => + Left(Failed(Exceptions.INVALID_URI_AUTHORITY, s"$uri has invalid authority.")) + case Some(authority) => + val parts = authority.split("@", 2) + if (parts.length < 2 || parts(0).isEmpty) { + Left( + Failed( + Exceptions.INVALID_URI, + s"Invalid URI '$uri' has a malformed authority, expected container name. " + + "Authority takes the form abfs://[@]")) + } else { + Right((parts(0), parts(1))) + } + } + } + + // One resolution's view of an AbfsConfiguration; records every credential it reads so a + // failure message can be redacted against them. Each password read may load a credential + // provider's keystore, so a successful resolution reads only what the mechanism needs. + private final class Session( + val h: Handles, + val abfsConf: AnyRef, + hadoopConf: Configuration, + account: String, + secrets: ArrayBuffer[String]) { + + def password(key: String): Option[String] = { + val value = h.getPasswordString(abfsConf, key) + value.foreach(v => recordSecret(v)) + value + } + + // Both spellings, since Hadoop quotes the raw value and the providers use the trimmed one. + private def recordSecret(value: String): Unit = + Seq(value, value.trim).filter(_.length >= MIN_REDACTED_LENGTH).foreach(secrets += _) + + /** `getTrimmedPasswordString`: blank means the default, then trimmed. */ + def trimmedPassword(key: String, default: String): String = + password(key).filterNot(StringUtils.isBlank).getOrElse(default).trim + + def storageAccountKey(): String = { + val key = h.getStorageAccountKey(abfsConf) + recordSecret(key) + key + } + + // Stops at the first key that is set, plain keys first since they need no password lookup. + def isConfigured: Boolean = + plainAuthKeys.exists(isSet(abfsConf, _, password = false)) || + passwordAuthKeys.exists(isSet(abfsConf, _, password = true)) || + hasKeyProviderAccountKey + + /** + * Reads every credential key the resolution may have skipped, so the failure message is + * redacted against all of them. A read that throws is ignored: the failure being reported is + * the one that matters. + */ + def recordSkippedSecrets(): Unit = { + passwordAuthKeys.foreach(key => bestEffort(isSet(abfsConf, key, password = true))) + bestEffort(hasKeyProviderAccountKey) + } + + private def bestEffort(read: => Boolean): Unit = + try { + read + () + } catch { + case _: LinkageError | NonFatal(_) => () + } + + // A value Hadoop cannot read still counts as set: Hadoop reads a key only when the mechanism + // it picks needs it, and then fails in its own words. + private def isSet(conf: AnyRef, key: String, password: Boolean): Boolean = + h.read(conf, key, password) match { + case Right(value) => + if (password) value.foreach(v => recordSecret(v)) + value.isDefined + case Left(_) => true + } + + // SimpleKeyProvider reads the key through its own two-argument AbfsConfiguration, whose + // container level on 3.4.2+ is the literal `key.null.`; a key placed there + // configures the account for Hadoop, so it counts here too. + private def hasKeyProviderAccountKey: Boolean = + h.hasContainerScopedConfiguration && isSet( + h.newKeyProviderConfiguration(hadoopConf, account), + Keys.ACCOUNT_KEY, + password = true) + } + + private def sharedKey(session: Session, uri: URI, account: String): Outcome = { + // Store.initializeClient checks the account name before fetching the key. + if (account.indexOf('.') <= 0) { + Failed(Exceptions.INVALID_URI, s"Invalid URI $uri - account name is not fully qualified.") + } else { + val key = session.storageAccountKey() + if (key.isEmpty) { + // SharedKeyCredentials rejects the empty key SimpleKeyProvider's Base64 check let through. + Failed(classOf[IllegalArgumentException].getName, "Invalid account key.") + } else { + Resolved(AuthTypes.SHARED_KEY, None, Map(Keys.ACCOUNT_KEY -> key)) + } + } + } + + private def sas(session: Session, authType: Enum[_]): Outcome = { + val h = session.h + h.getTokenProviderClass( + session.abfsConf, + authType, + Keys.SAS_TOKEN_PROVIDER_TYPE, + h.sasTokenProviderClass) match { + case Some(custom) => Declined(AuthTypes.SAS, Some(custom.getName)) + case None => + // Hadoop's own checks: 3.3.4 has no fixed token and fails here, 3.4.1+ requires one. + h.getSASTokenProvider(session.abfsConf) + val fixed = session.trimmedPassword(Keys.SAS_FIXED_TOKEN, "") + Resolved(AuthTypes.SAS, None, Map(Keys.SAS_FIXED_TOKEN -> fixed)) + } + } + + private def oauth(session: Session, authType: Enum[_]): Outcome = { + val h = session.h + val providerClass = h.getTokenProviderClass( + session.abfsConf, + authType, + Keys.OAUTH_PROVIDER_TYPE, + h.accessTokenProviderClass) + val className = providerClass.map(_.getName) + val assertionProvider = + if (h.hasClientAssertionProvider && className.contains(Providers.WORKLOAD_IDENTITY)) { + session.password(Keys.CLIENT_ASSERTION_PROVIDER_TYPE).map(_.trim).filter(_.nonEmpty) + } else { + None + } + assertionProvider match { + case Some(assertion) => + // 3.5.0 Workload Identity with a custom ClientAssertionProvider: user code per request. + Declined(AuthTypes.OAUTH, className, Some(assertion)) + case None => + // Hadoop validates the configuration and throws in its own words: an unresolved class, + // a class that is not one of its five built-ins (matched by `==`, never instantiated), + // a missing mandatory key. The built-in constructors do no I/O. + h.getTokenProvider(session.abfsConf) + Resolved(AuthTypes.OAUTH, className, oauthValues(session, className.getOrElse(""))) + } + } + + // The keys each built-in provider's constructor receives, in the form Hadoop hands them over: + // mandatory and plain keys verbatim, defaulted keys trimmed with the AuthConfigurations default. + private def oauthValues(session: Session, providerClass: String): Map[String, String] = { + val h = session.h + def plain(key: String): Map[String, String] = session.password(key).map(key -> _).toMap + def defaulted(key: String, defaultField: String): Map[String, String] = + Map(key -> session.trimmedPassword(key, h.authDefault(defaultField))) + def authority(): Map[String, String] = { + val value = session.trimmedPassword( + Keys.MSI_AUTHORITY, + h.authDefault("DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY")) + Map(Keys.MSI_AUTHORITY -> appendSlashIfNeeded(value)) + } + + providerClass match { + case Providers.CLIENT_CREDS => + plain(Keys.CLIENT_ENDPOINT) ++ plain(Keys.CLIENT_ID) ++ plain(Keys.CLIENT_SECRET) + case Providers.USER_PASSWORD => + plain(Keys.CLIENT_ENDPOINT) ++ plain(Keys.USER_NAME) ++ plain(Keys.USER_PASSWORD) + case Providers.MSI => + defaulted(Keys.MSI_ENDPOINT, "DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT") ++ + plain(Keys.MSI_TENANT) ++ plain(Keys.CLIENT_ID) ++ authority() + case Providers.REFRESH_TOKEN => + defaulted( + Keys.REFRESH_TOKEN_ENDPOINT, + "DEFAULT_FS_AZURE_ACCOUNT_OAUTH_REFRESH_TOKEN_ENDPOINT") ++ + plain(Keys.REFRESH_TOKEN) ++ plain(Keys.CLIENT_ID) + case Providers.WORKLOAD_IDENTITY => + authority() ++ plain(Keys.MSI_TENANT) ++ plain(Keys.CLIENT_ID) ++ + defaulted(Keys.TOKEN_FILE, "DEFAULT_FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE") + case _ => Map.empty + } + } + + private def appendSlashIfNeeded(authority: String): String = + if (authority.endsWith("/")) authority else authority + "/" + + /** + * The options for the native store: markers, plus the resolved values for `Resolved`. The + * provider marker is `comet.azure.sas.provider.class` for SAS and + * `comet.azure.oauth.provider.class` otherwise; a Workload Identity decline over a custom + * client assertion keeps the provider marker and adds + * `comet.azure.oauth.assertion.provider.class`. + */ + def toOptions(outcome: Outcome): Map[String, String] = outcome match { + case Resolved(authType, providerClass, values) => + Map(Markers.RESOLUTION -> Markers.RESOLVED, Markers.AUTH_TYPE -> authType) ++ + providerClass.map(providerMarker(authType) -> _) ++ values + case Unconfigured => + Map(Markers.RESOLUTION -> Markers.NONE) + case Declined(authType, providerClass, assertionProviderClass) => + Map(Markers.RESOLUTION -> Markers.DECLINED, Markers.AUTH_TYPE -> authType) ++ + providerClass.map(providerMarker(authType) -> _) ++ + assertionProviderClass.map(Markers.OAUTH_ASSERTION_PROVIDER_CLASS -> _) + case Failed(exceptionClass, message) => + Map( + Markers.RESOLUTION -> Markers.ERROR, + Markers.ERROR_CLASS -> exceptionClass, + Markers.ERROR_MESSAGE -> message) + } + + private def providerMarker(authType: String): String = + if (authType == AuthTypes.SAS) Markers.SAS_PROVIDER_CLASS else Markers.OAUTH_PROVIDER_CLASS + + /** + * Hides every `secret` of at least [[MIN_REDACTED_LENGTH]] characters in `message`, and the raw + * key a `SimpleKeyProvider` message quotes. + */ + private[objectstore] def redact(message: String, secrets: Iterable[String]): String = { + val withoutQuotedKey = QUOTED_KEY_PATTERN.replaceAllIn(message, "key =\"" + REDACTED + "\"") + secrets + .filter(_.length >= MIN_REDACTED_LENGTH) + .foldLeft(withoutQuotedKey)((text, secret) => text.replace(secret, REDACTED)) + } + + private def failed(uri: URI, e: Throwable, secrets: Iterable[String]): Failed = { + val chain = causeChain(e) + val thrown = chain.head + def describe(t: Throwable, verbatim: Boolean): String = { + val name = t.getClass.getName + if (verbatim) s"$name: ${redact(Option(t.getMessage).getOrElse(""), secrets)}" else name + } + logWarning( + s"Hadoop ABFS authentication for ${uri.getScheme}://${uri.getRawAuthority} failed on the " + + s"driver: ${chain.map(describe(_, verbatim = true)).mkString(", caused by ")}") + + val forwarded = chain.map(t => describe(t, isVerbatim(t))) + val hidden = chain.exists(t => !isVerbatim(t)) + val pointer = if (hidden) "; the full message is in the Spark driver log" else "" + val message = thrown match { + case _: ClassNotFoundException => + "hadoop-azure is not on the driver classpath; the native ABFS scan resolves " + + s"authentication through it (${forwarded.mkString(", caused by ")})" + case _ => forwarded.mkString(", caused by ") + pointer + } + Failed(thrown.getClass.getName, message) + } + + // KeyProviderException drops its cause, so its own message is all there is at that level. + private def causeChain(e: Throwable): Seq[Throwable] = { + val chain = ArrayBuffer[Throwable](e) + var cause = e.getCause + while (cause != null && chain.length < MAX_CAUSE_DEPTH && !chain.contains(cause)) { + chain += cause + cause = cause.getCause + } + chain.toSeq + } +} diff --git a/spark/src/main/scala/org/apache/comet/objectstore/AbfsReflection.scala b/spark/src/main/scala/org/apache/comet/objectstore/AbfsReflection.scala new file mode 100644 index 00000000000..9d9c5d0e2d0 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/objectstore/AbfsReflection.scala @@ -0,0 +1,299 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.lang.reflect.{Constructor, InvocationTargetException, Method} +import java.net.URI + +import scala.util.{Failure, Success, Try} +import scala.util.control.NonFatal + +import org.apache.hadoop.conf.Configuration + +import org.apache.comet.util.ClassLoaders + +/** + * Reflective access to the hadoop-azure on the classpath. Comet does not link against + * hadoop-azure; every version-dependent surface of `AbfsConfiguration` is resolved here, once per + * class loader, so the resolver in [[AbfsAuthResolver]] can ask Hadoop the same questions + * `AzureBlobFileSystemStore.initializeClient` asks. + */ +private[objectstore] object AbfsReflection { + + object ClassNames { + val ABFS_CONFIGURATION = "org.apache.hadoop.fs.azurebfs.AbfsConfiguration" + val ABFS_SERVICE_TYPE = "org.apache.hadoop.fs.azurebfs.constants.AbfsServiceType" + val AUTH_CONFIGURATIONS = "org.apache.hadoop.fs.azurebfs.constants.AuthConfigurations" + val AUTH_TYPE = "org.apache.hadoop.fs.azurebfs.services.AuthType" + val ACCESS_TOKEN_PROVIDER = "org.apache.hadoop.fs.azurebfs.oauth2.AccessTokenProvider" + val SAS_TOKEN_PROVIDER = "org.apache.hadoop.fs.azurebfs.extensions.SASTokenProvider" + // Present from hadoop-azure 3.4.1: the fixed SAS token path. + val FIXED_SAS_TOKEN_PROVIDER = "org.apache.hadoop.fs.azurebfs.services.FixedSASTokenProvider" + // Present from hadoop-azure 3.5.0: Workload Identity with a custom client assertion. + val CLIENT_ASSERTION_PROVIDER = "org.apache.hadoop.fs.azurebfs.oauth2.ClientAssertionProvider" + } + + // Store.getAbfsServiceTypeFromUrl: BLOB when the URI text contains this, DFS otherwise. + private val BLOB_DOMAIN_MARKER = ".blob." + + /** + * The operations the resolver needs from an `AbfsConfiguration`. A `Configuration` object is + * passed as `AnyRef` because its class is only known at run time. Hadoop's own exceptions + * propagate unwrapped from every method. + */ + trait Handles { + + /** True when hadoop-azure takes the container name (the 3.4.2+ four-argument constructor). */ + def hasContainerScopedConfiguration: Boolean + + /** True when `fs.azure.sas.fixed.token` exists (hadoop-azure 3.4.1+). */ + def hasFixedSasTokenProvider: Boolean + + /** True when a Workload Identity client assertion provider can be configured (3.5.0+). */ + def hasClientAssertionProvider: Boolean + + def accessTokenProviderClass: Class[_] + + def sasTokenProviderClass: Class[_] + + /** + * `new AbfsConfiguration(conf, account, container, serviceType)`, or the two-argument form. + */ + def newAbfsConfiguration( + conf: Configuration, + account: String, + container: String, + uri: URI): AnyRef + + /** + * `new AbfsConfiguration(conf, account)`: the form `SimpleKeyProvider` reads the key with. + */ + def newKeyProviderConfiguration(conf: Configuration, account: String): AnyRef + + /** `getAuthType(account)`: the `AuthType` enum constant. */ + def getAuthType(abfsConf: AnyRef, account: String): Enum[_] + + /** + * `getPasswordString(key)`: container (3.4.2+), account, then global, via credential + * providers. + */ + def getPasswordString(abfsConf: AnyRef, key: String): Option[String] + + /** + * `getPasswordString(key)` when `password` is set, otherwise `get(key)` (account-specific, + * then global). A value Hadoop cannot read (a `${...}` cycle, a failing credential provider) + * is returned as `Left` rather than thrown; fatal errors and linkage or reflection failures + * still propagate. + */ + def read(abfsConf: AnyRef, key: String, password: Boolean): Either[Throwable, Option[String]] + + /** + * `getTokenProviderClass(authType, key, null, xface)`: the provider class Hadoop would load. + */ + def getTokenProviderClass( + abfsConf: AnyRef, + authType: Enum[_], + key: String, + xface: Class[_]): Option[Class[_]] + + /** `getStorageAccountKey()`: runs the configured `KeyProvider`, as Hadoop does at FS init. */ + def getStorageAccountKey(abfsConf: AnyRef): String + + /** `getTokenProvider()`, called only for Hadoop's validation of a built-in OAuth provider. */ + def getTokenProvider(abfsConf: AnyRef): Unit + + /** `getSASTokenProvider()`, called only for Hadoop's validation of the fixed SAS token. */ + def getSASTokenProvider(abfsConf: AnyRef): Unit + + /** A `String` constant of `AuthConfigurations`, by field name. */ + def authDefault(fieldName: String): String + } + + // Keyed by the AbfsConfiguration class itself: a ClassValue holds nothing that pins a loader, + // and `computeValue` throwing caches nothing, so a loader that gains hadoop-azure later + // (ADD JAR) is probed again on the next call. + private val cache = new ClassValue[Handles] { + override protected def computeValue(abfsConfigurationClass: Class[_]): Handles = + new ReflectiveHandles(abfsConfigurationClass) + } + + /** The loader Hadoop's `FileSystem.get` would instantiate the ABFS classes with. */ + def loaderFor(hadoopConf: Configuration): ClassLoader = + Option(hadoopConf.getClassLoader) + .getOrElse(ClassLoaders.contextOrDefault(getClass.getClassLoader)) + + /** The handles for `loader`'s hadoop-azure, or the failure that keeps it from loading. */ + def handles(loader: ClassLoader): Try[Handles] = + attempt(cache.get(loadClass(loader, ClassNames.ABFS_CONFIGURATION))) + + // `Try` lets a `LinkageError` through, and a hadoop-azure whose own dependencies are missing + // fails with exactly that (`NoClassDefFoundError`), so it is caught here alongside `NonFatal`. + private def attempt[A](body: => A): Try[A] = { + try { + Success(body) + } catch { + case e @ (_: LinkageError | NonFatal(_)) => Failure(e) + } + } + + private def loadClass(loader: ClassLoader, name: String): Class[_] = + ClassLoaders.loadClass(name, loader) + + // Sibling classes come from the loader that defined AbfsConfiguration, the one hadoop-azure + // itself links against. + private final class ReflectiveHandles(abfsConfigurationClass: Class[_]) extends Handles { + private val loader = abfsConfigurationClass.getClassLoader + private val authTypeClass = loadClass(loader, ClassNames.AUTH_TYPE) + private val authConfigurationsClass = loadClass(loader, ClassNames.AUTH_CONFIGURATIONS) + + override val accessTokenProviderClass: Class[_] = + loadClass(loader, ClassNames.ACCESS_TOKEN_PROVIDER) + override val sasTokenProviderClass: Class[_] = + loadClass(loader, ClassNames.SAS_TOKEN_PROVIDER) + + private def isLoadable(name: String): Boolean = attempt(loadClass(loader, name)).isSuccess + + override val hasFixedSasTokenProvider: Boolean = isLoadable( + ClassNames.FIXED_SAS_TOKEN_PROVIDER) + override val hasClientAssertionProvider: Boolean = + isLoadable(ClassNames.CLIENT_ASSERTION_PROVIDER) + + // (Configuration, account, container, AbfsServiceType) exists from 3.4.2; the service type + // enum arrived with it, so a missing enum means the two-argument constructor. + private val serviceTypeClass: Option[Class[_]] = + attempt(loadClass(loader, ClassNames.ABFS_SERVICE_TYPE)).toOption + private val containerConstructor: Option[Constructor[_]] = serviceTypeClass.flatMap { st => + val ctor: Try[Constructor[_]] = attempt( + abfsConfigurationClass + .getConstructor(classOf[Configuration], classOf[String], classOf[String], st)) + ctor.toOption + } + private val accountConstructor: Constructor[_] = + abfsConfigurationClass.getConstructor(classOf[Configuration], classOf[String]) + + private val getAuthTypeMethod = + abfsConfigurationClass.getMethod("getAuthType", classOf[String]) + private val getMethod = abfsConfigurationClass.getMethod("get", classOf[String]) + private val getPasswordStringMethod = + abfsConfigurationClass.getMethod("getPasswordString", classOf[String]) + private val getTokenProviderClassMethod = abfsConfigurationClass.getMethod( + "getTokenProviderClass", + authTypeClass, + classOf[String], + classOf[Class[_]], + classOf[Class[_]]) + private val getStorageAccountKeyMethod = + abfsConfigurationClass.getMethod("getStorageAccountKey") + private val getTokenProviderMethod = abfsConfigurationClass.getMethod("getTokenProvider") + private val getSASTokenProviderMethod = + abfsConfigurationClass.getMethod("getSASTokenProvider") + + override def hasContainerScopedConfiguration: Boolean = containerConstructor.isDefined + + override def newAbfsConfiguration( + conf: Configuration, + account: String, + container: String, + uri: URI): AnyRef = { + (containerConstructor, serviceTypeClass) match { + case (Some(ctor), Some(st)) => + val serviceTypeName = if (uri.toString.contains(BLOB_DOMAIN_MARKER)) "BLOB" else "DFS" + construct(ctor, conf, account, container, enumConstant(st, serviceTypeName)) + case _ => + construct(accountConstructor, conf, account) + } + } + + override def newKeyProviderConfiguration(conf: Configuration, account: String): AnyRef = + construct(accountConstructor, conf, account) + + override def getAuthType(abfsConf: AnyRef, account: String): Enum[_] = + invoke(getAuthTypeMethod, abfsConf, account).asInstanceOf[Enum[_]] + + override def getPasswordString(abfsConf: AnyRef, key: String): Option[String] = + Option(invoke(getPasswordStringMethod, abfsConf, key)).map(_.asInstanceOf[String]) + + override def read( + abfsConf: AnyRef, + key: String, + password: Boolean): Either[Throwable, Option[String]] = { + val method = if (password) getPasswordStringMethod else getMethod + try { + Right(Option(method.invoke(abfsConf, key)).map(_.asInstanceOf[String])) + } catch { + // NonFatal excludes LinkageError, which rethrows below like any other failure. + case e: InvocationTargetException if e.getCause != null && NonFatal(e.getCause) => + Left(e.getCause) + case e: InvocationTargetException => throw unwrapped(e) + } + } + + override def getTokenProviderClass( + abfsConf: AnyRef, + authType: Enum[_], + key: String, + xface: Class[_]): Option[Class[_]] = { + Option(invoke(getTokenProviderClassMethod, abfsConf, authType, key, null, xface)) + .map(_.asInstanceOf[Class[_]]) + } + + override def getStorageAccountKey(abfsConf: AnyRef): String = + invoke(getStorageAccountKeyMethod, abfsConf).asInstanceOf[String] + + override def getTokenProvider(abfsConf: AnyRef): Unit = { + invoke(getTokenProviderMethod, abfsConf) + () + } + + override def getSASTokenProvider(abfsConf: AnyRef): Unit = { + invoke(getSASTokenProviderMethod, abfsConf) + () + } + + override def authDefault(fieldName: String): String = + authConfigurationsClass.getField(fieldName).get(null).asInstanceOf[String] + + private def enumConstant(enumClass: Class[_], name: String): AnyRef = + enumClass.getEnumConstants + .find(_.asInstanceOf[Enum[_]].name == name) + .getOrElse(throw new NoSuchFieldException(s"${enumClass.getName}.$name")) + .asInstanceOf[AnyRef] + + // Reflection wraps whatever the target throws; hand the resolver Hadoop's own exception. + private def invoke(method: Method, target: AnyRef, args: AnyRef*): AnyRef = { + try { + method.invoke(target, args: _*) + } catch { + case e: InvocationTargetException => throw unwrapped(e) + } + } + + private def construct(ctor: Constructor[_], args: AnyRef*): AnyRef = { + try { + ctor.newInstance(args: _*).asInstanceOf[AnyRef] + } catch { + case e: InvocationTargetException => throw unwrapped(e) + } + } + + private def unwrapped(e: InvocationTargetException): Throwable = + Option(e.getCause).getOrElse(e) + } +} diff --git a/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala b/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala index 75328649fea..b7c09da4694 100644 --- a/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala +++ b/spark/src/main/scala/org/apache/comet/objectstore/NativeConfig.scala @@ -37,12 +37,10 @@ object NativeConfig { "s3a" -> Seq("fs.s3a."), // Google Cloud Storage configurations "gs" -> Seq("fs.gs."), - // Azure Blob Storage configurations (can use both prefixes) + // Azure Blob Storage configurations (can use both prefixes). ABFS (abfs, abfss) is not + // forwarded by prefix: `AbfsAuthResolver` resolves its authentication for the one account. "wasb" -> Seq("fs.azure.", "fs.wasb."), - "wasbs" -> Seq("fs.azure.", "fs.wasb."), - // Azure Data Lake Storage Gen2 (ABFS) configurations. Hadoop ABFS authentication keys - "abfs" -> Seq("fs.azure.", "fs.abfs."), - "abfss" -> Seq("fs.azure.", "fs.abfss.", "fs.abfs.")) + "wasbs" -> Seq("fs.azure.", "fs.wasb.")) // Some alias filesystems report the literal authority "default" when the URI has none (e.g. // `scheme:///bucket/key`); the real bucket is then promoted from the URL path. Keys under this @@ -188,6 +186,10 @@ object NativeConfig { * Extract object store configs (S3, GCS, Azure, ...) from the Hadoop configuration for native * DataFusion. Captures global and per-bucket keys; native code prefers per-bucket. * + * For `abfs`/`abfss` the `fs.azure.*` namespace is not copied. [[AbfsAuthResolver]] asks the + * classpath's hadoop-azure which mechanism and values apply to the URI's account and forwards + * them as `comet.azure.*` markers plus global-named keys. + * * A scheme listed in `COMET_S3_COMPLIANT_SCHEMES_KEY` reuses the `fs.s3a.*` surface, and its * vendor `fs...*` keys are translated to `fs.s3a.*` (see * `translateVendorKeys`) AFTER the `fs.s3a.*` pass, so real vendor values override conflicting @@ -214,6 +216,13 @@ object NativeConfig { } val s3CompliantSchemes = parseSchemeSet(s3CompliantRaw) + // ABFS: Hadoop resolves the mechanism and values for this URI's account on the driver, and + // only those travel to native (never another account's keys, never the whole namespace). + if (scheme == "abfs" || scheme == "abfss") { + return options.toMap ++ AbfsAuthResolver.toOptions( + AbfsAuthResolver.resolve(hadoopConf, uri)) + } + // Prefixes for this scheme; a configured S3-compliant alias reuses the fs.s3a.* surface. val prefixes = objectStoreConfigPrefixes.get(scheme).orElse { if (s3CompliantSchemes.contains(scheme)) objectStoreConfigPrefixes.get("s3a") else None @@ -250,6 +259,41 @@ object NativeConfig { options.toMap } + /** + * The object store options for a scan over `uris`, or the reason the native scan cannot read + * them. The options are those of the first URI. Native reads every file of a partition through + * the store built for the partition's first file, and an ABFS store is bound to one container + * of one account, so a scan whose ABFS files span more than one container or account, or that + * mixes ABFS and other schemes, is declined. + */ + def extractScanObjectStoreOptions( + hadoopConf: Configuration, + uris: Iterable[URI]): Either[String, Map[String, String]] = { + val iter = uris.iterator + if (!iter.hasNext) { + return Right(Map.empty) + } + val first = iter.next() + val firstIsAbfs = isAbfs(first) + while (iter.hasNext) { + val uri = iter.next() + if (isAbfs(uri) != firstIsAbfs) { + return Left( + "Native scan does not support a scan that mixes ABFS and other file systems " + + s"(${first.getScheme} and ${uri.getScheme})") + } + if (firstIsAbfs && uri.getRawAuthority != first.getRawAuthority) { + return Left( + "Native scan does not support ABFS files from more than one container or account " + + s"(${first.getRawAuthority} and ${uri.getRawAuthority})") + } + } + Right(extractObjectStoreOptions(hadoopConf, first)) + } + + private def isAbfs(uri: URI): Boolean = + lowerScheme(uri).exists(scheme => scheme == "abfs" || scheme == "abfss") + /** * The value Hadoop's own consumers observe for `key`. `Configuration#get` expands any `${...}` * variable reference in the stored value against other conf entries and system properties, diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index a47f1c2a6b4..569ccc277f7 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -443,7 +443,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) convertToComet(scan, CometIcebergNativeScan).getOrElse(scan) case scan: CometBatchScanExec if scan.wrapped.scan.isInstanceOf[CSVScan] => - convertToComet(scan, CometCsvNativeScanExec).getOrElse(scan) + convertToComet(scan, CometCsvNativeScanExec).getOrElse(scan.wrapped) // Comet JVM + native scan for V1 and V2 case op if isCometScan(op) => diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index c7c9a31909d..176c4c595ae 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -585,7 +585,14 @@ case class CometScanRule(session: SparkSession) val effectiveUri = new java.net.URI(effectiveLocation) - val hadoopS3Options = NativeConfig.extractObjectStoreOptions(hadoopConf, effectiveUri) + // Only a scheme the native Iceberg reader can open is worth resolving credentials + // for: the ABFS resolver may run a KeyProvider, and the scan is rejected below anyway. + val hadoopS3Options = + if (CometScanRule.isIcebergReadableScheme(effectiveUri, s3CompliantSchemes)) { + NativeConfig.extractObjectStoreOptions(hadoopConf, effectiveUri) + } else { + Map.empty[String, String] + } // Promote the DATA bucket's per-bucket `fs.s3a.bucket..*` settings to global (what // the native FileIO reads), NOT the metadata bucket's: the FileIO opens data/delete diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala index 286937635b7..38e7cdbf2e4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala @@ -251,6 +251,23 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS } if (scanTypes.length == scan.output.length) { + // Collect S3/cloud storage configurations + val hadoopConf = scan.relation.sparkSession.sessionState + .newHadoopConfWithOptions(scan.relation.options) + + // Use selectedPartitions (static) instead of getFilePartitions() because at planning time + // DPP subqueries haven't been resolved yet. Object store options don't depend on DPP. + val fileUris = scan.selectedPartitions.view.flatMap(_.files).map(_.getPath.toUri) + val objectStoreOptions = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match { + case Right(options) => options + case Left(reason) => + // CometExecRule falls back to the wrapped Spark scan, so it carries the reason too. + withFallbackReason(scan, reason) + withFallbackReason(scan.wrapped, reason) + return None + } + commonBuilder.addAllFields(scanTypes.asJava) // Sink operators don't have children @@ -279,14 +296,6 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS return None } - // Extract object store options from first file (S3 configs apply to all files in scan). - // Use selectedPartitions (static) instead of getFilePartitions() because at planning time - // DPP subqueries haven't been resolved yet. Object store options don't depend on DPP. - val firstFileUri = scan.selectedPartitions - .flatMap(_.files.headOption) - .headOption - .map(_.getPath.toUri) - // Constant metadata columns (file_path, file_name, file_size, file_block_start, // file_block_length, file_modification_time) are known before opening the file and // constant for every row read from it, exactly like partition columns. Spark places @@ -356,18 +365,10 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS commonBuilder.setAllowTypePromotion(CometConf.COMET_SCHEMA_EVOLUTION_ENABLED) commonBuilder.setAllowTimestampLtzToNtz(CometConf.COMET_ALLOW_TIMESTAMP_LTZ_AS_NTZ) - // Collect S3/cloud storage configurations - val hadoopConf = scan.relation.sparkSession.sessionState - .newHadoopConfWithOptions(scan.relation.options) - commonBuilder.setEncryptionEnabled(CometParquetUtils.encryptionEnabled(hadoopConf)) - firstFileUri.foreach { uri => - val objectStoreOptions = - NativeConfig.extractObjectStoreOptions(hadoopConf, uri) - objectStoreOptions.foreach { case (key, value) => - commonBuilder.putObjectStoreOptions(key, value) - } + objectStoreOptions.foreach { case (key, value) => + commonBuilder.putObjectStoreOptions(key, value) } // Set common data in NativeScan (file_partition will be populated at execution time) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala index dc92648b8d5..3d9f3f773fc 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometCsvNativeScanExec.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan import com.google.common.base.Objects import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.objectstore.NativeConfig import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass} import org.apache.comet.serde.OperatorOuterClass.Operator @@ -93,6 +94,18 @@ object CometCsvNativeScanExec extends CometOperatorSerde[CometBatchScanExec] { new CSVOptions(csvScan.options.asScala.toMap, columnPruning, timeZone) } val filePartitions = op.inputPartitions.map(_.asInstanceOf[FilePartition]) + val hadoopConf = + sessionState.newHadoopConfWithOptions(op.session.sparkContext.conf.getAll.toMap) + val fileUris = filePartitions.view.flatMap(_.files).map(_.pathUri) + val objectStoreOptions = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match { + case Right(options) => options + case Left(reason) => + // CometExecRule falls back to the wrapped Spark scan, so it carries the reason too. + withFallbackReason(op, reason) + withFallbackReason(op.wrapped, reason) + return None + } val csvOptionsProto = csvOptions2Proto(options) val dataSchemaProto = schema2Proto(csvScan.dataSchema) val readSchemaFieldNames = csvScan.readDataSchema.fieldNames @@ -104,15 +117,6 @@ object CometCsvNativeScanExec extends CometOperatorSerde[CometBatchScanExec] { val partitionSchemaProto = schema2Proto(csvScan.readPartitionSchema) val partitionsProto = filePartitions.map(partition2Proto(_, csvScan.readPartitionSchema)) - val objectStoreOptions = filePartitions.headOption - .flatMap { partitionFile => - val hadoopConf = sessionState - .newHadoopConfWithOptions(op.session.sparkContext.conf.getAll.toMap) - partitionFile.files.headOption - .map(file => NativeConfig.extractObjectStoreOptions(hadoopConf, file.pathUri)) - } - .getOrElse(Map.empty) - csvScanBuilder.putAllObjectStoreOptions(objectStoreOptions.asJava) csvScanBuilder.setCsvOptions(csvOptionsProto) csvScanBuilder.addAllFilePartitions(partitionsProto.asJava) diff --git a/spark/src/test/java/org/apache/comet/hadoop/fs/FakeAbfsFileSystem.java b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeAbfsFileSystem.java new file mode 100644 index 00000000000..fe032e0d9db --- /dev/null +++ b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeAbfsFileSystem.java @@ -0,0 +1,59 @@ +/* + * 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. + */ + +package org.apache.comet.hadoop.fs; + +import java.io.IOException; +import java.net.URI; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.RawLocalFileSystem; + +/** + * A local-disk-backed FileSystem for the {@code abfss} scheme: {@code abfss://c@a/} is the + * local file at {@code }, so a test can plan a scan over several ABFS containers and + * accounts, and run its Spark fallback, without Azure. + */ +public class FakeAbfsFileSystem extends RawLocalFileSystem { + + private URI uri; + + public FakeAbfsFileSystem() { + // Avoid `URI scheme is not "file"` error on + // RawLocalFileSystem$DeprecatedRawLocalFileStatus.getOwner + RawLocalFileSystem.useStatIfAvailable(); + } + + @Override + public void initialize(URI name, Configuration conf) throws IOException { + super.initialize(name, conf); + uri = URI.create(name.getScheme() + "://" + name.getRawAuthority()); + } + + @Override + public String getScheme() { + return "abfss"; + } + + @Override + public URI getUri() { + // The superclass constructor qualifies its working directory before initialize runs. + return uri != null ? uri : URI.create("abfss://uninitialized@fake.dfs.core.windows.net"); + } +} diff --git a/spark/src/test/resources/abfs-auth-parity/expected-3.3.4.json b/spark/src/test/resources/abfs-auth-parity/expected-3.3.4.json new file mode 100644 index 00000000000..6cc13e20926 --- /dev/null +++ b/spark/src/test/resources/abfs-auth-parity/expected-3.3.4.json @@ -0,0 +1,876 @@ +{ + "cases": { + "at-account-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-oauth-global-sharedkey": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-account-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-account-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-global-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-global-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-unset-no-keys": { + "outcome": "unconfigured", + "values": {} + }, + "at-unset-other-account-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-base64-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-base64-non-ascii": { + "errorClass": "java.lang.ArrayIndexOutOfBoundsException", + "outcome": "failed", + "values": {} + }, + "cal-container-account-key-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-container-oauth-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-empty-account-key": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-empty-mandatory-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + } + }, + "cal-endpoint-with-path-prefix": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.example.net/prefix/tenant-abc/oauth2/v2.0/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-global-enum-parsed-first": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-msi-no-tenant": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-msi-no-tenant-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-null-fsname-key": { + "outcome": "unconfigured", + "values": {} + }, + "cal-null-fsname-key-beats-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cal-oauth-keys-without-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-port-key-matches": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "cal-port-key-missing": { + "outcome": "unconfigured", + "values": {} + }, + "cal-sas-container-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-guard-with-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-guard-without-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-untrimmed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-userbound-sas": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-wi-client-assertion-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "cal-wi-on-334": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "cp-jceks-account-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-beats-clear-text": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-client-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cp-jceks-no-fallback": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-no-fallback-clear-text-only": { + "outcome": "unconfigured", + "values": {} + }, + "cp-no-fallback-with-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-account-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-account-type-global-provider": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-global-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-global-type-global-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-sas-account-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-account-type-global-provider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "guard-sas-global-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-global-type-global-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "oauth-clientcreds": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-container-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-global": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-mixed-scopes": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-secret": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-other-account-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-padded-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": " https://login.microsoftonline.com/tenant-abc/oauth2/token ", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-custom-no-provider-type": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-custom-not-adaptee": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-msi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-blank-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-container-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-custom-endpoint-authority": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.example.net/", + "fs.azure.account.oauth2.msi.endpoint": "http://msi.example.net/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-no-provider-type": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-empty": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-built-in": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-refreshtoken": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.microsoftonline.com/Common/oauth2/token" + } + }, + "oauth-refreshtoken-custom-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.example.net/token" + } + }, + "oauth-refreshtoken-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-refreshtoken-no-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.user.name": "user@example.com", + "fs.azure.account.oauth2.user.password": "sha256:c55bf0309fee4de8125a00ee97699a19a269caa35afcc8d1058bbd6283167fc8" + } + }, + "oauth-userpassword-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-password": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-user": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-blank-assertion-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-container-token-file": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-no-client-id": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-no-tenant": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-padded-token-file": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-blank-fixed": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "sas-custom-and-fixed": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-custom-only": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-fixed-global": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "sas-fixed-only": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "sas-neither": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "sas-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-question-mark": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-beats-global": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-account-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-container-ignored-with-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-global-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-global-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-custom": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:5be524205bd6aee5594ca7d1288c421a4cca089d63453278d1f75277c6d382bc" + } + }, + "sk-keyprovider-custom-null": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.ConfigurationPropertyNotFoundException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-empty": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-not-a-keyprovider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-simple": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-unknown-class": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-echo": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-shell-script-missing": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-script-unset": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "subst-cross-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:41800564f8313f5a8d3f468f205a7c39b55f63fa6efdb2dddd0005441f3c44bb" + } + }, + "subst-env-default": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "fallback-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "subst-env-home": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-account-without-dot": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-blob-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-blob-in-path-only": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-dfs-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-empty-container": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-no-at": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-null-authority": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-percent-encoded-container": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-port-in-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "uri-second-at": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account-lower-case-key": { + "outcome": "unconfigured", + "values": {} + } + }, + "hadoopAzureVersion": "3.3.4" +} diff --git a/spark/src/test/resources/abfs-auth-parity/expected-3.4.1.json b/spark/src/test/resources/abfs-auth-parity/expected-3.4.1.json new file mode 100644 index 00000000000..98ef579fe81 --- /dev/null +++ b/spark/src/test/resources/abfs-auth-parity/expected-3.4.1.json @@ -0,0 +1,936 @@ +{ + "cases": { + "at-account-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-oauth-global-sharedkey": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-account-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-account-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-global-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-global-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-unset-no-keys": { + "outcome": "unconfigured", + "values": {} + }, + "at-unset-other-account-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-base64-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-base64-non-ascii": { + "errorClass": "java.lang.ArrayIndexOutOfBoundsException", + "outcome": "failed", + "values": {} + }, + "cal-container-account-key-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-container-oauth-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "cal-empty-account-key": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-empty-mandatory-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + } + }, + "cal-endpoint-with-path-prefix": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.example.net/prefix/tenant-abc/oauth2/v2.0/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-global-enum-parsed-first": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-msi-no-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-msi-no-tenant-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-null-fsname-key": { + "outcome": "unconfigured", + "values": {} + }, + "cal-null-fsname-key-beats-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cal-oauth-keys-without-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-port-key-matches": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "cal-port-key-missing": { + "outcome": "unconfigured", + "values": {} + }, + "cal-sas-container-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-guard-with-fixed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-sas-guard-without-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-untrimmed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-userbound-sas": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-wi-client-assertion-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "cal-wi-on-334": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "cp-jceks-account-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-beats-clear-text": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-client-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cp-jceks-no-fallback": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-no-fallback-clear-text-only": { + "outcome": "unconfigured", + "values": {} + }, + "cp-no-fallback-with-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-account-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-account-type-global-provider": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-global-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-global-type-global-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-sas-account-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-account-type-global-provider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "guard-sas-global-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-global-type-global-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "oauth-clientcreds": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-container-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-global": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-mixed-scopes": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-secret": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-other-account-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-padded-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": " https://login.microsoftonline.com/tenant-abc/oauth2/token ", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-custom-no-provider-type": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-custom-not-adaptee": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-msi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-blank-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-container-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-custom-endpoint-authority": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.example.net/", + "fs.azure.account.oauth2.msi.endpoint": "http://msi.example.net/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-no-provider-type": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-empty": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-built-in": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-refreshtoken": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.microsoftonline.com/Common/oauth2/token" + } + }, + "oauth-refreshtoken-custom-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.example.net/token" + } + }, + "oauth-refreshtoken-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-refreshtoken-no-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.user.name": "user@example.com", + "fs.azure.account.oauth2.user.password": "sha256:c55bf0309fee4de8125a00ee97699a19a269caa35afcc8d1058bbd6283167fc8" + } + }, + "oauth-userpassword-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-password": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-user": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-blank-assertion-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-container-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-no-tenant": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-padded-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/mnt/tokens/token" + } + }, + "sas-blank-fixed": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-custom-and-fixed": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-custom-only": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-fixed-global": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-fixed-only": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-neither": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-question-mark": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:4f629ccc08252cde45726db17f66f06819423bba3254103f893dce6a1e626fc8" + } + }, + "sk-account-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-beats-global": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-account-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-container-ignored-with-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-global-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-global-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-custom": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:5be524205bd6aee5594ca7d1288c421a4cca089d63453278d1f75277c6d382bc" + } + }, + "sk-keyprovider-custom-null": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.ConfigurationPropertyNotFoundException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-empty": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-not-a-keyprovider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-simple": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-unknown-class": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-echo": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-shell-script-missing": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-script-unset": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "subst-cross-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:41800564f8313f5a8d3f468f205a7c39b55f63fa6efdb2dddd0005441f3c44bb" + } + }, + "subst-env-default": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "fallback-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "subst-env-home": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-account-without-dot": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-blob-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-blob-in-path-only": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-dfs-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-empty-container": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-no-at": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-null-authority": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-percent-encoded-container": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-port-in-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "uri-second-at": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account-lower-case-key": { + "outcome": "unconfigured", + "values": {} + } + }, + "hadoopAzureVersion": "3.4.1" +} diff --git a/spark/src/test/resources/abfs-auth-parity/expected-3.4.2.json b/spark/src/test/resources/abfs-auth-parity/expected-3.4.2.json new file mode 100644 index 00000000000..d2a368a7998 --- /dev/null +++ b/spark/src/test/resources/abfs-auth-parity/expected-3.4.2.json @@ -0,0 +1,947 @@ +{ + "cases": { + "at-account-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-oauth-global-sharedkey": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-account-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-account-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-global-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-global-userbound": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-unset-no-keys": { + "outcome": "unconfigured", + "values": {} + }, + "at-unset-other-account-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-base64-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-base64-non-ascii": { + "errorClass": "java.lang.ArrayIndexOutOfBoundsException", + "outcome": "failed", + "values": {} + }, + "cal-container-account-key-only": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-container-oauth-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-empty-account-key": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-empty-mandatory-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + } + }, + "cal-endpoint-with-path-prefix": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.example.net/prefix/tenant-abc/oauth2/v2.0/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-global-enum-parsed-first": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-msi-no-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-msi-no-tenant-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-null-fsname-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:bb99e5bb3a756453730056b900774c14cb31358abd79a0903bf0e0349d0ade7a" + } + }, + "cal-null-fsname-key-beats-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:bb99e5bb3a756453730056b900774c14cb31358abd79a0903bf0e0349d0ade7a" + } + }, + "cal-oauth-keys-without-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-port-key-matches": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "cal-port-key-missing": { + "outcome": "unconfigured", + "values": {} + }, + "cal-sas-container-fixed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-sas-guard-with-fixed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-sas-guard-without-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-untrimmed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-userbound-sas": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-wi-client-assertion-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "cal-wi-on-334": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "cp-jceks-account-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-beats-clear-text": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-client-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cp-jceks-no-fallback": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-no-fallback-clear-text-only": { + "outcome": "unconfigured", + "values": {} + }, + "cp-no-fallback-with-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-account-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-account-type-global-provider": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-global-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-global-type-global-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-sas-account-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-account-type-global-provider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "guard-sas-global-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-global-type-global-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "oauth-clientcreds": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-container-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:1c008528aad647fbd3246a750e3f57dd6877398fbf25416d9269571ee4f35129" + } + }, + "oauth-clientcreds-global": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-mixed-scopes": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-secret": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-other-account-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-padded-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": " https://login.microsoftonline.com/tenant-abc/oauth2/token ", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-custom-no-provider-type": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-custom-not-adaptee": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-msi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-blank-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-container-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "container-tenant" + } + }, + "oauth-msi-custom-endpoint-authority": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.example.net/", + "fs.azure.account.oauth2.msi.endpoint": "http://msi.example.net/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-no-provider-type": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-empty": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-built-in": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-refreshtoken": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.microsoftonline.com/Common/oauth2/token" + } + }, + "oauth-refreshtoken-custom-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.example.net/token" + } + }, + "oauth-refreshtoken-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-refreshtoken-no-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.user.name": "user@example.com", + "fs.azure.account.oauth2.user.password": "sha256:c55bf0309fee4de8125a00ee97699a19a269caa35afcc8d1058bbd6283167fc8" + } + }, + "oauth-userpassword-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-password": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-user": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-blank-assertion-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-container-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/mnt/tokens/container-token" + } + }, + "oauth-wi-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-no-tenant": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-padded-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/mnt/tokens/token" + } + }, + "sas-blank-fixed": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-custom-and-fixed": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-custom-only": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-fixed-global": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-fixed-only": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-neither": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-question-mark": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:4f629ccc08252cde45726db17f66f06819423bba3254103f893dce6a1e626fc8" + } + }, + "sk-account-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-beats-global": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-account-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-container-ignored-with-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-global-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-global-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-custom": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:5be524205bd6aee5594ca7d1288c421a4cca089d63453278d1f75277c6d382bc" + } + }, + "sk-keyprovider-custom-null": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.ConfigurationPropertyNotFoundException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-empty": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-not-a-keyprovider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-simple": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-unknown-class": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-echo": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-shell-script-missing": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-script-unset": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "subst-cross-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:41800564f8313f5a8d3f468f205a7c39b55f63fa6efdb2dddd0005441f3c44bb" + } + }, + "subst-env-default": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "fallback-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "subst-env-home": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-account-without-dot": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-blob-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-blob-in-path-only": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-dfs-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-empty-container": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-no-at": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-null-authority": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-percent-encoded-container": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "container-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-port-in-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "uri-second-at": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account-lower-case-key": { + "outcome": "unconfigured", + "values": {} + } + }, + "hadoopAzureVersion": "3.4.2" +} diff --git a/spark/src/test/resources/abfs-auth-parity/expected-3.5.0.json b/spark/src/test/resources/abfs-auth-parity/expected-3.5.0.json new file mode 100644 index 00000000000..7197ec2fb7a --- /dev/null +++ b/spark/src/test/resources/abfs-auth-parity/expected-3.5.0.json @@ -0,0 +1,943 @@ +{ + "cases": { + "at-account-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-account-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-oauth-global-sharedkey": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-account-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-account-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-account-userbound": { + "authType": "UserboundSASWithOAuth", + "outcome": "declined", + "values": {} + }, + "at-global-custom": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "at-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-lower-case": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "at-global-oauth": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "at-global-sas-custom": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "at-global-sharedkey": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "at-global-userbound": { + "authType": "UserboundSASWithOAuth", + "outcome": "declined", + "values": {} + }, + "at-unset-no-keys": { + "outcome": "unconfigured", + "values": {} + }, + "at-unset-other-account-only": { + "outcome": "unconfigured", + "values": {} + }, + "cal-base64-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-base64-non-ascii": { + "errorClass": "java.lang.ArrayIndexOutOfBoundsException", + "outcome": "failed", + "values": {} + }, + "cal-container-account-key-only": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-container-oauth-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-empty-account-key": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-empty-mandatory-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + } + }, + "cal-endpoint-with-path-prefix": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.example.net/prefix/tenant-abc/oauth2/v2.0/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cal-global-enum-parsed-first": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "cal-msi-no-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-msi-no-tenant-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token" + } + }, + "cal-null-fsname-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:bb99e5bb3a756453730056b900774c14cb31358abd79a0903bf0e0349d0ade7a" + } + }, + "cal-null-fsname-key-beats-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:bb99e5bb3a756453730056b900774c14cb31358abd79a0903bf0e0349d0ade7a" + } + }, + "cal-oauth-keys-without-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "cal-port-key-matches": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "cal-port-key-missing": { + "outcome": "unconfigured", + "values": {} + }, + "cal-sas-container-fixed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-sas-guard-with-fixed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-sas-guard-without-fixed-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "cal-sas-untrimmed-token": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "cal-userbound-sas": { + "authType": "UserboundSASWithOAuth", + "outcome": "declined", + "values": {} + }, + "cal-wi-client-assertion-provider": { + "assertionProviderClass": "com.example.comet.NeverLoadedAssertionProvider", + "authType": "OAuth", + "outcome": "declined", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": {} + }, + "cal-wi-on-334": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "cp-jceks-account-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-beats-clear-text": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-jceks-client-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "cp-jceks-no-fallback": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "cp-no-fallback-clear-text-only": { + "outcome": "unconfigured", + "values": {} + }, + "cp-no-fallback-with-auth-type": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-account-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-account-type-global-provider": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "guard-oauth-global-type-account-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-oauth-global-type-global-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "guard-sas-account-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-account-type-global-provider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "guard-sas-global-type-account-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "guard-sas-global-type-global-provider": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "oauth-clientcreds": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-container-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:1c008528aad647fbd3246a750e3f57dd6877398fbf25416d9269571ee4f35129" + } + }, + "oauth-clientcreds-global": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-mixed-scopes": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-no-secret": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-clientcreds-other-account-secret": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-clientcreds-padded-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": " https://login.microsoftonline.com/tenant-abc/oauth2/token ", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-custom-no-provider-type": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-custom-not-adaptee": { + "authType": "Custom", + "outcome": "declined", + "values": {} + }, + "oauth-msi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-blank-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-container-tenant": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "container-tenant" + } + }, + "oauth-msi-custom-endpoint-authority": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.example.net/", + "fs.azure.account.oauth2.msi.endpoint": "http://msi.example.net/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-msi-no-client-id": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider", + "values": { + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.endpoint": "http://169.254.169.254/metadata/identity/oauth2/token", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc" + } + }, + "oauth-no-provider-type": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-empty": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-not-built-in": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "oauth-provider-padded": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "oauth-refreshtoken": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.microsoftonline.com/Common/oauth2/token" + } + }, + "oauth-refreshtoken-custom-endpoint": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.refresh.token": "sha256:e65009f6e0ae9fc204adcb73a208f63c582b7839bde7fe836da36a3df6ae7a85", + "fs.azure.account.oauth2.refresh.token.endpoint": "https://login.example.net/token" + } + }, + "oauth-refreshtoken-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-refreshtoken-no-token": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.user.name": "user@example.com", + "fs.azure.account.oauth2.user.password": "sha256:c55bf0309fee4de8125a00ee97699a19a269caa35afcc8d1058bbd6283167fc8" + } + }, + "oauth-userpassword-no-endpoint": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-password": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-userpassword-no-user": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-blank-assertion-provider": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/var/run/secrets/azure/tokens/azure-identity-token" + } + }, + "oauth-wi-container-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/mnt/tokens/container-token" + } + }, + "oauth-wi-no-client-id": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-no-tenant": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException", + "outcome": "failed", + "values": {} + }, + "oauth-wi-padded-token-file": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", + "values": { + "fs.azure.account.oauth2.client.id": "client-123", + "fs.azure.account.oauth2.msi.authority": "https://login.microsoftonline.com/", + "fs.azure.account.oauth2.msi.tenant": "tenant-abc", + "fs.azure.account.oauth2.token.file": "/mnt/tokens/token" + } + }, + "sas-blank-fixed": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-custom-and-fixed": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-custom-only": { + "authType": "SAS", + "outcome": "declined", + "providerClass": "org.apache.comet.objectstore.ParitySasTokenProvider", + "values": {} + }, + "sas-fixed-global": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-fixed-only": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:332a89b8fbd21ed8631f462437d07dcc560375e25159ff1db89cdb7945c0bd1b" + } + }, + "sas-neither": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException", + "outcome": "failed", + "values": {} + }, + "sas-provider-missing-class": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-provider-not-a-provider": { + "errorClass": "java.lang.RuntimeException", + "outcome": "failed", + "values": {} + }, + "sas-question-mark": { + "authType": "SAS", + "outcome": "resolved", + "values": { + "fs.azure.sas.fixed.token": "sha256:4f629ccc08252cde45726db17f66f06819423bba3254103f893dce6a1e626fc8" + } + }, + "sk-account-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-beats-global": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-account-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-account-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-account-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-container-ignored-with-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-global-bad-char": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-bad-length": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-empty": { + "errorClass": "java.lang.IllegalArgumentException", + "outcome": "failed", + "values": {} + }, + "sk-global-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-global-valid": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-custom": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:5be524205bd6aee5594ca7d1288c421a4cca089d63453278d1f75277c6d382bc" + } + }, + "sk-keyprovider-custom-null": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.ConfigurationPropertyNotFoundException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-empty": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-not-a-keyprovider": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-padded": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-keyprovider-simple": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-keyprovider-unknown-class": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-echo": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "sk-shell-script-missing": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "sk-shell-script-unset": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException", + "outcome": "failed", + "values": {} + }, + "subst-cross-key": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:41800564f8313f5a8d3f468f205a7c39b55f63fa6efdb2dddd0005441f3c44bb" + } + }, + "subst-env-default": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "fallback-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "subst-env-home": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-account-without-dot": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-blob-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-blob-in-path-only": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-dfs-host": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-empty-container": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException", + "outcome": "failed", + "values": {} + }, + "uri-no-at": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-null-authority": { + "errorClass": "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException", + "outcome": "failed", + "values": {} + }, + "uri-percent-encoded-container": { + "authType": "OAuth", + "outcome": "resolved", + "providerClass": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider", + "values": { + "fs.azure.account.oauth2.client.endpoint": "https://login.microsoftonline.com/tenant-abc/oauth2/token", + "fs.azure.account.oauth2.client.id": "container-client", + "fs.azure.account.oauth2.client.secret": "sha256:c1a39329c3e1799f02b7067ae57bbece8f4accbc5b952cc94fe97799332f1b27" + } + }, + "uri-port-in-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:e1e799c979ec65dcc29a55d9fe8d6dba055a00b83ababed9600d2ea0462f4fe3" + } + }, + "uri-second-at": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account": { + "authType": "SharedKey", + "outcome": "resolved", + "values": { + "fs.azure.account.key": "sha256:1c1185e02ff3e23b3e5a1c5bc86cf15d4126caa3dcde0fdb6e93adc4deec119e" + } + }, + "uri-upper-case-account-lower-case-key": { + "outcome": "unconfigured", + "values": {} + } + }, + "hadoopAzureVersion": "3.5.0" +} diff --git a/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParityCases.scala b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParityCases.scala new file mode 100644 index 00000000000..73652d0cee4 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParityCases.scala @@ -0,0 +1,1014 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.net.URI +import java.util.{Date, Locale} + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.azurebfs.extensions.{CustomTokenProviderAdaptee, SASTokenProvider} +import org.apache.hadoop.fs.azurebfs.services.KeyProvider + +import org.apache.comet.objectstore.AbfsAuthResolver.{Keys, Providers} + +/** + * The configurations [[AbfsAuthParitySuite]] runs through the resolver and through the real + * `AzureBlobFileSystem`. Built in code so nothing checked in can drift from the cases; the + * recorded expectation per hadoop-azure version lives in `abfs-auth-parity/expected-*.json`. + * + * A case may carry `pins`: the outcome kind and error class that reading the hadoop-azure source + * predicts per version. Every `calibration` case has one for all four versions, so the recorded + * files are checked against that reading rather than trusted. + */ +object AbfsAuthParityCases { + + val VERSIONS: Seq[String] = Seq("3.3.4", "3.4.1", "3.4.2", "3.5.0") + + object Outcomes { + val RESOLVED = "resolved" + val UNCONFIGURED = "unconfigured" + val DECLINED = "declined" + val FAILED = "failed" + } + + object ExceptionClasses { + private val pkg = "org.apache.hadoop.fs.azurebfs.contracts.exceptions." + val KEY_PROVIDER: String = pkg + "KeyProviderException" + val INVALID_URI: String = pkg + "InvalidUriException" + val INVALID_URI_AUTHORITY: String = pkg + "InvalidUriAuthorityException" + val TOKEN_ACCESS_PROVIDER: String = pkg + "TokenAccessProviderException" + val SAS_TOKEN_PROVIDER: String = pkg + "SASTokenProviderException" + val CONFIGURATION_PROPERTY_NOT_FOUND: String = pkg + "ConfigurationPropertyNotFoundException" + val ILLEGAL_ARGUMENT: String = classOf[IllegalArgumentException].getName + val RUNTIME: String = classOf[RuntimeException].getName + val ARRAY_INDEX: String = classOf[ArrayIndexOutOfBoundsException].getName + } + + /** The outcome the hadoop-azure source predicts for one version. */ + final case class Pin(outcome: String, errorClass: Option[String] = None) + + object Pin { + val resolved: Pin = Pin(Outcomes.RESOLVED) + val unconfigured: Pin = Pin(Outcomes.UNCONFIGURED) + val declined: Pin = Pin(Outcomes.DECLINED) + def failed(exceptionClass: String): Pin = Pin(Outcomes.FAILED, Some(exceptionClass)) + } + + /** + * @param liveValues + * forwarded values that depend on the machine (an environment variable); the suite checks + * them directly and records a placeholder instead. + */ + final case class ParityCase( + id: String, + group: String, + uri: URI, + conf: Map[String, String], + note: String, + pins: Map[String, Pin] = Map.empty, + liveValues: Map[String, String] = Map.empty) + + def pins(v334: Pin, v341: Pin, v342: Pin, v350: Pin): Map[String, Pin] = + Map("3.3.4" -> v334, "3.4.1" -> v341, "3.4.2" -> v342, "3.5.0" -> v350) + + def allVersions(pin: Pin): Map[String, Pin] = pins(pin, pin, pin, pin) + + /** `before` on 3.3.4, `after` from 3.4.1 (fixed SAS token, Workload Identity, optional MSI). */ + def from341(before: Pin, after: Pin): Map[String, Pin] = pins(before, after, after, after) + + /** `before` up to 3.4.1, `after` from 3.4.2 (container-scoped keys). */ + def from342(before: Pin, after: Pin): Map[String, Pin] = pins(before, before, after, after) + + /** `before` up to 3.4.2, `after` on 3.5.0 (user-bound SAS, client assertion provider). */ + def from350(before: Pin, after: Pin): Map[String, Pin] = pins(before, before, before, after) + + val ACCOUNT = "myacct.dfs.core.windows.net" + val CONTAINER = "data" + val URI_DFS = new URI(s"abfss://$CONTAINER@$ACCOUNT/path/file.parquet") + + // Base64 of "secret", "global", "container", "null", "custom", "port", "env". + val KEY_ACCOUNT = "c2VjcmV0" + val KEY_GLOBAL = "Z2xvYmFs" + val KEY_CONTAINER = "Y29udGFpbmVy" + val KEY_NULL_LITERAL = "bnVsbA==" + val KEY_CUSTOM = "Y3VzdG9t" + val KEY_PORT = "cG9ydA==" + val KEY_ENV = "ZW52" + + val ENDPOINT = "https://login.microsoftonline.com/tenant-abc/oauth2/token" + val CLIENT_ID = "client-123" + val CLIENT_SECRET = "client-secret-value" + val TENANT = "tenant-abc" + val FIXED_SAS = "sv=2020-08-04&sig=fixed-signature" + val DEFAULT_MSI_ENDPOINT = "http://169.254.169.254/metadata/identity/oauth2/token" + val DEFAULT_AUTHORITY = "https://login.microsoftonline.com/" + val DEFAULT_TOKEN_FILE = "/var/run/secrets/azure/tokens/azure-identity-token" + + val SIMPLE_KEY_PROVIDER = "org.apache.hadoop.fs.azurebfs.services.SimpleKeyProvider" + val SHELL_KEY_PROVIDER = "org.apache.hadoop.fs.azurebfs.services.ShellDecryptionKeyProvider" + val MISSING_CLASS = "com.example.comet.MissingProvider" + val NOT_A_PROVIDER: String = classOf[String].getName + val CUSTOM_ADAPTEE: String = classOf[ParityCustomTokenAdaptee].getName + val SAS_PROVIDER: String = classOf[ParitySasTokenProvider].getName + val FOREIGN_ACCESS_TOKEN_PROVIDER: String = + classOf[NeverInstantiatedAccessTokenProvider].getName + + /** The aliases the suite stores in its JCEKS keystore, with no clear-text counterpart. */ + val credentialEntries: Map[String, String] = + Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT, acct(Keys.CLIENT_SECRET) -> CLIENT_SECRET) + + /** + * The variable `${env.COMET_ABFS_PARITY_UNSET}` falls back to; the suite asserts it is unset. + */ + val UNSET_ENV_VAR = "COMET_ABFS_PARITY_UNSET" + + def acct(key: String, account: String = ACCOUNT): String = s"$key.$account" + + def cont(key: String, container: String = CONTAINER, account: String = ACCOUNT): String = + s"$key.$container.$account" + + private def uriFor(container: String, account: String): URI = + new URI(s"abfss://$container@$account/path/file.parquet") + + private val clientCreds: Map[String, String] = Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.CLIENT_CREDS, + acct(Keys.CLIENT_ENDPOINT) -> ENDPOINT, + acct(Keys.CLIENT_ID) -> CLIENT_ID, + acct(Keys.CLIENT_SECRET) -> CLIENT_SECRET) + + private val msi: Map[String, String] = Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.MSI, + acct(Keys.MSI_TENANT) -> TENANT, + acct(Keys.CLIENT_ID) -> CLIENT_ID) + + private val workloadIdentity: Map[String, String] = Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.WORKLOAD_IDENTITY, + acct(Keys.MSI_TENANT) -> TENANT, + acct(Keys.CLIENT_ID) -> CLIENT_ID) + + private val refreshToken: Map[String, String] = Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.REFRESH_TOKEN, + acct(Keys.REFRESH_TOKEN) -> "refresh-token-value", + acct(Keys.CLIENT_ID) -> CLIENT_ID) + + private val userPassword: Map[String, String] = Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.USER_PASSWORD, + acct(Keys.CLIENT_ENDPOINT) -> ENDPOINT, + acct(Keys.USER_NAME) -> "user@example.com", + acct(Keys.USER_PASSWORD) -> "user-password-value") + + private val sasCustom: Map[String, String] = + Map(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> SAS_PROVIDER) + + // 3.3.4 has no fixed SAS token, so a configuration that names no provider class fails there in + // TokenAccessProviderException and in SASTokenProviderException from 3.4.1. + private val fixedSasOnly: Map[String, Pin] = + from341(Pin.failed(ExceptionClasses.TOKEN_ACCESS_PROVIDER), Pin.resolved) + private val noSasSource: Map[String, Pin] = from341( + Pin.failed(ExceptionClasses.TOKEN_ACCESS_PROVIDER), + Pin.failed(ExceptionClasses.SAS_TOKEN_PROVIDER)) + + // WorkloadIdentityTokenProvider is absent from 3.3.4; Configuration.getClass raises the + // RuntimeException the resolver reports (Hadoop itself would wrap it once more). + private val fromWorkloadIdentity: Map[String, Pin] = + from341(Pin.failed(ExceptionClasses.RUNTIME), Pin.resolved) + + def cases(credentialProviderPath: String): Seq[ParityCase] = + uriCases ++ authTypeCases ++ sharedKeyCases ++ oauthCases ++ guardCases ++ sasCases ++ + credentialProviderCases(credentialProviderPath) ++ substitutionCases ++ calibrationCases + + private def c( + id: String, + group: String, + note: String, + conf: Map[String, String], + uri: URI = URI_DFS, + pins: Map[String, Pin] = Map.empty, + liveValues: Map[String, String] = Map.empty): ParityCase = + ParityCase(id, group, uri, conf, note, pins, liveValues) + + // ---- URI shapes ---- + + private def uriCases: Seq[ParityCase] = { + val blobAccount = "myacct.blob.core.windows.net" + val portAccount = s"$ACCOUNT:8080" + val upperAccount = "MyAcct.DFS.core.windows.net" + val encodedContainer = "my%2Ddata" + val secondAt = s"user@$ACCOUNT" + Seq( + c("uri-dfs-host", "uri", "plain dfs host", Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT)), + // hns.enabled=false: 3.4.2+ reject a Blob endpoint for an HNS account after authentication. + c( + "uri-blob-host", + "uri", + ".blob. host selects the BLOB service type on 3.4.2+", + Map( + acct(Keys.ACCOUNT_KEY, blobAccount) -> KEY_ACCOUNT, + "fs.azure.account.hns.enabled" -> "false"), + uri = uriFor(CONTAINER, blobAccount)), + c( + "uri-blob-in-path-only", + "uri", + ".blob. in the path also selects BLOB (Store checks the whole URI text)", + Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT, "fs.azure.account.hns.enabled" -> "false"), + uri = new URI(s"abfss://$CONTAINER@$ACCOUNT/x.blob.y/file.parquet")), + c( + "uri-port-in-account", + "uri", + "the port is part of the account name Hadoop looks keys up under", + Map( + acct(Keys.ACCOUNT_KEY, portAccount) -> KEY_PORT, + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, portAccount)), + c( + "uri-upper-case-account", + "uri", + "account name case is kept when matching keys", + Map(acct(Keys.ACCOUNT_KEY, upperAccount) -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, upperAccount)), + c( + "uri-upper-case-account-lower-case-key", + "uri", + "a lower-case key does not configure a mixed-case account", + Map(acct(Keys.ACCOUNT_KEY, upperAccount.toLowerCase(Locale.ROOT)) -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, upperAccount), + pins = allVersions(Pin.unconfigured)), + c( + "uri-percent-encoded-container", + "uri", + "container-scoped keys match the raw (still encoded) container name", + clientCreds ++ Map(cont(Keys.CLIENT_ID, encodedContainer) -> "container-client"), + uri = uriFor(encodedContainer, ACCOUNT)), + c( + "uri-second-at", + "uri", + "everything after the first '@' is the account name", + Map(acct(Keys.ACCOUNT_KEY, secondAt) -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, secondAt)), + c( + "uri-no-at", + "uri", + "no '@' in the authority", + Map(Keys.ACCOUNT_KEY -> KEY_GLOBAL), + uri = new URI(s"abfss://$ACCOUNT/path/file.parquet"), + pins = allVersions(Pin.failed(ExceptionClasses.INVALID_URI_AUTHORITY))), + c( + "uri-empty-container", + "uri", + "empty container name", + Map(Keys.ACCOUNT_KEY -> KEY_GLOBAL), + uri = new URI(s"abfss://@$ACCOUNT/path/file.parquet"), + pins = allVersions(Pin.failed(ExceptionClasses.INVALID_URI))), + c( + "uri-null-authority", + "uri", + "no authority at all", + Map(Keys.ACCOUNT_KEY -> KEY_GLOBAL), + uri = new URI("abfss:///path/file.parquet"), + pins = allVersions(Pin.failed(ExceptionClasses.INVALID_URI_AUTHORITY))), + c( + "uri-account-without-dot", + "uri", + "SharedKey needs a fully qualified account name", + Map(acct(Keys.ACCOUNT_KEY, "myacct") -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, "myacct"), + pins = allVersions(Pin.failed(ExceptionClasses.INVALID_URI)))) + } + + // ---- auth type values and scopes ---- + + private def authTypeCases: Seq[ParityCase] = { + val oauthKeys = clientCreds - acct(Keys.AUTH_TYPE) + val sharedKey = Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT) + val badEnum = Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT) + def scoped(scope: String, key: String): String = + if (scope == "global") key else acct(key) + Seq("global", "account").flatMap { scope => + val at = scoped(scope, Keys.AUTH_TYPE) + Seq( + c( + s"at-$scope-sharedkey", + "auth-type", + s"SharedKey ($scope)", + sharedKey + (at -> "SharedKey")), + c(s"at-$scope-oauth", "auth-type", s"OAuth ($scope)", oauthKeys + (at -> "OAuth")), + c( + s"at-$scope-sas-custom", + "auth-type", + s"SAS with a custom provider ($scope)", + Map(at -> "SAS", acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> SAS_PROVIDER), + pins = allVersions(Pin.declined)), + c( + s"at-$scope-custom", + "auth-type", + s"Custom ($scope)", + Map(at -> "Custom", acct(Keys.OAUTH_PROVIDER_TYPE) -> CUSTOM_ADAPTEE), + pins = allVersions(Pin.declined)), + c( + s"at-$scope-lower-case", + "auth-type", + s"'oauth' is not an enum constant ($scope)", + oauthKeys + (at -> "oauth"), + pins = allVersions(badEnum)), + c( + s"at-$scope-padded", + "auth-type", + s"' OAuth ' is trimmed by getEnum ($scope)", + oauthKeys + (at -> " OAuth "), + pins = allVersions(Pin.resolved)), + c( + s"at-$scope-empty", + "auth-type", + s"'' is not unset: Enum.valueOf fails ($scope)", + oauthKeys + (at -> ""), + pins = allVersions(badEnum)), + c( + s"at-$scope-userbound", + "auth-type", + s"UserboundSASWithOAuth exists from 3.5.0 ($scope)", + Map(at -> "UserboundSASWithOAuth", acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> SAS_PROVIDER), + pins = from350(badEnum, Pin.declined))) + } ++ Seq( + c( + "at-account-oauth-global-sharedkey", + "auth-type", + "the account value wins over a valid global one", + clientCreds + (Keys.AUTH_TYPE -> "SharedKey") + (Keys.ACCOUNT_KEY -> KEY_GLOBAL), + pins = allVersions(Pin.resolved)), + c( + "at-unset-no-keys", + "auth-type", + "nothing configured for the account", + Map("fs.azure.account.hns.enabled" -> "true"), + pins = allVersions(Pin.unconfigured)), + c( + "at-unset-other-account-only", + "auth-type", + "another account's key does not configure this one", + Map(acct(Keys.ACCOUNT_KEY, "other.dfs.core.windows.net") -> KEY_GLOBAL), + pins = allVersions(Pin.unconfigured))) + } + + // ---- SharedKey ---- + + private def sharedKeyCases: Seq[ParityCase] = { + val kpe = Pin.failed(ExceptionClasses.KEY_PROVIDER) + val emptyKey = Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT) + val badKeys = Seq( + ("bad-length", "c2VjcmV", "length % 4 != 0", kpe), + ("bad-char", "c2Vj!mV0", "a character outside the Base64 alphabet", kpe), + ("padded", " c2VjcmV0 ", "whitespace is not trimmed, so the length is wrong", kpe), + ("empty", "", "'' passes the Base64 check and fails in SharedKeyCredentials", emptyKey)) + val placements = Seq(("account", acct(Keys.ACCOUNT_KEY)), ("global", Keys.ACCOUNT_KEY)) + val validity = placements.flatMap { case (placement, key) => + c( + s"sk-$placement-valid", + "sharedkey", + s"valid key ($placement)", + Map(key -> KEY_ACCOUNT)) +: + badKeys.map { case (suffix, value, note, pin) => + c( + s"sk-$placement-$suffix", + "sharedkey", + s"$note ($placement)", + Map(key -> value), + pins = allVersions(pin)) + } + } + val providers = Seq( + c( + "sk-keyprovider-simple", + "sharedkey", + "explicit SimpleKeyProvider", + Map( + acct(Keys.KEY_PROVIDER) -> SIMPLE_KEY_PROVIDER, + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT)), + c( + "sk-keyprovider-unknown-class", + "sharedkey", + "keyprovider names a class that does not exist", + Map(acct(Keys.KEY_PROVIDER) -> MISSING_CLASS, acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-keyprovider-not-a-keyprovider", + "sharedkey", + "keyprovider names a class that is not a KeyProvider", + Map(acct(Keys.KEY_PROVIDER) -> NOT_A_PROVIDER, acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-keyprovider-empty", + "sharedkey", + "keyprovider '' is a class name Hadoop cannot load", + Map(acct(Keys.KEY_PROVIDER) -> "", acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-keyprovider-padded", + "sharedkey", + "keyprovider is read with get, which does not trim", + Map( + acct(Keys.KEY_PROVIDER) -> s" $SIMPLE_KEY_PROVIDER ", + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-keyprovider-custom", + "sharedkey", + "a custom KeyProvider supplies the key; no fs.azure.account.key needed", + Map(acct(Keys.KEY_PROVIDER) -> classOf[ParityKeyProvider].getName), + pins = allVersions(Pin.resolved)), + c( + "sk-keyprovider-custom-null", + "sharedkey", + "a KeyProvider returning null is reported against the account name", + Map(acct(Keys.KEY_PROVIDER) -> classOf[ParityNullKeyProvider].getName), + pins = allVersions(Pin.failed(ExceptionClasses.CONFIGURATION_PROPERTY_NOT_FOUND))), + c( + "sk-shell-echo", + "sharedkey", + "ShellDecryptionKeyProvider runs /bin/echo, which returns the envelope", + Map( + acct(Keys.KEY_PROVIDER) -> SHELL_KEY_PROVIDER, + acct(Keys.SHELL_KEY_PROVIDER_SCRIPT) -> "/bin/echo", + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(Pin.resolved)), + c( + "sk-shell-script-unset", + "sharedkey", + "ShellDecryptionKeyProvider without a script", + Map(acct(Keys.KEY_PROVIDER) -> SHELL_KEY_PROVIDER, acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-shell-script-missing", + "sharedkey", + "ShellDecryptionKeyProvider with a script that does not exist", + Map( + acct(Keys.KEY_PROVIDER) -> SHELL_KEY_PROVIDER, + acct(Keys.SHELL_KEY_PROVIDER_SCRIPT) -> "/nonexistent/comet-decrypt.sh", + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe)), + c( + "sk-account-beats-global", + "sharedkey", + "the account key wins over the global one", + Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT, Keys.ACCOUNT_KEY -> KEY_GLOBAL)), + c( + "sk-container-ignored-with-account", + "sharedkey", + "SimpleKeyProvider never reads the container-scoped key", + Map(cont(Keys.ACCOUNT_KEY) -> KEY_CONTAINER, acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(Pin.resolved))) + validity ++ providers + } + + // ---- OAuth ---- + + private def oauthCases: Seq[ParityCase] = { + val tape = Pin.failed(ExceptionClasses.TOKEN_ACCESS_PROVIDER) + def missing(base: Map[String, String], key: String, id: String, pin: Map[String, Pin]) = + c( + s"oauth-$id", + "oauth", + s"${key.stripPrefix("fs.azure.account.oauth2.")} missing", + base - acct(key), + pins = pin) + + Seq( + // ClientCreds + c("oauth-clientcreds", "oauth", "ClientCreds, all keys set", clientCreds), + missing(clientCreds, Keys.CLIENT_ENDPOINT, "clientcreds-no-endpoint", allVersions(tape)), + missing(clientCreds, Keys.CLIENT_ID, "clientcreds-no-client-id", allVersions(tape)), + missing(clientCreds, Keys.CLIENT_SECRET, "clientcreds-no-secret", allVersions(tape)), + c( + "oauth-clientcreds-padded-endpoint", + "oauth", + "mandatory keys are forwarded untrimmed", + clientCreds + (acct(Keys.CLIENT_ENDPOINT) -> s" $ENDPOINT ")), + c( + "oauth-clientcreds-container-secret", + "oauth", + "container-scoped secret wins on 3.4.2+, is ignored before", + clientCreds + (cont(Keys.CLIENT_SECRET) -> "container-secret-value")), + c( + "oauth-clientcreds-global", + "oauth", + "ClientCreds configured globally", + clientCreds.map { case (k, v) => k.stripSuffix(s".$ACCOUNT") -> v }), + c( + "oauth-clientcreds-mixed-scopes", + "oauth", + "account auth type and provider, global endpoint and secret", + Map( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.CLIENT_CREDS, + Keys.CLIENT_ENDPOINT -> ENDPOINT, + acct(Keys.CLIENT_ID) -> CLIENT_ID, + Keys.CLIENT_SECRET -> CLIENT_SECRET)), + c( + "oauth-clientcreds-other-account-secret", + "oauth", + "another account's secret is never forwarded", + clientCreds + (acct(Keys.CLIENT_SECRET, "other.dfs.core.windows.net") -> "other-secret")), + // UserPassword + c("oauth-userpassword", "oauth", "UserPassword, all keys set", userPassword), + missing(userPassword, Keys.USER_NAME, "userpassword-no-user", allVersions(tape)), + missing(userPassword, Keys.USER_PASSWORD, "userpassword-no-password", allVersions(tape)), + missing(userPassword, Keys.CLIENT_ENDPOINT, "userpassword-no-endpoint", allVersions(tape)), + // MSI + c("oauth-msi", "oauth", "MSI with tenant and client id, defaults for the rest", msi), + c( + "oauth-msi-custom-endpoint-authority", + "oauth", + "MSI endpoint and authority are trimmed; the authority gets a trailing slash", + msi ++ Map( + acct(Keys.MSI_ENDPOINT) -> " http://msi.example.net/token ", + acct(Keys.MSI_AUTHORITY) -> " https://login.example.net ")), + c( + "oauth-msi-blank-endpoint", + "oauth", + "a blank MSI endpoint means the default", + msi + (acct(Keys.MSI_ENDPOINT) -> " ")), + c( + "oauth-msi-no-client-id", + "oauth", + "client id: mandatory on 3.3.4, optional from 3.4.1", + msi - acct(Keys.CLIENT_ID), + pins = from341(tape, Pin.resolved)), + c( + "oauth-msi-container-tenant", + "oauth", + "container-scoped MSI tenant on 3.4.2+", + msi + (cont(Keys.MSI_TENANT) -> "container-tenant")), + // RefreshToken + c("oauth-refreshtoken", "oauth", "RefreshToken, all keys set", refreshToken), + missing(refreshToken, Keys.REFRESH_TOKEN, "refreshtoken-no-token", allVersions(tape)), + missing(refreshToken, Keys.CLIENT_ID, "refreshtoken-no-client-id", allVersions(tape)), + c( + "oauth-refreshtoken-custom-endpoint", + "oauth", + "refresh token endpoint is trimmed", + refreshToken + (acct( + Keys.REFRESH_TOKEN_ENDPOINT) -> " https://login.example.net/token ")), + // WorkloadIdentity + c( + "oauth-wi", + "oauth", + "Workload Identity with the default token file", + workloadIdentity, + pins = fromWorkloadIdentity), + c( + "oauth-wi-no-tenant", + "oauth", + "Workload Identity tenant is mandatory", + workloadIdentity - acct(Keys.MSI_TENANT), + pins = from341(Pin.failed(ExceptionClasses.RUNTIME), tape)), + c( + "oauth-wi-no-client-id", + "oauth", + "Workload Identity client id is mandatory", + workloadIdentity - acct(Keys.CLIENT_ID), + pins = from341(Pin.failed(ExceptionClasses.RUNTIME), tape)), + c( + "oauth-wi-padded-token-file", + "oauth", + "token file is trimmed", + workloadIdentity + (acct(Keys.TOKEN_FILE) -> " /mnt/tokens/token "), + pins = fromWorkloadIdentity), + c( + "oauth-wi-container-token-file", + "oauth", + "container-scoped token file on 3.4.2+", + workloadIdentity + (cont(Keys.TOKEN_FILE) -> "/mnt/tokens/container-token"), + pins = fromWorkloadIdentity), + c( + "oauth-wi-blank-assertion-provider", + "oauth", + "a blank client assertion provider is ignored on 3.5.0", + workloadIdentity + (acct(Keys.CLIENT_ASSERTION_PROVIDER_TYPE) -> " "), + pins = fromWorkloadIdentity), + // provider class shapes + c( + "oauth-no-provider-type", + "oauth", + "OAuth without a provider class", + clientCreds - acct(Keys.OAUTH_PROVIDER_TYPE), + pins = allVersions(Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT))), + c( + "oauth-provider-not-built-in", + "oauth", + "an AccessTokenProvider outside the five built-ins is rejected by ==", + clientCreds + (acct(Keys.OAUTH_PROVIDER_TYPE) -> FOREIGN_ACCESS_TOKEN_PROVIDER), + pins = allVersions(Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT))), + c( + "oauth-provider-not-a-provider", + "oauth", + "a class that is not an AccessTokenProvider", + clientCreds + (acct(Keys.OAUTH_PROVIDER_TYPE) -> NOT_A_PROVIDER), + pins = allVersions(Pin.failed(ExceptionClasses.RUNTIME))), + c( + "oauth-provider-missing-class", + "oauth", + "a provider class that does not exist", + clientCreds + (acct(Keys.OAUTH_PROVIDER_TYPE) -> MISSING_CLASS), + pins = allVersions(Pin.failed(ExceptionClasses.RUNTIME))), + c( + "oauth-provider-empty", + "oauth", + "provider type '' is an unloadable class name", + clientCreds + (acct(Keys.OAUTH_PROVIDER_TYPE) -> ""), + pins = allVersions(Pin.failed(ExceptionClasses.RUNTIME))), + c( + "oauth-provider-padded", + "oauth", + "provider type is trimmed by Configuration.getClass", + clientCreds + (acct(Keys.OAUTH_PROVIDER_TYPE) -> s" ${Providers.CLIENT_CREDS} ")), + c( + "oauth-custom-not-adaptee", + "oauth", + "Custom with a class that is not a CustomTokenProviderAdaptee: declined, Hadoop fails", + Map( + acct(Keys.AUTH_TYPE) -> "Custom", + acct(Keys.OAUTH_PROVIDER_TYPE) -> FOREIGN_ACCESS_TOKEN_PROVIDER), + pins = allVersions(Pin.declined)), + c( + "oauth-custom-no-provider-type", + "oauth", + "Custom without a provider class: declined, Hadoop fails", + Map(acct(Keys.AUTH_TYPE) -> "Custom"), + pins = allVersions(Pin.declined))) + } + + // ---- auth-type guard: a global provider class needs a matching global auth type ---- + + private def guardCases: Seq[ParityCase] = { + def scopedKey(scope: String, key: String): String = if (scope == "global") key else acct(key) + val oauthValues = Map( + acct(Keys.CLIENT_ENDPOINT) -> ENDPOINT, + acct(Keys.CLIENT_ID) -> CLIENT_ID, + acct(Keys.CLIENT_SECRET) -> CLIENT_SECRET) + val unresolvedProviderClass = Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT) + val scopes = Seq( + ("global", "global"), + ("account", "account"), + ("account", "global"), + ("global", "account")) + scopes.flatMap { case (authScope, providerScope) => + val guardFails = authScope == "account" && providerScope == "global" + Seq( + c( + s"guard-oauth-$authScope-type-$providerScope-provider", + "guard", + s"$authScope auth type, $providerScope OAuth provider class", + oauthValues ++ Map( + scopedKey(authScope, Keys.AUTH_TYPE) -> "OAuth", + scopedKey(providerScope, Keys.OAUTH_PROVIDER_TYPE) -> Providers.CLIENT_CREDS), + pins = allVersions(if (guardFails) unresolvedProviderClass else Pin.resolved)), + c( + s"guard-sas-$authScope-type-$providerScope-provider", + "guard", + s"$authScope auth type, $providerScope SAS provider class", + Map( + scopedKey(authScope, Keys.AUTH_TYPE) -> "SAS", + scopedKey(providerScope, Keys.SAS_TOKEN_PROVIDER_TYPE) -> SAS_PROVIDER), + pins = if (guardFails) noSasSource else allVersions(Pin.declined))) + } + } + + // ---- SAS ---- + + private def sasCases: Seq[ParityCase] = { + val fixed = Map(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_FIXED_TOKEN) -> FIXED_SAS) + Seq( + c("sas-fixed-only", "sas", "fixed token only", fixed, pins = fixedSasOnly), + c( + "sas-custom-only", + "sas", + "custom provider only", + sasCustom, + pins = allVersions(Pin.declined)), + c( + "sas-custom-and-fixed", + "sas", + "the custom provider wins over the fixed token", + sasCustom + (acct(Keys.SAS_FIXED_TOKEN) -> FIXED_SAS), + pins = allVersions(Pin.declined)), + c( + "sas-neither", + "sas", + "neither provider nor token", + Map(acct(Keys.AUTH_TYPE) -> "SAS"), + pins = noSasSource), + c( + "sas-blank-fixed", + "sas", + "a blank fixed token counts as unset", + fixed + (acct(Keys.SAS_FIXED_TOKEN) -> " "), + pins = noSasSource), + c( + "sas-question-mark", + "sas", + "a leading '?' is kept at init; the client strips it per request", + fixed + (acct(Keys.SAS_FIXED_TOKEN) -> s"?$FIXED_SAS"), + pins = fixedSasOnly), + c( + "sas-fixed-global", + "sas", + "global auth type and fixed token", + Map(Keys.AUTH_TYPE -> "SAS", Keys.SAS_FIXED_TOKEN -> FIXED_SAS), + pins = fixedSasOnly), + c( + "sas-provider-not-a-provider", + "sas", + "a class that is not a SASTokenProvider", + Map(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> NOT_A_PROVIDER), + pins = allVersions(Pin.failed(ExceptionClasses.RUNTIME))), + c( + "sas-provider-missing-class", + "sas", + "a provider class that does not exist", + Map(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> MISSING_CLASS), + pins = allVersions(Pin.failed(ExceptionClasses.RUNTIME)))) + } + + // ---- credential providers (JCEKS) and the clear-text fallback ---- + + private def credentialProviderCases(providerPath: String): Seq[ParityCase] = { + val provider = Map("hadoop.security.credential.provider.path" -> providerPath) + val noFallback = Map("hadoop.security.credential.clear-text-fallback" -> "false") + val kpe = Pin.failed(ExceptionClasses.KEY_PROVIDER) + Seq( + c( + "cp-jceks-account-key", + "credential-provider", + "the account key comes from the keystore; no clear-text key", + provider, + pins = allVersions(Pin.resolved)), + c( + "cp-jceks-beats-clear-text", + "credential-provider", + "the keystore is consulted before the clear-text value", + provider + (acct(Keys.ACCOUNT_KEY) -> KEY_GLOBAL), + pins = allVersions(Pin.resolved)), + c( + "cp-jceks-no-fallback", + "credential-provider", + "keystore with the clear-text fallback off", + provider ++ noFallback, + pins = allVersions(Pin.resolved)), + c( + "cp-jceks-client-secret", + "credential-provider", + "the OAuth client secret comes from the keystore", + provider ++ (clientCreds - acct(Keys.CLIENT_SECRET)), + pins = allVersions(Pin.resolved)), + c( + "cp-no-fallback-clear-text-only", + "credential-provider", + "fallback off, no keystore: a clear-text key is invisible, so nothing is configured", + noFallback + (acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(Pin.unconfigured)), + c( + "cp-no-fallback-with-auth-type", + "credential-provider", + "fallback off, no keystore, SharedKey named: the key provider finds nothing", + noFallback ++ Map( + acct(Keys.AUTH_TYPE) -> "SharedKey", + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(kpe))) + } + + // ---- ${...} substitution ---- + + private def substitutionCases: Seq[ParityCase] = { + // Built by concatenation so scalac does not read them as interpolations missing their prefix. + val crossKey = "${" + "comet.parity.account.key}" + val envWithDefault = "${env." + UNSET_ENV_VAR + ":-fallback-client}" + val envHome = "${" + "env.HOME}" + Seq( + c( + "subst-cross-key", + "substitution", + "a reference to another key resolves through Configuration.get", + Map(acct(Keys.ACCOUNT_KEY) -> crossKey, "comet.parity.account.key" -> KEY_ENV), + pins = allVersions(Pin.resolved)), + c( + "subst-env-default", + "substitution", + "an environment reference with a default falls back when the variable is unset", + clientCreds + (acct(Keys.CLIENT_ID) -> envWithDefault), + pins = allVersions(Pin.resolved)), + c( + "subst-env-home", + "substitution", + "an environment reference reads the variable", + clientCreds + (acct(Keys.CLIENT_ID) -> envHome), + pins = allVersions(Pin.resolved), + liveValues = Map(Keys.CLIENT_ID -> sys.env.getOrElse("HOME", envHome)))) + } + + // ---- calibration: one case per version difference and per surprising Hadoop rule ---- + + private def calibrationCases: Seq[ParityCase] = { + val tape = Pin.failed(ExceptionClasses.TOKEN_ACCESS_PROVIDER) + val kpe = Pin.failed(ExceptionClasses.KEY_PROVIDER) + val badEnum = Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT) + val portAccount = s"$ACCOUNT:8080" + val prefixedEndpoint = "https://login.example.net/prefix/tenant-abc/oauth2/v2.0/token" + val assertionProvider = "com.example.comet.NeverLoadedAssertionProvider" + Seq( + c( + "cal-sas-guard-with-fixed-token", + "calibration", + "account SAS type + global SAS provider class: the guard drops the class, the token wins", + Map( + acct(Keys.AUTH_TYPE) -> "SAS", + Keys.SAS_TOKEN_PROVIDER_TYPE -> SAS_PROVIDER, + acct(Keys.SAS_FIXED_TOKEN) -> FIXED_SAS), + pins = fixedSasOnly), + c( + "cal-sas-guard-without-fixed-token", + "calibration", + "account SAS type + global SAS provider class and no token: nothing is left", + Map(acct(Keys.AUTH_TYPE) -> "SAS", Keys.SAS_TOKEN_PROVIDER_TYPE -> SAS_PROVIDER), + pins = noSasSource), + c( + "cal-sas-container-fixed-token", + "calibration", + "container-scoped fixed token: honoured from 3.4.2, invisible on 3.4.1", + Map(acct(Keys.AUTH_TYPE) -> "SAS", cont(Keys.SAS_FIXED_TOKEN) -> FIXED_SAS), + pins = pins( + Pin.failed(ExceptionClasses.TOKEN_ACCESS_PROVIDER), + Pin.failed(ExceptionClasses.SAS_TOKEN_PROVIDER), + Pin.resolved, + Pin.resolved)), + c( + "cal-sas-untrimmed-token", + "calibration", + "the fixed token is trimmed before use", + Map(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_FIXED_TOKEN) -> s" $FIXED_SAS "), + pins = fixedSasOnly), + c( + "cal-endpoint-with-path-prefix", + "calibration", + "the client endpoint is forwarded verbatim, whatever its path shape", + clientCreds + (acct(Keys.CLIENT_ENDPOINT) -> prefixedEndpoint), + pins = allVersions(Pin.resolved)), + c( + "cal-port-key-matches", + "calibration", + "with a port in the URI, only the port-qualified key configures the account", + Map(acct(Keys.ACCOUNT_KEY, portAccount) -> KEY_PORT), + uri = uriFor(CONTAINER, portAccount), + pins = allVersions(Pin.resolved)), + c( + "cal-port-key-missing", + "calibration", + "with a port in the URI, the port-less key belongs to another account", + Map(acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + uri = uriFor(CONTAINER, portAccount), + pins = allVersions(Pin.unconfigured)), + c( + "cal-container-oauth-endpoint", + "calibration", + "container-scoped client endpoint: read from 3.4.2, missing before", + (clientCreds - acct(Keys.CLIENT_ENDPOINT)) + (cont(Keys.CLIENT_ENDPOINT) -> ENDPOINT), + pins = from342(tape, Pin.resolved)), + c( + "cal-container-account-key-only", + "calibration", + "a container-scoped account key alone: unread by SimpleKeyProvider on every version", + Map(cont(Keys.ACCOUNT_KEY) -> KEY_CONTAINER), + pins = from342(Pin.unconfigured, kpe)), + c( + "cal-global-enum-parsed-first", + "calibration", + "a bad global auth type fails even when the account one is valid", + clientCreds + (Keys.AUTH_TYPE -> "oauth"), + pins = allVersions(badEnum)), + c( + "cal-empty-mandatory-secret", + "calibration", + "'' passes the mandatory getter and is forwarded", + clientCreds + (acct(Keys.CLIENT_SECRET) -> ""), + pins = allVersions(Pin.resolved)), + c( + "cal-msi-no-tenant", + "calibration", + "MSI tenant: mandatory on 3.3.4, optional from 3.4.1", + msi - acct(Keys.MSI_TENANT), + pins = from341(tape, Pin.resolved)), + c( + "cal-msi-no-tenant-no-client-id", + "calibration", + "MSI with neither tenant nor client id", + msi - acct(Keys.MSI_TENANT) - acct(Keys.CLIENT_ID), + pins = from341(tape, Pin.resolved)), + c( + "cal-null-fsname-key", + "calibration", + "SimpleKeyProvider's 2-arg configuration probes the literal key.null. on 3.4.2+", + Map(s"${Keys.ACCOUNT_KEY}.null.$ACCOUNT" -> KEY_NULL_LITERAL), + pins = from342(Pin.unconfigured, Pin.resolved)), + c( + "cal-null-fsname-key-beats-account", + "calibration", + "key.null. is the first level on 3.4.2+, so it wins over the account key", + Map( + s"${Keys.ACCOUNT_KEY}.null.$ACCOUNT" -> KEY_NULL_LITERAL, + acct(Keys.ACCOUNT_KEY) -> KEY_ACCOUNT), + pins = allVersions(Pin.resolved)), + c( + "cal-base64-bad-length", + "calibration", + "Base64 validation: length % 4 != 0", + Map(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0c2V"), + pins = allVersions(kpe)), + c( + "cal-base64-non-ascii", + "calibration", + "Base64 validation: a char >= U+0080 escapes as ArrayIndexOutOfBoundsException", + Map(acct(Keys.ACCOUNT_KEY) -> "c2Vjcm\u00e90"), + pins = allVersions(Pin.failed(ExceptionClasses.ARRAY_INDEX))), + c( + "cal-empty-account-key", + "calibration", + "'' passes Base64 validation and fails in SharedKeyCredentials", + Map(acct(Keys.ACCOUNT_KEY) -> ""), + pins = allVersions(Pin.failed(ExceptionClasses.ILLEGAL_ARGUMENT))), + c( + "cal-wi-on-334", + "calibration", + "WorkloadIdentityTokenProvider does not exist on 3.3.4", + workloadIdentity, + pins = fromWorkloadIdentity), + c( + "cal-wi-client-assertion-provider", + "calibration", + "3.5.0 Workload Identity with a client assertion provider is declined, never loaded", + workloadIdentity + (acct(Keys.CLIENT_ASSERTION_PROVIDER_TYPE) -> s" $assertionProvider "), + pins = + pins(Pin.failed(ExceptionClasses.RUNTIME), Pin.resolved, Pin.resolved, Pin.declined)), + c( + "cal-userbound-sas", + "calibration", + "UserboundSASWithOAuth: an enum constant only from 3.5.0", + Map( + acct(Keys.AUTH_TYPE) -> "UserboundSASWithOAuth", + acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> SAS_PROVIDER), + pins = from350(badEnum, Pin.declined)), + c( + "cal-oauth-keys-without-auth-type", + "calibration", + "OAuth keys without an auth type are a SharedKey misconfiguration, not 'unconfigured'", + clientCreds - acct(Keys.AUTH_TYPE), + pins = allVersions(kpe))) + } +} + +/** A `KeyProvider` that supplies a fixed key; stands in for a vault-backed provider. */ +class ParityKeyProvider extends KeyProvider { + override def getStorageAccountKey(accountName: String, conf: Configuration): String = + AbfsAuthParityCases.KEY_CUSTOM +} + +/** A `KeyProvider` that finds nothing: Hadoop reports the account as a missing property. */ +class ParityNullKeyProvider extends KeyProvider { + override def getStorageAccountKey(accountName: String, conf: Configuration): String = null +} + +/** A `CustomTokenProviderAdaptee` that initializes fine, so a Custom FileSystem init succeeds. */ +class ParityCustomTokenAdaptee extends CustomTokenProviderAdaptee { + override def initialize(configuration: Configuration, accountName: String): Unit = () + override def getAccessToken(): String = "parity-token" + override def getExpiryTime(): Date = new Date(Long.MaxValue) +} + +/** A `SASTokenProvider` that initializes fine, so a SAS FileSystem init succeeds. */ +class ParitySasTokenProvider extends SASTokenProvider { + override def initialize(configuration: Configuration, accountName: String): Unit = () + override def getSASToken( + account: String, + fileSystem: String, + path: String, + operation: String): String = AbfsAuthParityCases.FIXED_SAS +} diff --git a/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParitySuite.scala b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParitySuite.scala new file mode 100644 index 00000000000..3d95864130a --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthParitySuite.scala @@ -0,0 +1,422 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.io.InputStream +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Path, Paths} +import java.security.MessageDigest +import java.util.Comparator + +import scala.collection.immutable.SortedMap +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.FileSystem +import org.apache.hadoop.security.alias.CredentialProviderFactory + +import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} + +import org.apache.comet.objectstore.AbfsAuthParityCases._ +import org.apache.comet.objectstore.AbfsAuthResolver._ +import org.apache.comet.objectstore.AbfsReflection.Handles + +/** + * Pins the resolver's behaviour per hadoop-azure version and cross-checks it against the real + * `AzureBlobFileSystem`, offline. + * + * Every [[AbfsAuthParityCases]] case runs through [[AbfsAuthResolver.resolve]] and is compared + * with `abfs-auth-parity/expected-.json`, where the version is the one the capability + * probes identify (never the Spark profile), and with the case's pins. Every case is then handed + * to `FileSystem.newInstance` with `fs.azure.account.hns.enabled=true`, which keeps 3.4.2+ + * initialization off the network. + * + * The FileSystem cross-check compares outcome classes, not values: `Resolved` needs a successful + * initialization, `Failed` an initialization that threw with the resolver's exception class in + * the cause chain, `Unconfigured` a `KeyProviderException`, and `Declined` either success or a + * failure while Hadoop loaded the provider the resolver left alone. Forwarded values are covered + * by the pins, the live-value checks and the digests in the recorded files, not by this check. + * + * To regenerate the expectation files, run the suite once per hadoop-azure line with + * `-Dcomet.abfs.parity.record=` (spark-3.5 for 3.3.4, spark-4.0 for 3.4.1, spark-4.1 for + * 3.4.2, spark-4.2 for 3.5.0) and copy the written files into + * `spark/src/test/resources/abfs-auth-parity/`. The pins and the FileSystem cross-check still + * assert while recording, so a recorded file never contradicts the hadoop-azure source. The file + * is written only when every case ran, so a filtered run cannot leave a partial file behind. + */ +class AbfsAuthParitySuite extends AnyFunSuite with Matchers with BeforeAndAfterAll { + import AbfsAuthParitySuite._ + + private lazy val handles: Handles = AbfsReflection.handles(getClass.getClassLoader).get + private lazy val version: String = versionLabel(handles) + + private val recordDir: Option[Path] = Option(System.getProperty(RECORD_PROPERTY)) + // An undefined Maven property reaches the forked JVM as "null" or as the literal + // "${comet.abfs.parity.record}"; neither is a directory. + .map(_.trim) + .filter(p => p.nonEmpty && p != "null" && !p.startsWith("$")) + .map(p => Paths.get(p)) + + // The build points java.io.tmpdir under target/, which a fresh checkout may not have yet. + private val tempDir: Path = Files.createTempDirectory( + Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))), + "abfs-auth-parity") + private val keystore: Path = tempDir.resolve("credentials.jceks") + private val credentialProviderPath = "jceks://file" + keystore.toUri.getPath + + private lazy val cases: Seq[ParityCase] = AbfsAuthParityCases.cases(credentialProviderPath) + private lazy val expected: Map[String, Recorded] = + if (recordDir.isDefined) Map.empty else loadExpected(version) + private val recorded = mutable.Map[String, Recorded]() + + override def beforeAll(): Unit = { + super.beforeAll() + val conf = new Configuration() + conf.set(CredentialProviderFactory.CREDENTIAL_PROVIDER_PATH, credentialProviderPath) + val provider = CredentialProviderFactory.getProviders(conf).get(0) + credentialEntries.foreach { case (alias, value) => + provider.createCredentialEntry(alias, value.toCharArray) + } + provider.flush() + } + + override def afterAll(): Unit = { + try { + recordDir.foreach { dir => + require( + recorded.size == cases.size, + s"only ${recorded.size} of ${cases.size} cases ran; not writing a partial expectation " + + s"file for hadoop-azure $version") + Files.createDirectories(dir) + val file = dir.resolve(s"expected-$version.json") + Files.write( + file, + renderExpected(version, recorded.toMap).getBytes(StandardCharsets.UTF_8)) + announce(s"AbfsAuthParitySuite recorded ${recorded.size} cases to $file") + } + } finally { + deleteRecursively(tempDir) + super.afterAll() + } + } + + // scalastyle:off println + private def announce(message: String): Unit = println(message) + // scalastyle:on println + + private def hadoopConf(c: ParityCase): Configuration = { + val conf = new Configuration() + c.conf.foreach { case (k, v) => conf.set(k, v) } + conf + } + + test("case ids are unique and every calibration case is pinned on all four versions") { + val ids = cases.map(_.id) + ids.distinct should have size ids.size.toLong + val unpinned = + cases.filter(_.group == "calibration").filterNot(_.pins.keySet == VERSIONS.toSet) + assert( + unpinned.isEmpty, + s"calibration cases without pins for every version: ${unpinned.map(_.id)}") + val unknownVersions = cases.flatMap(_.pins.keySet).distinct.filterNot(VERSIONS.contains) + assert(unknownVersions.isEmpty, s"pins for unknown versions: $unknownVersions") + assert(!sys.env.contains(UNSET_ENV_VAR), s"$UNSET_ENV_VAR must not be set in the environment") + info( + s"hadoop-azure $version, ${cases.size} cases, " + + cases + .groupBy(_.group) + .toSeq + .sortBy(_._1) + .map { case (g, cs) => s"$g=${cs.size}" } + .mkString(", ")) + } + + test("the expectation file matches the hadoop-azure version and the case list") { + assume(recordDir.isEmpty, "recording") + expected.keySet shouldBe cases.map(_.id).toSet + } + + cases.foreach { c => + test(s"${c.group}/${c.id}: ${c.note}") { + val outcome = resolve(hadoopConf(c), c.uri) + c.liveValues.foreach { case (key, value) => + withClue(s"live value of $key: ") { + outcome shouldBe a[Resolved] + outcome.asInstanceOf[Resolved].values.get(key) shouldBe Some(value) + } + } + val actual = Recorded.of(outcome, c.liveValues.keySet) + c.pins.get(version).foreach { pin => + withClue(s"pin for $version: ") { + actual.outcome shouldBe pin.outcome + actual.errorClass shouldBe pin.errorClass + } + } + recorded(c.id) = actual + if (recordDir.isEmpty) { + withClue(s"recorded expectation for hadoop-azure $version: ") { + actual shouldBe expected(c.id) + } + } + } + } + + test("FileSystem.newInstance agrees with the resolver on every case") { + val timings = mutable.ArrayBuffer[(String, Long)]() + val disagreements = cases.flatMap { c => + val outcome = resolve(hadoopConf(c), c.uri) + val started = System.nanoTime() + val fsResult = initializeFileSystem(c) + timings += c.id -> (System.nanoTime() - started) / 1000000L + checkAgreement(outcome, fsResult).map(reason => s"${c.id}: $reason") + } + info(s"${cases.size - disagreements.size} of ${cases.size} cases agree") + // A slow initialization would point at a network attempt, which this suite must never make. + val slowest = timings.sortBy(-_._2).take(3).map { case (id, ms) => s"$id ${ms}ms" } + info(s"slowest FileSystem initializations: ${slowest.mkString(", ")}") + assert(disagreements.isEmpty, disagreements.mkString("\n", "\n", "")) + } + + /** The real initialization, offline: HNS known, caches off, closed whatever happens. */ + private def initializeFileSystem(c: ParityCase): Either[Throwable, Unit] = { + val conf = hadoopConf(c) + if (Option(conf.get(HNS_ENABLED)).isEmpty) conf.set(HNS_ENABLED, "true") + conf.set("fs.abfs.impl.disable.cache", "true") + conf.set("fs.abfss.impl.disable.cache", "true") + var fs: Option[FileSystem] = None + try { + fs = Some(FileSystem.newInstance(c.uri, conf)) + Right(()) + } catch { + // The resolver reports LinkageErrors too, so the cross-check must see them. + case e @ (_: LinkageError | NonFatal(_)) => Left(e) + } finally { + fs.foreach(_.close()) + } + } + + private def checkAgreement(outcome: Outcome, fs: Either[Throwable, Unit]): Option[String] = { + def chain(e: Throwable): Seq[String] = causeChain(e).map(_.getClass.getName) + (outcome, fs) match { + case (_: Resolved, Right(())) => None + case (_: Resolved, Left(e)) => Some(s"resolver resolved, FileSystem threw ${chain(e)}") + case (Failed(cls, _), Right(())) => + Some(s"resolver failed with $cls, FileSystem initialized") + case (Failed(cls, _), Left(e)) => + if (chain(e).contains(cls)) None + else Some(s"resolver failed with $cls, FileSystem threw ${chain(e)}") + case (Unconfigured, Right(())) => Some("resolver unconfigured, FileSystem initialized") + case (Unconfigured, Left(e)) => + if (chain(e).contains(ExceptionClasses.KEY_PROVIDER)) None + else Some(s"resolver unconfigured, FileSystem threw ${chain(e)} not KeyProviderException") + case (_: Declined, Right(())) => None + case (_: Declined, Left(e)) => + // Hadoop instantiates the provider the resolver declined to; whatever that does is fine. + if (chain(e).exists(PROVIDER_FAILURES.contains)) None + else Some(s"resolver declined, FileSystem threw ${chain(e)} outside provider loading") + } + } +} + +object AbfsAuthParitySuite { + val RECORD_PROPERTY = "comet.abfs.parity.record" + val RESOURCE_DIR = "abfs-auth-parity" + private val HNS_ENABLED = "fs.azure.account.hns.enabled" + private val MAX_CAUSE_DEPTH = 8 + + private val PROVIDER_FAILURES: Set[String] = Set( + ExceptionClasses.TOKEN_ACCESS_PROVIDER, + ExceptionClasses.SAS_TOKEN_PROVIDER, + ExceptionClasses.ILLEGAL_ARGUMENT) + + /** The forwarded values whose recorded form is a SHA-256 digest. */ + private val SECRET_KEYS: Set[String] = Set( + Keys.ACCOUNT_KEY, + Keys.CLIENT_SECRET, + Keys.USER_PASSWORD, + Keys.REFRESH_TOKEN, + Keys.SAS_FIXED_TOKEN) + private val ENVIRONMENT_PLACEHOLDER = "" + + /** The hadoop-azure line on the classpath, from what it can do rather than from a version. */ + def versionLabel(handles: Handles): String = + if (handles.hasClientAssertionProvider) "3.5.0" + else if (handles.hasContainerScopedConfiguration) "3.4.2" + else if (handles.hasFixedSasTokenProvider) "3.4.1" + else "3.3.4" + + /** + * One case's expectation; `values` holds digests for secrets and a placeholder for live values. + */ + final case class Recorded( + outcome: String, + authType: Option[String], + providerClass: Option[String], + assertionProviderClass: Option[String], + errorClass: Option[String], + values: SortedMap[String, String]) + + object Recorded { + def of(outcome: Outcome, liveKeys: Set[String]): Recorded = outcome match { + case Resolved(authType, providerClass, values) => + val normalized = values.map { case (k, v) => + val recorded = + if (liveKeys.contains(k)) ENVIRONMENT_PLACEHOLDER + else if (SECRET_KEYS.contains(k)) sha256(v) + else v + k -> recorded + } + Recorded( + Outcomes.RESOLVED, + Some(authType), + providerClass, + None, + None, + SortedMap(normalized.toSeq: _*)) + case Unconfigured => + Recorded(Outcomes.UNCONFIGURED, None, None, None, None, SortedMap.empty) + case Declined(authType, providerClass, assertionProviderClass) => + Recorded( + Outcomes.DECLINED, + Some(authType), + providerClass, + assertionProviderClass, + None, + SortedMap.empty) + case Failed(exceptionClass, _) => + Recorded(Outcomes.FAILED, None, None, None, Some(exceptionClass), SortedMap.empty) + } + } + + private def sha256(value: String): String = { + val digest = + MessageDigest.getInstance("SHA-256").digest(value.getBytes(StandardCharsets.UTF_8)) + "sha256:" + digest.map(b => f"$b%02x").mkString + } + + private def causeChain(e: Throwable): Seq[Throwable] = { + val chain = mutable.ArrayBuffer[Throwable](e) + var cause = e.getCause + while (cause != null && chain.length < MAX_CAUSE_DEPTH && !chain.contains(cause)) { + chain += cause + cause = cause.getCause + } + chain.toSeq + } + + private def loadExpected(version: String): Map[String, Recorded] = { + val resource = s"/$RESOURCE_DIR/expected-$version.json" + val stream: InputStream = getClass.getResourceAsStream(resource) + require( + stream != null, + s"$resource is missing; record it with -D$RECORD_PROPERTY= on a classpath with " + + s"hadoop-azure $version") + try { + val root = new ObjectMapper().readTree(stream) + require( + root.get("hadoopAzureVersion").asText == version, + s"$resource records ${root.get("hadoopAzureVersion").asText}, classpath has $version") + val cases = root.get("cases") + // fieldNames() rather than fields(): the latter is deprecated on the Jackson Spark 4 ships + // and its replacement is missing on the one Spark 3 ships. + cases.fieldNames().asScala.map(id => id -> parseRecorded(cases.get(id))).toMap + } finally { + stream.close() + } + } + + private def parseRecorded(node: JsonNode): Recorded = { + def optional(field: String): Option[String] = Option(node.get(field)).map(_.asText) + val values = Option(node.get("values")) + .map(v => v.fieldNames().asScala.map(k => k -> v.get(k).asText).toSeq) + .getOrElse(Seq.empty) + Recorded( + node.get("outcome").asText, + optional("authType"), + optional("providerClass"), + optional("assertionProviderClass"), + optional("errorClass"), + SortedMap(values: _*)) + } + + /** Sorted keys and a fixed layout, so a regeneration diff shows only behaviour changes. */ + private def renderExpected(version: String, recorded: Map[String, Recorded]): String = { + val cases: SortedMap[String, Any] = SortedMap(recorded.toSeq: _*).map { case (id, r) => + val fields: Seq[(String, Any)] = Seq( + Some("outcome" -> r.outcome), + r.authType.map("authType" -> _), + r.providerClass.map("providerClass" -> _), + r.assertionProviderClass.map("assertionProviderClass" -> _), + r.errorClass.map("errorClass" -> _), + Some("values" -> r.values)).flatten + id -> SortedMap[String, Any](fields: _*) + } + val root = SortedMap[String, Any]("cases" -> cases, "hadoopAzureVersion" -> version) + val out = new StringBuilder + renderJson(root, 0, out) + out.append('\n').toString + } + + private def renderJson(value: Any, indent: Int, out: StringBuilder): Unit = value match { + case s: String => out.append(quote(s)) + case m: scala.collection.Map[_, _] if m.isEmpty => out.append("{}") + case m: scala.collection.Map[_, _] => + val pad = " " * (indent + 1) + out.append("{\n") + val entries = m.toSeq.map { case (k, v) => k.toString -> v }.sortBy(_._1) + entries.zipWithIndex.foreach { case ((k, v), i) => + out.append(pad).append(quote(k)).append(": ") + renderJson(v, indent + 1, out) + if (i < entries.size - 1) out.append(',') + out.append('\n') + } + out.append(" " * indent).append('}') + case other => throw new IllegalArgumentException(s"unsupported JSON value: $other") + } + + private def quote(s: String): String = { + val out = new StringBuilder("\"") + s.foreach { + case '"' => out.append("\\\"") + case '\\' => out.append("\\\\") + case '\n' => out.append("\\n") + case '\r' => out.append("\\r") + case '\t' => out.append("\\t") + case ch if ch < 0x20 || ch > 0x7e => out.append('\\').append('u').append(f"${ch.toInt}%04x") + case ch => out.append(ch) + } + out.append('"').toString + } + + private def deleteRecursively(dir: Path): Unit = { + if (Files.exists(dir)) { + Files.walk(dir).sorted(Comparator.reverseOrder[Path]()).forEach { (p: Path) => + Files.deleteIfExists(p) + () + } + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthResolverSuite.scala b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthResolverSuite.scala new file mode 100644 index 00000000000..4d69c2e01a3 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/objectstore/AbfsAuthResolverSuite.scala @@ -0,0 +1,703 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.net.URI + +import scala.collection.mutable.ArrayBuffer +import scala.util.{Failure, Success, Try} + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.azurebfs.AbfsConfiguration +import org.apache.hadoop.fs.azurebfs.extensions.SASTokenProvider +import org.apache.hadoop.fs.azurebfs.oauth2.{AccessTokenProvider, AzureADToken} +import org.apache.hadoop.fs.azurebfs.services.KeyProvider + +import org.apache.comet.objectstore.AbfsAuthResolver._ +import org.apache.comet.objectstore.AbfsReflection.Handles + +/** + * Runs the resolver against the hadoop-azure on the test classpath (3.3.4 on the spark-3.x + * profiles, 3.4.1, 3.4.2 and 3.5.0 on spark-4.0, 4.1 and 4.2). Cases that depend on a feature one + * line lacks are gated on the capability probes rather than on a version number. + */ +class AbfsAuthResolverSuite extends AnyFunSuite with Matchers { + + private val account = "myacct.dfs.core.windows.net" + private val container = "data" + private val uri = new URI(s"abfss://$container@$account/path/file.parquet") + + private val keyProviderException = + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.KeyProviderException" + private val invalidUriException = + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriException" + private val invalidUriAuthorityException = + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.InvalidUriAuthorityException" + private val tokenAccessProviderException = + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.TokenAccessProviderException" + + private lazy val handles: Handles = AbfsReflection.handles(getClass.getClassLoader).get + + private def acct(key: String): String = s"$key.$account" + + private def cont(key: String): String = s"$key.$container.$account" + + private def conf(pairs: (String, String)*): Configuration = { + val hadoopConf = new Configuration() + pairs.foreach { case (k, v) => hadoopConf.set(k, v) } + hadoopConf + } + + private def resolveWith(pairs: (String, String)*): Outcome = resolve(conf(pairs: _*), uri) + + /** + * `key` set to a `${...}` cycle through `other`, which `Configuration#get` cannot expand and + * throws on. Built by concatenation so scalac does not read it as a missing interpolator. + */ + private def cyclic(key: String, other: String): Seq[(String, String)] = + Seq(key -> ("${" + other + "}"), other -> ("${" + key + "}")) + + private def assertNoAzureKeys(options: Map[String, String]): Unit = { + val leaked = options.keys.filter(_.startsWith("fs.")) + assert(leaked.isEmpty, s"fs.* keys forwarded: $leaked") + } + + // ---- authorityParts ---- + + test("authorityParts - raw authority split once at the first '@', port and case kept") { + authorityParts(new URI("abfss://data@myacct.dfs.core.windows.net:8080/p")) shouldBe + Right(("data", "myacct.dfs.core.windows.net:8080")) + authorityParts(new URI("abfss://Data@MyAcct.DFS.core.windows.net/p")) shouldBe + Right(("Data", "MyAcct.DFS.core.windows.net")) + authorityParts(new URI("abfss://my%20data@myacct.dfs.core.windows.net/p")) shouldBe + Right(("my%20data", "myacct.dfs.core.windows.net")) + authorityParts(new URI("abfss://data@user@myacct.dfs.core.windows.net/p")) shouldBe + Right(("data", "user@myacct.dfs.core.windows.net")) + } + + test("authorityParts - no '@', empty container and null authority fail as Hadoop does") { + val noAt = "abfss://myacct.dfs.core.windows.net/p" + authorityParts(new URI(noAt)) shouldBe + Left(Failed(invalidUriAuthorityException, s"$noAt has invalid authority.")) + + val emptyContainer = "abfss://@myacct.dfs.core.windows.net/p" + authorityParts(new URI(emptyContainer)) shouldBe Left( + Failed( + invalidUriException, + s"Invalid URI '$emptyContainer' has a malformed authority, expected container name. " + + "Authority takes the form abfs://[@]")) + + val nullAuthority = "abfss:///p" + authorityParts(new URI(nullAuthority)) shouldBe + Left(Failed(invalidUriAuthorityException, s"$nullAuthority has invalid authority.")) + } + + // ---- reflection ---- + + test("handles - the constructor choice matches the classpath's AbfsConfiguration") { + // Hadoop 3.4.2 added `containerConf` together with the constructor that takes the container. + val hasContainerConf = + Try(classOf[AbfsConfiguration].getMethod("containerConf", classOf[String])).isSuccess + handles.hasContainerScopedConfiguration shouldBe hasContainerConf + info( + s"container-scoped=${handles.hasContainerScopedConfiguration} " + + s"fixedSas=${handles.hasFixedSasTokenProvider} " + + s"clientAssertion=${handles.hasClientAssertionProvider}") + } + + test("handles - a failed lookup is not cached; the same class yields the same handles") { + val loader = new FlippingLoader(getClass.getClassLoader) + AbfsReflection.handles(loader).isFailure shouldBe true + AbfsReflection.handles(loader).isFailure shouldBe true + loader.blocked = false + // The delegating loader hands back the parent's AbfsConfiguration class, so the cache entry + // is the parent's too. + AbfsReflection.handles(loader) shouldBe Success(handles) + } + + // ---- Unconfigured ---- + + test("Unconfigured - no auth-bearing key for the account, even with other Azure keys set") { + resolveWith() shouldBe Unconfigured + resolveWith("fs.azure.account.hns.enabled" -> "true") shouldBe Unconfigured + // Another account's key does not configure this one. + resolveWith("fs.azure.account.key.other.dfs.core.windows.net" -> "c2VjcmV0") shouldBe + Unconfigured + toOptions(Unconfigured) shouldBe Map(Markers.RESOLUTION -> Markers.NONE) + } + + test("not Unconfigured - OAuth keys without an auth type fail in Hadoop's SharedKey words") { + // Hadoop defaults to SharedKey, whose key provider finds no account key. The environment must + // not step in for a misconfigured account. + val outcome = resolveWith(acct(Keys.CLIENT_ID) -> "client-123") + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe keyProviderException + assertNoAzureKeys(toOptions(outcome)) + } + + test("not Unconfigured - a fixed SAS token without an auth type, on every hadoop-azure") { + val outcome = resolveWith(acct(Keys.SAS_FIXED_TOKEN) -> "sv=2020-08-04&sig=fixed-signature") + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe keyProviderException + assertNoAzureKeys(toOptions(outcome)) + } + + test("not Unconfigured - every auth-bearing key set alone hands the decision to Hadoop") { + val plainValues = Map( + Keys.AUTH_TYPE -> "SharedKey", + Keys.KEY_PROVIDER -> "org.apache.hadoop.fs.azurebfs.services.SimpleKeyProvider", + Keys.SHELL_KEY_PROVIDER_SCRIPT -> "/bin/echo", + Keys.OAUTH_PROVIDER_TYPE -> Providers.CLIENT_CREDS, + Keys.SAS_TOKEN_PROVIDER_TYPE -> classOf[NeverInstantiatedSasTokenProvider].getName) + plainValues.keySet shouldBe plainAuthKeys.toSet + val cases = plainValues.toSeq ++ passwordAuthKeys.map(_ -> "some-value") + cases.foreach { case (key, value) => + withClue(s"$key alone: ") { + val outcome = resolveWith(acct(key) -> value) + outcome should not be Unconfigured + outcome shouldBe a[Failed] + assertNoAzureKeys(toOptions(outcome)) + } + } + } + + // ---- Declined ---- + + test("Declined - Custom auth type is forwarded by name, never instantiated") { + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "Custom", + acct(Keys.OAUTH_PROVIDER_TYPE) -> classOf[NeverInstantiatedAccessTokenProvider].getName) + outcome shouldBe Declined("Custom", None) + toOptions(outcome) shouldBe + Map(Markers.RESOLUTION -> Markers.DECLINED, Markers.AUTH_TYPE -> "Custom") + } + + test("Declined - a custom SAS token provider class is named, never constructed") { + val providerClass = classOf[NeverInstantiatedSasTokenProvider].getName + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "SAS", + acct(Keys.SAS_TOKEN_PROVIDER_TYPE) -> providerClass, + acct(Keys.SAS_FIXED_TOKEN) -> "sv=2020-08-04&sig=ignored-when-a-class-is-set") + outcome shouldBe Declined("SAS", Some(providerClass)) + val options = toOptions(outcome) + options(Markers.SAS_PROVIDER_CLASS) shouldBe providerClass + assertNoAzureKeys(options) + } + + test("Failed - an AccessTokenProvider that is not one of the five built-ins fails in Hadoop") { + // Hadoop matches the class by `==` and rejects anything else without constructing it. + val providerClass = classOf[NeverInstantiatedAccessTokenProvider].getName + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> providerClass, + acct(Keys.CLIENT_ID) -> "client-123") + outcome shouldBe Failed( + classOf[IllegalArgumentException].getName, + s"java.lang.IllegalArgumentException: Failed to initialize class $providerClass") + assertNoAzureKeys(toOptions(outcome)) + } + + test("Declined - Workload Identity with a client assertion provider (hadoop-azure 3.5.0+)") { + assume(handles.hasClientAssertionProvider, "no ClientAssertionProvider on this classpath") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.WORKLOAD_IDENTITY, + acct(Keys.MSI_TENANT) -> "tenant-abc", + acct(Keys.CLIENT_ID) -> "client-123", + acct(Keys.CLIENT_ASSERTION_PROVIDER_TYPE) -> " com.example.NeverLoadedAssertionProvider ") + outcome shouldBe Declined( + "OAuth", + Some(Providers.WORKLOAD_IDENTITY), + Some("com.example.NeverLoadedAssertionProvider")) + val options = toOptions(outcome) + options(Markers.OAUTH_PROVIDER_CLASS) shouldBe Providers.WORKLOAD_IDENTITY + options(Markers.OAUTH_ASSERTION_PROVIDER_CLASS) shouldBe + "com.example.NeverLoadedAssertionProvider" + assertNoAzureKeys(options) + } + + test("client assertion key is ignored where Hadoop ignores it (3.4.1 and 3.4.2)") { + assume( + handles.hasFixedSasTokenProvider && !handles.hasClientAssertionProvider, + "only on a hadoop-azure with Workload Identity but without ClientAssertionProvider") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.WORKLOAD_IDENTITY, + acct(Keys.MSI_TENANT) -> "tenant-abc", + acct(Keys.CLIENT_ID) -> "client-123", + acct(Keys.CLIENT_ASSERTION_PROVIDER_TYPE) -> "com.example.NeverLoadedAssertionProvider") + outcome shouldBe a[Resolved] + outcome.asInstanceOf[Resolved].providerClass shouldBe Some(Providers.WORKLOAD_IDENTITY) + } + + // ---- Failed ---- + + test("Failed - a bad Base64 account key carries Hadoop's class and a redacted message") { + val rawKey = "this-is-not-base64!!" + val outcome = resolveWith(acct(Keys.ACCOUNT_KEY) -> rawKey) + outcome shouldBe a[Failed] + val failed = outcome.asInstanceOf[Failed] + failed.exceptionClass shouldBe keyProviderException + failed.message should include("KeyProviderException") + failed.message should include("driver log") + failed.message should not include rawKey + val options = toOptions(outcome) + options(Markers.RESOLUTION) shouldBe Markers.ERROR + options(Markers.ERROR_CLASS) shouldBe keyProviderException + options.values.foreach(v => v should not include rawKey) + assertNoAzureKeys(options) + } + + test("Failed - a KeyProvider's own IllegalArgumentException is reported by class only") { + // Hadoop's IllegalArgumentExceptions name keys and classes; a user KeyProvider's may quote a + // secret the resolver never read, so only its origin decides whether the text is forwarded. + val outcome = + resolveWith(acct(Keys.KEY_PROVIDER) -> classOf[SecretLeakingKeyProvider].getName) + outcome shouldBe Failed( + classOf[IllegalArgumentException].getName, + "java.lang.IllegalArgumentException; the full message is in the Spark driver log") + val options = toOptions(outcome) + options(Markers.ERROR_CLASS) shouldBe classOf[IllegalArgumentException].getName + options.values.foreach(v => v should not include SecretLeakingKeyProvider.marker) + assertNoAzureKeys(options) + } + + test("Failed - a credential the resolution never read is still redacted") { + // A client secret pasted into the auth type: Hadoop's enum error quotes it, and resolution + // stops at the auth type, so only the failure path reads the secret it must hide. + val secret = "oauth-client-secret-4f9a2c" + val outcome = resolveWith(acct(Keys.AUTH_TYPE) -> secret, acct(Keys.CLIENT_SECRET) -> secret) + outcome shouldBe a[Failed] + val failed = outcome.asInstanceOf[Failed] + failed.exceptionClass shouldBe classOf[IllegalArgumentException].getName + failed.message should include("No enum constant") + failed.message should include("") + failed.message should not include secret + } + + test("Failed - an unknown auth type value is Hadoop's enum error") { + val outcome = resolveWith(acct(Keys.AUTH_TYPE) -> "oauth") + outcome shouldBe a[Failed] + val failed = outcome.asInstanceOf[Failed] + failed.exceptionClass shouldBe classOf[IllegalArgumentException].getName + failed.message should include("No enum constant") + failed.message should include("oauth") + } + + test("Failed - account-specific OAuth with only a global provider class (auth-type guard)") { + // getTokenProviderClass reads the global provider class only when the GLOBAL auth type + // matches, so Hadoop sees no class and fails to initialize the provider. + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + Keys.OAUTH_PROVIDER_TYPE -> Providers.CLIENT_CREDS, + acct(Keys.CLIENT_ENDPOINT) -> "https://login.microsoftonline.com/tenant/oauth2/token", + acct(Keys.CLIENT_ID) -> "client-123", + acct(Keys.CLIENT_SECRET) -> "client-secret-value") + outcome shouldBe Failed( + classOf[IllegalArgumentException].getName, + "java.lang.IllegalArgumentException: Failed to initialize null") + } + + test("Failed - a missing mandatory OAuth key surfaces Hadoop's cause chain") { + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.CLIENT_CREDS, + acct(Keys.CLIENT_ENDPOINT) -> "https://login.microsoftonline.com/tenant/oauth2/token", + acct(Keys.CLIENT_ID) -> "client-123") + outcome shouldBe a[Failed] + val failed = outcome.asInstanceOf[Failed] + failed.exceptionClass shouldBe tokenAccessProviderException + failed.message should include(s"Configuration property ${Keys.CLIENT_SECRET} not found.") + assertNoAzureKeys(toOptions(outcome)) + } + + test("Failed - SharedKey with an account name that is not fully qualified") { + val shortUri = new URI("abfss://data@myacct/path") + val outcome = resolve(conf("fs.azure.account.key.myacct" -> "c2VjcmV0"), shortUri) + outcome shouldBe + Failed(invalidUriException, s"Invalid URI $shortUri - account name is not fully qualified.") + } + + test("Failed - a LinkageError from the reflection layer is an error marker, not a crash") { + val outcome = resolve(conf(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0"), uri, Success(LinkageStub)) + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe classOf[NoSuchMethodError].getName + assertNoAzureKeys(toOptions(outcome)) + } + + test("Failed - hadoop-azure missing from the classpath") { + val missing = new ClassNotFoundException(AbfsReflection.ClassNames.ABFS_CONFIGURATION) + val outcome = resolve(conf(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0"), uri, Failure(missing)) + outcome shouldBe a[Failed] + val failed = outcome.asInstanceOf[Failed] + failed.exceptionClass shouldBe classOf[ClassNotFoundException].getName + failed.message should include("hadoop-azure is not on the driver classpath") + failed.message should include(AbfsReflection.ClassNames.ABFS_CONFIGURATION) + } + + test("redact - hides quoted keys and long secrets, leaves short values alone") { + val message = "Failure for acct key =\"c2VjcmV0c2VjcmV0\": token=the-secret-value port=8080" + redact(message, Seq("the-secret-value", "8080")) shouldBe + "Failure for acct key =\"\": token= port=8080" + } + + test("redact - a padded secret is hidden in its trimmed form too") { + val padded = " the-padded-secret " + redact("provider said the-padded-secret was rejected", Seq(padded, padded.trim)) shouldBe + "provider said was rejected" + } + + // ---- Resolved ---- + + test("Resolved - SharedKey from the account-scoped key, forwarded under the global name") { + val outcome = resolveWith(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0") + outcome shouldBe Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "c2VjcmV0")) + toOptions(outcome) shouldBe Map( + Markers.RESOLUTION -> Markers.RESOLVED, + Markers.AUTH_TYPE -> "SharedKey", + Keys.ACCOUNT_KEY -> "c2VjcmV0") + } + + test("Resolved - SharedKey reads no more credentials than Hadoop needs for the key") { + // Every getPassword can load a credential provider's keystore, so the resolver must not + // read the credentials of mechanisms Hadoop did not select. + def sharedKeyConf(): CountingConfiguration = { + val hadoopConf = new CountingConfiguration + hadoopConf.set(Keys.AUTH_TYPE, "SharedKey") + hadoopConf.set(acct(Keys.ACCOUNT_KEY), "c2VjcmV0") + hadoopConf + } + val direct = sharedKeyConf() + handles.getStorageAccountKey( + handles.newAbfsConfiguration(direct, account, container, uri)) shouldBe "c2VjcmV0" + direct.passwordReads should not be empty + + val viaResolver = sharedKeyConf() + resolve(viaResolver, uri) shouldBe + Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "c2VjcmV0")) + val reads = s"Hadoop read ${direct.passwordReads}, the resolver ${viaResolver.passwordReads}" + withClue(reads) { + viaResolver.passwordReads.size should be <= direct.passwordReads.size + } + } + + test("Resolved - SharedKey with an unreadable OAuth client secret it never reads") { + val outcome = resolveWith( + Seq(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0") ++ + cyclic(Keys.CLIENT_SECRET, "unused.secret"): _*) + outcome shouldBe Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "c2VjcmV0")) + } + + test("Resolved - SharedKey with an unreadable OAuth provider type it never reads") { + val outcome = resolveWith( + Seq(acct(Keys.ACCOUNT_KEY) -> "c2VjcmV0") ++ + cyclic(Keys.OAUTH_PROVIDER_TYPE, "unused.provider"): _*) + outcome shouldBe Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "c2VjcmV0")) + } + + test("not Unconfigured - an unreadable key alone fails in Hadoop's SharedKey words") { + Seq(Keys.CLIENT_SECRET, Keys.OAUTH_PROVIDER_TYPE).foreach { key => + withClue(s"$key alone: ") { + val outcome = resolveWith(cyclic(key, "unused.value"): _*) + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe keyProviderException + } + } + } + + test("Failed - an unreadable account key fails whether or not SharedKey is named") { + Seq(Seq.empty[(String, String)], Seq(acct(Keys.AUTH_TYPE) -> "SharedKey")).foreach { + authType => + withClue(s"auth type $authType: ") { + val outcome = resolveWith(authType ++ cyclic(acct(Keys.ACCOUNT_KEY), "cycle.key"): _*) + outcome shouldBe a[Failed] + assertNoAzureKeys(toOptions(outcome)) + } + } + } + + test("Resolved - SharedKey from the global key, with a port-scoped account") { + resolveWith(Keys.ACCOUNT_KEY -> "Z2xvYmFs") shouldBe + Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "Z2xvYmFs")) + + val portUri = new URI("abfss://data@myacct.dfs.core.windows.net:8080/path") + val portConf = conf( + "fs.azure.account.key.myacct.dfs.core.windows.net:8080" -> "cG9ydA==", + acct(Keys.ACCOUNT_KEY) -> "bm9wb3J0") + resolve(portConf, portUri) shouldBe + Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "cG9ydA==")) + } + + test("Resolved - mixed-case account names are matched exactly, as Hadoop does") { + val mixedUri = new URI("abfss://data@MyAcct.dfs.core.windows.net/path") + resolve( + conf("fs.azure.account.key.MyAcct.dfs.core.windows.net" -> "bWl4ZWQ="), + mixedUri) shouldBe + Resolved("SharedKey", None, Map(Keys.ACCOUNT_KEY -> "bWl4ZWQ=")) + // The lower-case key belongs to a different account name. + resolve( + conf("fs.azure.account.key.myacct.dfs.core.windows.net" -> "bWl4ZWQ="), + mixedUri) shouldBe + Unconfigured + } + + test("Resolved - ClientCreds forwards only its keys, never another account's") { + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.CLIENT_CREDS, + acct(Keys.CLIENT_ENDPOINT) -> "https://login.microsoftonline.com/tenant/oauth2/token", + acct(Keys.CLIENT_ID) -> "client-123", + acct(Keys.CLIENT_SECRET) -> "client-secret-value", + // Keys MSI would read, and another account's secret: none of them may travel. + acct(Keys.MSI_TENANT) -> "tenant-abc", + "fs.azure.account.oauth2.client.secret.other.dfs.core.windows.net" -> "other-secret-value") + outcome shouldBe Resolved( + "OAuth", + Some(Providers.CLIENT_CREDS), + Map( + Keys.CLIENT_ENDPOINT -> "https://login.microsoftonline.com/tenant/oauth2/token", + Keys.CLIENT_ID -> "client-123", + Keys.CLIENT_SECRET -> "client-secret-value")) + val options = toOptions(outcome) + options(Markers.OAUTH_PROVIDER_CLASS) shouldBe Providers.CLIENT_CREDS + options.values.foreach(v => v should not include "other-secret-value") + options.keys.filter(_.startsWith("fs.")) should contain theSameElementsAs + Seq(Keys.CLIENT_ENDPOINT, Keys.CLIENT_ID, Keys.CLIENT_SECRET) + } + + test("Resolved - global OAuth configuration applies when the global auth type matches") { + val globalConf = conf( + Keys.AUTH_TYPE -> "OAuth", + Keys.OAUTH_PROVIDER_TYPE -> Providers.CLIENT_CREDS, + Keys.CLIENT_ENDPOINT -> "https://login.microsoftonline.com/tenant/oauth2/token", + Keys.CLIENT_ID -> "global-client", + // A variable reference resolves through Configuration#get, as it does for Hadoop. Built by + // concatenation so scalac does not read it as a string missing its interpolator. + "my.client.secret" -> "substituted-secret", + Keys.CLIENT_SECRET -> ("${" + "my.client.secret}")) + resolve(globalConf, uri) shouldBe Resolved( + "OAuth", + Some(Providers.CLIENT_CREDS), + Map( + Keys.CLIENT_ENDPOINT -> "https://login.microsoftonline.com/tenant/oauth2/token", + Keys.CLIENT_ID -> "global-client", + Keys.CLIENT_SECRET -> "substituted-secret")) + } + + test("Resolved - MSI applies Hadoop's defaults, trimming and the authority slash") { + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.MSI, + acct(Keys.MSI_TENANT) -> "tenant-abc", + acct(Keys.CLIENT_ID) -> "client-123", + acct(Keys.MSI_AUTHORITY) -> " https://login.example.net ") + outcome shouldBe Resolved( + "OAuth", + Some(Providers.MSI), + Map( + Keys.MSI_ENDPOINT -> "http://169.254.169.254/metadata/identity/oauth2/token", + Keys.MSI_TENANT -> "tenant-abc", + Keys.CLIENT_ID -> "client-123", + Keys.MSI_AUTHORITY -> "https://login.example.net/")) + } + + test("Resolved - MSI without tenant and client id (optional from hadoop-azure 3.4.1)") { + assume(handles.hasFixedSasTokenProvider, "hadoop-azure 3.3.4 makes both keys mandatory") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.MSI, + acct(Keys.MSI_AUTHORITY) -> "") + outcome shouldBe Resolved( + "OAuth", + Some(Providers.MSI), + Map( + Keys.MSI_ENDPOINT -> "http://169.254.169.254/metadata/identity/oauth2/token", + Keys.MSI_AUTHORITY -> "https://login.microsoftonline.com/")) + } + + test("Resolved - Workload Identity with the default token file (hadoop-azure 3.4.1+)") { + assume(handles.hasFixedSasTokenProvider, "no WorkloadIdentityTokenProvider on 3.3.4") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.WORKLOAD_IDENTITY, + acct(Keys.MSI_TENANT) -> "tenant-abc", + acct(Keys.CLIENT_ID) -> "client-123") + outcome shouldBe Resolved( + "OAuth", + Some(Providers.WORKLOAD_IDENTITY), + Map( + Keys.MSI_AUTHORITY -> "https://login.microsoftonline.com/", + Keys.MSI_TENANT -> "tenant-abc", + Keys.CLIENT_ID -> "client-123", + Keys.TOKEN_FILE -> "/var/run/secrets/azure/tokens/azure-identity-token")) + } + + test("Resolved - container-scoped keys win on 3.4.2+ and are ignored before") { + assume(handles.hasFixedSasTokenProvider, "no WorkloadIdentityTokenProvider on 3.3.4") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "OAuth", + acct(Keys.OAUTH_PROVIDER_TYPE) -> Providers.WORKLOAD_IDENTITY, + acct(Keys.MSI_TENANT) -> "tenant-abc", + acct(Keys.CLIENT_ID) -> "account-client", + cont(Keys.CLIENT_ID) -> "container-client", + cont(Keys.TOKEN_FILE) -> " /mnt/tokens/container-token ") + val expectedClient = + if (handles.hasContainerScopedConfiguration) "container-client" else "account-client" + val expectedTokenFile = + if (handles.hasContainerScopedConfiguration) "/mnt/tokens/container-token" + else "/var/run/secrets/azure/tokens/azure-identity-token" + outcome shouldBe Resolved( + "OAuth", + Some(Providers.WORKLOAD_IDENTITY), + Map( + Keys.MSI_AUTHORITY -> "https://login.microsoftonline.com/", + Keys.MSI_TENANT -> "tenant-abc", + Keys.CLIENT_ID -> expectedClient, + Keys.TOKEN_FILE -> expectedTokenFile)) + } + + test("Resolved - fixed SAS token is trimmed and forwarded (hadoop-azure 3.4.1+)") { + assume(handles.hasFixedSasTokenProvider, "no fixed SAS token on 3.3.4") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "SAS", + acct(Keys.SAS_FIXED_TOKEN) -> " ?sv=2020-08-04&sig=fixed-signature ") + outcome shouldBe + Resolved("SAS", None, Map(Keys.SAS_FIXED_TOKEN -> "?sv=2020-08-04&sig=fixed-signature")) + toOptions(outcome) should not contain key(Markers.SAS_PROVIDER_CLASS) + } + + test("Failed - SAS without a provider class on hadoop-azure 3.3.4, which has no fixed token") { + assume(!handles.hasFixedSasTokenProvider, "3.4.1+ accept the fixed token") + val outcome = resolveWith( + acct(Keys.AUTH_TYPE) -> "SAS", + acct(Keys.SAS_FIXED_TOKEN) -> "sv=2020-08-04&sig=fixed-signature") + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe tokenAccessProviderException + assertNoAzureKeys(toOptions(outcome)) + } + + test("Failed - SAS with neither class nor token (hadoop-azure 3.4.1+)") { + assume(handles.hasFixedSasTokenProvider, "3.3.4 fails differently, covered above") + val outcome = resolveWith(acct(Keys.AUTH_TYPE) -> "SAS", acct(Keys.SAS_FIXED_TOKEN) -> " ") + outcome shouldBe a[Failed] + outcome.asInstanceOf[Failed].exceptionClass shouldBe + "org.apache.hadoop.fs.azurebfs.contracts.exceptions.SASTokenProviderException" + } +} + +/** Records the key of every `getPassword` lookup, each of which may load a keystore. */ +class CountingConfiguration extends Configuration { + private val reads = ArrayBuffer[String]() + + def passwordReads: Seq[String] = reads.synchronized(reads.toList) + + override def getPassword(name: String): Array[Char] = { + reads.synchronized(reads += name) + super.getPassword(name) + } +} + +/** A delegating loader that refuses `AbfsConfiguration` until `blocked` is cleared. */ +class FlippingLoader(parent: ClassLoader) extends ClassLoader(parent) { + @volatile var blocked = true + + override def loadClass(name: String, resolve: Boolean): Class[_] = { + if (blocked && name == AbfsReflection.ClassNames.ABFS_CONFIGURATION) { + throw new ClassNotFoundException(name) + } + super.loadClass(name, resolve) + } +} + +object SecretLeakingKeyProvider { + val marker = "vault-secret-marker-7b3e9c" +} + +/** A KeyProvider whose failure quotes a secret the resolver never read from the conf. */ +class SecretLeakingKeyProvider extends KeyProvider { + override def getStorageAccountKey(accountName: String, conf: Configuration): String = + throw new IllegalArgumentException(s"bad key ${SecretLeakingKeyProvider.marker}") +} + +/** Fails loudly if anything constructs it: the resolver must only ever name this class. */ +class NeverInstantiatedSasTokenProvider extends SASTokenProvider { + NeverInstantiated.fail("NeverInstantiatedSasTokenProvider") + + override def initialize(configuration: Configuration, accountName: String): Unit = () + + override def getSASToken( + account: String, + fileSystem: String, + path: String, + operation: String): String = + throw new UnsupportedOperationException("not reached") +} + +/** An `AccessTokenProvider` that is not one of Hadoop's built-ins; never constructed. */ +class NeverInstantiatedAccessTokenProvider extends AccessTokenProvider { + NeverInstantiated.fail("NeverInstantiatedAccessTokenProvider") + + override protected def refreshToken(): AzureADToken = + throw new UnsupportedOperationException("not reached") +} + +/** Handles whose first call fails with a `LinkageError`, as a mismatched hadoop-azure would. */ +object LinkageStub extends Handles { + // Built here so the file has no `throw new ...Error(` (scalastyle); thrown below. + private val linkageFailure = new NoSuchMethodError("stub: AbfsConfiguration.") + + private def notReached: Nothing = throw new UnsupportedOperationException("not reached") + + override def hasContainerScopedConfiguration: Boolean = false + override def hasFixedSasTokenProvider: Boolean = false + override def hasClientAssertionProvider: Boolean = false + override def accessTokenProviderClass: Class[_] = classOf[AnyRef] + override def sasTokenProviderClass: Class[_] = classOf[AnyRef] + override def newAbfsConfiguration( + conf: Configuration, + account: String, + container: String, + uri: URI): AnyRef = throw linkageFailure + override def newKeyProviderConfiguration(conf: Configuration, account: String): AnyRef = + notReached + override def getAuthType(abfsConf: AnyRef, account: String): Enum[_] = notReached + override def getPasswordString(abfsConf: AnyRef, key: String): Option[String] = notReached + override def read( + abfsConf: AnyRef, + key: String, + password: Boolean): Either[Throwable, Option[String]] = notReached + override def getTokenProviderClass( + abfsConf: AnyRef, + authType: Enum[_], + key: String, + xface: Class[_]): Option[Class[_]] = notReached + override def getStorageAccountKey(abfsConf: AnyRef): String = notReached + override def getTokenProvider(abfsConf: AnyRef): Unit = notReached + override def getSASTokenProvider(abfsConf: AnyRef): Unit = notReached + override def authDefault(fieldName: String): String = notReached +} + +/** Throws from a constructor without leaving the compiler a dead-code warning to report. */ +object NeverInstantiated { + def fail(name: String): Unit = throw new IllegalStateException(s"$name was constructed") +} diff --git a/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala new file mode 100644 index 00000000000..016491d8af8 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/objectstore/AbfsScanFallbackSuite.scala @@ -0,0 +1,149 @@ +/* + * 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. + */ + +package org.apache.comet.objectstore + +import java.io.File +import java.nio.file.{Files, Paths} + +import org.apache.commons.io.FileUtils +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.comet.CometNativeScanExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StructType} + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.hadoop.fs.FakeAbfsFileSystem + +/** + * Plans scans over `abfss://` paths served by [[FakeAbfsFileSystem]], with authentication + * resolved offline by the hadoop-azure on the test classpath. A scan Comet declines runs on + * Spark's own scan against the local files. + */ +class AbfsScanFallbackSuite extends CometTestBase { + + private val multipleAuthoritiesReason = "more than one container or account" + + private var rootDir: File = _ + + override protected def sparkConf: SparkConf = + super.sparkConf + .set("spark.hadoop.fs.abfss.impl", classOf[FakeAbfsFileSystem].getName) + .set("spark.hadoop.fs.abfss.impl.disable.cache", "true") + + override def beforeAll(): Unit = { + // java.io.tmpdir is under target/, which a fresh checkout may not have yet. + val tmp = Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + rootDir = Files.createTempDirectory(tmp, "comet-abfs-scan").toFile + super.beforeAll() + } + + protected override def afterAll(): Unit = { + if (rootDir != null) FileUtils.deleteDirectory(rootDir) + super.afterAll() + } + + private def localDir(name: String): String = new File(rootDir, name).getAbsolutePath + + private def abfss(container: String, account: String, dir: String): String = + s"abfss://$container@$account.dfs.core.windows.net${localDir(dir)}" + + private def accountKeys(keys: (String, String)*): Seq[(String, String)] = + keys.map { case (account, key) => + s"fs.azure.account.key.$account.dfs.core.windows.net" -> key + } + + private def writeParquet(dir: String, from: Int): Unit = + spark.range(from.toLong, from + 10L).toDF("id").write.parquet(localDir(dir)) + + private def writeCsv(dir: String, from: Int): Unit = + spark + .range(from.toLong, from + 10L) + .selectExpr("cast(id as int) as a", "cast(id * 2 as int) as b") + .write + .csv(localDir(dir)) + + private val csvConfs = + Seq(CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", SQLConf.USE_V1_SOURCE_LIST.key -> "") + + private def readCsv(paths: String*): DataFrame = + spark.read + .schema(new StructType().add("a", IntegerType).add("b", IntegerType)) + .csv(paths: _*) + + test("native Parquet scan over two accounts with different keys falls back to Spark") { + writeParquet("parquet-a", 0) + writeParquet("parquet-b", 100) + withSQLConf(accountKeys("accta" -> "a2V5LWE=", "acctb" -> "a2V5LWI="): _*) { + val df = + spark.read.parquet( + abfss("data", "accta", "parquet-a"), + abfss("data", "acctb", "parquet-b")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } + + test("native Parquet scan over two containers of one account falls back to Spark") { + writeParquet("parquet-c", 0) + writeParquet("parquet-d", 100) + withSQLConf(accountKeys("accta" -> "a2V5LWE="): _*) { + val df = + spark.read.parquet( + abfss("data", "accta", "parquet-c"), + abfss("logs", "accta", "parquet-d")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } + + test("native Parquet scan over two directories of one container stays native") { + writeParquet("parquet-e", 0) + writeParquet("parquet-f", 100) + val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") +: + accountKeys("accta" -> "a2V5LWE=") + withSQLConf(confs: _*) { + // Plan only: running the native scan would send its requests to Azure. + val plan = spark.read + .parquet(abfss("data", "accta", "parquet-e"), abfss("data", "accta", "parquet-f")) + .queryExecution + .executedPlan + assert(plan.collect { case scan: CometNativeScanExec => scan }.size == 1, plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(!reasons.exists(_.contains("ABFS")), reasons) + } + } + + test("native CSV scan over two accounts with different keys falls back to Spark") { + writeCsv("csv-a", 0) + writeCsv("csv-b", 100) + withSQLConf(csvConfs ++ accountKeys("accta" -> "a2V5LWE=", "acctb" -> "a2V5LWI="): _*) { + val df = readCsv(abfss("data", "accta", "csv-a"), abfss("data", "acctb", "csv-b")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } + + test("native CSV scan over two containers of one account falls back to Spark") { + writeCsv("csv-c", 0) + writeCsv("csv-d", 100) + withSQLConf(csvConfs ++ accountKeys("accta" -> "a2V5LWE="): _*) { + val df = readCsv(abfss("data", "accta", "csv-c"), abfss("logs", "accta", "csv-d")) + checkSparkAnswerAndFallbackReason(df, multipleAuthoritiesReason) + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/objectstore/NativeConfigSuite.scala b/spark/src/test/scala/org/apache/comet/objectstore/NativeConfigSuite.scala index 1cfd1d36b92..d62733aa421 100644 --- a/spark/src/test/scala/org/apache/comet/objectstore/NativeConfigSuite.scala +++ b/spark/src/test/scala/org/apache/comet/objectstore/NativeConfigSuite.scala @@ -26,7 +26,7 @@ import org.scalatest.matchers.should.Matchers import org.apache.hadoop.conf.Configuration -import org.apache.comet.CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY +import org.apache.comet.CometConf.{COMET_LIBHDFS_SCHEMES_KEY, COMET_S3_COMPLIANT_SCHEMES_KEY} class NativeConfigSuite extends AnyFunSuite with Matchers { @@ -81,38 +81,44 @@ class NativeConfigSuite extends AnyFunSuite with Matchers { assert(unsupportedOptions.isEmpty, "Unsupported scheme should return empty options") } - test("extractObjectStoreOptions - ABFS forwards Hadoop fs.azure.* auth keys") { - // ABFS auth (account keys, OAuth, MSI/Workload Identity, SAS) lives under fs.azure.*, not - // fs.abfs.*. Verify abfs[s] forwards fs.azure.* (earlier versions dropped these credentials). + test("extractObjectStoreOptions - ABFS forwards the auth Hadoop resolves for the account") { + // ABFS auth lives under fs.azure.*; Hadoop (AbfsConfiguration) picks the mechanism and the + // values for the URI's account, and only those travel, under the global key names, with the + // comet.azure.* markers. Account-scoped keys and other accounts' keys stay on the driver. val hadoopConf = new Configuration() hadoopConf.set("fs.azure.account.auth.type.myacct.dfs.core.windows.net", "OAuth") hadoopConf.set( "fs.azure.account.oauth.provider.type.myacct.dfs.core.windows.net", - "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider") + "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider") + hadoopConf.set( + "fs.azure.account.oauth2.client.endpoint.myacct.dfs.core.windows.net", + "https://login.microsoftonline.com/tenant-abc/oauth2/token") hadoopConf.set("fs.azure.account.oauth2.client.id.myacct.dfs.core.windows.net", "client-123") - hadoopConf.set("fs.azure.account.oauth2.msi.tenant.myacct.dfs.core.windows.net", "tenant-abc") hadoopConf.set( - "fs.azure.account.oauth2.token.file.myacct.dfs.core.windows.net", - "/var/run/secrets/azure/tokens/azure-identity-token") + "fs.azure.account.oauth2.client.secret.myacct.dfs.core.windows.net", + "client-secret-value") + hadoopConf.set( + "fs.azure.account.oauth2.client.secret.other.dfs.core.windows.net", + "other-secret-value") Seq( "abfs://data@myacct.dfs.core.windows.net/path/file.parquet", "abfss://data@myacct.dfs.core.windows.net/path/file.parquet").foreach { path => val opts = NativeConfig.extractObjectStoreOptions(hadoopConf, new URI(path)) - assert( - opts("fs.azure.account.oauth2.client.id.myacct.dfs.core.windows.net") == "client-123", - s"client id should be forwarded for $path") - assert( - opts("fs.azure.account.oauth2.msi.tenant.myacct.dfs.core.windows.net") == "tenant-abc", - s"tenant id should be forwarded for $path") - assert( - opts("fs.azure.account.oauth2.token.file.myacct.dfs.core.windows.net") == - "/var/run/secrets/azure/tokens/azure-identity-token", - s"federated token file should be forwarded for $path") - assert( - opts("fs.azure.account.oauth.provider.type.myacct.dfs.core.windows.net") == - "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider", - s"oauth provider type should be forwarded for $path") + withClue(s"$path: $opts") { + assert(opts("comet.azure.resolution") == "resolved") + assert(opts("comet.azure.auth.type") == "OAuth") + assert( + opts("comet.azure.oauth.provider.class") == + "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider") + assert( + opts("fs.azure.account.oauth2.client.endpoint") == + "https://login.microsoftonline.com/tenant-abc/oauth2/token") + assert(opts("fs.azure.account.oauth2.client.id") == "client-123") + assert(opts("fs.azure.account.oauth2.client.secret") == "client-secret-value") + assert(!opts.keys.exists(_.endsWith(".myacct.dfs.core.windows.net"))) + assert(!opts.values.exists(_ == "other-secret-value")) + } } } @@ -354,4 +360,133 @@ class NativeConfigSuite extends AnyFunSuite with Matchers { NativeConfig.resolveS3CompliantSchemes(conf) == Set("blob", "minio", "r2"), "schemes must be split on commas, trimmed, lowercased, with blanks dropped") } + + test("extractObjectStoreOptions - abfs without Azure auth keys is 'none' with no fs.* keys") { + // The one case where native may read the AZURE_* environment. A key set under the comet + // marker name in the Hadoop conf is never forwarded; the resolver's verdict replaces it. + val hadoopConf = new Configuration() + hadoopConf.set(COMET_LIBHDFS_SCHEMES_KEY, "hdfs") + hadoopConf.set("comet.azure.resolution", "resolved") + hadoopConf.set("fs.azure.account.hns.enabled", "true") + + Seq("abfs", "abfss").foreach { scheme => + val opts = NativeConfig.extractObjectStoreOptions( + hadoopConf, + new URI(s"$scheme://data@myacct.dfs.core.windows.net/path/file.parquet")) + withClue(s"$scheme: $opts") { + assert( + opts == Map(COMET_LIBHDFS_SCHEMES_KEY -> "hdfs", "comet.azure.resolution" -> "none")) + } + } + } + + test("extractObjectStoreOptions - abfs resolution failures travel as error markers only") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key.myacct.dfs.core.windows.net", "not-base64-at-all!!!") + + val opts = NativeConfig.extractObjectStoreOptions( + hadoopConf, + new URI("abfss://data@myacct.dfs.core.windows.net/path/file.parquet")) + assert(opts("comet.azure.resolution") == "error", opts.toString) + assert(opts("comet.azure.error.class").endsWith("KeyProviderException"), opts.toString) + assert(!opts.keys.exists(_.startsWith("fs.")), opts.toString) + assert(!opts.values.exists(_.contains("not-base64-at-all!!!")), opts.toString) + } + + test("extractObjectStoreOptions - wasb still forwards fs.azure.* by prefix, without markers") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key.myacct.blob.core.windows.net", "azure-key") + hadoopConf.set("fs.azure.sas.fixed.token", "sv=2020-08-04&sig=fixed") + + Seq( + "wasb://data@myacct.blob.core.windows.net/key", + "wasbs://data@myacct.blob.core.windows.net/key").foreach { path => + val opts = NativeConfig.extractObjectStoreOptions(hadoopConf, new URI(path)) + withClue(s"$path: $opts") { + assert(opts("fs.azure.account.key.myacct.blob.core.windows.net") == "azure-key") + assert(opts("fs.azure.sas.fixed.token") == "sv=2020-08-04&sig=fixed") + assert(!opts.keys.exists(_.startsWith("comet.azure."))) + } + } + Seq("s3a://bucket/key", "gs://bucket/key").foreach { path => + val opts = NativeConfig.extractObjectStoreOptions(hadoopConf, new URI(path)) + assert(!opts.keys.exists(_.startsWith("comet.azure.")), s"$path: $opts") + } + } + + private def abfsFile(container: String, account: String, file: String): URI = + new URI(s"abfss://$container@$account.dfs.core.windows.net/path/$file") + + private def assertDeclined(hadoopConf: Configuration, uris: Seq[URI]): Unit = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, uris) match { + case Left(reason) => + reason should include(uris.head.getRawAuthority) + reason should include(uris.last.getRawAuthority) + case Right(options) => fail(s"expected a decline, got $options") + } + + test("extractScanObjectStoreOptions - abfs accounts with different keys decline the scan") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key.accta.dfs.core.windows.net", "a2V5LWE=") + hadoopConf.set("fs.azure.account.key.acctb.dfs.core.windows.net", "a2V5LWI=") + assertDeclined( + hadoopConf, + Seq( + abfsFile("data", "accta", "a1.parquet"), + abfsFile("data", "accta", "a2.parquet"), + abfsFile("data", "acctb", "b1.parquet"))) + } + + test("extractScanObjectStoreOptions - containers of one account decline the scan") { + // Native reads a whole partition through the store of its first file's container. + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key.accta.dfs.core.windows.net", "a2V5LWE=") + assertDeclined( + hadoopConf, + Seq(abfsFile("data", "accta", "a1.parquet"), abfsFile("logs", "accta", "a2.parquet"))) + } + + test("extractScanObjectStoreOptions - accounts sharing a global key decline the scan") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key", "Z2xvYmFs") + assertDeclined( + hadoopConf, + Seq(abfsFile("data", "accta", "a.parquet"), abfsFile("data", "acctb", "b.parquet"))) + } + + test("extractScanObjectStoreOptions - files in one container keep the first file's options") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key.accta.dfs.core.windows.net", "a2V5LWE=") + val uris = (1 to 5).map(i => abfsFile("data", "accta", s"a$i.parquet")) + + val single = NativeConfig.extractObjectStoreOptions(hadoopConf, uris.head) + assert(single("fs.azure.account.key") == "a2V5LWE=", single.toString) + NativeConfig.extractScanObjectStoreOptions(hadoopConf, uris) shouldBe Right(single) + } + + test("extractScanObjectStoreOptions - abfs and non-abfs files in one scan decline it") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.azure.account.key", "Z2xvYmFs") + val abfs = abfsFile("data", "accta", "a.parquet") + val s3 = new URI("s3a://bucket/path/b.parquet") + + Seq(Seq(abfs, s3), Seq(s3, abfs)).foreach { uris => + val result = NativeConfig.extractScanObjectStoreOptions(hadoopConf, uris) + assert(result.isLeft, s"$uris: $result") + } + } + + test("extractScanObjectStoreOptions - non-abfs scans keep the first file's options") { + val hadoopConf = new Configuration() + hadoopConf.set("fs.s3a.access.key", "s3-access-key") + hadoopConf.set("fs.s3a.bucket.other.access.key", "other-access-key") + val uris = Seq( + new URI("s3a://bucket/path/a.parquet"), + new URI("s3a://other/path/b.parquet"), + new URI("s3://bucket/path/c.parquet")) + + NativeConfig.extractScanObjectStoreOptions(hadoopConf, uris) shouldBe + Right(NativeConfig.extractObjectStoreOptions(hadoopConf, uris.head)) + NativeConfig.extractScanObjectStoreOptions(hadoopConf, Seq.empty) shouldBe Right(Map.empty) + } } From 79dd56f3660e4016e290b05635670f1f1033cca2 Mon Sep 17 00:00:00 2001 From: Feng Zhang Date: Sun, 11 Oct 2026 19:08:42 +0000 Subject: [PATCH 68/89] fix: fall back from native V2 scans when the plan reads input_file_name (#6754) * fix: fall back from native V2 scans when the plan reads input_file_name The native Iceberg and CSV V2 scans do not set InputFileBlockHolder, so input_file_name, input_file_block_start and input_file_block_length returned empty values. Fall back to Spark's reader as the V1 path does. Closes #6707 * refactor: reuse readsInputFileBlock in the V2 scan check #6703 added the shared helper for the same predicate. * test: cover the native CSV V2 scan's input_file_name fallback Document the fallback in the CSV section of datasources.md. * test: move the Iceberg input_file_name test to a SQL fixture Match the full fallback reason there and in the CSV test. --- docs/source/user-guide/latest/datasources.md | 4 +- docs/source/user-guide/latest/iceberg.md | 1 + .../apache/comet/rules/CometScanRule.scala | 14 ++++- .../sql-tests/iceberg/input_file_name.sql | 61 +++++++++++++++++++ .../comet/csv/CometCsvNativeReadSuite.scala | 29 +++++++++ 5 files changed, 106 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/iceberg/input_file_name.sql diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index 0a5fe4aded4..458538a102d 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -39,7 +39,9 @@ See the [Iceberg Guide] and [Iceberg Writes](iceberg-writes.md) for more informa Comet provides experimental Rust-based CSV scan support. When `spark.comet.scan.csv.v2.enabled` is enabled, CSV files are read in Rust for improved performance. This feature is experimental and performance benefits are workload-dependent. Only Spark's DataSource V2 CSV scan is accelerated, and Spark reads CSV through the V1 API by -default, so also remove `csv` from `spark.sql.sources.useV1SourceList`. +default, so also remove `csv` from `spark.sql.sources.useV1SourceList`. Queries that use +`input_file_name()`, `input_file_block_start()` or `input_file_block_length()` fall back to Spark's CSV +reader, which sets the values these functions read. Alternatively, when `spark.comet.convert.csv.enabled` is enabled, data from Spark's CSV reader is immediately converted into Arrow format, allowing the Comet pipeline to take over after that. diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 4c3abe00216..fbdf81c3712 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -186,6 +186,7 @@ The following scenarios will fall back to the JVM Iceberg reader: - Encrypted tables with 192-bit data keys (no AES-192-GCM in the underlying crypto) - Delete files in a format other than Parquet or Puffin (Avro or ORC positional/equality deletes) - Tables backed by Avro or ORC data files (only Parquet is accelerated) +- Queries that use `input_file_name()`, `input_file_block_start()` or `input_file_block_length()` - Scans whose data or delete files span more than one S3 bucket (the native reader uses one object-store configuration per scan) - Tables partitioned on `BINARY` or `DECIMAL` (with precision >28) columns diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index 176c4c595ae..edda7d3fabd 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -111,7 +111,7 @@ case class CometScanRule(session: SparkSession) // is ever offered it. `transformV2Scan` applies the guard right after its contrib hook // declines. case scanExec: BatchScanExec => - transformV2Scan(scanExec) + transformV2Scan(fullPlan, scanExec) } plan.transform { @@ -375,7 +375,7 @@ case class CometScanRule(session: SparkSession) Some(CometScanExec(scanExec, session)) } - private def transformV2Scan(scanExec: BatchScanExec): SparkPlan = { + private def transformV2Scan(plan: SparkPlan, scanExec: BatchScanExec): SparkPlan = { // Give any optional, out-of-tree scan contrib (e.g. Lance) first crack at this V2 scan. On a // default build no contrib is registered, so this returns None and we proceed with Comet's @@ -394,6 +394,16 @@ case class CometScanRule(session: SparkSession) return withFallbackReason(scanExec, "Iceberg Metadata tables are not supported") } + // As in transformV1Scan: these expressions read InputFileBlockHolder, which the source's own + // reader sets per file. Comet's native V2 scans (Iceberg, CSV) do not, so they would return + // empty/default values (https://github.com/apache/datafusion-comet/issues/6707). + if (CometScanRule.readsInputFileBlock(plan)) { + return withFallbackReason( + scanExec, + "Native V2 scan is not compatible with input_file_name, " + + "input_file_block_start, or input_file_block_length") + } + // NOTE: there is no blanket metadata-column guard here. Comet's built-in V2 paths handle // metadata columns individually -- the Iceberg path supports the ones in // `CometIcebergNativeScan.MetadataFieldIds` and rejects the rest, the CSV path rejects all -- diff --git a/spark/src/test/resources/sql-tests/iceberg/input_file_name.sql b/spark/src/test/resources/sql-tests/iceberg/input_file_name.sql new file mode 100644 index 00000000000..a473e720ae9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/iceberg/input_file_name.sql @@ -0,0 +1,61 @@ +-- 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. + +-- input_file_name, input_file_block_start and input_file_block_length read InputFileBlockHolder, +-- which Iceberg's Spark reader sets as it opens each data file. The native Iceberg scan does not, +-- so a plan that uses them falls back to Spark's reader +-- (https://github.com/apache/datafusion-comet/issues/6707). + +-- Native Iceberg scan is unsupported on Spark 4.2: no Iceberg spark-runtime is published for +-- 4.2 yet. See https://github.com/apache/datafusion-comet/issues/4969. +-- MaxSparkVersion: 4.1 + +-- Config: spark.sql.catalog.test_cat=org.apache.iceberg.spark.SparkCatalog +-- Config: spark.sql.catalog.test_cat.type=hadoop +-- Config: spark.sql.catalog.test_cat.warehouse=/tmp/comet-iceberg-sql-test +-- Config: spark.comet.enabled=true +-- Config: spark.comet.exec.enabled=true +-- Config: spark.comet.scan.icebergNative.enabled=true + +statement +DROP TABLE IF EXISTS test_cat.db.input_file + +statement +CREATE TABLE test_cat.db.input_file (id BIGINT) USING iceberg + +-- Three inserts write three data files +statement +INSERT INTO test_cat.db.input_file SELECT id FROM range(0, 1000) + +statement +INSERT INTO test_cat.db.input_file SELECT id FROM range(1000, 2000) + +statement +INSERT INTO test_cat.db.input_file SELECT id FROM range(2000, 3000) + +query expect_fallback(Native V2 scan is not compatible with input_file_name) +SELECT input_file_name(), input_file_block_start(), input_file_block_length(), id +FROM test_cat.db.input_file + +-- A Comet filter above the scan +query expect_fallback(Native V2 scan is not compatible with input_file_name) +SELECT input_file_name(), input_file_block_start(), input_file_block_length(), id +FROM test_cat.db.input_file WHERE id >= 0 + +-- Without these functions the scan stays native +query +SELECT id FROM test_cat.db.input_file WHERE id >= 0 diff --git a/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala b/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala index 6e743cfc90d..eeb809764be 100644 --- a/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/csv/CometCsvNativeReadSuite.scala @@ -147,4 +147,33 @@ class CometCsvNativeReadSuite extends CometTestBase { } } } + + // https://github.com/apache/datafusion-comet/issues/6707. These functions read + // InputFileBlockHolder, which the native CSV scan does not set. + test("Native csv read - input_file_name falls back to Spark's reader") { + withTempPath { dir => + spark.range(30).repartition(3).write.csv(dir.toString) + withSQLConf( + CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> "") { + val read = spark.read.schema("id LONG").csv(dir.toString) + val df = read.selectExpr( + "input_file_name()", + "input_file_block_start()", + "input_file_block_length()", + "id") + val (_, plan) = checkSparkAnswerAndFallbackReason( + df, + "Native V2 scan is not compatible with input_file_name") + assert( + collect(plan) { case s: CometCsvNativeScanExec => s }.isEmpty, + s"Expected Spark's CSV reader but found CometCsvNativeScanExec. Plan:\n$plan") + // Without these functions the scan stays native + val (_, nativePlan) = checkSparkAnswer(read.where("id >= 0")) + assert( + collect(nativePlan) { case s: CometCsvNativeScanExec => s }.nonEmpty, + s"Expected CometCsvNativeScanExec. Plan:\n$nativePlan") + } + } + } } From 512d05fb7f6834ad0a87ec2eea7b94fb9b269eda Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 11 Oct 2026 21:44:57 +0000 Subject: [PATCH 69/89] fix: accept all BMP Unicode decimal digits in cast(string as decimal) (#6841) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: accept all BMP Unicode decimal digits in cast(string as decimal) Spark casts a string to decimal with `new java.math.BigDecimal(s.trim)`, which accepts any char for which `Character.isDigit(char)` is true, in both the mantissa and the exponent. Comet only normalized fullwidth digits, so strings such as '٣' (U+0663 ARABIC-INDIC DIGIT THREE) or '1e३' returned NULL, or raised CAST_INVALID_INPUT under ANSI, where Spark returns a value. Normalize every non-ASCII Nd digit in the Basic Multilingual Plane to ASCII before parsing. Supplementary-plane digits stay rejected, since BigDecimal sees them as surrogate pairs. The pure-ASCII fast path is unchanged. Integral, floating-point and date casts are left alone: Spark only accepts ASCII digits there. * perf: keep ASCII and fullwidth digits off the table lookup Normalizing every char through the Nd table slowed down already supported fullwidth and mixed ASCII/fullwidth input. Scan bytes instead, copying ASCII directly and mapping the fullwidth UTF-8 encoding in place as before, and only decode a char for the table lookup otherwise. The output buffer is preallocated, since it is never longer than the input. Add fullwidth, mixed ASCII/fullwidth and Arabic-Indic cases to the cast_from_string benchmark. * fix: drop removed SparkCastOptions argument in the new decimal benchmark #6135 removed the third argument of SparkCastOptions::new on main. --- .../expressions/_category_template/cast.md | 7 +- native/spark-expr/benches/cast_from_string.rs | 46 ++++ .../spark-expr/src/conversion_funcs/string.rs | 200 ++++++++++++++++-- .../cast_string_to_decimal_unicode_digits.sql | 88 ++++++++ ..._string_to_decimal_unicode_digits_ansi.sql | 76 +++++++ 5 files changed, 393 insertions(+), 24 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits_ansi.sql diff --git a/docs/source/user-guide/latest/compatibility/expressions/_category_template/cast.md b/docs/source/user-guide/latest/compatibility/expressions/_category_template/cast.md index c023539f36e..f55199125c3 100644 --- a/docs/source/user-guide/latest/compatibility/expressions/_category_template/cast.md +++ b/docs/source/user-guide/latest/compatibility/expressions/_category_template/cast.md @@ -64,8 +64,11 @@ including: `java.lang.String.trim` that Spark applies before handing the value to `BigDecimal`. - Null bytes (`\u0000`) at the start or end of a string are therefore trimmed. Null bytes embedded in the middle of a string produce `NULL`. -- Fullwidth Unicode digits (U+FF10–U+FF19, e.g. `１２３.４５`) are treated as their ASCII - equivalents, so `CAST('１２３.４５' AS DECIMAL(10,2))` returns `123.45`. +- Non-ASCII decimal digits that Java's `BigDecimal` accepts, i.e. every Unicode `Nd` digit in + the Basic Multilingual Plane (fullwidth `１２３.４５`, Arabic-Indic `١٢٣`, Devanagari `१२३`, + ...), are treated as their ASCII equivalents in both the mantissa and the exponent, so + `CAST('١٢٣.٤٥' AS DECIMAL(10,2))` returns `123.45`. Digits outside the Basic Multilingual + Plane, such as U+1D7CE MATHEMATICAL BOLD DIGIT ZERO, are rejected as they are by Spark. - Scientific notation (e.g. `1.23E+5`) is supported. - Special values (`inf`, `infinity`, `nan`) produce `NULL`. diff --git a/native/spark-expr/benches/cast_from_string.rs b/native/spark-expr/benches/cast_from_string.rs index cd344670127..a64d762a6b7 100644 --- a/native/spark-expr/benches/cast_from_string.rs +++ b/native/spark-expr/benches/cast_from_string.rs @@ -136,6 +136,29 @@ fn criterion_benchmark(c: &mut Criterion) { group.finish(); } + // str -> decimal with non-ASCII digits, which take the normalization path in + // `parse_string_to_decimal` instead of the pure-ASCII fast path + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, ""); + let cast = Cast::new( + expr.clone(), + DataType::Decimal128(38, 10), + spark_cast_options, + None, + None, + ); + let mut group = c.benchmark_group("cast_string_to_decimal_unicode/legacy"); + for (name, zero, every_other) in [ + ("fullwidth", '\u{FF10}', false), + ("mixed_ascii_fullwidth", '\u{FF10}', true), + ("arabic_indic", '\u{0660}', false), + ] { + let batch = create_unicode_decimal_cast_string_batch(zero, every_other); + group.bench_function(name, |b| { + b.iter(|| cast.evaluate(&batch).unwrap()); + }); + } + group.finish(); + // str -> boolean and str -> float benchmarks, with and without the leading/trailing // whitespace that exercises the trim helpers in `conversion_funcs::trim` let bool_batch = create_boolean_string_batch(false); @@ -252,6 +275,29 @@ fn create_decimal_cast_string_batch() -> RecordBatch { }) } +/// Create batch of decimal strings whose digits are re-spelled in the script whose zero is +/// `zero`, either all of them or every other one (mixing with ASCII) +fn create_unicode_decimal_cast_string_batch(zero: char, every_other: bool) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(42); + string_batch(8192, 10, move |_| { + let ascii = format!( + "-{}.{}", + rng.random_range(0..1_000_000_000u32), + rng.random_range(0..1_000_000_000u32) + ); + ascii + .chars() + .enumerate() + .map(|(i, c)| match c.to_digit(10) { + Some(d) if !every_other || i % 2 == 0 => { + char::from_u32(u32::from(zero) + d).unwrap() + } + _ => c, + }) + .collect() + }) +} + fn config() -> Criterion { Criterion::default() } diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index f278bafda68..89c18255d0f 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -461,37 +461,110 @@ fn cast_string_to_decimal256_impl( )) } -/// Normalize fullwidth Unicode digits (U+FF10–U+FF19) to their ASCII equivalents. +/// Code point of the zero of each run of ten non-ASCII decimal digits (Unicode general +/// category `Nd`) in the Basic Multilingual Plane, in ascending order. Every run is contiguous, +/// so the digit at `zero + n` has value `n`. /// -/// Spark's UTF8String parser treats fullwidth digits as numerically equivalent to -/// ASCII digits, e.g. "１２３.４５" parses as 123.45. Each fullwidth digit encodes -/// to exactly three UTF-8 bytes: [0xEF, 0xBC, 0x90+n] for digit n. The ASCII -/// equivalent is 0x30+n, so the conversion is: third_byte - 0x60. +/// This is the set of `char`s for which Java's `Character.isDigit(char)` is true, minus ASCII. +/// It is identical under JDK 17 (Unicode 13) and JDK 21 (Unicode 15); all `Nd` digits added +/// since lie outside the BMP. +const BMP_NON_ASCII_DIGIT_ZEROS: [u16; 36] = [ + 0x0660, // Arabic-Indic + 0x06F0, // Extended Arabic-Indic + 0x07C0, // NKo + 0x0966, // Devanagari + 0x09E6, // Bengali + 0x0A66, // Gurmukhi + 0x0AE6, // Gujarati + 0x0B66, // Oriya + 0x0BE6, // Tamil + 0x0C66, // Telugu + 0x0CE6, // Kannada + 0x0D66, // Malayalam + 0x0DE6, // Sinhala Lith + 0x0E50, // Thai + 0x0ED0, // Lao + 0x0F20, // Tibetan + 0x1040, // Myanmar + 0x1090, // Myanmar Shan + 0x17E0, // Khmer + 0x1810, // Mongolian + 0x1946, // Limbu + 0x19D0, // New Tai Lue + 0x1A80, // Tai Tham Hora + 0x1A90, // Tai Tham Tham + 0x1B50, // Balinese + 0x1BB0, // Sundanese + 0x1C40, // Lepcha + 0x1C50, // Ol Chiki + 0xA620, // Vai + 0xA8D0, // Saurashtra + 0xA900, // Kayah Li + 0xA9D0, // Javanese + 0xA9F0, // Myanmar Tai Laing + 0xAA50, // Cham + 0xABF0, // Meetei Mayek + 0xFF10, // Fullwidth +]; + +/// The value of `c` if it is a non-ASCII decimal digit that Java's `BigDecimal` accepts. /// -/// All other bytes (ASCII or other multi-byte sequences) are passed through -/// unchanged, so the output is valid UTF-8 whenever the input is. -fn normalize_fullwidth_digits(s: &str) -> String { - let bytes = s.as_bytes(); +/// `BigDecimal` scans UTF-16 code units and tests each with `Character.isDigit(char)`, so only +/// BMP digits count: a supplementary digit such as U+1D7CE MATHEMATICAL BOLD DIGIT ZERO is a +/// surrogate pair, neither half of which is a digit, and the parse fails. +#[inline] +fn non_ascii_digit_value(c: char) -> Option { + let cp = u32::from(c); + if cp > 0xFFFF { + return None; + } + let idx = BMP_NON_ASCII_DIGIT_ZEROS.partition_point(|&zero| u32::from(zero) <= cp); + let zero = u32::from(BMP_NON_ASCII_DIGIT_ZEROS[idx.checked_sub(1)?]); + let value = cp - zero; + (value < 10).then_some(value as u8) +} + +/// Replace every non-ASCII decimal digit that Java's `BigDecimal` accepts with its ASCII +/// equivalent, e.g. "١٢٣.٤٥" and "１２３.４５" both become "123.45". +/// +/// Spark parses decimals with `new java.math.BigDecimal(str)`, which reads each digit through +/// `Character.digit(c, 10)` in both the mantissa and the exponent, so digits from any script +/// can be mixed freely ("1٢e३" is 12000). All other characters are kept unchanged and are +/// rejected later by [`parse_decimal_str`]. +fn normalize_unicode_digits(s: &str) -> String { + // Every non-ASCII digit is at least two UTF-8 bytes and becomes one, so the output is + // never longer than the input. let mut out = Vec::with_capacity(s.len()); + let bytes = s.as_bytes(); let mut i = 0; while i < bytes.len() { - if i + 2 < bytes.len() - && bytes[i] == 0xEF - && bytes[i + 1] == 0xBC - && bytes[i + 2] >= 0x90 - && bytes[i + 2] <= 0x99 + let b = bytes[i]; + // ASCII (including mixed ASCII/fullwidth input) and fullwidth digits, encoded as + // [0xEF, 0xBC, 0x90 + n], are the common cases, so they skip the table lookup. + if b.is_ascii() { + out.push(b); + i += 1; + } else if b == 0xEF + && bytes.get(i + 1) == Some(&0xBC) + && matches!(bytes.get(i + 2), Some(0x90..=0x99)) { // e.g. 0x91 - 0x60 = 0x31 = b'1' out.push(bytes[i + 2] - 0x60); i += 3; } else { - out.push(bytes[i]); - i += 1; + // `i` is on a char boundary: every branch above advances by a whole char. + let mut rest = s[i..].chars(); + let c = rest.next().unwrap(); + let next = bytes.len() - rest.as_str().len(); + match non_ascii_digit_value(c) { + Some(d) => out.push(b'0' + d), + None => out.extend_from_slice(&bytes[i..next]), + } + i = next; } } - // SAFETY: we only replace valid 3-byte UTF-8 sequences [EF BC 9X] with a - // single ASCII byte; all other bytes are copied unchanged, preserving the - // UTF-8 invariant of the input. + // SAFETY: every branch either copies a whole UTF-8 encoded char unchanged or replaces it + // with a single ASCII byte, so the output is valid UTF-8 because the input is. unsafe { String::from_utf8_unchecked(out) } } @@ -585,13 +658,13 @@ fn parse_string_to_decimal(input_str: &str, precision: u8, scale: i8) -> SparkRe // producing NULL. let trimmed = trim_java_string(input_str); - // Normalize fullwidth digits to ASCII. Fast path skips the allocation for - // pure-ASCII strings, which is the common case. + // Normalize non-ASCII digits (Arabic-Indic, Devanagari, fullwidth, ...) to ASCII. Fast path + // skips the allocation for pure-ASCII strings, which is the common case. let normalized; let trimmed = if trimmed.is_ascii() { trimmed } else { - normalized = normalize_fullwidth_digits(trimmed); + normalized = normalize_unicode_digits(trimmed); normalized.as_str() }; @@ -2335,6 +2408,89 @@ mod tests { assert!(parse_string_to_decimal(&over_long, 38, 10).is_err()); } + #[test] + fn test_non_ascii_digit_value() { + for &zero in BMP_NON_ASCII_DIGIT_ZEROS.iter() { + for n in 0..10u8 { + let c = char::from_u32(u32::from(zero) + u32::from(n)).unwrap(); + assert_eq!(non_ascii_digit_value(c), Some(n), "{c:?}"); + } + // The code point just past each run is not a digit of that run. + let after = char::from_u32(u32::from(zero) + 10).unwrap(); + assert_eq!(non_ascii_digit_value(after), None, "{after:?}"); + } + for c in [ + '0', + '9', + 'a', + '.', + '\u{065F}', // just before Arabic-Indic zero + '\u{00B2}', // SUPERSCRIPT TWO (No) + '\u{2163}', // ROMAN NUMERAL FOUR (Nl) + '\u{2460}', // CIRCLED DIGIT ONE (No) + '\u{066B}', // ARABIC DECIMAL SEPARATOR + '\u{FFFF}', + '\u{1D7CE}', // MATHEMATICAL BOLD DIGIT ZERO: Nd, but outside the BMP + '\u{104A3}', // OSMANYA DIGIT THREE: Nd, but outside the BMP + ] { + assert_eq!(non_ascii_digit_value(c), None, "{c:?}"); + } + } + + #[test] + fn test_parse_string_to_decimal_unicode_digits() { + // Expected values come from `new java.math.BigDecimal(s)` on JDK 17, which is what + // Spark's Decimal.fromString / fromStringANSI call. + let ok = |s: &str, precision: u8, scale: i8| { + parse_string_to_decimal(s, precision, scale) + .unwrap() + .unwrap() + }; + assert_eq!(ok("٣", 10, 2), 300); // Arabic-Indic + assert_eq!(ok("۳", 10, 2), 300); // Extended Arabic-Indic + assert_eq!(ok("३", 10, 2), 300); // Devanagari + assert_eq!(ok("৩", 10, 2), 300); // Bengali + assert_eq!(ok("๓", 10, 2), 300); // Thai + assert_eq!(ok("߃", 10, 2), 300); // NKo + assert_eq!(ok("꧓", 10, 2), 300); // Javanese + assert_eq!(ok("１２３.４５", 10, 2), 12345); // Fullwidth + assert_eq!(ok("1٣", 10, 2), 1300); + assert_eq!(ok("-٣", 10, 2), -300); + assert_eq!(ok("+٣.٣٣", 10, 2), 333); + assert_eq!(ok(".٣", 10, 2), 30); + assert_eq!(ok("١٢٣.٤٥٥", 10, 2), 12346); // HALF_UP rounding + assert_eq!(ok("1e٣", 10, 2), 100000); // non-ASCII exponent digits + assert_eq!(ok("1E-٣", 10, 3), 1); + assert_eq!(ok("1e+٣", 10, 2), 100000); + assert_eq!(ok("1٢e३", 10, 0), 12000); // mixed scripts + assert_eq!(ok(" \t٣\n", 10, 2), 300); // ASCII trim still applies + assert_eq!(ok(&"٣".repeat(20), 38, 0), 33333333333333333333); + + // Rejected by BigDecimal: not Nd, or Nd outside the BMP. + for s in [ + "\u{1D7D0}", // MATHEMATICAL BOLD DIGIT TWO + "1\u{1D7D0}", + "1e\u{1D7D0}", + "\u{2163}", // ROMAN NUMERAL FOUR + "\u{00B2}", // SUPERSCRIPT TWO + "١\u{066B}٥", // ARABIC DECIMAL SEPARATOR is not '.' + "\u{00A0}٣", // non-ASCII whitespace is not trimmed + ] { + // The error reports the original string, not the normalized one. + match parse_string_to_decimal(s, 10, 2) { + Err(SparkError::CastInvalidValue { value, .. }) => assert_eq!(value, s), + other => panic!("{s:?}: {other:?}"), + } + } + + // Overflow is reported as such, not as an invalid input. + let err = parse_string_to_decimal("٣٣٣٣٣٣٣٣٣", 10, 2).unwrap_err(); + assert!( + matches!(err, SparkError::NumericValueOutOfRange { .. }), + "{err}" + ); + } + #[test] #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_cast_string_to_timestamp() { diff --git a/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits.sql b/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits.sql new file mode 100644 index 00000000000..f8d22a1d553 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits.sql @@ -0,0 +1,88 @@ +-- 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. + +-- cast(string as decimal) with non-ASCII digits, non-ANSI mode. Spark parses the trimmed string +-- with `new java.math.BigDecimal(s)`, which accepts every char for which +-- `Character.isDigit(char)` is true: all Unicode `Nd` digits in the Basic Multilingual Plane, +-- in the mantissa and in the exponent, mixed freely across scripts. Supplementary-plane digits +-- (e.g. U+1D7D0 MATHEMATICAL BOLD DIGIT TWO) are surrogate pairs and are rejected. +-- +-- The integral, floating-point and date casts of the same strings are ASCII-only in Spark +-- (UTF8String.toLong, Double.parseDouble, SparkDateTimeUtils.stringToDate) and must stay NULL. + +-- Config: spark.sql.ansi.enabled=false + +statement +CREATE TABLE test_cast_str_dec_unicode(id int, s string) USING parquet + +-- 1: U+0663 ARABIC-INDIC DIGIT THREE +-- 2: U+06F3 EXTENDED ARABIC-INDIC DIGIT THREE +-- 3: U+0969 DEVANAGARI DIGIT THREE +-- 4: U+09E9 BENGALI DIGIT THREE +-- 5: U+0E53 THAI DIGIT THREE +-- 6: U+07C3 NKO DIGIT THREE +-- 7: U+A9D3 JAVANESE DIGIT THREE +-- 8: fullwidth digits +-- 9-17: signs, '.', HALF_UP rounding, exponent digits, mixed scripts, ASCII trim +-- 18: 20 Arabic-Indic digits, too many for the cast targets below +-- 19-25: rejected by BigDecimal: supplementary-plane Nd digits, ROMAN NUMERAL FOUR (Nl), +-- SUPERSCRIPT TWO (No), ARABIC DECIMAL SEPARATOR, leading NBSP (not trimmed) +statement +INSERT INTO test_cast_str_dec_unicode VALUES + (1, '٣'), + (2, '۳'), + (3, '३'), + (4, '৩'), + (5, '๓'), + (6, '߃'), + (7, '꧓'), + (8, '１２３.４５'), + (9, '1٣'), + (10, '-٣'), + (11, '+٣.٣٣'), + (12, '.٣'), + (13, '١٢٣.٤٥٥'), + (14, '1e٣'), + (15, '1E-٣'), + (16, '1٢e३'), + (17, concat(' ٣', chr(9))), + (18, '٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣٣'), + (19, '𝟐'), + (20, '1𝟐'), + (21, '1e𝟐'), + (22, 'Ⅳ'), + (23, '²'), + (24, '١٫٥'), + (25, concat(chr(160), '٣')), + (26, NULL) + +query +SELECT id, cast(s as decimal(10,2)), cast(s as decimal(38,10)), cast(s as decimal(5,0)) +FROM test_cast_str_dec_unicode ORDER BY id + +query +SELECT id, try_cast(s as decimal(10,2)), try_cast(s as decimal(38,0)) +FROM test_cast_str_dec_unicode ORDER BY id + +-- Spark's integral, floating-point and date parsers only accept ASCII digits +query +SELECT id, cast(s as int), cast(s as bigint), cast(s as double), cast(s as date) +FROM test_cast_str_dec_unicode ORDER BY id + +-- literal arguments +query +SELECT cast('٣' as decimal(10,2)), cast('१२३.४५' as decimal(10,2)), cast('𝟐' as decimal(10,2)) diff --git a/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits_ansi.sql b/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits_ansi.sql new file mode 100644 index 00000000000..f546893dfab --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/cast/cast_string_to_decimal_unicode_digits_ansi.sql @@ -0,0 +1,76 @@ +-- 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. + +-- cast(string as decimal) with non-ASCII digits, ANSI mode. Spark parses the trimmed string +-- with `new java.math.BigDecimal(s)`, which accepts every Unicode `Nd` digit in the Basic +-- Multilingual Plane, so these must cast to a value rather than raise CAST_INVALID_INPUT. +-- Supplementary-plane digits are rejected by BigDecimal and must still raise. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_cast_str_dec_unicode_ansi(id int, s string) USING parquet + +-- Arabic-Indic, Devanagari, Thai, fullwidth, signs, '.', rounding, exponent digits, mixed +-- scripts and ASCII trim; every value fits decimal(10,2) +statement +INSERT INTO test_cast_str_dec_unicode_ansi VALUES + (1, '٣'), + (2, '३'), + (3, '๓'), + (4, '１２３.４５'), + (5, '-٣'), + (6, '+٣.٣٣'), + (7, '١٢٣.٤٥٥'), + (8, '1e٣'), + (9, '1E-٣'), + (10, '1٢e३'), + (11, concat(' ٣', chr(9))), + (12, NULL) + +statement +CREATE TABLE test_cast_str_dec_unicode_ansi_bad(s string) USING parquet + +-- U+1D7D0 MATHEMATICAL BOLD DIGIT TWO is Nd but outside the BMP +statement +INSERT INTO test_cast_str_dec_unicode_ansi_bad VALUES ('𝟐') + +statement +CREATE TABLE test_cast_str_dec_unicode_ansi_big(s string) USING parquet + +statement +INSERT INTO test_cast_str_dec_unicode_ansi_big VALUES ('٣٣٣٣٣٣٣٣٣') + +query +SELECT id, cast(s as decimal(10,2)), cast(s as decimal(38,10)) +FROM test_cast_str_dec_unicode_ansi ORDER BY id + +query +SELECT id, try_cast(s as decimal(10,2)) FROM test_cast_str_dec_unicode_ansi ORDER BY id + +query expect_error(CAST_INVALID_INPUT) +SELECT cast(s as decimal(10,2)) FROM test_cast_str_dec_unicode_ansi_bad + +-- the value is valid but does not fit decimal(10,2) +query expect_error(NUMERIC_VALUE_OUT_OF_RANGE) +SELECT cast(s as decimal(10,2)) FROM test_cast_str_dec_unicode_ansi_big + +-- try_cast suppresses both errors +query +SELECT try_cast(s as decimal(10,2)) FROM test_cast_str_dec_unicode_ansi_bad +UNION ALL +SELECT try_cast(s as decimal(10,2)) FROM test_cast_str_dec_unicode_ansi_big From cf6673808dc92d0be7ed331431fe9a2e97067a7d Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Sun, 11 Oct 2026 22:31:50 +0000 Subject: [PATCH 70/89] fix: preserve CASE WHEN struct field names from the first THEN branch (#6827) * fix: preserve CASE WHEN struct field names from the first THEN branch * test: keep CASE WHEN tests inline for focused review * test: cover COALESCE struct field names * test: cover nested and null CASE WHEN struct branches --- .../src/conditional_funcs/case_when.rs | 104 ++++++++++-------- .../conditional/case_when_field_names.sql | 67 +++++++++++ .../conditional/coalesce_field_names.sql | 36 ++++++ 3 files changed, 163 insertions(+), 44 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 74df26c977c..192e6046f61 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -31,7 +31,6 @@ use datafusion::common::{ exec_err, internal_datafusion_err, internal_err, DataFusionError, Result, ScalarValue, }; use datafusion::logical_expr::type_coercion::binary::type_union_coercion; -use datafusion::logical_expr::type_coercion::other::get_coerce_type_for_case_expression; use datafusion::logical_expr::{ColumnarValue, Operator}; use datafusion::physical_expr::expressions::{ BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, @@ -60,8 +59,17 @@ pub fn create_case_when( .transpose()? .unwrap_or(DataType::Null); - let Some(coerce_type) = get_coerce_type_for_case_expression(&then_types, Some(&else_type)) - else { + // Spark merges THEN types from left to right, with ELSE last, retaining the first THEN's + // struct field names. DataFusion's CASE coercion starts with ELSE and can match fields by name. + let coerce_type = then_types.split_first().and_then(|(first, remaining)| { + remaining + .iter() + .chain([&else_type]) + .try_fold(first.clone(), |common, next| { + merge_branch_types(&common, next) + }) + }); + let Some(coerce_type) = coerce_type else { return Ok(Arc::new(CaseWhenExpr::try_new(when_then, else_expr)?)); }; let when_then = when_then @@ -86,7 +94,7 @@ pub fn create_if_expr( ) -> Result> { let true_type = true_expr.data_type(input_schema)?; let false_type = false_expr.data_type(input_schema)?; - let Some(common_type) = if_common_type(&true_type, &false_type) else { + let Some(common_type) = merge_branch_types(&true_type, &false_type) else { return Ok(Arc::new(IfExpr::new(if_expr, true_expr, false_expr))); }; Ok(Arc::new(IfExpr::new( @@ -96,34 +104,30 @@ pub fn create_if_expr( ))) } -/// Reconciles Spark IF branches positionally, retaining THEN names and merging nullability. +/// Merges conditional branch types positionally, retaining left names and merging nullability. /// Spark has already coerced the branches to the same SQL type. DataFusion's struct union may /// instead match by name, pairing different positions when names differ only in case. -fn if_common_type(then_type: &DataType, else_type: &DataType) -> Option { +fn merge_branch_types(left_type: &DataType, right_type: &DataType) -> Option { use arrow::datatypes::FieldRef; - fn field(then_field: &FieldRef, else_field: &FieldRef) -> Option { + fn field(left: &FieldRef, right: &FieldRef) -> Option { Some(Arc::new( - then_field - .as_ref() + left.as_ref() .clone() - .with_data_type(if_common_type( - then_field.data_type(), - else_field.data_type(), - )?) - .with_nullable(then_field.is_nullable() || else_field.is_nullable()), + .with_data_type(merge_branch_types(left.data_type(), right.data_type())?) + .with_nullable(left.is_nullable() || right.is_nullable()), )) } - match (then_type, else_type) { - (DataType::Struct(then_fields), DataType::Struct(else_fields)) => { - if then_fields.len() != else_fields.len() { + match (left_type, right_type) { + (DataType::Struct(left_fields), DataType::Struct(right_fields)) => { + if left_fields.len() != right_fields.len() { return None; } Some(DataType::Struct( - then_fields + left_fields .iter() - .zip(else_fields) + .zip(right_fields) .map(|(t, e)| field(t, e)) .collect::>>()? .into(), @@ -133,7 +137,7 @@ fn if_common_type(then_type: &DataType, else_type: &DataType) -> Option { Some(DataType::Map(field(t, e)?, *t_sorted && *e_sorted)) } - _ => type_union_coercion(then_type, else_type), + _ => type_union_coercion(left_type, right_type), } } @@ -1338,11 +1342,17 @@ mod tests { /// CASE branches that share a Spark timestamp type but carry different Arrow timezone labels /// are reconciled by relabelling them, which used to panic because the cast had no timezone. + /// Like IF, CASE retains the first THEN's label when present. #[test] fn case_reconciles_timestamp_timezone_labels() { let utc = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); for then_label in [Some("Etc/UTC"), None] { let then_type = DataType::Timestamp(TimeUnit::Microsecond, then_label.map(Into::into)); + let expected_type = if then_label.is_some() { + then_type.clone() + } else { + utc.clone() + }; let schema = Schema::new(vec![ Field::new("b", DataType::Boolean, true), Field::new("t", then_type, true), @@ -1362,16 +1372,17 @@ mod tests { let c = |name: &str| col(name, &schema).unwrap(); let case = create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap(); let result = case.evaluate(&batch).unwrap().into_array(2).unwrap(); - assert_eq!(result.data_type(), &utc, "{then_label:?}"); + assert_eq!(case.data_type(&schema).unwrap(), expected_type); + assert_eq!(result.data_type(), &expected_type, "{then_label:?}"); let values = result.as_primitive::(); assert_eq!(values.values().to_vec(), vec![1, 20], "{then_label:?}"); } } - /// Evaluates an IF over `batch`, whose first row takes the THEN branch and second row the ELSE - /// branch, and over each row alone, and checks that every result has the type the IF reports. + /// Evaluates a conditional over `batch`, whose first row takes THEN and second row ELSE, + /// and over each row alone, checking that every result has the type the expression reports. /// Returns the result for the whole batch. - fn evaluate_if(expr: &Arc, batch: &RecordBatch) -> ArrayRef { + fn evaluate_conditional(expr: &Arc, batch: &RecordBatch) -> ArrayRef { let data_type = expr.data_type(&batch.schema()).unwrap(); // Both branches, only THEN, and only ELSE, whose array `IfExpr` returns unchanged for (offset, len) in [(0, 2), (0, 1), (1, 1)] { @@ -1389,7 +1400,7 @@ mod tests { /// Spark matches struct fields by position even when their case-distinct names are reordered. /// Name-based union pairs the first INT with the second DOUBLE and silently widens it. #[test] - fn if_reconciles_case_variant_fields_positionally() { + fn conditional_branches_reconcile_case_variant_fields_positionally() { use arrow::array::Float64Array; let then_fields: arrow::datatypes::Fields = vec![ @@ -1436,12 +1447,6 @@ mod tests { ) .unwrap(); let c = |name: &str| col(name, &schema).unwrap(); - let expr = create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(); - assert_eq!( - expr.data_type(&schema).unwrap(), - DataType::Struct(expected_fields.clone()) - ); - let result = evaluate_if(&expr, &batch); let expected = StructArray::new( expected_fields, vec![ @@ -1450,7 +1455,14 @@ mod tests { ], None, ); - assert_eq!(result.as_ref(), &expected); + for expr in [ + create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(), + create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap(), + ] { + assert_eq!(expr.data_type(&schema).unwrap(), *expected.data_type()); + let result = evaluate_conditional(&expr, &batch); + assert_eq!(result.as_ref(), &expected, "{expr}"); + } } /// IF branches that share a Spark timestamp type but carry different Arrow timezone labels are @@ -1493,17 +1505,17 @@ mod tests { timestamp(label), "{labels}" ); - let result = evaluate_if(&expr, &batch); + let result = evaluate_conditional(&expr, &batch); let values = result.as_primitive::(); assert_eq!(values.values().to_vec(), vec![1, 20], "{labels}"); } } - /// IF branches whose struct field differs only in whether it can be NULL, or in the case of its + /// Conditional branches whose struct field differs only in nullability or in the case of its /// name, which Spark adds no cast for. The result's field can be NULL if either branch's can, - /// and has the THEN branch's name, as in Spark's `If.dataType`. + /// and has the THEN branch's name, as in Spark's `If.dataType` and `CaseWhen.dataType`. #[test] - fn if_reconciles_struct_field_nullability_and_names() { + fn conditional_branches_reconcile_struct_field_nullability_and_names() { let field = |name: &str, nullable: bool| Arc::new(Field::new(name, DataType::Int32, nullable)); let struct_type = |field: &FieldRef| DataType::Struct(vec![Arc::clone(field)].into()); @@ -1542,16 +1554,20 @@ mod tests { ) .unwrap(); let c = |name: &str| col(name, &schema).unwrap(); - let expr = create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(); let fields = format!("{then_field:?} {else_field:?}"); - assert_eq!( - expr.data_type(&schema).unwrap(), - struct_type(&result_field), - "{fields}" - ); - let result = evaluate_if(&expr, &batch); let expected = column(&result_field, vec![Some(1), else_values[1]]); - assert_eq!(result.as_ref(), expected.as_ref(), "{fields}"); + for expr in [ + create_if_expr(c("b"), c("t"), c("e"), &schema).unwrap(), + create_case_when(vec![(c("b"), c("t"))], Some(c("e")), &schema).unwrap(), + ] { + assert_eq!( + expr.data_type(&schema).unwrap(), + struct_type(&result_field), + "{fields} {expr}" + ); + let result = evaluate_conditional(&expr, &batch); + assert_eq!(result.as_ref(), expected.as_ref(), "{fields} {expr}"); + } } } diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql b/spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql new file mode 100644 index 00000000000..f8106cb2a25 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/conditional/case_when_field_names.sql @@ -0,0 +1,67 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Spark preserves the first THEN's struct field names when branches differ only in case. +-- Native to_json exposes Arrow field names that the row comparison would otherwise ignore. +-- https://github.com/apache/datafusion-comet/issues/6482 +-- Config: spark.sql.caseSensitive=false +-- Config: spark.comet.expression.StructsToJson.allowIncompatible=true + +statement +CREATE TABLE test_case_field_names(p boolean, q boolean, i int, d double) USING parquet + +statement +INSERT INTO test_case_field_names VALUES + (true, false, 1, 1.5), (false, true, 2, 2.5), + (NULL, false, 3, 3.5), (false, NULL, NULL, NULL) + +-- The result keeps the THEN name for both selected branches, in either name order. +query expect_native(casewhen) +SELECT + to_json(CASE WHEN p THEN named_struct('x', i) ELSE named_struct('X', i) END), + to_json(CASE WHEN p THEN named_struct('X', i) ELSE named_struct('x', i) END) +FROM test_case_field_names + +-- Preserve the first THEN's names at both levels of a nested struct, in either branch order. +query expect_native(casewhen) +SELECT + to_json(CASE WHEN p THEN named_struct('a', named_struct('x', i)) + ELSE named_struct('A', named_struct('X', i)) END), + to_json(CASE WHEN p THEN named_struct('A', named_struct('X', i)) + ELSE named_struct('a', named_struct('x', i)) END) +FROM test_case_field_names + +-- Later THEN branches and an omitted ELSE must also preserve the first THEN's name. +query expect_native(casewhen) +SELECT + to_json(CASE WHEN p THEN named_struct('x', i) + WHEN q THEN named_struct('X', i + 10) ELSE named_struct('X', i + 20) END), + to_json(CASE WHEN p THEN named_struct('x', i) + WHEN q THEN named_struct('X', i + 10) END) +FROM test_case_field_names + +-- A typed NULL in the first THEN still determines the result's struct field name. +query expect_native(casewhen) +SELECT to_json(CASE WHEN p THEN CAST(NULL AS STRUCT) + WHEN q THEN named_struct('x', i) ELSE named_struct('X', i) END) +FROM test_case_field_names + +-- Fields align by position. Matching these names would incorrectly combine INT with DOUBLE. +query expect_native(casewhen) +SELECT to_json(CASE WHEN p THEN named_struct('x', i, 'X', CAST(5.5 AS DOUBLE)) + ELSE named_struct('X', 0, 'x', d) END) +FROM test_case_field_names diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql new file mode 100644 index 00000000000..45596f3cc26 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/conditional/coalesce_field_names.sql @@ -0,0 +1,36 @@ +-- 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. + +-- COALESCE lowers to CASE WHEN with its last argument as ELSE, but must retain the first +-- argument's struct field names. Native to_json exposes names that row comparison ignores. +-- https://github.com/apache/datafusion-comet/issues/6482 +-- Config: spark.sql.caseSensitive=false +-- Config: spark.comet.expression.StructsToJson.allowIncompatible=true + +statement +CREATE TABLE test_coalesce_field_names(p boolean, q boolean, i int) USING parquet + +statement +INSERT INTO test_coalesce_field_names VALUES + (true, false, 1), (false, true, 2), (NULL, false, 3), (false, NULL, NULL) + +query expect_native(coalesce) +SELECT to_json(coalesce( + CASE WHEN p THEN named_struct('x', i) END, + CASE WHEN q THEN named_struct('X', i + 100) END, + named_struct('X', 7))) +FROM test_coalesce_field_names From 71ae2284841dbcab0efc2d6523bbfcbe8ddde8df Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:17:36 +0000 Subject: [PATCH 71/89] test: migrate required-child hash regressions to SQL fixtures (#6864) --- .../expressions/hash/hash_required_child.sql | 49 ++++++++++ .../comet/CometHashExpressionSuite.scala | 97 +------------------ 2 files changed, 50 insertions(+), 96 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/hash/hash_required_child.sql diff --git a/spark/src/test/resources/sql-tests/expressions/hash/hash_required_child.sql b/spark/src/test/resources/sql-tests/expressions/hash/hash_required_child.sql new file mode 100644 index 00000000000..4f95e40cfa4 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/hash/hash_required_child.sql @@ -0,0 +1,49 @@ +-- 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. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- CTAS preserves REQUIRED children. CREATE + INSERT can rebuild them as nullable, +-- which would hide #5754: a child can contain a value underneath a NULL parent. +-- Write with Spark so these tests exercise the native reader and hash kernels. +statement +SET spark.comet.enabled=false + +statement +CREATE TABLE hash_required_child USING parquet AS SELECT /*+ COALESCE(1) */ c FROM VALUES (named_struct('a', 1)), (NULL), (named_struct('a', 3)), (NULL) AS v(c) + +-- The parent's null mask must also reach grandchildren through a REQUIRED struct. +statement +CREATE TABLE hash_required_nested_child USING parquet AS SELECT /*+ COALESCE(1) */ c FROM VALUES (named_struct('b', named_struct('x', 1))), (NULL), (named_struct('b', named_struct('x', 3))), (NULL) AS v(c) + +statement +SET spark.comet.enabled=true + +query expect_native(hash,xxhash64) +SELECT hash(c), xxhash64(c) FROM hash_required_child ORDER BY 1, 2 + +query expect_native(hash,xxhash64) +SELECT hash(c), xxhash64(c) FROM hash_required_nested_child ORDER BY 1, 2 + +-- A valid array element wrapping the NULL struct preserves its hidden child values. +-- A NULL array element itself is rebuilt and would not exercise the same path. +query expect_native(hash,xxhash64) +SELECT hash(array(named_struct('tag', 1, 'b', c))), xxhash64(array(named_struct('tag', 1, 'b', c))) FROM hash_required_child ORDER BY 1, 2 + +-- Two elements exercise chaining the second element's hash onto the first. +query expect_native(hash,xxhash64) +SELECT hash(array(named_struct('tag', 1, 'b', c), named_struct('tag', 2, 'b', c))), xxhash64(array(named_struct('tag', 1, 'b', c), named_struct('tag', 2, 'b', c))) FROM hash_required_child ORDER BY 1, 2 diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index bbf6c4289ad..c20e5899b37 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -21,9 +21,8 @@ package org.apache.comet import scala.util.Random -import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} @@ -259,100 +258,6 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } - test("hash - null struct with a required child field") { - // `c` is nullable but its child is REQUIRED, so Spark writes - // `optional group c { required int32 a; }`. On read the child leaf has nowhere to record a - // null of its own, so its buffer holds a value at exactly the rows where the struct is null. - // That is the shape where the parent's null mask has to reach the children; a struct built in - // the plan, or one whose child is also nullable, has null children there and hides the bug. - withTempPath { dir => - val schema = StructType( - Seq( - StructField( - "c", - StructType(Seq(StructField("a", IntegerType, nullable = false))), - nullable = true))) - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val rows = Seq(Row(Row(1)), Row(null), Row(Row(3)), Row(null)) - spark - .createDataFrame(spark.sparkContext.parallelize(rows), schema) - .coalesce(1) - .write - .parquet(dir.toString) - } - spark.read.parquet(dir.toString).createOrReplaceTempView("null_struct_t") - checkSparkAnswerAndOperator("SELECT hash(c), xxhash64(c) FROM null_struct_t ORDER BY 1, 2") - } - } - - test("hash - null struct whose child is itself a struct") { - // The union has to recurse: the outer struct's nulls reach the inner struct, whose own - // children are required and so carry values under the null. - withTempPath { dir => - val inner = StructType(Seq(StructField("x", IntegerType, nullable = false))) - val schema = StructType( - Seq( - StructField( - "c", - StructType(Seq(StructField("b", inner, nullable = false))), - nullable = true))) - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val rows = Seq(Row(Row(Row(1))), Row(null), Row(Row(Row(3))), Row(null)) - spark - .createDataFrame(spark.sparkContext.parallelize(rows), schema) - .coalesce(1) - .write - .parquet(dir.toString) - } - spark.read.parquet(dir.toString).createOrReplaceTempView("null_nested_struct_t") - checkSparkAnswerAndOperator( - "SELECT hash(c), xxhash64(c) FROM null_nested_struct_t ORDER BY 1, 2") - } - } - - test("hash - list element wrapping a null struct with a required child field") { - // The per-element path in `hash_list_array!`, which #5567 made usable as a shuffle key. - // - // Wrapping the struct rather than using it as the element directly is what makes this - // reachable: `array(named_struct('tag', 1, 'b', c))` produces an element that is itself - // valid, so it is copied rather than rebuilt, and the null `c` inside it keeps the child - // values Parquet wrote under it. Using `c` as the element directly does not reproduce, - // because a null element is rebuilt on the way in and the hidden values go with it. - // - // So this covers a null struct nested inside a valid element. The unit test in murmur3.rs - // covers the complementary shape, where the element itself is the null struct. - withTempPath { dir => - val schema = StructType( - Seq( - StructField( - "c", - StructType(Seq(StructField("a", IntegerType, nullable = false))), - nullable = true))) - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val rows = Seq(Row(Row(1)), Row(null), Row(Row(3)), Row(null)) - spark - .createDataFrame(spark.sparkContext.parallelize(rows), schema) - .coalesce(1) - .write - .parquet(dir.toString) - } - spark.read.parquet(dir.toString).createOrReplaceTempView("wrapped_null_struct_t") - - checkSparkAnswerAndOperator(""" - SELECT - hash(array(named_struct('tag', 1, 'b', c))), - xxhash64(array(named_struct('tag', 1, 'b', c))) - FROM wrapped_null_struct_t ORDER BY 1, 2""") - - // Two elements, so the hash of the second chains onto the first. - checkSparkAnswerAndOperator(""" - SELECT - hash(array(named_struct('tag', 1, 'b', c), named_struct('tag', 2, 'b', c))), - xxhash64(array(named_struct('tag', 1, 'b', c), named_struct('tag', 2, 'b', c))) - FROM wrapped_null_struct_t ORDER BY 1, 2""") - } - } - test("hash - struct with array field") { withTable("t") { sql("CREATE TABLE t(c STRUCT>) USING parquet") From 46f63116febd8e3269e1ce2e1c91ca7669718f75 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 11:18:00 +0000 Subject: [PATCH 72/89] fix: keep one row per key in a native PartialMerge that runs short of memory (#6867) * fix: keep one row per key in a native PartialMerge that runs short of memory Comet ran Spark's PartialMerge aggregate as DataFusion Partial, which emits the groups it holds when it runs short of memory and starts again. In Spark's single-distinct plan the stage above the PartialMerge counts its rows, so count(DISTINCT) counted every copy of a key. Run a PartialMerge as DataFusion Single instead, which spills. MergeAsPartial evaluates to a struct of its state, and a projection unpacks it into the state columns that the Partial plan emitted. SpillReplayPool also covers SingleHashAggregateStream, whose spill replay starves the same way as the final aggregates' do (#6254). Closes #6857. * Emit PartialMerge states a column at a time for row-only aggregates MergeAsPartial now wraps an inner aggregate that has no GroupsAccumulator (first, last and any_value without an ordering) in a GroupsAccumulatorAdapter itself. Under Single mode, DataFusion otherwise wrapped the wrapper's row accumulators and evaluated every group to a one-row struct. Also write the planner's two output wrappers the same way, check the planner test's output with assert_batches_sorted_eq and cover first there, and drop two settings that the memory-pressure tests don't read. --- native/core/src/execution/jni_api.rs | 2 +- .../execution/memory_pools/spill_replay.rs | 19 +- native/core/src/execution/merge_as_partial.rs | 179 ++++++++++++--- native/core/src/execution/planner.rs | 206 ++++++++++++++++-- .../comet/exec/CometAggregateSuite.scala | 117 +++++++++- 5 files changed, 453 insertions(+), 70 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 0812b9f7087..a2646dd8daf 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -963,7 +963,7 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato if count.children.len() == 1) }) } - // PartialMerge is represented as native Partial, but consumes states, not rows. + // PartialMerge consumes states, not rows. _ => false, }, _ => true, diff --git a/native/core/src/execution/memory_pools/spill_replay.rs b/native/core/src/execution/memory_pools/spill_replay.rs index e7d947e38c1..086e8d87374 100644 --- a/native/core/src/execution/memory_pools/spill_replay.rs +++ b/native/core/src/execution/memory_pools/spill_replay.rs @@ -21,7 +21,8 @@ //! replays them through an `OrderedFinalAggregateStream` that has no way to spill, so a refused //! memory request there fails the task. `FinalHashAggregateStream` does this, and so does //! `OrderedFinalAggregateStream` itself, which DataFusion uses when the input is sorted on some of -//! the grouping keys. The merge reserves read buffers for as many spill files as fit, and those +//! the grouping keys. So does `SingleHashAggregateStream`, which Comet runs for a Spark +//! PartialMerge (#6857). The merge reserves read buffers for as many spill files as fit, and those //! buffers belong to the same consumer, so the replay often finds the consumer's share already //! taken. //! @@ -31,9 +32,9 @@ //! //! This relies on the following in DataFusion 55.1, which a DataFusion upgrade has to re-check: //! -//! - The two aggregates name their consumers `FinalHashAggregateStream[{partition}]` and -//! `OrderedFinalAggregateStream[{partition}]`, and their merge and replay reserve through sibling -//! reservations of that consumer. +//! - The three aggregates name their consumers `FinalHashAggregateStream[{partition}]`, +//! `OrderedFinalAggregateStream[{partition}]` and `SingleHashAggregateStream[{partition}]`, and +//! their merge and replay reserve through sibling reservations of that consumer. //! - A final aggregate grows one reservation while another of its reservations holds memory only //! during the replay, while the merge holds its read buffers. Before that, the aggregate's table //! is its only reservation holding memory, so a refusal stands and makes it spill. The merge @@ -45,7 +46,8 @@ //! //! Remove this, as #6583 describes, once Comet's DataFusion has apache/datafusion#25383, which //! leaves the replay room when the merge picks its files. The #6254 tests in `CometAggregateSuite` -//! fail without this on DataFusion 55.1, so they show whether the replay still needs it. +//! fail without this on DataFusion 55.1, and its PartialMerge read-back test for #6857 fails +//! without it on 55.2, so they show whether the replay still needs it. use std::collections::HashMap; use std::fmt; @@ -174,11 +176,12 @@ pub(super) fn unwrap_spill_replay(pool: &Arc) -> Option<&Arc bool { let name = consumer.name(); name.starts_with("FinalHashAggregateStream[") || name.starts_with("OrderedFinalAggregateStream[") + || name.starts_with("SingleHashAggregateStream[") } #[cfg(test)] @@ -353,14 +356,14 @@ mod tests { for name in [ "FinalHashAggregateStream[3]", "OrderedFinalAggregateStream[0]", + "SingleHashAggregateStream[0]", ] { assert!(is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); } - // Partial aggregates emit early instead of spilling, and Comet never plans single ones. + // Partial aggregates emit early instead of spilling. for name in [ "PartialHashAggregateStream[0]", "OrderedPartialAggregateStream[0]", - "SingleHashAggregateStream[0]", "ExternalSorterMerge[0]", ] { assert!(!is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs index 7f483a761c5..4dc4c4809f7 100644 --- a/native/core/src/execution/merge_as_partial.rs +++ b/native/core/src/execution/merge_as_partial.rs @@ -20,18 +20,21 @@ //! Spark's PartialMerge mode merges intermediate state buffers and outputs intermediate //! state (not final values). DataFusion has no equivalent mode — `Partial` calls //! `update_batch` and outputs state, while `Final` calls `merge_batch` and outputs -//! evaluated results. +//! evaluated results. `PartialReduce` merges state and outputs state, but from DataFusion 56 +//! it emits a group more than once when it runs short of memory, as `Partial` does. //! -//! This wrapper bridges the gap: it operates under DataFusion's `Partial` mode (which -//! outputs state) but redirects `update_batch` calls to `merge_batch`, giving merge -//! semantics with state output. +//! This wrapper bridges the gap: it redirects `update_batch` calls to `merge_batch`, and it +//! evaluates to its state, packed into a struct. Under DataFusion's `Partial` mode, which +//! outputs state, that gives merge semantics with state output. Under `Single` mode, which +//! spills rather than emit a group twice, the planner unpacks the struct into the state columns. use std::fmt::Debug; use std::hash::{Hash, Hasher}; +use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray}; -use arrow::datatypes::{DataType, Field, FieldRef}; -use datafusion::common::Result; +use arrow::array::{ArrayRef, BooleanArray, StructArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema}; +use datafusion::common::{internal_err, Result}; use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::function::StateFieldsArgs; use datafusion::logical_expr::{ @@ -39,16 +42,23 @@ use datafusion::logical_expr::{ Signature, Volatility, }; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; +use datafusion::physical_expr::{ + expressions::Column, GroupsAccumulatorAdapter, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion::physical_plan::{projection::ProjectionExec, ExecutionPlan}; use datafusion::scalar::ScalarValue; +use datafusion_comet_spark_expr::GetStructField; use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; -/// An AggregateUDF wrapper that gives merge semantics in Partial mode. +/// An AggregateUDF wrapper that merges state and outputs state in Partial or Single mode. /// /// When DataFusion runs an AggregateExec in Partial mode, it calls `update_batch` /// on each accumulator and outputs `state()`. This wrapper intercepts `update_batch` /// and redirects it to `merge_batch` on the inner accumulator, effectively -/// implementing PartialMerge: merge inputs, output state. +/// implementing PartialMerge: merge inputs, output state. In Single mode, DataFusion +/// outputs `evaluate()` instead, which returns the inner accumulator's state as a struct +/// of the state fields. /// /// We store the inner AggregateUDF (not the AggregateFunctionExpr) to avoid keeping /// references to UnboundColumn expressions that would panic if evaluated. @@ -56,10 +66,12 @@ use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; pub struct MergeAsPartialUDF { /// The inner aggregate UDF, cloned from the original expression. inner_udf: AggregateUDF, - /// Pre-computed return type from the original expression. - return_type: DataType, - /// Pre-computed state fields from the original expression. - cached_state_fields: Vec, + /// The original expression's return field, which the inner accumulator is built with, since + /// the wrapper returns its state instead. + inner_return_field: FieldRef, + /// Pre-computed state fields from the original expression, which are also the fields of + /// the struct the wrapper evaluates to. + cached_state_fields: Fields, /// Cached signature that accepts state field types. signature: Signature, /// Override argument fields used only when constructing the inner accumulator. @@ -87,8 +99,7 @@ impl Hash for MergeAsPartialUDF { impl MergeAsPartialUDF { pub fn new(inner_expr: &AggregateFunctionExpr) -> Result { let name = format!("merge_as_partial_{}", inner_expr.name()); - let return_type = inner_expr.field().data_type().clone(); - let cached_state_fields = inner_expr.state_fields()?; + let cached_state_fields = Fields::from(inner_expr.state_fields()?); let accumulator_expr_fields = Self::accumulator_expr_fields(inner_expr, &cached_state_fields); let state_decoder = PartialMergeStateDecoder::try_new(inner_expr, &cached_state_fields)?; @@ -99,7 +110,7 @@ impl MergeAsPartialUDF { Ok(Self { inner_udf: inner_expr.fun().clone(), - return_type, + inner_return_field: inner_expr.field(), cached_state_fields, signature, accumulator_expr_fields, @@ -136,6 +147,7 @@ impl MergeAsPartialUDF { .as_deref() .unwrap_or(args.expr_fields); AccumulatorArgs { + return_field: Arc::clone(&self.inner_return_field), expr_fields, ..args } @@ -153,14 +165,20 @@ impl AggregateUDFImpl for MergeAsPartialUDF { fn return_type(&self, _arg_types: &[DataType]) -> Result { // In Partial mode, return_type isn't used for output schema (state_fields is). - // Return the inner function's return type for consistency. - Ok(self.return_type.clone()) + // In Single mode, the wrapper evaluates to its state. + Ok(DataType::Struct(self.cached_state_fields.clone())) + } + + fn is_nullable(&self) -> bool { + // The state struct itself is never null, so each state field unpacked from it keeps its + // own nullability. + false } fn state_fields(&self, _args: StateFieldsArgs) -> Result> { // Cached at construction: state schema depends on the inner aggregate's // return type, not on StateFieldsArgs. - Ok(self.cached_state_fields.clone()) + Ok(self.cached_state_fields.to_vec()) } fn accumulator(&self, args: AccumulatorArgs) -> Result> { @@ -170,24 +188,34 @@ impl AggregateUDFImpl for MergeAsPartialUDF { Ok(Box::new(MergeAsPartialAccumulator { inner: inner_acc, state_decoder: self.state_decoder.clone(), + state_fields: self.cached_state_fields.clone(), })) } - fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - self.inner_udf - .groups_accumulator_supported(self.inner_accumulator_args(args)) + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + // An inner aggregate without a GroupsAccumulator gets an adapter below, so that Single + // mode emits its states a column at a time rather than evaluating each group to a struct. + true } fn create_groups_accumulator( &self, args: AccumulatorArgs, ) -> Result> { - let inner_acc = self + let inner_args = self.inner_accumulator_args(args); + let inner_acc: Box = if self .inner_udf - .create_groups_accumulator(self.inner_accumulator_args(args))?; + .groups_accumulator_supported(inner_args.clone()) + { + self.inner_udf.create_groups_accumulator(inner_args)? + } else { + let factory = InnerAccumulatorFactory::new(&self.inner_udf, &inner_args); + Box::new(GroupsAccumulatorAdapter::new(move || factory.create())) + }; Ok(Box::new(MergeAsPartialGroupsAccumulator { inner: inner_acc, state_decoder: self.state_decoder.clone(), + state_fields: self.cached_state_fields.clone(), })) } @@ -204,10 +232,88 @@ impl AggregateUDFImpl for MergeAsPartialUDF { } } -/// Accumulator wrapper that redirects update_batch to merge_batch. +/// Owned copies of the arguments that build an inner row accumulator, which a +/// `GroupsAccumulatorAdapter` does for each new group, after `AccumulatorArgs` has gone. +struct InnerAccumulatorFactory { + udf: AggregateUDF, + return_field: FieldRef, + schema: Schema, + ignore_nulls: bool, + order_bys: Vec, + is_reversed: bool, + name: String, + is_distinct: bool, + exprs: Vec>, + expr_fields: Vec, +} + +impl InnerAccumulatorFactory { + fn new(udf: &AggregateUDF, args: &AccumulatorArgs) -> Self { + Self { + udf: udf.clone(), + return_field: Arc::clone(&args.return_field), + schema: args.schema.clone(), + ignore_nulls: args.ignore_nulls, + order_bys: args.order_bys.to_vec(), + is_reversed: args.is_reversed, + name: args.name.to_string(), + is_distinct: args.is_distinct, + exprs: args.exprs.to_vec(), + expr_fields: args.expr_fields.to_vec(), + } + } + + fn create(&self) -> Result> { + self.udf.accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + ignore_nulls: self.ignore_nulls, + order_bys: &self.order_bys, + is_reversed: self.is_reversed, + name: &self.name, + is_distinct: self.is_distinct, + exprs: &self.exprs, + expr_fields: &self.expr_fields, + }) + } +} + +/// Projects the output of a Single-mode `aggregate` whose aggregates are all [`MergeAsPartialUDF`]s +/// into the output of a Partial-mode one: its `num_grouping` grouping columns, then one column per +/// state field, unpacked from the struct each aggregate evaluates to. +pub(crate) fn unpack_state_structs( + aggregate: Arc, + num_grouping: usize, +) -> Result> { + let schema = aggregate.schema(); + let mut exprs: Vec<(Arc, String)> = Vec::new(); + for (idx, field) in schema.fields().iter().enumerate() { + let column: Arc = Arc::new(Column::new(field.name(), idx)); + if idx < num_grouping { + exprs.push((column, field.name().clone())); + continue; + } + let DataType::Struct(state_fields) = field.data_type() else { + return internal_err!( + "Expected a struct of aggregate state in column {idx}, got {}", + field.data_type() + ); + }; + for (ordinal, state_field) in state_fields.iter().enumerate() { + exprs.push(( + Arc::new(GetStructField::new(Arc::clone(&column), ordinal)), + state_field.name().clone(), + )); + } + } + Ok(Arc::new(ProjectionExec::try_new(exprs, aggregate)?)) +} + +/// Accumulator wrapper that redirects update_batch to merge_batch and evaluates to its state. struct MergeAsPartialAccumulator { inner: Box, state_decoder: PartialMergeStateDecoder, + state_fields: Fields, } impl Debug for MergeAsPartialAccumulator { @@ -229,7 +335,17 @@ impl Accumulator for MergeAsPartialAccumulator { } fn evaluate(&mut self) -> Result { - self.inner.evaluate() + let state = self + .inner + .state()? + .iter() + .map(ScalarValue::to_array) + .collect::>>()?; + Ok(ScalarValue::Struct(Arc::new(StructArray::try_new( + self.state_fields.clone(), + state, + None, + )?))) } fn state(&mut self) -> Result> { @@ -241,10 +357,12 @@ impl Accumulator for MergeAsPartialAccumulator { } } -/// GroupsAccumulator wrapper that redirects update_batch to merge_batch. +/// GroupsAccumulator wrapper that redirects update_batch to merge_batch and evaluates to its +/// state. struct MergeAsPartialGroupsAccumulator { inner: Box, state_decoder: PartialMergeStateDecoder, + state_fields: Fields, } impl Debug for MergeAsPartialGroupsAccumulator { @@ -290,7 +408,12 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator { } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - self.inner.evaluate(emit_to) + let state = self.inner.state(emit_to)?; + Ok(Arc::new(StructArray::try_new( + self.state_fields.clone(), + state, + None, + )?)) } fn state(&mut self, emit_to: EmitTo) -> Result> { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ae5916092b1..cc96e847c6d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1575,13 +1575,17 @@ impl PhysicalPlanner { let mode = match proto_mode { ProtoAggregateMode::Partial => DFAggregateMode::Partial, ProtoAggregateMode::Final => DFAggregateMode::Final, - // PartialMerge: Partial + MergeAsPartial - ProtoAggregateMode::PartialMerge => DFAggregateMode::Partial, + // Spark's distinct-aggregate plan counts the rows of a PartialMerge, so it must + // emit each key once. DataFusion's Partial emits what it holds when it runs + // short of memory and starts again, while its Single spills. So run the + // PartialMerge as Single, whose MergeAsPartial aggregates merge their input and + // evaluate to their states, and unpack those states below. + ProtoAggregateMode::PartialMerge => DFAggregateMode::Single, }; // Check if any expression uses PartialMerge mode. When present, // those expressions are wrapped with MergeAsPartial to get merge - // semantics inside a Partial-mode AggregateExec. + // semantics with state output. let partial_merge_value = ProtoAggregateMode::PartialMerge as i32; let has_partial_merge = proto_mode == ProtoAggregateMode::PartialMerge || agg.expr_modes.contains(&partial_merge_value); @@ -1680,6 +1684,16 @@ impl PhysicalPlanner { )?, ); + let mut output = aggregate; + let mut additional_native_plans = vec![]; + if proto_mode == ProtoAggregateMode::PartialMerge { + additional_native_plans.push(Arc::clone(&output)); + output = crate::execution::merge_as_partial::unpack_state_structs( + output, + agg.grouping_exprs.len(), + )?; + } + // Spark's SortAggregateExec reports its output as ordered by the grouping keys, // and Spark may have removed a sort above it on that basis. DataFusion emits // groups in input order only when it sees the input sorted on them. When that @@ -1687,11 +1701,11 @@ impl PhysicalPlanner { // relation that reaches native code as an unordered ScanExec, the hash table can // emit groups out of order, so sort the aggregate output instead. if agg.ordered_by_grouping_keys && !agg.grouping_exprs.is_empty() { - let aggregate_schema = aggregate.schema(); + let output_schema = output.schema(); let ordering: Vec = (0..agg.grouping_exprs.len()) .map(|idx| { PhysicalSortExpr::new( - Arc::new(Column::new(aggregate_schema.field(idx).name(), idx)), + Arc::new(Column::new(output_schema.field(idx).name(), idx)), SortOptions { descending: false, nulls_first: true, @@ -1699,31 +1713,25 @@ impl PhysicalPlanner { ) }) .collect(); - if !aggregate + if !output .equivalence_properties() .ordering_satisfy(ordering.clone())? { - let sort: Arc = Arc::new(SortExec::new( - LexOrdering::new(ordering).unwrap(), - Arc::clone(&aggregate), - )); - return Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new_with_additional( - spark_plan.plan_id, - sort, - vec![child], - vec![aggregate], - )), - )); + additional_native_plans.push(Arc::clone(&output)); + output = + Arc::new(SortExec::new(LexOrdering::new(ordering).unwrap(), output)); } } Ok(( scans, shuffle_scans, - Arc::new(SparkPlan::new(spark_plan.plan_id, aggregate, vec![child])), + Arc::new(SparkPlan::new_with_additional( + spark_plan.plan_id, + output, + vec![child], + additional_native_plans, + )), )) } @@ -6316,6 +6324,162 @@ mod tests { assert_eq!("AggregateExec", planned.native_plan.name()); } + #[tokio::test] + async fn partial_merge_runs_as_single_and_emits_merged_states() { + use arrow::array::{BooleanArray, Int64Array}; + use datafusion::assert_batches_sorted_eq; + use datafusion::common::tree_node::{Transformed, TreeNode}; + use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; + use spark_expression::agg_expr::ExprStruct as AggExprStruct; + + // A Spark PartialMerge of max(v), count(v) and first(v) grouped by k, over k and their + // states. DataFusion's first has no GroupsAccumulator without an ordering. + let int64 = spark_expression::DataType { + type_id: 4, + type_info: None, + }; + let boolean = spark_expression::DataType { + type_id: 0, + type_info: None, + }; + let v = Expr { + expr_struct: Some(Unbound(spark_expression::UnboundReference { + name: "v".to_string(), + datatype: Some(int64.clone()), + })), + ..Default::default() + }; + let partial_merge = spark_operator::AggregateMode::PartialMerge as i32; + let aggregate = |mode: i32| Operator { + plan_id: 1, + children: vec![Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields: vec![ + int64.clone(), + int64.clone(), + int64.clone(), + int64.clone(), + boolean.clone(), + ], + source: String::new(), + })), + ..Default::default() + }], + op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate { + grouping_exprs: vec![Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(int64.clone()), + })), + ..Default::default() + }], + agg_exprs: vec![ + spark_expression::AggExpr { + expr_struct: Some(AggExprStruct::Max(spark_expression::Max { + child: Some(v.clone()), + datatype: Some(int64.clone()), + })), + ..Default::default() + }, + spark_expression::AggExpr { + expr_struct: Some(AggExprStruct::Count(spark_expression::Count { + children: vec![v.clone()], + })), + ..Default::default() + }, + spark_expression::AggExpr { + expr_struct: Some(AggExprStruct::First(spark_expression::First { + child: Some(v.clone()), + datatype: Some(int64.clone()), + ignore_nulls: false, + })), + ..Default::default() + }, + ], + mode, + expr_modes: vec![partial_merge; 3], + initial_input_buffer_offset: 1, + ..Default::default() + })), + ..Default::default() + }; + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + + // DataFusion's Single spills rather than emit a key twice, and the projection above it + // unpacks the states into the columns that the Partial it replaces emitted. + let (_, _, planned) = planner + .create_plan(&aggregate(partial_merge), &mut vec![], 1) + .unwrap(); + assert_eq!("ProjectionExec", planned.native_plan.name()); + assert_eq!(1, planned.additional_native_plans.len()); + let single = planned.additional_native_plans[0] + .downcast_ref::() + .unwrap(); + assert_eq!(&AggregateMode::Single, single.mode()); + let (_, _, partial) = planner + .create_plan( + &aggregate(spark_operator::AggregateMode::Partial as i32), + &mut vec![], + 1, + ) + .unwrap(); + assert_eq!("AggregateExec", partial.native_plan.name()); + assert_eq!(partial.native_plan.schema(), planned.native_plan.schema()); + + let batch = RecordBatch::try_new( + planned.children[0].schema(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 1, 2, 3])), + Arc::new(Int64Array::from(vec![ + Some(10), + Some(5), + Some(7), + None, + Some(1), + ])), + Arc::new(Int64Array::from(vec![1, 2, 3, 0, 4])), + Arc::new(Int64Array::from(vec![ + Some(100), + None, + Some(101), + Some(200), + None, + ])), + Arc::new(BooleanArray::from(vec![true, false, true, true, false])), + ], + ) + .unwrap(); + let input: Arc = + MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap(); + let native_plan = Arc::clone(&planned.native_plan) + .transform_up(|node| { + Ok(if node.name() == "ScanExec" { + Transformed::yes(Arc::clone(&input)) + } else { + Transformed::no(node) + }) + }) + .unwrap() + .data; + let results = collect(native_plan.execute(0, task_ctx).unwrap()) + .await + .unwrap(); + assert_batches_sorted_eq!( + [ + "+-------+------------+--------------+--------------------+---------------------------+", + "| col_0 | max[value] | count[count] | first[first_value] | first[first_value_is_set] |", + "+-------+------------+--------------+--------------------+---------------------------+", + "| 1 | 10 | 4 | 100 | true |", + "| 2 | 5 | 2 | 200 | true |", + "| 3 | 1 | 4 | | false |", + "+-------+------------+--------------+--------------------+---------------------------+", + ], + &results + ); + } + #[tokio::test] async fn projection_prunes_filter_output_and_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; 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 a8bd439c188..2790838a310 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{CometListenerBusUtils, SparkConf} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{Column, CometTestBase, DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.Cast -import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortAggregateExec, CometSortExec} @@ -156,19 +156,21 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } /** - * Returns the native final aggregates in the plan that `df` last ran, after checking that there - * is one and that each spilled. + * Returns the native aggregates in the plan that `df` last ran whose aggregate functions all + * run in `mode`, after checking that there is one and that each spilled. */ - private def spilledFinalAggregates(df: DataFrame): Seq[CometHashAggregateExec] = { + private def spilledAggregates( + df: DataFrame, + mode: AggregateMode): Seq[CometHashAggregateExec] = { val plan = df.queryExecution.executedPlan - val finalAggregates = collect(plan) { - case aggregate: CometHashAggregateExec if aggregate.modes.contains(Final) => aggregate + val aggregates = collect(plan) { + case aggregate: CometHashAggregateExec if aggregate.modes == Seq(mode) => aggregate } - assert(finalAggregates.nonEmpty, s"Expected a native final aggregate:\n$plan") + assert(aggregates.nonEmpty, s"Expected a native $mode aggregate:\n$plan") assert( - finalAggregates.forall(_.metrics("spill_count").value > 0), - "The final aggregate did not spill") - finalAggregates + aggregates.forall(_.metrics("spill_count").value > 0), + s"The $mode aggregate did not spill") + aggregates } test("final aggregate that has spilled reads its spill files back (issue #6254)") { @@ -193,7 +195,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } // checkCometAnswer runs the query once, so the plan read below is the one that ran. checkCometAnswer(df, expected) - spilledFinalAggregates(df) + spilledAggregates(df, Final) } } } @@ -233,7 +235,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Only a sort in the same native plan keeps the final aggregate's input sorted. assert( - spilledFinalAggregates(df).exists(_.child match { + spilledAggregates(df, Final).exists(_.child match { case partial: CometHashAggregateExec => partial.child.isInstanceOf[CometSortExec] case _ => false }), @@ -243,6 +245,97 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + /** + * Runs `f` with a native memory pool small enough to make the PartialMerge stage of a distinct + * aggregate run short of memory: about 4 MiB, 0.002 of the suite's 2g off-heap. + */ + private def withDistinctMergeMemoryPressure(f: => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002")(f) + } + + test("PartialMerge that runs short of memory keeps one row per key (issue #6857)") { + // Spark plans count(DISTINCT d) next to sum(v) with a PartialMerge keyed by (g, d), and the + // stage above it counts the PartialMerge's rows, so the PartialMerge has to spill rather than + // emit a key twice. + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 400000, 1, 1) + .selectExpr("id % 100 AS g", "id % 200000 AS d", "id AS v") + .write + .parquet(path) + withDistinctMergeMemoryPressure { + val df = spark.read.parquet(path).groupBy("g").agg(sum("v"), count_distinct(col("d"))) + // Each g holds the 4000 ids g + 100 * i, whose d values pair up into 2000 distinct ones. + val expected = (0 until 100).map(g => Row(g.toLong, 4000L * g + 799800000L, 2000L)) + checkCometAnswer(df, expected) + spilledAggregates(df, PartialMerge) + } + } + } + + test("PartialMerge that has spilled reads its spill files back (issue #6857)") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 125000, 1, 1) + .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v") + .write + .parquet(path) + // A 2 MiB pool, shared with the aggregate above it, makes the PartialMerge keyed by (g, k) + // spill several runs of its 62,500 groups. As in the #6254 tests, merging the runs then + // takes nearly all of its share, and reading them back must not fail the task. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (2.0 / 2048).toString) { + val df = + spark.read.parquet(path).groupBy("g").agg(expr("max(v)"), count_distinct(col("k"))) + checkCometAnswer(df, Seq(Row(0L, 124999L, 62500L))) + spilledAggregates(df, PartialMerge) + } + } + } + + test("collect_list PartialMerge short of memory keeps one row per key (issue #6857)") { + // The collect_list states fill the PartialMerge's memory too. Over a Spark partial, the + // PartialMerge decodes Spark's buffers as it merges them, but reads back the native states + // it spilled. + def padded(i: Int): String = "x" * (128 - i.toString.length) + i + val rows = (0 until 32768).map(i => (i % 4096, i % 8192, padded(i))) + // Each k holds the ids k + 4096 * j, whose d values are k and k + 4096. + val expected = (0 until 4096).map { k => + Row(k, 2L, (0 until 8).map(j => padded(k + 4096 * j)).sorted) + } + val localRows = rows.toDF("k", "d", "s") + withTempPath { dir => + val path = dir.getCanonicalPath + localRows.write.parquet(path) + for (sparkPartial <- Seq(false, true)) { + withDistinctMergeMemoryPressure { + val source = if (sparkPartial) localRows else spark.read.parquet(path) + val df = + source.groupBy("k").agg(count_distinct(col("d")), sort_array(collect_list("s"))) + checkCometAnswer(df, expected) + spilledAggregates(df, PartialMerge) + val sparkPartials = collect(df.queryExecution.executedPlan) { + case agg: ObjectHashAggregateExec + if agg.aggregateExpressions.forall(_.mode == Partial) => + agg + } + assert( + sparkPartials.nonEmpty == sparkPartial, + "Expected a Spark partial aggregate only over the local rows:\n" + + df.queryExecution.executedPlan) + } + } + } + } + test("collect_list over struct with non-nullable fields") { // Building a struct from non-nullable columns yields non-nullable struct fields. Native // collect_list derives its declared output element type from the argument's declared type, From 66e0c7009ba05d58832f30a1b4606a4acea448a5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 11:18:21 +0000 Subject: [PATCH 73/89] test: fail Comet tests when a native plan mislabels a timestamp's timezone (#6868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Native plans must label every TimestampType value Timestamp(µs, "UTC") and every TimestampNTZType value with no timezone. The label never changes the values, so checkSparkAnswer cannot see a wrong one, and the casts at ScanExec and the shuffle writer's SchemaAlignExec repair it silently. Add the internal test-only config spark.comet.testing.strictTimestampLabels.enabled, which CometTestBase turns on. With it, CometExecIterator checks each batch a native plan returns against the plan's Spark output types, ScanExec fails instead of casting a relabelled input timestamp, and the native shuffle writer fails planning instead of aligning one. --- native/common/src/lib.rs | 4 +- native/common/src/schema.rs | 134 ++++++++++++++++++ native/core/src/execution/jni_api.rs | 9 +- native/core/src/execution/operators/scan.rs | 70 ++++++++- native/core/src/execution/planner.rs | 16 ++- native/core/src/execution/planner/shuffle.rs | 63 +++++++- native/core/src/execution/spark_config.rs | 2 + .../scala/org/apache/comet/CometConf.scala | 17 +++ .../org/apache/comet/CometExecIterator.scala | 14 +- .../org/apache/comet/vector/NativeUtil.scala | 86 ++++++++++- .../apache/spark/sql/comet/CometExecRDD.scala | 13 +- .../spark/sql/comet/CometExecUtils.scala | 3 +- .../comet/CometIcebergNativeScanExec.scala | 3 +- .../spark/sql/comet/CometNativeScanExec.scala | 3 +- .../CometTakeOrderedAndProjectExec.scala | 10 +- .../apache/spark/sql/comet/operators.scala | 9 +- .../apache/comet/exec/CometExecSuite.scala | 6 + .../apache/comet/vector/NativeUtilSuite.scala | 104 +++++++++++++- .../org/apache/spark/sql/CometTestBase.scala | 4 + 19 files changed, 542 insertions(+), 28 deletions(-) diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index e68613053f1..95a4007efca 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -30,7 +30,9 @@ mod utils; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; pub use ffi_offsets::zero_offsets; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; -pub use schema::{cast_and_stamp_schema, widen_nested_nullability}; +pub use schema::{ + cast_and_stamp_schema, describe_timestamp_label_mismatch, widen_nested_nullability, +}; pub use struct_nulls::{child_with_parent_nulls, children_with_parent_nulls}; pub use utf8::decode_string_arrays; pub use utils::{bytes_to_i128, decode_utf8_spark_lossy}; diff --git a/native/common/src/schema.rs b/native/common/src/schema.rs index 312861999f7..bfee814e5a8 100644 --- a/native/common/src/schema.rs +++ b/native/common/src/schema.rs @@ -154,6 +154,55 @@ fn nullability(field: &FieldRef) -> String { format!("{qualifier} {}", field.data_type()) } +/// Describes the first timestamp whose timezone label differs between `expected` and `actual`, as +/// a dotted path rooted at `path`. Returns `None` when every timestamp in `actual` carries the +/// label `expected` declares, whatever else differs. A dictionary on either side is looked through, +/// since boundaries unpack dictionaries anyway. +/// +/// A label never changes the values, which are UTC microseconds whatever it says, so the cast at a +/// boundary repairs it and nothing past the boundary can tell. Inside a native plan the label does +/// matter: any other zone breaks comparisons and type coercion against Comet's `"UTC"`, and a +/// missing one makes kernels read the values as wall-clock time. Strict test runs +/// (`spark.comet.testing.strictTimestampLabels.enabled`) fail at the boundary instead of casting, +/// so the test that produced the drift reports it. +pub fn describe_timestamp_label_mismatch( + path: &str, + expected: &DataType, + actual: &DataType, +) -> Option { + match (expected, actual) { + (DataType::Timestamp(_, e), DataType::Timestamp(_, a)) if e != a => { + Some(format!("{path}: expected {expected}, found {actual}")) + } + (_, DataType::Dictionary(_, a)) => describe_timestamp_label_mismatch(path, expected, a), + (DataType::Dictionary(_, e), _) => describe_timestamp_label_mismatch(path, e, actual), + (DataType::List(e), DataType::List(a)) + | (DataType::LargeList(e), DataType::LargeList(a)) + | (DataType::FixedSizeList(e, _), DataType::FixedSizeList(a, _)) => { + describe_timestamp_label_mismatch( + &format!("{path}.element"), + e.data_type(), + a.data_type(), + ) + } + (DataType::Map(e, _), DataType::Map(a, _)) => describe_timestamp_label_mismatch( + &format!("{path}.entries"), + e.data_type(), + a.data_type(), + ), + (DataType::Struct(e), DataType::Struct(a)) if e.len() == a.len() => { + e.iter().zip(a.iter()).find_map(|(e, a)| { + describe_timestamp_label_mismatch( + &format!("{path}.{}", e.name()), + e.data_type(), + a.data_type(), + ) + }) + } + _ => None, + } +} + /// Returns `base` with nested field nullability widened to also cover `other`, so that arrays of /// either type can be stamped with the result. Shapes that do not line up (different struct field /// counts, different list flavours, ...) are left as `base`; those are real type differences and @@ -416,6 +465,91 @@ mod tests { assert!(msg.contains("payload: expected Struct"), "{msg}"); } + fn timestamp(tz: Option<&str>) -> DataType { + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, tz.map(Into::into)) + } + + #[test] + fn describes_timestamp_label_mismatches() { + let utc = timestamp(Some("UTC")); + assert_eq!( + describe_timestamp_label_mismatch("ts", &utc, ×tamp(Some("Etc/UTC"))).unwrap(), + r#"ts: expected Timestamp(µs, "UTC"), found Timestamp(µs, "Etc/UTC")"# + ); + // The worst case: an LTZ value labelled as NTZ is read as wall-clock time downstream. + assert!(describe_timestamp_label_mismatch("ts", &utc, ×tamp(None)).is_some()); + assert!(describe_timestamp_label_mismatch("ts", ×tamp(None), &utc).is_some()); + assert_eq!(describe_timestamp_label_mismatch("ts", &utc, &utc), None); + assert_eq!( + describe_timestamp_label_mismatch("ts", ×tamp(None), ×tamp(None)), + None + ); + } + + #[test] + fn describes_nested_timestamp_label_paths() { + let wrap = |tz: Option<&str>| { + DataType::Struct(Fields::from(vec![ + Field::new("id", DataType::Int64, true), + Field::new( + "events", + DataType::List(Arc::new(Field::new_list_field(timestamp(tz), true))), + true, + ), + ])) + }; + assert_eq!( + describe_timestamp_label_mismatch("c0", &wrap(Some("UTC")), &wrap(None)).unwrap(), + r#"c0.events.element: expected Timestamp(µs, "UTC"), found Timestamp(µs)"# + ); + + let map = |tz: Option<&str>| { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", timestamp(tz), true), + ])), + false, + )), + false, + ) + }; + assert_eq!( + describe_timestamp_label_mismatch("m", &map(Some("UTC")), &map(Some("+00:00"))) + .unwrap(), + r#"m.entries.value: expected Timestamp(µs, "UTC"), found Timestamp(µs, "+00:00")"# + ); + } + + /// Only the label is checked. A dictionary, a nullability flag or a different leaf type is + /// left for the boundary cast, and a dictionary is looked through to the timestamp inside. + #[test] + fn timestamp_label_check_ignores_other_differences() { + let utc = timestamp(Some("UTC")); + let dict = + |value: DataType| DataType::Dictionary(Box::new(DataType::Int32), Box::new(value)); + assert_eq!( + describe_timestamp_label_mismatch("ts", &utc, &dict(utc.clone())), + None + ); + assert!( + describe_timestamp_label_mismatch("ts", &utc, &dict(timestamp(Some("Etc/UTC")))) + .is_some() + ); + assert_eq!( + describe_timestamp_label_mismatch("c0", &DataType::Int64, &DataType::Int32), + None + ); + let list = + |nullable: bool| DataType::List(Arc::new(Field::new_list_field(utc.clone(), nullable))); + assert_eq!( + describe_timestamp_label_mismatch("c0", &list(true), &list(false)), + None + ); + } + #[test] fn error_on_column_count_mismatch() { let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, true)])); diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index a2646dd8daf..f8e9e497f6f 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -135,7 +135,8 @@ use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::CometObjectStoreRegistry; @@ -555,6 +556,8 @@ struct ExecutionContext { pub session_ctx: Arc, /// Whether to enable additional debugging checks & messages pub debug_native: bool, + /// Whether boundaries fail on a relabelled timestamp instead of casting it (test-only) + pub strict_timestamp_labels: bool, /// Whether to write native plans with metrics to stdout pub explain_native: bool, /// Whether to log memory usage on each call to execute_plan @@ -706,6 +709,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( struct PlanSettings { spark_config: HashMap, debug_native: bool, + strict_timestamp_labels: bool, explain_native: bool, tracing_enabled: bool, max_temp_directory_size: u64, @@ -726,6 +730,7 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult, pub schema: SchemaRef, + /// Fail on an input timestamp whose timezone label differs from `data_types` instead of + /// casting it. Test-only; see `spark.comet.testing.strictTimestampLabels.enabled`. + pub strict_timestamp_labels: bool, /// Used in unit tests to mock the input batch; otherwise written by `pull_next` on each /// poll. pub batch: Arc>>, @@ -98,9 +101,15 @@ impl ScanExec { metrics: metrics_set, baseline_metrics, schema, + strict_timestamp_labels: false, }) } + pub fn with_strict_timestamp_labels(mut self, strict: bool) -> Self { + self.strict_timestamp_labels = strict; + self + } + /// Unpack all dictionary types because some DataFusion operators /// and expressions do not support dictionary types fn unpack_dictionary_type(dt: &DataType) -> DataType { @@ -311,21 +320,45 @@ impl ScanStream<'_> { let new_columns: Vec = columns .iter() .zip(schema_fields.iter()) - .map(|(column, f)| { + .enumerate() + .map(|(idx, (column, f))| { if column.data_type() != f.data_type() { + self.check_timestamp_labels(idx, f, column)?; let mut timer = self.cast_time.timer(); let cast_array = cast_with_options(column, f.data_type(), &self.cast_options); timer.stop(); - cast_array + cast_array.map_err(DataFusionError::from) } else { Ok(Arc::clone(column)) } }) - .collect::, _>>()?; + .collect::>>()?; let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); RecordBatch::try_new_with_options(Arc::clone(&self.schema), new_columns, &options) .map_err(|e| arrow_datafusion_err!(e)) } + + /// In strict test runs, rejects an input column whose timestamps are labelled with a + /// different timezone than the declared type, rather than letting the cast hide it. + fn check_timestamp_labels( + &self, + idx: usize, + field: &Field, + column: &ArrayRef, + ) -> DataFusionResult<()> { + if !self.scan.strict_timestamp_labels { + return Ok(()); + } + match describe_timestamp_label_mismatch(field.name(), field.data_type(), column.data_type()) + { + Some(detail) => Err(DataFusionError::Execution(format!( + "ScanExec [{}] col[{idx}] labels a timestamp differently from the type Spark \ + declared, at {detail}", + self.scan.input_source_description + ))), + None => Ok(()), + } + } } impl Stream for ScanStream<'_> { @@ -534,4 +567,33 @@ mod import_tests { let s = out.as_any().downcast_ref::().unwrap(); assert_eq!(s.value(0), "\u{FFFD}"); } + + /// The cast repairs an input timestamp labelled with another zone, which is what hides the + /// drift; a strict test run rejects it instead. + #[test] + fn strict_scan_rejects_relabelled_timestamps() { + use arrow::array::TimestampMicrosecondArray; + use arrow::datatypes::TimeUnit; + use futures::StreamExt; + + let declared = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let scan_one_batch = |strict: bool| { + let mut scan = + ScanExec::new(TEST_EXEC_CONTEXT_ID, None, "test", vec![declared.clone()]) + .unwrap() + .with_strict_timestamp_labels(strict); + let mut stream = scan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let column: ArrayRef = + Arc::new(TimestampMicrosecondArray::from(vec![1_i64, 2]).with_timezone("Etc/UTC")); + scan.set_input_batch(InputBatch::new(vec![column], Some(2))); + futures::executor::block_on(stream.next()).unwrap() + }; + + let batch = scan_one_batch(false).unwrap(); + assert_eq!(batch.column(0).data_type(), &declared); + + let err = scan_one_batch(true).unwrap_err().to_string(); + assert!(err.contains("ScanExec [test] col[0]"), "{err}"); + assert!(err.contains(r#"found Timestamp(µs, "Etc/UTC")"#), "{err}"); + } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index cc96e847c6d..ceca9e657af 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -350,6 +350,9 @@ pub struct PhysicalPlanner { /// How comparisons treat floating-point operands. `Raw` only while planning a scan's data /// filters; see [`Self::create_data_filter`]. float_operands: FloatOperands, + /// Whether a boundary that would cast a timestamp's timezone label to the declared one fails + /// instead. Test-only; see `spark.comet.testing.strictTimestampLabels.enabled`. + strict_timestamp_labels: bool, } impl Default for PhysicalPlanner { @@ -370,6 +373,7 @@ impl PhysicalPlanner { class_loader: None, shuffle_partition_pusher: None, float_operands: FloatOperands::Normalize, + strict_timestamp_labels: false, } } @@ -448,6 +452,15 @@ impl PhysicalPlanner { self } + pub fn with_strict_timestamp_labels(mut self, strict: bool) -> Self { + self.strict_timestamp_labels = strict; + self + } + + pub(crate) fn strict_timestamp_labels(&self) -> bool { + self.strict_timestamp_labels + } + /// Attach a propagated Spark `TaskContext` global reference. Called by the JNI `executePlan` /// entry with whatever was captured at `createPlan` time. The planner clones this `Option` /// into every `JvmScalarUdfExpr` it builds. @@ -2136,7 +2149,8 @@ impl PhysicalPlanner { // The `ScanExec` operator will take actual arrays from Spark during execution let scan = - ScanExec::new(self.exec_context_id, input_source, &scan.source, data_types)?; + ScanExec::new(self.exec_context_id, input_source, &scan.source, data_types)? + .with_strict_timestamp_labels(self.strict_timestamp_labels); Ok(( vec![scan.clone()], diff --git a/native/core/src/execution/planner/shuffle.rs b/native/core/src/execution/planner/shuffle.rs index 43dbd97047f..eb2fc6c8119 100644 --- a/native/core/src/execution/planner/shuffle.rs +++ b/native/core/src/execution/planner/shuffle.rs @@ -27,8 +27,9 @@ use arrow::{ use datafusion::{ common::ScalarValue, physical_expr::{expressions::Literal, LexOrdering, PhysicalExpr, PhysicalSortExpr}, - physical_plan::ExecutionPlan, + physical_plan::{display::DisplayableExecutionPlan, ExecutionPlan}, }; +use datafusion_comet_common::describe_timestamp_label_mismatch; use datafusion_comet_proto::spark_operator::{ self, CompressionCodec as SparkCompressionCodec, Operator, }; @@ -79,6 +80,7 @@ impl OperatorBuilder for ShuffleWriterBuilder { let writer_input = align_shuffle_writer_input( Arc::clone(&child.native_plan), &writer.expected_output_schema, + planner.strict_timestamp_labels(), )?; let partitioning = planner @@ -284,14 +286,37 @@ impl OperatorBuilder for ShuffleScanBuilder { /// Wrap `child` in a `SchemaAlignExec` when its output drifts from what Spark catalyst /// declared. See . +/// +/// With `strict_timestamp_labels`, a timestamp labelled with another timezone than declared fails +/// planning instead of being cast. The child is native, so the label is wrong inside its plan too. fn align_shuffle_writer_input( child: Arc, expected_proto: &[spark_operator::SparkStructField], + strict_timestamp_labels: bool, ) -> Result, ExecutionError> { if expected_proto.is_empty() { return Ok(child); } let expected = convert_spark_types_to_arrow_schema(expected_proto); + if strict_timestamp_labels { + let actual = child.schema(); + let mismatch = expected + .fields() + .iter() + .zip(actual.fields()) + .enumerate() + .find_map(|(idx, (e, a))| { + describe_timestamp_label_mismatch(e.name(), e.data_type(), a.data_type()) + .map(|detail| (idx, detail)) + }); + if let Some((idx, detail)) = mismatch { + return Err(GeneralError(format!( + "Shuffle writer input col[{idx}] labels a timestamp differently from the type \ + Spark declared, at {detail}. Native plan:\n{}", + DisplayableExecutionPlan::new(child.as_ref()).indent(true) + ))); + } + } SchemaAlignExec::try_new_or_passthrough(child, &expected) .map_err(|e| ExecutionError::DataFusionError(e.to_string())) } @@ -790,4 +815,40 @@ mod tests { assert!(error.to_string().contains("No input for shuffle scan")); } + + /// The writer's alignment casts a timestamp its child labelled with another zone; a strict + /// test run fails planning instead, naming the column. + #[test] + fn strict_shuffle_writer_rejects_relabelled_timestamps() { + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion::physical_plan::empty::EmptyExec; + + let child: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("Etc/UTC".into())), + true, + )])))); + let expected = [spark_operator::SparkStructField { + name: "ts".to_string(), + data_type: Some(spark_expression::DataType { + type_id: spark_expression::data_type::DataTypeId::Timestamp as i32, + type_info: None, + }), + nullable: true, + ..Default::default() + }]; + + let aligned = align_shuffle_writer_input(Arc::clone(&child), &expected, false).unwrap(); + assert!(aligned.is::()); + + let error = align_shuffle_writer_input(child, &expected, true) + .unwrap_err() + .to_string(); + assert!(error.contains("Shuffle writer input col[0]"), "{error}"); + assert!( + error.contains(r#"found Timestamp(µs, "Etc/UTC")"#), + "{error}" + ); + } } diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 1de88ccbeec..7755f6d5dc4 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -28,6 +28,8 @@ pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str = "spark.comet.exec.aggregate.skipPartial.enabled"; pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = "spark.comet.exec.spill.compression.codec"; +pub(crate) const COMET_STRICT_TIMESTAMP_LABELS: &str = + "spark.comet.testing.strictTimestampLabels.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; /// Comet configs read through this trait must be resolved by the JVM first: diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index c6ac50606f3..9bd32816c64 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1082,6 +1082,23 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_STRICT_TIMESTAMP_LABELS: ConfigEntry[Boolean] = + conf("spark.comet.testing.strictTimestampLabels.enabled") + .category(CATEGORY_TESTING) + .doc( + "Test-only. When enabled, Comet fails a query whose native plan labels a timestamp " + + "with a different timezone than Comet's convention, which is `UTC` for " + + "`TimestampType` and no timezone for `TimestampNTZType`. Every native output batch " + + "is checked against the Spark output types, and the casts that scans and the native " + + "shuffle writer would otherwise apply to repair a label fail instead. A label never " + + "changes the values, so comparing results cannot detect it, but inside a native plan " + + "a wrong label breaks comparisons and type coercion, and a missing one makes kernels " + + "read the values as wall-clock time. Enabled for all Comet test suites via " + + "`CometTestBase`.") + .internal() + .booleanConf + .createWithDefault(false) + val COMET_SCAN_CONTRIB_DETECT_CONFLICTS: ConfigEntry[Boolean] = conf("spark.comet.scan.contrib.detectConflicts.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 27f880f9132..8997be2d7e9 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -36,6 +36,7 @@ import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.comet.CometMetricNode import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized._ import org.apache.spark.util.SerializableConfiguration @@ -76,6 +77,9 @@ import org.apache.comet.vector.NativeUtil * Whether to read the shuffle writer's partition offsets when the plan reaches the end of its * output, for a plan rooted at a native shuffle writer with a local destination. Remote shuffle * reports its partition lengths through its pusher instead, so it leaves this false. + * @param outputTypes + * The Spark types of the plan's output columns, if the caller knows them. Used only to check + * timestamp labels under `spark.comet.testing.strictTimestampLabels.enabled`. */ class CometExecIterator( val id: Long, @@ -90,11 +94,14 @@ class CometExecIterator( shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, taskFilePaths: Seq[String] = Seq.empty, shufflePartitionPusher: Option[ShufflePartitionPusher] = None, - capturePartitionOffsets: Boolean = false) + capturePartitionOffsets: Boolean = false, + outputTypes: Seq[DataType] = Seq.empty) extends Iterator[ColumnarBatch] with Logging { private val tracingEnabled = CometConf.COMET_TRACING_ENABLED.get() + private val checkTimestampLabels = + outputTypes.nonEmpty && CometConf.COMET_STRICT_TIMESTAMP_LABELS.get() private val memoryMXBean = ManagementFactory.getMemoryMXBean private val nativeLib = new Native() private val nativeUtil = new NativeUtil() @@ -258,6 +265,10 @@ class CometExecIterator( traceMemoryUsage() } + if (checkTimestampLabels) { + result.foreach(nativeUtil.checkTimestampLabels(_, outputTypes)) + } + result } catch { // Native saw only the text of what a JVM input threw, so rethrow the exception itself. @@ -879,6 +890,7 @@ object CometExecIterator extends Logging { CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, + CometConf.COMET_STRICT_TIMESTAMP_LABELS, CometConf.COMET_TRACING_ENABLED).foreach { entry => builder.putEntries(entry.key, entry.get(SQLConf.get).toString) } diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 84554134c76..8dd49d8fb1f 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -20,18 +20,21 @@ package org.apache.comet.vector import scala.collection.mutable +import scala.jdk.CollectionConverters._ import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.util.AutoCloseables import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.dictionary.DictionaryProvider +import org.apache.arrow.vector.types.pojo.{ArrowType, Field} import org.apache.spark.SparkException -import org.apache.spark.sql.comet.execution.arrow.ConstantColumnVectors +import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, ConstantColumnVectors} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.ConstantColumnVector +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator} +import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometRuntimeException} /** * Provides functionality for importing Arrow vectors from native code and wrapping them as @@ -292,6 +295,36 @@ class NativeUtil extends AutoCloseable { } } + /** + * Fails, closing `batch`, if a column that this instance imported from native labels a + * timestamp with a different timezone than Comet's convention for its Spark type in + * `expectedTypes`. Nothing past this point reads the label, so this is the last place a wrong + * one can be seen. Test-only; see `CometConf.COMET_STRICT_TIMESTAMP_LABELS`. + */ + def checkTimestampLabels(batch: ColumnarBatch, expectedTypes: Seq[DataType]): Unit = { + val mismatch = expectedTypes.indices + .take(batch.numCols()) + .iterator + .map { i => + batch.column(i) match { + case vector: CometVector => + NativeUtil.timestampLabelMismatch( + vector.getValueVector.getField, + expectedTypes(i), + dictionaryProvider, + s"col[$i]") + case _ => None + } + } + .collectFirst { case Some(detail) => detail } + mismatch.foreach { detail => + batch.close() + throw new CometRuntimeException( + "Native plan output labels a timestamp with a different timezone than the type Spark " + + s"declared, at $detail") + } + } + override def close(): Unit = { // closing the dictionary provider also closes the dictionary arrays dictionaryProvider.close() @@ -299,6 +332,55 @@ class NativeUtil extends AutoCloseable { } object NativeUtil { + + /** + * Describes the first timestamp in `field` whose timezone label is not Comet's convention for + * the Spark type at the same position in `dataType`, as a path rooted at `path`. That is + * `CometArrowStream.NATIVE_TIMEZONE` for `TimestampType` and no timezone for + * `TimestampNTZType`. Returns `None` when every label is right. Nothing else about `field` is + * checked. A dictionary-encoded field is checked by its values, looked up in `provider`. + */ + private[comet] def timestampLabelMismatch( + field: Field, + dataType: DataType, + provider: DictionaryProvider, + path: String): Option[String] = { + val valueField = Option(field.getDictionary) + .map(encoding => provider.lookup(encoding.getId).getVector.getField) + .getOrElse(field) + val children = valueField.getChildren.asScala + def child(index: Int, childType: DataType, childPath: String): Option[String] = + children + .lift(index) + .flatMap(timestampLabelMismatch(_, childType, provider, childPath)) + (dataType, valueField.getType) match { + case (TimestampType, ts: ArrowType.Timestamp) + if ts.getTimezone != CometArrowStream.NATIVE_TIMEZONE => + Some( + s"$path: expected TimestampType labelled ${CometArrowStream.NATIVE_TIMEZONE}, found $ts") + case (TimestampNTZType, ts: ArrowType.Timestamp) if ts.getTimezone != null => + Some(s"$path: expected TimestampNTZType with no timezone, found $ts") + case (StructType(fields), _: ArrowType.Struct) => + fields.indices.iterator + .map(i => child(i, fields(i).dataType, s"$path.${fields(i).name}")) + .collectFirst { case Some(detail) => detail } + case (ArrayType(elementType, _), _: ArrowType.List) => + child(0, elementType, s"$path.element") + case (MapType(keyType, valueType, _), _: ArrowType.Map) => + children.headOption.flatMap { entries => + val keyValue = entries.getChildren.asScala + keyValue + .lift(0) + .flatMap(timestampLabelMismatch(_, keyType, provider, s"$path.key")) + .orElse( + keyValue + .lift(1) + .flatMap(timestampLabelMismatch(_, valueType, provider, s"$path.value"))) + } + case _ => None + } + } + def rootAsBatch(arrowRoot: VectorSchemaRoot): ColumnarBatch = { rootAsBatch(arrowRoot, null) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index 1d876dfb83f..db1ca44c200 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -25,6 +25,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.rdd.RDD import org.apache.spark.sql.comet.execution.shuffle.CometShuffledBatchRDD import org.apache.spark.sql.execution.ScalarSubquery +import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.SerializableConfiguration @@ -67,7 +68,8 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]] = None, encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, - @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty) + @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty, + outputTypes: Seq[DataType] = Seq.empty) extends RDD[ColumnarBatch](sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { // Determine partition count: from inputs if available, otherwise from parameter @@ -136,7 +138,8 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption, encryptedFilePaths, shuffleBlockIters, - taskFilePaths = partition.filePaths) + taskFilePaths = partition.filePaths, + outputTypes = outputTypes) // Register ScalarSubqueries so native code can look them up subqueries.foreach(sub => CometScalarSubquery.setSubquery(it.id, sub)) @@ -229,7 +232,8 @@ object CometExecRDD { broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]] = None, encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, - perPartitionFilePaths: Array[Seq[String]] = Array.empty): CometExecRDD = { + perPartitionFilePaths: Array[Seq[String]] = Array.empty, + outputTypes: Seq[DataType] = Seq.empty): CometExecRDD = { // scalastyle:on new CometExecRDD( @@ -246,6 +250,7 @@ object CometExecRDD { broadcastedHadoopConfForEncryption, encryptedFilePaths, shuffleScanIndices, - perPartitionFilePaths) + perPartitionFilePaths, + outputTypes) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala index 5628ec039b3..5c97bebf271 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala @@ -78,7 +78,8 @@ object CometExecUtils { outputAttribute.length, serializedPlan, numParts, - idx) + idx, + outputAttribute.map(_.dataType)) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index e058a5c0916..0cb7b885b6e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -293,7 +293,8 @@ case class CometIcebergNativeScanExec( defaultNumPartitions = perPartitionData.length, numOutputCols = output.length, nativeMetrics = nativeMetrics, - subqueries = Seq.empty) { + subqueries = Seq.empty, + outputTypes = output.map(_.dataType)) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { // Register before super.compute creates the CometExecIterator, so this listener runs // after the iterator's close has published the final scan metrics. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 37694992211..bff854a2eee 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -285,7 +285,8 @@ case class CometNativeScanExec( Seq.empty, broadcastedHadoopConfForEncryption, encryptedFilePaths, - perPartitionFilePaths = perPartitionFilePaths) { + perPartitionFilePaths = perPartitionFilePaths, + outputTypes = output.map(_.dataType)) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { // Register before super.compute creates the CometExecIterator, so this listener runs // after the iterator's close has published the final scan metrics. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala index 6e2ed501ecb..88a797d123d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala @@ -170,7 +170,8 @@ case class CometTakeOrderedAndProjectExec( numOutputCols, serializedTopK, numParts, - idx) + idx, + child.output.map(_.dataType)) } } @@ -189,16 +190,17 @@ case class CometTakeOrderedAndProjectExec( // Serialize the plan once before mapping to avoid repeated serialization per partition val serializedFinalPlan = CometExec.serializeNativePlan(finalNativePlan(childRDD.getNumPartitions).get) - val finalOutputLength = output.length + val finalOutputTypes = output.map(_.dataType) val finalInputSchema = CometUtils.fromAttributes(child.output) singlePartitionRDD.mapPartitionsInternal { iter => val it = CometExec.getCometIterator( CometArrowStream .inputObjects(iter, finalInputSchema, "CometTakeOrderedAndProject-final"), - finalOutputLength, + finalOutputTypes.length, serializedFinalPlan, 1, - 0) + 0, + finalOutputTypes) setSubqueries(it.id, this) Option(TaskContext.get()).foreach { context => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index bcae614556b..a712f44796a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -801,7 +801,8 @@ object CometExec { numOutputCols: Int, serializedPlan: Array[Byte], numParts: Int, - partitionIdx: Int): CometExecIterator = { + partitionIdx: Int, + outputTypes: Seq[DataType] = Seq.empty): CometExecIterator = { new CometExecIterator( newIterId, inputObjects, @@ -811,7 +812,8 @@ object CometExec { numParts, partitionIdx, broadcastedHadoopConfForEncryption = None, - encryptedFilePaths = Seq.empty) + encryptedFilePaths = Seq.empty, + outputTypes = outputTypes) } def getCometIterator( @@ -950,7 +952,8 @@ abstract class CometNativeExec extends CometExec { ctx.broadcastedHadoopConfForEncryption, ctx.encryptedFilePaths, ctx.shuffleScanIndices, - perPartitionFilePaths = ctx.perPartitionFilePaths) { + perPartitionFilePaths = ctx.perPartitionFilePaths, + outputTypes = output.map(_.dataType)) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { // Register before super.compute creates the CometExecIterator, so this listener runs // after the iterator's close has published the final scan metrics. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 478c62f4eda..a59eff0e02a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -110,6 +110,12 @@ class CometExecSuite extends CometTestBase { assert(resolved.get(CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE.key) == "10737418240") flags.foreach(flag => assert(resolved.get(flag.key) == "true", flag.key)) } + + // CometTestBase turns this one on, so it has no default to check here. An explicitly set + // value crosses verbatim, so only the resolved list turns `TRUE` into `true`. + withSQLConf(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key -> "TRUE") { + assert(entries.get(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key) == "true") + } } test("native sort spill files are compressed with the default spill codec") { diff --git a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala index c60cf5cc624..a3046b7f99b 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -27,19 +27,20 @@ import scala.util.Using import org.apache.arrow.c.{ArrowArray, ArrowSchema, Data} import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{IntVector, UInt4Vector, VarCharVector} +import org.apache.arrow.vector.{IntVector, TimeStampMicroTZVector, UInt4Vector, VarCharVector} import org.apache.arrow.vector.complex.StructVector import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider +import org.apache.arrow.vector.types.TimeUnit import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.CometExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.ConstantColumnVector -import org.apache.spark.sql.types.{BinaryType, IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, MapType, StringType, StructField, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometConf} +import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometConf, CometRuntimeException} import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} class NativeUtilSuite extends CometTestBase { @@ -535,4 +536,101 @@ class NativeUtilSuite extends CometTestBase { assert(actual.sameElements(Array(variantType))) } } + + private def timestampField(name: String, timeZone: String): Field = + new Field( + name, + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, timeZone)), + null) + + test("timestampLabelMismatch finds a wrong timezone label at any depth") { + def mismatch(field: Field, dataType: DataType): Option[String] = + NativeUtil.timestampLabelMismatch(field, dataType, null, "col[0]") + + assert(mismatch(timestampField("ts", "UTC"), TimestampType).isEmpty) + assert(mismatch(timestampField("ts", null), TimestampNTZType).isEmpty) + assert( + mismatch(timestampField("ts", "Etc/UTC"), TimestampType).contains( + "col[0]: expected TimestampType labelled UTC, found Timestamp(MICROSECOND, Etc/UTC)")) + // An LTZ value labelled as NTZ, which kernels would read as wall-clock time. + assert(mismatch(timestampField("ts", null), TimestampType).isDefined) + assert(mismatch(timestampField("ts", "UTC"), TimestampNTZType).isDefined) + + val struct = new Field( + "s", + FieldType.nullable(ArrowType.Struct.INSTANCE), + Seq( + new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "events", + FieldType.nullable(ArrowType.List.INSTANCE), + Seq(timestampField("element", null)).asJava)).asJava) + val structType = + StructType( + Seq(StructField("id", IntegerType), StructField("events", ArrayType(TimestampType)))) + assert(mismatch(struct, structType).exists(_.startsWith("col[0].events.element: "))) + + val entries = new Field( + "entries", + FieldType.notNullable(ArrowType.Struct.INSTANCE), + Seq( + new Field("key", FieldType.notNullable(ArrowType.Utf8.INSTANCE), null), + timestampField("value", "+00:00")).asJava) + val map = new Field("m", FieldType.nullable(new ArrowType.Map(false)), Seq(entries).asJava) + assert( + mismatch(map, MapType(StringType, TimestampType)).exists(_.startsWith("col[0].value: "))) + } + + test("checkTimestampLabels closes a batch it rejects and keeps one it accepts") { + Using.resource(new RootAllocator(Long.MaxValue)) { allocator => + def labelled(timeZone: String): ColumnarBatch = { + val vector = new TimeStampMicroTZVector("ts", allocator, timeZone) + vector.allocateNew(1) + vector.setSafe(0, 1L) + vector.setValueCount(1) + new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), 1) + } + Using.resource(new NativeUtil) { nativeUtil => + val accepted = labelled("UTC") + nativeUtil.checkTimestampLabels(accepted, Seq(TimestampType)) + assert(allocator.getAllocatedMemory > 0) + accepted.close() + + val error = intercept[CometRuntimeException] { + nativeUtil.checkTimestampLabels(labelled("Etc/UTC"), Seq(TimestampType)) + } + assert(error.getMessage.contains("at col[0]: expected TimestampType labelled UTC")) + assert(allocator.getAllocatedMemory == 0) + } + } + } + + test("checkTimestampLabels reads a dictionary column's label from its imported values") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + val indexType = new ArrowType.Int(32, true) + val encoding = new DictionaryEncoding(0L, false, indexType) + val indexField = new Field("ts", new FieldType(true, indexType, encoding), null) + val (arrays, schemas) = nativeUtil.allocateArrowStructs(1) + + Using.resources( + new TimeStampMicroTZVector("dictionary", allocator, "Etc/UTC"), + new IntVector(indexField, allocator)) { (values, indices) => + values.allocateNew(1) + values.setSafe(0, 1L) + values.setValueCount(1) + indices.allocateNew(1) + indices.setSafe(0, 0) + indices.setValueCount(1) + val provider = new MapDictionaryProvider(new Dictionary(values, encoding)) + Data.exportVector(allocator, indices, provider, arrays(0), schemas(0)) + } + + val batch = new ColumnarBatch(nativeUtil.importVector(arrays, schemas).toArray, 1) + assert(batch.column(0).isInstanceOf[CometDictionaryVector]) + val error = intercept[CometRuntimeException] { + nativeUtil.checkTimestampLabels(batch, Seq(TimestampType)) + } + assert(error.getMessage.contains("found Timestamp(MICROSECOND, Etc/UTC)"), error.getMessage) + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 60b646fd74e..5ff3f7ed08f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -126,6 +126,10 @@ abstract class CometTestBase // generic " is not supported" message mask the missing reason. // See https://github.com/apache/datafusion-comet/issues/5230. conf.set(CometConf.COMET_STRICT_FALLBACK_REASONS.key, "true") + // Fail a native plan that labels a timestamp with another timezone than Comet's convention. + // The label never changes the values, so result comparisons cannot catch it, and the casts + // at scans and shuffle writes would otherwise repair it without a trace. + conf.set(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key, "true") // SortOrder is incompatible for mixed zero and negative zero floating point values, but // this is an edge case, and we expect most users to allow sorts on floating point, so we // enable this for the tests From 019f3be9d106a8076e8ab22ba31466195bdf3e5b Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 12 Oct 2026 11:19:48 +0000 Subject: [PATCH 74/89] fix: use Spark-compatible shuffle for wide decimal hash keys (#6005) * fix: route wide decimal hash shuffle keys through Spark-compatible hashing * test: retain safe aggregate buffers after wide decimal shuffle fallback * Preserve one-partition decimal shuffle and cover mixed joins * docs: fix native shuffle table formatting * test: avoid implicit widening in wide decimal benchmark * refactor: drop the wide-decimal shuffle guards that this rule makes unreachable #6564 and #6607 added narrower versions of the same restriction: the typed Dataset check in nativeShuffleFailureReasons and the decimal half of hashesDifferentlyFromSpark. supportedHashPartitioningDataType now declines every multi-partition hash key with a decimal wider than 18 digits, so neither can fire. The typed Dataset check only ran when no other reason had been found, and convertsInputForNativeShuffle checks nativeShuffleFailureReasons before hashesDifferentlyFromSpark. The tests those PRs added in CometTypedDatasetSuite and CometShuffleInputConversionSuite assert the routing and still pass. * test: cover a decimal join with a Spark-bucketed input A table bucketed on a DECIMAL(38, 0) key needs no exchange in a join, so only the other input is shuffled, and its partitions have to match the buckets Spark wrote. Without the hash key restriction, native shuffle put the rows in other partitions and the join returned 128 of 1000 rows, with both scans native and no Spark shuffle in the plan. DECIMAL(18, 0) stays on native shuffle and matches. --------- Co-authored-by: Andy Grove --- docs/source/contributor-guide/jvm_shuffle.md | 4 +- .../contributor-guide/native_shuffle.md | 11 +- .../user-guide/latest/tuning/shuffle.md | 11 + .../scala/org/apache/comet/serde/hash.scala | 1 + .../shuffle/CometShuffleExchangeExec.scala | 54 +---- .../approved-plans-v1_4/q49/extended.txt | 2 +- .../q70a/extended.txt | 2 +- .../q14a/extended.txt | 2 +- .../approved-plans-v2_7/q14a/extended.txt | 2 +- .../approved-plans-v2_7/q36a/extended.txt | 2 +- .../approved-plans-v2_7/q49/extended.txt | 2 +- .../approved-plans-v2_7/q5a/extended.txt | 2 +- .../approved-plans-v2_7/q70a/extended.txt | 2 +- .../approved-plans-v2_7/q77a/extended.txt | 2 +- .../approved-plans-v2_7/q80a/extended.txt | 2 +- .../approved-plans-v2_7/q86a/extended.txt | 2 +- .../comet/CometFuzzAggregateSuite.scala | 28 ++- .../org/apache/comet/CometFuzzTestSuite.scala | 24 +- .../comet/exec/CometNativeShuffleSuite.scala | 215 +++++++++++++++++- .../CometWideDecimalShuffleBenchmark.scala | 95 ++++++++ .../CometCelebornShufflePlanningSuite.scala | 32 ++- 21 files changed, 412 insertions(+), 85 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometWideDecimalShuffleBenchmark.scala diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 474764b90e2..1d8552b0f0f 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -52,7 +52,9 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE range key always falls back here. `HashPartitioning` keys must be primitive only by default: setting `spark.comet.shuffle.native.partitioning.hash.nested.enabled` to `true` keeps struct and array keys, and map keys on Spark 4.0 and later, on the native path. The config defaults to - `false`, so a complex hash key falls back to JVM columnar shuffle unless it is enabled. See + `false`, so a complex hash key falls back to JVM columnar shuffle unless it is enabled. + Decimal hash keys with precision greater than 18, including nested decimal leaves, also use + JVM shuffle when there is more than one partition and the shuffle mode allows it. See [Supported partition key types](native_shuffle.md#when-native-shuffle-is-used) for the exact rules. Complex types are fully supported as data columns in both implementations. diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 704fa2a0ea7..80d57fe8081 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -83,12 +83,9 @@ Native shuffle (`CometExchange`) is selected when all of the following condition Spark's `mapsort` normalization makes physical entry order irrelevant. A collated string at any depth still disqualifies the key. The config defaults to `false` pending measurement of the nested hashing paths, so by default a complex hash key falls back to JVM shuffle. - - A hash key that is or contains a decimal wider than 18 digits stays on JVM shuffle when the - shuffle's stage starts at a typed `Dataset` conversion - (`spark.comet.convert.typedDataset.enabled`) and the shuffle has more than one partition. - Native shuffle hashes such decimals differently from Spark - ([#5994](https://github.com/apache/datafusion-comet/issues/5994)). Without the conversion - this shuffle would have used JVM shuffle, and a join partner may still use it. + Decimal leaves with precision greater than 18 also disqualify a hash key when there is more + than one output partition: native decimal hashing differs from Spark's. A one-partition hash + exchange stays native because the writer serializes it as `SinglePartition` without hashing. ## Architecture @@ -531,7 +528,7 @@ independently compressed, allowing parallel decompression during reads. | Input format | Columnar (direct from Comet operators) | Row-based (via ColumnarToRowExec) | | Partitioning logic | Rust implementation | Spark's partitioner | | Supported schemes | Hash, Range, Single, RoundRobin | Hash, Range, Single, RoundRobin | -| Partition key types | Primitives only (Hash, Range) | Any type | +| Partition key types | See the hash and range rules above | Any type | | Performance | Higher (no format conversion) | Lower (columnar→row→columnar) | | Writer variants | Single path | Bypass (hash) and sort-based | diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index b7ad3d2f090..31aca7bd024 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -53,6 +53,17 @@ scalar types. Hash partitioning keys must be scalar types unless and later) map keys. That setting is disabled by default until the performance of the nested hashing paths has been measured. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays. +Hash partitioning into multiple partitions on decimal keys with precision greater than 18 falls back because native +hashing does not match Spark's partition assignments. Mixing the two partitioners can silently lose join rows. +The restriction also applies to decimal leaves in nested hash keys. Wider decimals remain supported as payload +columns, range partitioning keys, and in single-partition shuffles, including one-partition hash exchanges. + +With `spark.comet.shuffle.mode=auto`, Comet uses Columnar Shuffle when eligible, adding a columnar-to-row-to-columnar +conversion. With `native`, it uses Spark shuffle, which can also restore aggregates such as `collect_list` and +`collect_set` to Spark to preserve buffer compatibility. Celeborn uses Spark/Celeborn shuffle for unsupported native +keys because it has no Comet columnar fallback. Spark-compatible native decimal hashing is tracked in +[#5994](https://github.com/apache/datafusion-comet/issues/5994). + ### Columnar (JVM) Shuffle Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index ee3e80059d5..760bb3c2829 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -134,6 +134,7 @@ private object HashUtils { } private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { + // Keep in sync with CometShuffleExchangeExec's hash-key restriction until #5994 is fixed. case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason) case s: StructType => s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index c5e5c7641e5..4113fa920b3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -36,11 +36,11 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Exp import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.catalyst.plans.physical._ -import org.apache.spark.sql.comet.{CometFilterExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometPlan, CometProjectExec, CometScanWrapper, CometSinkPlaceHolder, CometSparkToColumnarExec, NativeExecContext} +import org.apache.spark.sql.comet.{CometFilterExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometPlan, CometProjectExec, CometSinkPlaceHolder, CometSparkToColumnarExec, NativeExecContext} import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec -import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, Exchange, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} +import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} @@ -579,38 +579,20 @@ object CometShuffleExchangeExec * Whether the shuffle hashes into more than one partition a key that native shuffle, reading * converted rows, may put in a different partition from Spark's partitioner. Native shuffle * must not take over such a shuffle from the JVM columnar shuffle: a join with an input that is - * still partitioned by Spark would put matching keys in different partitions. The keys are: - * - * - A decimal wider than 18 digits, which native shuffle hashes differently (#5994). - * - A string, or any value computed from one, such as `hash(s)`. Spark reads the string's - * bytes as they are, but the import of the converted batch into native replaces invalid - * UTF-8 before native shuffle evaluates the key. - * - * TODO: allow wide decimals once native hashing matches Spark for them. + * still partitioned by Spark would put matching keys in different partitions. Such a key is a + * string, or any value computed from one, such as `hash(s)`. Spark reads the string's bytes as + * they are, but the import of the converted batch into native replaces invalid UTF-8 before + * native shuffle evaluates the key. */ private def hashesDifferentlyFromSpark(s: ShuffleExchangeExec): Boolean = s.outputPartitioning match { case HashPartitioning(expressions, numPartitions) => numPartitions > 1 && expressions.exists { key => - key.dataType.existsRecursively(DecimalType.isByteArrayDecimalType) || key.exists(_.dataType.existsRecursively(_.isInstanceOf[StringType])) } case _ => false } - /** - * Whether the stage feeding a shuffle starts at a typed Dataset conversion (see - * [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]). `CometExecRule` decides the shuffle - * before it removes its placeholders, so the conversion is still inside its `CometScanWrapper`. - */ - private def readsTypedDatasetConversion(plan: SparkPlan): Boolean = plan match { - case _: Exchange => false - case CometScanWrapper(_, wrapped) => readsTypedDatasetConversion(wrapped) - case conversion: CometSparkToColumnarExec => - conversion.child.isInstanceOf[SerializeFromObjectExec] - case other => other.children.exists(readsTypedDatasetConversion) - } - /** * Reasons the native shuffle path cannot handle this shuffle. Empty means native is supported. * Pure: does not tag the node. @@ -642,12 +624,11 @@ object CometShuffleExchangeExec _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | _: DateType => true - case _: DecimalType => - // TODO enforce this check - // https://github.com/apache/datafusion-comet/issues/3079 - // Decimals with precision > 18 require Java BigDecimal conversion before hashing - // d.precision <= 18 - true + case d: DecimalType => + // Match the SQL hash restriction in serde/HashUtils until #5994 fixes native encoding. + // Different partition assignments break mixed native/Spark joins. A single partition + // does not hash the key: CometNativeShuffleWriter serializes it as SinglePartition. + d.precision <= 18 || s.outputPartitioning.numPartitions == 1 case dt if isTimeType(dt) => true case StructType(fields) if nestedHashPartitioningEnabled => @@ -709,19 +690,6 @@ object CometShuffleExchangeExec reasons += s"unsupported hash partitioning data type for native shuffle: $dt" } } - // A typed Dataset conversion moves the shuffle above it from Comet's columnar shuffle, - // which partitions with Spark's hash, to native shuffle. Native shuffle hashes a decimal - // wider than 18 digits differently from Spark, so a join with an input that is still on - // the columnar shuffle would put matching keys in different partitions. Leave such a - // shuffle where it was. A single partition hashes nothing. - // TODO: remove once native hashing matches Spark for wide decimals (#5994). - if (CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf) && reasons.isEmpty && - partitioning.numPartitions > 1 && - expressions.exists(_.dataType.existsRecursively(DecimalType.isByteArrayDecimalType)) && - readsTypedDatasetConversion(s.child)) { - reasons += "a shuffle above a typed Dataset conversion that hashes a decimal wider " + - "than 18 digits stays on Comet's columnar shuffle, which hashes it as Spark does" - } case SinglePartition => // we already checked that the input types are supported case RangePartitioning(orderings, _) => diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt index 6deb75bea15..1f9d12fc934 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt index 185ce8a5247..09918cf843c 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt index 9a11b4895a7..2a3100957f1 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt index a4ef99efa8e..8f5b9835de5 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt index f11c15706a4..04405d9b84e 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt index 6deb75bea15..1f9d12fc934 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt index f42b9795c47..4b41ce8a173 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt index 594ae116008..6b6725c38f5 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt index daccf86c3f7..a3fbdd2c329 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt index f4db0be1188..ea972a9b91d 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt index 6269f5ddc51..b2fe83f4b36 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala index 3674887f80f..32b21cb92ad 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala @@ -19,6 +19,10 @@ package org.apache.comet +import org.apache.spark.sql.execution.aggregate.HashAggregateExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.types.DecimalType + import org.apache.comet.DataTypeSupport.isComplexType class CometFuzzAggregateSuite extends CometFuzzTestBase { @@ -31,7 +35,16 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - checkSparkAnswerAndOperator(sql) + val hasWideDecimalKey = df.schema(col).dataType match { + case d: DecimalType => d.precision > 18 + case _ => false + } + // Wide decimal hash keys require Spark shuffle when columnar shuffle is disabled. + if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { + checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) + } else { + checkSparkAnswerAndOperator(sql) + } } } @@ -55,7 +68,18 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - checkSparkAnswerAndOperator(sql) + val hasWideDecimalKey = Seq("c1", "c2", "c3", col).exists { key => + df.schema(key).dataType match { + case d: DecimalType => d.precision > 18 + case _ => false + } + } + // Check both GROUP BY and DISTINCT keys, not unrelated decimal payload columns. + if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { + checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) + } else { + checkSparkAnswerAndOperator(sql) + } } } diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index 55fb6010c99..91bd098c1f3 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -160,6 +160,17 @@ class CometFuzzTestSuite extends CometFuzzTestBase { } test("distribute by single column (complex types)") { + // Inspect only the key schema: any wide decimal leaf requires Spark's hash partition + // assignments, even when native hashing of nested keys is enabled. + def hasWideDecimal(dataType: DataType): Boolean = dataType match { + case decimal: DecimalType => decimal.precision > 18 + case StructType(fields) => fields.exists(field => hasWideDecimal(field.dataType)) + case ArrayType(elementType, _) => hasWideDecimal(elementType) + case MapType(keyType, valueType, _) => + hasWideDecimal(keyType) || hasWideDecimal(valueType) + case _ => false + } + val df = spark.read.parquet(filename) df.createOrReplaceTempView("t1") val columns = df.schema.fields.filter(f => isComplexType(f.dataType)).map(_.name) @@ -180,15 +191,20 @@ class CometFuzzTestSuite extends CometFuzzTestBase { } assert(cometShuffleExchanges.length == expectedNumCometShuffles) - // With the config enabled these keys do run through native shuffle. This is the widest - // nested-type coverage in the repo, so it is worth asserting that they are admitted rather - // than only that they fall back. + // Enabling nested keys admits supported types, but wide decimal leaves still require + // Spark's hash partition assignments. JVM shuffle supports both kinds of key. withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") { val enabledDf = spark.sql(sql) enabledDf.collect() val enabledPlan = enabledDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan - assert(collectCometShuffleExchanges(enabledPlan).length == 1) + val expectedEnabledShuffles = + if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && + hasWideDecimal(df.schema(col).dataType)) 0 + else 1 + assert( + collectCometShuffleExchanges(enabledPlan).length == expectedEnabledShuffles, + s"Unexpected shuffle for ${df.schema(col)}") } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 0c56e11b9ab..8d91ec71765 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -41,15 +41,16 @@ import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.expressions.aggregate.Final import org.apache.spark.sql.catalyst.plans.logical.LocalRelation import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning -import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource} -import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, LocalTableScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} -import org.apache.spark.sql.functions.{broadcast, col, count, countDistinct, sum} +import org.apache.spark.sql.functions.{broadcast, col, count, countDistinct, spark_partition_id, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructField, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} @@ -440,7 +441,209 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val shuffled = df .select($"_1") .repartition(10, col(c)) - checkShuffleAnswer(shuffled, 1, checkNativeOperators = true) + val nativeHashSupported = df.schema(c).dataType match { + case d: DecimalType => d.precision <= 18 + case _ => true + } + checkShuffleAnswer( + shuffled, + if (nativeHashSupported) 1 else 0, + checkNativeOperators = nativeHashSupported) + } + } + } + } + } + } + + for (precision <- Seq(18, 19, 38)) { + test(s"decimal hash shuffle preserves Spark partitions at precision $precision") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTable("decimal_shuffle") { + sql(s"CREATE TABLE decimal_shuffle(id INT, k DECIMAL($precision, 0)) USING parquet") + val maximum = "9" * precision + sql(s"""INSERT INTO decimal_shuffle VALUES + |(0, NULL), (1, 0), (2, 1), (3, -1), (4, $maximum), (5, -$maximum) + |""".stripMargin) + for (mode <- Seq("native", "auto", "jvm")) { + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val shuffled = spark.table("decimal_shuffle").repartition(7, $"k") + val native = precision <= 18 && mode != "jvm" + val sparkShuffle = precision > 18 && mode == "native" + checkCometExchange(shuffled, if (sparkShuffle) 0 else 1, native) + assert(shuffled.queryExecution.executedPlan.collect { case _: ShuffleExchangeExec => + 1 + }.sum == (if (sparkShuffle) 1 else 0)) + // Result equality alone cannot detect a different hash partition assignment. + checkSparkAnswer(shuffled.withColumn("partition", spark_partition_id())) + } + } + } + } + } + } + + test("wide decimals remain supported in shuffle payloads, ranges and single partitions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_NATIVE_RANGE_PARTITIONING_ENABLED.key -> "true") { + withTable("decimal_shuffle") { + sql("CREATE TABLE decimal_shuffle(id INT, k DECIMAL(38, 0)) USING parquet") + sql( + "INSERT INTO decimal_shuffle VALUES (0, NULL), (1, 1), (2, -1), " + + "(3, 99999999999999999999999999999999999999)") + for (mode <- Seq("native", "auto", "jvm")) { + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val input = spark.table("decimal_shuffle") + Seq( + input.repartition(7, $"id"), + input.repartitionByRange(7, $"k"), + input.repartition(1)).foreach { shuffled => + checkCometExchange(shuffled, 1, native = mode != "jvm") + checkSparkAnswer(shuffled) + } + } + } + } + } + } + + test( + "wide decimal hash shuffle keeps native aggregates unless multiple partitions need fallback") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i % 3, i % 7)), "decimal_shuffle") { + for ((precision, partitions, mode) <- Seq( + (18, 2, "native"), + (38, 2, "native"), + (38, 1, "native"), + (38, 1, "auto")); + function <- Seq("collect_list", "collect_set")) { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString, + CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val key = s"CAST(_1 AS DECIMAL($precision, 0))" + val df = + sql(s"SELECT $key, sort_array($function(_2)) FROM decimal_shuffle GROUP BY $key") + val plan = df.queryExecution.executedPlan + val nativeExpected = precision <= 18 || partitions == 1 + assert( + plan.collect { case _: CometHashAggregateExec => 1 }.sum == + (if (nativeExpected) 2 else 0), + plan.treeString) + assert( + plan.collect { case _: ObjectHashAggregateExec => 1 }.sum == + (if (nativeExpected) 0 else 2), + plan.treeString) + // Restoring Spark's aggregate buffers must retain the accelerated input scan. + assert(plan.collect { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) + val exchanges = checkCometExchange(df, if (nativeExpected) 1 else 0, native = true) + // GROUP BY retains HashPartitioning even when there is only one partition. + assert(exchanges.forall(_.outputPartitioning.isInstanceOf[HashPartitioning])) + checkSparkAnswer(df) + } + } + } + } + } + + test("wide decimal join keeps native and Spark inputs copartitioned") { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "auto", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "false", + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet,json", + SQLConf.SHUFFLE_PARTITIONS.key -> "7", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + withTable("decimal_parquet", "decimal_json") { + for ((table, format) <- Seq("decimal_parquet" -> "parquet", "decimal_json" -> "json")) { + sql(s"CREATE TABLE $table(id INT, k DECIMAL(38, 0)) USING $format") + sql(s"""INSERT INTO $table VALUES + |(1, 1), (2, -1), (3, 123456789012345678901234567890), + |(4, 99999999999999999999999999999999999999) + |""".stripMargin) + } + for (adaptive <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val df = sql("""SELECT p.id, j.id FROM decimal_parquet p + |JOIN decimal_json j ON p.k = j.k""".stripMargin) + checkAnswer(df, Seq(Row(1, 1), Row(2, 2), Row(3, 3), Row(4, 4))) + val plan = df.queryExecution.executedPlan + val exchanges = collect(plan) { case e: CometShuffleExchangeExec => e } + assert(exchanges.size == 2, plan.treeString) + assert(exchanges.forall(_.shuffleType == CometColumnarShuffle), plan.treeString) + assert(collect(plan) { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) + assert(collect(plan) { case _: FileSourceScanExec => 1 }.sum == 1, plan.treeString) + } + } + } + } + } + + test("decimal join keeps a Spark-bucketed input copartitioned") { + withSQLConf( + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + SQLConf.BUCKETING_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "8", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + for (precision <- Seq(18, 38)) { + withTable("decimal_bucketed", "decimal_plain") { + // Spark writes the buckets with its own hash. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val input = + spark.range(0, 1000).selectExpr(s"CAST(id AS DECIMAL($precision, 0)) AS k", "id") + input.write.bucketBy(8, "k").saveAsTable("decimal_bucketed") + input.write.saveAsTable("decimal_plain") + } + for (adaptive <- Seq(false, true); mode <- Seq("auto", "native")) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val df = sql("""SELECT b.id, p.id FROM decimal_bucketed b + |JOIN decimal_plain p ON b.k = p.k""".stripMargin) + val (_, plan) = checkSparkAnswer(df) + assert(collect(plan) { case _: CometNativeScanExec => 1 }.sum == 2, plan.treeString) + // Only the plain input is shuffled, so its partitions have to match the buckets. + val shuffles = collect(plan) { + case e: CometShuffleExchangeExec => Some(e.shuffleType) + case _: ShuffleExchangeExec => None + } + val expected = + if (precision <= 18) Some(CometNativeShuffle) + else if (mode == "auto") Some(CometColumnarShuffle) + else None + assert(shuffles == Seq(expected), plan.treeString) + } + } + } + } + } + } + + test("decimal hash shuffle checks nested keys recursively") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withNestedHashPartitioning { + for (precision <- Seq(18, 38)) { + withTable("decimal_shuffle") { + sql(s"""CREATE TABLE decimal_shuffle( + |id INT, s STRUCT>, + |a ARRAY>) USING parquet + |""".stripMargin) + sql("""INSERT INTO decimal_shuffle VALUES + |(0, NULL, NULL), + |(1, named_struct('a', array(1, -1)), array(named_struct('d', 1))), + |(2, named_struct('a', array(2, NULL)), array(named_struct('d', NULL))) + |""".stripMargin) + for (key <- Seq("s", "a")) { + val shuffled = spark.table("decimal_shuffle").repartition(7, col(key)) + checkCometExchange(shuffled, if (precision <= 18) 1 else 0, native = true) + checkSparkAnswer(shuffled.withColumn("partition", spark_partition_id())) } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWideDecimalShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWideDecimalShuffleBenchmark.scala new file mode 100644 index 00000000000..6801213425d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWideDecimalShuffleBenchmark.scala @@ -0,0 +1,95 @@ +/* + * 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. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.Row +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * Measures wide-decimal hash shuffle routing with COUNT, independently of decimal AVG support. + * Run this benchmark on both revisions to compare the routing change. Fixture creation, result + * validation and route reporting are untimed. The session uses local[1]; `--reverse` reverses + * case order, and `--validate-only` checks the fixture and results without collecting timings. + */ +object CometWideDecimalShuffleBenchmark extends CometBenchmarkBase { + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val rows = 1024 * 1024 + val groups = 10000 + val partitions = 4 + val filePartitionBytes = 16 * 1024 * 1024 + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString, + // Keep each file in one split, and prevent Spark from combining files into one task. + SQLConf.FILES_MAX_PARTITION_BYTES.key -> filePartitionBytes.toString, + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> filePartitionBytes.toString, + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + withTempPath { dir => + withTempTable("parquetV1Table") { + val query = "SELECT k, COUNT(v) FROM parquetV1Table GROUP BY k" + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + prepareTable( + dir, + spark + .range(0L, rows.toLong, 1L, partitions) + .selectExpr( + s"CAST(id % $groups AS DECIMAL(38, 2)) AS k", + "CAST(id % 97 AS INT) AS v")) + assert(spark.table("parquetV1Table").rdd.getNumPartitions == partitions) + expected = spark.sql(query).collect().toSeq.sortBy(_.getDecimal(0)) + } + val benchmark = new Benchmark("wide_decimal_hash_shuffle", rows.toLong, output = output) + val modes = Seq("Spark", "native", "auto", "jvm") + for (mode <- (if (mainArgs.contains("--reverse")) modes.reverse else modes)) { + val configs = Seq( + CometConf.COMET_ENABLED.key -> (mode != "Spark").toString, + CometConf.COMET_SHUFFLE_MODE.key -> (if (mode == "Spark") "native" else mode)) + withSQLConf(configs: _*) { + val df = spark.sql(query) + assert(df.collect().toSeq.sortBy(_.getDecimal(0)) == expected) + val plan = df.queryExecution.executedPlan + val routes = plan.collect { + case exchange: CometShuffleExchangeExec => exchange.shuffleType.toString + case _: ShuffleExchangeExec => "Spark shuffle" + } + assert(routes.size == 1, plan.treeString) + benchmark.out.println(s"Wide-decimal shuffle mode=$mode: ${routes.head}") + benchmark.out.println(plan.treeString) + } + benchmark.addCase(mode) { _ => + withSQLConf(configs: _*) { + spark.sql(query).collect() + } + } + } + if (!mainArgs.contains("--validate-only")) benchmark.run() + } + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index 1d217bf6a25..5be64f696dd 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -703,17 +703,27 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { } } - test(s"unsupported native repartition executes Spark fallback with AQE=$adaptive") { - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - CometConf.COMET_SHUFFLE_MODE.key -> "native", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "false") { - val nativeRegistrations = manager.nativeRegistrations.get() - val query = input.repartition(2) - assertSparkExchange(query.queryExecution.executedPlan) - checkAnswer(query, (1L to 32L).map(Row(_))) - assert(cometExchanges(query.queryExecution.executedPlan).isEmpty) - assert(manager.nativeRegistrations.get() == nativeRegistrations) + for (wideDecimal <- Seq(false, true)) { + test( + s"unsupported repartition executes Spark fallback: wideDecimal=$wideDecimal, " + + s"AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "false") { + val nativeRegistrations = manager.nativeRegistrations.get() + val sparkRegistrations = manager.sparkRegistrations.get() + val query = if (wideDecimal) { + input.repartition(2, col("value").cast("decimal(38, 0)")) + } else { + input.repartition(2) + } + assertSparkExchange(query.queryExecution.executedPlan) + checkAnswer(query, (1L to 32L).map(Row(_))) + assert(cometExchanges(query.queryExecution.executedPlan).isEmpty) + assert(manager.nativeRegistrations.get() == nativeRegistrations) + assert(manager.sparkRegistrations.get() > sparkRegistrations) + } } } From 489632ab1ea3d4af3dfac0c9f34d5e9b45da3bd0 Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Mon, 12 Oct 2026 11:22:03 +0000 Subject: [PATCH 75/89] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning (#6814) * fix: preserve empty Iceberg partitions after dynamic pruning * test: simplify empty Iceberg scan credential regression * test: cover single-split Iceberg joins after DPP --- .../src/execution/operators/iceberg_scan.rs | 8 +- .../operator/CometIcebergNativeScan.scala | 17 +- .../comet/CometIcebergNativeSuite.scala | 280 +++++++++++++++++- 3 files changed, 297 insertions(+), 8 deletions(-) diff --git a/native/core/src/execution/operators/iceberg_scan.rs b/native/core/src/execution/operators/iceberg_scan.rs index ad69cac4f03..7b18c7e3bea 100644 --- a/native/core/src/execution/operators/iceberg_scan.rs +++ b/native/core/src/execution/operators/iceberg_scan.rs @@ -35,7 +35,8 @@ use datafusion::physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, }; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, Partitioning, + PlanProperties, }; use futures::{Stream, StreamExt, TryStreamExt}; use iceberg::arrow::ScanMetrics; @@ -178,6 +179,11 @@ impl IcebergScanExec { context: Arc, ) -> DFResult { let output_schema = Arc::clone(&self.output_schema); + // DPP can leave an empty execution partition so that global aggregates still run. + // No files need reading, so do not initialize storage or request catalog credentials. + if tasks.is_empty() { + return Ok(Box::pin(EmptyRecordBatchStream::new(output_schema))); + } let file_io = load_file_io( &self.catalog_properties, &self.metadata_location, diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index 335f02ff034..22895e5145a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -1280,17 +1280,24 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit perPartitionBuilders += partitionBuilder.build() } case other if other.getClass.getName == "org.apache.spark.rdd.ParallelCollectionRDD" => - // Spark's BatchScanExec.inputRDD returns sparkContext.parallelize(empty, 1) when - // DPP filtering removes all input partitions. That ParallelCollectionRDD is the only - // non-DataSourceRDD shape its inputRDD produces, so reaching this branch means "DPP - // pruned everything"; emit no per-partition data and let native execution return empty. + // Spark 3.4's BatchScanExec.inputRDD returns sparkContext.parallelize(empty, 1) when + // DPP removes every split of a scan that reports SinglePartition. Among supported + // versions, only Spark 3.4 infers SinglePartition from one execution partition, + // which may be a single split or a group of splits. Other fully pruned scans take + // the DataSourceRDD branch above: ungrouped scans have no partitions, while grouped + // scans may retain empty partitions. Preserve Spark's empty + // execution partition so global aggregates still run and scans in the same native block + // retain matching partition counts when no exchange separates them. // Re-querying scan.toBatch.planInputPartitions() to verify is unreliable because // Iceberg's Scan state after filter() doesn't always reflect post-DPP partitions on // a re-call (V2 scan state is one-shot for the materialized inputRDD). Matched by class // name because ParallelCollectionRDD is private[spark]. logDebug( "BatchScanExec.inputRDD is ParallelCollectionRDD (DPP pruned all partitions); " + - "skipping per-partition serialization") + "preserving empty execution partitions") + other.partitions.foreach { _ => + perPartitionBuilders += OperatorOuterClass.IcebergScan.getDefaultInstance + } case other => throw new IllegalStateException( "Expected DataSourceRDD or ParallelCollectionRDD from BatchScanExec, " + diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index f271f7d6ec2..02d2640d10d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -22,6 +22,7 @@ package org.apache.comet import java.io.File import java.net.URI import java.nio.charset.StandardCharsets.UTF_8 +import java.util.UUID import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ @@ -41,11 +42,12 @@ import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, StringType, StructType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, IntegerType, StringType, StructType, TimestampType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.cloud.s3.MinioLocationScopedCredentialProvider import org.apache.comet.iceberg.{IcebergReflection, RESTCatalogHelper} -import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} import org.apache.comet.testing.{FuzzDataGenerator, SchemaGenOptions} /** @@ -5150,6 +5152,280 @@ class CometIcebergNativeSuite } } + Seq(false, true).foreach { aqeEnabled => + test(s"DPP - global aggregate after pruning a single split (AQE=$aqeEnabled)") { + assume(icebergAvailable, "Iceberg not available") + withTempIcebergDir { warehouseDir => + val dimDir = new File(warehouseDir, "dim_parquet") + withSQLConf( + "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.aqe_cat.type" -> "hadoop", + "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString, + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + val factTable = "aqe_cat.db.single_split_dpp_fact" + spark.sql(s""" + CREATE TABLE $factTable (amount INT, store_id INT) + USING iceberg PARTITIONED BY (store_id) + """) + try { + spark + .range(100) + .selectExpr("CAST(id AS INT) AS amount", "5 AS store_id") + .coalesce(1) + .writeTo(factTable) + .append() + + // The regression requires one original split: Spark 3.4 can then omit the + // exchange between the partial and final global aggregates. + val unprunedPlan = spark.table(factTable).queryExecution.executedPlan + assertSingleNativeScan(unprunedPlan) + val unprunedScan = collectIcebergNativeScans(unprunedPlan).head + assert(unprunedScan.numPartitions == 1, s"Expected one split:\n$unprunedPlan") + assert( + unprunedScan.perPartitionData + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) + .sum == 1, + s"Expected one file scan task before pruning:\n$unprunedPlan") + + spark + .range(1) + .selectExpr("99 AS store_id", "'X' AS country") + .write + .parquet(dimDir.getAbsolutePath) + withTempView("single_split_dim") { + spark.read + .parquet(dimDir.getAbsolutePath) + .createOrReplaceTempView("single_split_dim") + val query = s"""SELECT /*+ BROADCAST(d) */ COUNT(*), SUM(f.amount) + |FROM $factTable f JOIN single_split_dim d ON f.store_id = d.store_id + |WHERE d.country = 'X'""".stripMargin + val expected = Seq(Row(0L, null)) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + checkAnswer(spark.sql(query), expected) + } + val df = spark.sql(query) + checkCometAnswer(df, expected) + + val plan = df.queryExecution.executedPlan + assertSingleNativeScan(plan) + val scan = collectIcebergNativeScans(plan).head + if (!isSpark35Plus) { + assert( + collect(plan) { case exchange: ShuffleExchangeLike => exchange }.isEmpty, + s"Expected the single-split aggregate to run without a shuffle:\n$plan") + assert(scan.numPartitions == 1, s"Expected one empty execution partition:\n$plan") + } + assert( + collectIcebergDPPSubqueries(plan).nonEmpty, + s"Expected executable DPP on the native scan:\n$plan") + assert( + scan.perPartitionData + .forall(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount == 0), + s"Expected DPP to prune all file scan tasks:\n$plan") + assert(scan.metrics("num_splits").value == 0, "Pruned files must not be read") + } + } finally { + spark.sql(s"DROP TABLE $factTable") + } + } + } + } + } + + for { + aqeEnabled <- Seq(false, true) + prunedSide <- Seq("a", "b") + } { + test(s"DPP - single-split sort-merge join (AQE=$aqeEnabled, pruned=$prunedSide)") { + assume(icebergAvailable, "Iceberg not available") + withTempIcebergDir { warehouseDir => + val dimDir = new File(warehouseDir, "dim_parquet") + withSQLConf( + "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.aqe_cat.type" -> "hadoop", + "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true", + CometConf.COMET_EXEC_SORT_MERGE_JOIN_ENABLED.key -> "true", + CometConf.COMET_FORCE_SHJ.key -> "false") { + val factA = "aqe_cat.db.single_split_join_a" + val factB = "aqe_cat.db.single_split_join_b" + val tables = Seq(factA, factB) + def fileTaskCount(scan: CometIcebergNativeScanExec): Int = + scan.perPartitionData + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) + .sum + + try { + tables.foreach { table => + spark.sql(s""" + CREATE TABLE $table (id INT, v INT, store_id INT) + USING iceberg PARTITIONED BY (store_id) + """) + spark + .range(100) + .selectExpr("CAST(id AS INT) AS id", "CAST(id AS INT) AS v", "5 AS store_id") + .coalesce(1) + .writeTo(table) + .append() + + val unprunedPlan = spark.table(table).queryExecution.executedPlan + assertSingleNativeScan(unprunedPlan) + val scan = collectIcebergNativeScans(unprunedPlan).head + assert(scan.numPartitions == 1, s"Expected one split for $table:\n$unprunedPlan") + assert( + fileTaskCount(scan) == 1, + s"Expected one file task for $table:\n$unprunedPlan") + } + + spark + .range(1) + .selectExpr("99 AS store_id", "'X' AS country") + .write + .parquet(dimDir.getAbsolutePath) + withTempView("single_split_join_dim") { + spark.read + .parquet(dimDir.getAbsolutePath) + .createOrReplaceTempView("single_split_join_dim") + val query = s"""SELECT /*+ BROADCAST(d), MERGE(a, b) */ COUNT(*), SUM(b.v) + |FROM $factA a JOIN $factB b ON a.id = b.id + |JOIN single_split_join_dim d ON $prunedSide.store_id = d.store_id + |WHERE d.country = 'X'""".stripMargin + val expected = Seq(Row(0L, null)) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + checkAnswer(spark.sql(query), expected) + } + val df = spark.sql(query) + val initialPlan = df.queryExecution.executedPlan + assert( + collectIcebergNativeScans(initialPlan).size == 2, + s"Expected two native Iceberg scans before execution:\n$initialPlan") + assert( + collectIcebergDPPSubqueries(initialPlan).size == 1, + s"Expected DPP on exactly one scan before execution:\n$initialPlan") + assert( + collect(initialPlan) { case join: CometSortMergeJoinExec => join }.size == 1, + s"Expected a native sort-merge join before execution:\n$initialPlan") + checkCometAnswer(df, expected) + + val plan = df.queryExecution.executedPlan + val scans = collectIcebergNativeScans(plan) + if (isSpark35Plus && aqeEnabled && scans.isEmpty) { + // Newer Spark versions shuffle the join inputs, allowing AQE to replace + // the join and its scans with an empty relation after DPP prunes one input. + // This checks result compatibility; the SinglePartition regression assertions + // below target Spark 3.4, where both scans must remain in one native block. + assert( + collect(plan) { case empty: CometEmptyRelationExec => empty }.nonEmpty, + s"Expected AQE to replace the pruned join with an empty relation:\n$plan") + } else { + assert(scans.size == 2, s"Expected two native Iceberg scans:\n$plan") + val prunedScans = scans.filter(scan => collectIcebergDPPSubqueries(scan).nonEmpty) + assert(prunedScans.size == 1, s"Expected DPP on exactly one scan:\n$plan") + val prunedScan = prunedScans.head + assert( + fileTaskCount(prunedScan) == 0, + s"Expected all file tasks to be pruned:\n$plan") + assert( + prunedScan.metrics("num_splits").value == 0, + "Pruned files must not be read") + assert( + scans.filterNot(_ eq prunedScan).map(fileTaskCount) == Seq(1), + s"Expected the other scan to retain its file task:\n$plan") + } + + if (!isSpark35Plus) { + val join = collectFirst(plan) { case j: CometSortMergeJoinExec => j } + .getOrElse(fail(s"Expected a native sort-merge join:\n$plan")) + assert( + collect(join) { case exchange: ShuffleExchangeLike => exchange }.isEmpty, + s"Expected no shuffle between the join and its scans:\n$plan") + val inputs = scala.collection.mutable.ArrayBuffer.empty[SparkPlan] + join.foreachUntilCometInput(join)(input => inputs += input) + val nativeScans = inputs.collect { case scan: CometIcebergNativeScanExec => scan } + assert(nativeScans.size == 2, s"Expected both scans in one native block:\n$plan") + assert( + nativeScans.forall(_.numPartitions == 1), + s"Expected both scans to retain one execution partition:\n$plan") + // Check execution input order, not just SQL aliases: pruning either input used + // to fail differently when its partition array was empty but the other was not. + val expectedTaskCounts = if (prunedSide == "a") Seq(0, 1) else Seq(1, 0) + assert( + nativeScans.map(fileTaskCount).toSeq == expectedTaskCounts, + s"Unexpected file tasks in native input order:\n$plan") + } + } + } finally { + tables.foreach(table => spark.sql(s"DROP TABLE IF EXISTS $table")) + } + } + } + } + } + + test("DPP - empty Iceberg scan does not request S3 policy locations") { + // Exercise the empty executor scan emitted when DPP preserves Spark's last partition. + // A fresh catalog/bucket prevents a previously initialized FileIO from hiding the call. + val identity = UUID.randomUUID().toString + val common = OperatorOuterClass.IcebergScanCommon + .newBuilder() + .setMetadataLocation(s"s3://empty-scan-$identity/table/metadata/v1.metadata.json") + .setCatalogName(s"empty_scan_$identity") + .setDataFileConcurrencyLimit(1) + .putCatalogProperties( + "s3.comet.credential.provider.class", + classOf[MinioLocationScopedCredentialProvider].getName) + .addRequiredSchema( + OperatorOuterClass.SparkStructField + .newBuilder() + .setName("amount") + .setNullable(true) + .setDataType(QueryPlanSerde.serializeDataType(IntegerType).get)) + .build() + val planBytes = OperatorOuterClass.Operator + .newBuilder() + .setIcebergScan(OperatorOuterClass.IcebergScan.newBuilder().setCommon(common)) + .build() + .toByteArray + + val result = spark.sparkContext + .parallelize(Seq(1), 1) + .mapPartitions { _ => + MinioLocationScopedCredentialProvider.resetCounters() + val iterator = new CometExecIterator( + CometExec.newIterId, + Array.empty[Object], + 1, + planBytes, + CometMetricNode(Map.empty), + 1, + 0) + try { + var rows = 0L + while (iterator.hasNext) { + val batch = iterator.next() + rows += batch.numRows() + batch.close() + } + Iterator.single((rows, MinioLocationScopedCredentialProvider.locationCallCount())) + } finally { + iterator.close() + } + } + .collect() + + assert(result.toSeq == Seq((0L, 0)), s"Expected no rows or policy requests: ${result.toSeq}") + } + test("AQE DPP - empty broadcast result prunes all partitions") { assume(icebergAvailable, "Iceberg not available") withTempIcebergDir { warehouseDir => From 16dcf844c5d71a5a704116c0d378f277b4b2da9a Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:52:18 +0000 Subject: [PATCH 76/89] test: migrate regex extraction coverage to SQL fixtures (#6882) Part of #6627. Preserve Java regex semantics and dispatcher routing while replacing 18 Scala tests, including multi-batch extraction coverage. --- .../expressions/string/regexp_extract.sql | 35 +++-- .../expressions/string/regexp_extract_all.sql | 34 ++++- .../string/regexp_extract_batches.sql | 32 +++++ .../expressions/string/regexp_instr.sql | 23 ++- .../apache/comet/CometRegExpJvmSuite.scala | 136 +----------------- 5 files changed, 103 insertions(+), 157 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/string/regexp_extract_batches.sql diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql index 81f01248444..987c457cb35 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql @@ -22,35 +22,54 @@ statement CREATE TABLE test_regexp_extract(s string) USING parquet statement -INSERT INTO test_regexp_extract VALUES ('abc123def'), ('no match'), (NULL), ('xyz789'), ('hello world'), ('aa') +INSERT INTO test_regexp_extract VALUES + ('abc123def'), ('no match'), (NULL), ('xyz789'), ('hello world'), ('aa'), + ('foo123bar'), ('abc'), ('def'), ('ab'), ('bb'), ('foobar'), ('foobaz'), + ('FOO123'), ('foo456'), ('bar789') -- group 0: entire match query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, r'\d+', 0) FROM test_regexp_extract +SELECT s, regexp_extract(s, r'\d+', 0) FROM test_regexp_extract -- group 1: first capturing group query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract +SELECT s, regexp_extract(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract -- group 2: second capturing group query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract +SELECT s, regexp_extract(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract -- no match returns empty string query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, 'NOMATCH', 0) FROM test_regexp_extract +SELECT s, regexp_extract(s, 'NOMATCH', 0) FROM test_regexp_extract -- backreference pattern (Java-only) query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, r'(\w)\1', 0) FROM test_regexp_extract +SELECT s, regexp_extract(s, r'(\w)\1', 0) FROM test_regexp_extract -- lookahead (Java-only) query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, r'abc(?=\d)', 0) FROM test_regexp_extract +SELECT s, regexp_extract(s, r'abc(?=\d)', 0) FROM test_regexp_extract + +query expect_dispatch(regexp_extract) +SELECT s, regexp_extract(s, 'foo(?=bar)', 0) FROM test_regexp_extract -- embedded flags (Java-only) query expect_dispatch(regexp_extract) -SELECT regexp_extract(s, '(?i)HELLO', 0) FROM test_regexp_extract +SELECT s, regexp_extract(s, '(?i)HELLO', 0) FROM test_regexp_extract + +query expect_dispatch(regexp_extract) +SELECT s, regexp_extract(s, r'(?i)(foo)(\d+)', 2) FROM test_regexp_extract + +-- Keep a separate all-null input batch, rather than only nulls mixed with matching rows. +statement +CREATE TABLE test_regexp_extract_nulls(s string) USING parquet + +statement +INSERT INTO test_regexp_extract_nulls VALUES (NULL), (NULL), (NULL) + +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'(\d+)', 1) FROM test_regexp_extract_nulls -- literal arguments query expect_dispatch(regexp_extract) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql index af3ada87c9c..060a6411194 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql @@ -22,31 +22,51 @@ statement CREATE TABLE test_regexp_extract_all(s string) USING parquet statement -INSERT INTO test_regexp_extract_all VALUES ('abc123def456'), ('no match'), (NULL), ('100-200-300'), ('hello world') +INSERT INTO test_regexp_extract_all VALUES + ('abc123def456'), ('no match'), (NULL), ('100-200-300'), ('hello world'), + ('x1y2z3'), ('cat bat hat'), ('no vowels'), ('a1b2c3'), ('x9y8'), ('abc'), ('def'), + ('foobar foobaz fooqux'), ('abc123def'), ('xyz789') -- group 0: all entire matches query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, r'\d+', 0) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, r'\d+', 0) FROM test_regexp_extract_all + +query expect_dispatch(regexp_extract_all) +SELECT s, regexp_extract_all(s, '[a-z]at', 0) FROM test_regexp_extract_all -- group 1: first capturing group from each match query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract_all + +query expect_dispatch(regexp_extract_all) +SELECT s, regexp_extract_all(s, r'(\d+)', 1) FROM test_regexp_extract_all -- group 2: second capturing group from each match query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract_all + +-- Single-character groups exercise a capture from each adjacent match. +query expect_dispatch(regexp_extract_all) +SELECT s, regexp_extract_all(s, r'([a-z])(\d)', 1) FROM test_regexp_extract_all + +query expect_dispatch(regexp_extract_all) +SELECT s, regexp_extract_all(s, r'([a-z])(\d)', 2) FROM test_regexp_extract_all -- no match returns empty array query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, 'NOMATCH', 0) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, 'NOMATCH', 0) FROM test_regexp_extract_all -- backreference pattern (Java-only) query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, r'(\d)\1', 0) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, r'(\d)\1', 0) FROM test_regexp_extract_all + +-- Multiple Java lookahead matches, followed by a non-match. +query expect_dispatch(regexp_extract_all) +SELECT s, regexp_extract_all(s, 'foo(?=ba[rz])', 0) FROM test_regexp_extract_all -- embedded flags (Java-only) query expect_dispatch(regexp_extract_all) -SELECT regexp_extract_all(s, '(?i)[A-Z]+', 0) FROM test_regexp_extract_all +SELECT s, regexp_extract_all(s, '(?i)[A-Z]+', 0) FROM test_regexp_extract_all -- literal arguments query expect_dispatch(regexp_extract_all) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_batches.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_batches.sql new file mode 100644 index 00000000000..884989ec770 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_batches.sql @@ -0,0 +1,32 @@ +-- 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. + +-- Config: spark.comet.batchSize=64 + +-- One Parquet file keeps enough rows in one task to cross Arrow batch boundaries. +statement +CREATE TABLE test_regexp_extract_batches(id int, s string) USING parquet + +statement +INSERT INTO test_regexp_extract_batches +SELECT CAST(id AS INT), + CASE WHEN id % 7 = 0 THEN NULL ELSE concat('item_', CAST(id AS STRING), '_value') END +FROM range(0, 5000, 1, 1) + +-- Keep the row ID so a shifted extraction or null bitmap cannot pass as the same result set. +query expect_dispatch(regexp_extract) +SELECT id, s, regexp_extract(s, r'item_(\d+)_value', 1) FROM test_regexp_extract_batches diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql index 71cf1c9dab1..dd88cb75109 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql @@ -22,27 +22,36 @@ statement CREATE TABLE test_regexp_instr(s string) USING parquet statement -INSERT INTO test_regexp_instr VALUES ('abc123def'), ('no match'), (NULL), ('123xyz'), ('hello world'), ('aa') +INSERT INTO test_regexp_instr VALUES + ('abc123def'), ('no match'), (NULL), ('123xyz'), ('hello world'), ('aa'), + ('456xyz'), ('abc123def456'), ('xyz'), ('abc'), ('def'), ('foobar'), ('foobaz'), ('xyz789') -- basic: position of first digit sequence query expect_dispatch(regexp_instr) -SELECT regexp_instr(s, r'\d+', 0) FROM test_regexp_instr +SELECT s, regexp_instr(s, r'\d+', 0) FROM test_regexp_instr --- group 1 (still returns position of entire match per Spark semantics) +-- Spark returns the position of the entire match for both third-argument values. query expect_dispatch(regexp_instr) -SELECT regexp_instr(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_instr +SELECT s, regexp_instr(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_instr + +query expect_dispatch(regexp_instr) +SELECT s, regexp_instr(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_instr -- no match returns 0 query expect_dispatch(regexp_instr) -SELECT regexp_instr(s, 'NOMATCH', 0) FROM test_regexp_instr +SELECT s, regexp_instr(s, 'NOMATCH', 0) FROM test_regexp_instr -- backreference pattern (Java-only) query expect_dispatch(regexp_instr) -SELECT regexp_instr(s, r'(\w)\1', 0) FROM test_regexp_instr +SELECT s, regexp_instr(s, r'(\w)\1', 0) FROM test_regexp_instr + +-- lookahead (Java-only) +query expect_dispatch(regexp_instr) +SELECT s, regexp_instr(s, 'foo(?=bar)', 0) FROM test_regexp_instr -- embedded flags (Java-only) query expect_dispatch(regexp_instr) -SELECT regexp_instr(s, '(?i)HELLO', 0) FROM test_regexp_instr +SELECT s, regexp_instr(s, '(?i)HELLO', 0) FROM test_regexp_instr -- literal arguments query expect_dispatch(regexp_instr) diff --git a/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala b/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala index b8db2171143..d90a299bed9 100644 --- a/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala @@ -677,90 +677,6 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - // ========== regexp_extract tests ========== - - test("regexp_extract: basic group extraction") { - withSubjects("abc123def", "no match", null, "xyz789") { - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract(s, '([a-z]+)(\\\\d+)', 1) FROM t")) - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract(s, '([a-z]+)(\\\\d+)', 2) FROM t")) - } - } - - test("regexp_extract: group 0 returns entire match") { - withSubjects("hello world", "foo123bar", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract(s, '\\\\d+', 0) FROM t")) - } - } - - test("regexp_extract: no match returns empty string") { - withSubjects("abc", "def", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract(s, '\\\\d+', 0) FROM t")) - } - } - - test("regexp_extract: backreference pattern (Java-only)") { - withSubjects("aa", "ab", "bb", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract(s, '(\\\\w)\\\\1', 0) FROM t")) - } - } - - test("regexp_extract: lookahead pattern (Java-only)") { - withSubjects("foobar", "foobaz", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract(s, 'foo(?=bar)', 0) FROM t")) - } - } - - test("regexp_extract: embedded flags (Java-only)") { - withSubjects("FOO123", "foo456", "bar789") { - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract(s, '(?i)(foo)(\\\\d+)', 2) FROM t")) - } - } - - test("regexp_extract: all-null column") { - withSubjects(null, null, null) { - checkSparkAnswerAndOperator(sql("SELECT regexp_extract(s, '(\\\\d+)', 1) FROM t")) - } - } - - // ========== regexp_extract_all tests ========== - - test("regexp_extract_all: basic extraction of all matches") { - withSubjects("abc123def456", "no match", null, "x1y2z3") { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract_all(s, '(\\\\d+)', 1) FROM t")) - } - } - - test("regexp_extract_all: group 0 returns full matches") { - withSubjects("cat bat hat", "no vowels", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract_all(s, '[a-z]at', 0) FROM t")) - } - } - - test("regexp_extract_all: multiple groups") { - withSubjects("a1b2c3", "x9y8", null) { - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract_all(s, '([a-z])(\\\\d)', 1) FROM t")) - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract_all(s, '([a-z])(\\\\d)', 2) FROM t")) - } - } - - test("regexp_extract_all: no matches returns empty array") { - withSubjects("abc", "def") { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract_all(s, '\\\\d+', 0) FROM t")) - } - } - - test("regexp_extract_all: lookahead pattern (Java-only)") { - withSubjects("foobar foobaz fooqux") { - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract_all(s, 'foo(?=ba[rz])', 0) FROM t")) - } - } - // ========== regexp_replace tests ========== test("regexp_replace: basic replacement") { @@ -847,33 +763,6 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - // ========== regexp_instr tests ========== - - test("regexp_instr: basic position finding") { - withSubjects("abc123def", "no match", null, "456xyz") { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, '\\\\d+', 0) FROM t")) - } - } - - test("regexp_instr: specific group position") { - withSubjects("abc123def456", "xyz", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, '([a-z]+)(\\\\d+)', 1) FROM t")) - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, '([a-z]+)(\\\\d+)', 2) FROM t")) - } - } - - test("regexp_instr: no match returns 0") { - withSubjects("abc", "def", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, '\\\\d+', 0) FROM t")) - } - } - - test("regexp_instr: lookahead (Java-only)") { - withSubjects("foobar", "foobaz", null) { - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, 'foo(?=bar)', 0) FROM t")) - } - } - // ========== split tests ========== test("split: basic regex split") { @@ -912,19 +801,7 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - // ========== multi-batch and combined tests ========== - - test("regexp_extract: many rows spanning multiple batches") { - withTable("t") { - sql("CREATE TABLE t (s STRING) USING parquet") - val values = (0 until 5000) - .map(i => if (i % 7 == 0) "(NULL)" else s"('item_${i}_value')") - .mkString(", ") - sql(s"INSERT INTO t VALUES $values") - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract(s, 'item_(\\\\d+)_value', 1) FROM t")) - } - } + // ========== combined tests ========== test("all regexp expressions combined in one query") { withSubjects("abc123def456", "hello world", null, "aa") { @@ -939,15 +816,4 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { |""".stripMargin)) } } - - test("expressions with no native path always run via the JVM dispatcher") { - withSubjects("abc123def", "no match", null, "xyz789") { - // regexp_extract / regexp_extract_all / regexp_instr have no native rust path, so they - // always run on Comet via the JVM codegen dispatcher rather than falling back to Spark. - checkSparkAnswerAndOperator( - sql("SELECT s, regexp_extract(s, '([a-z]+)(\\\\d+)', 2) FROM t")) - checkSparkAnswerAndOperator(sql("SELECT s, regexp_extract_all(s, '\\\\d+', 0) FROM t")) - checkSparkAnswerAndOperator(sql("SELECT s, regexp_instr(s, '\\\\d+', 0) FROM t")) - } - } } From 031070adc40b7f40491ddfd512296b44aaef3ade Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 12:28:31 +0000 Subject: [PATCH 77/89] refactor: stop normalizing float divisors now that percentile_approx orders NaN as Spark does (#6866) * refactor: stop normalizing float divisors now that percentile_approx orders NaN as Spark does CometExecRule wrapped the divisor of a FLOAT or DOUBLE Divide or Remainder in a Project or Filter in NormalizeNaNAndZero. Comparisons don't need it, since the native comparison builder follows Spark's ordering, including the check for a zero divisor. It stayed only because it also canonicalized a NaN quotient: 1.0D / (-d) of a NaN d is a NaN with the sign bit set, and percentile_approx ordered that NaN below every other value (#6519). #6829 fixed percentile_approx, so the rule can go. nan_divisor.sql keeps its percentile_approx query, which now checks the sign-bit NaN quotient directly, along with the hash and extrema queries for the earlier consumers. * test: regenerate TPC-DS plans without the divisor normalization Without NormalizeNaNAndZero around float divisors, the 11 approved plans that hold a float division in a Project or Filter count one native expression fewer, and q78 no longer lists knownfloatingpointnormalized and normalizenanandzero among its dispatched expressions. Regenerated on Spark 3.4, 3.5, 4.0, 4.1 and 4.2; 4.1 and 4.2 need no override. --- .../apache/comet/rules/CometExecRule.scala | 44 ++----------------- .../expressions/math/nan_divisor.sql | 5 +-- .../q83/extended.txt | 2 +- .../q83.ansi/extended.txt | 2 +- .../approved-plans-v1_4/q21/extended.txt | 2 +- .../approved-plans-v1_4/q34/extended.txt | 2 +- .../approved-plans-v1_4/q39a/extended.txt | 2 +- .../approved-plans-v1_4/q39b/extended.txt | 2 +- .../approved-plans-v1_4/q73/extended.txt | 2 +- .../approved-plans-v1_4/q78/extended.txt | 4 +- .../approved-plans-v1_4/q83/extended.txt | 2 +- .../approved-plans-v2_7/q34/extended.txt | 2 +- .../approved-plans-v2_7/q78/extended.txt | 4 +- 13 files changed, 18 insertions(+), 57 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 569ccc277f7..32faea60e65 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,9 +25,8 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, Expression, FloatLiteral, KnownFloatingPointNormalized, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Expression, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} -import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.catalyst.util.sideBySide @@ -937,41 +936,6 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) newPlan } - /** - * Wraps the divisor of a floating-point `Divide` or `Remainder` in a Project or Filter in - * `NormalizeNaNAndZero`. Native comparisons already follow Spark's ordering, including the - * check for a zero divisor, so the quotient does not need this. It also canonicalizes a NaN - * divisor, though, and with it the NaN quotient: `1.0D / (-d)` of a NaN `d` is otherwise a NaN - * with the sign bit set, which `percentile_approx` still orders below every other value - * (https://github.com/apache/datafusion-comet/issues/6519). Remove this once that is fixed. - */ - private def normalizeDivisors(plan: SparkPlan): SparkPlan = { - def normalize(expr: Expression): Expression = expr.transformUp { - case Divide(left, right, evalMode) => Divide(left, normalizeNaNAndZero(right), evalMode) - case Remainder(left, right, evalMode) => - Remainder(left, normalizeNaNAndZero(right), evalMode) - } - plan.transformUp { - case p: ProjectExec => - ProjectExec(p.projectList.map(normalize(_).asInstanceOf[NamedExpression]), p.child) - case f: FilterExec => FilterExec(normalize(f.condition), f.child) - } - } - - private def normalizeNaNAndZero(expr: Expression): Expression = { - expr match { - case _: KnownFloatingPointNormalized => expr - case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr - case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr - case _ => - expr.dataType match { - case _: FloatType | _: DoubleType => - KnownFloatingPointNormalized(NormalizeNaNAndZero(expr)) - case _ => expr - } - } - } - private def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) { @@ -992,17 +956,15 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { - val normalizedPlan = normalizeDivisors(plan) - // The rewrite removes the sorts of each join it converts, so put back any sort a kept // operator above still needs. This runs before transform(), which converts the sorts it // adds like any other SortExec. val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - RewriteJoin.restoreRequiredOrdering(normalizedPlan.transformUp { case p => + RewriteJoin.restoreRequiredOrdering(plan.transformUp { case p => RewriteJoin.rewrite(p, conf) }) } else { - normalizedPlan + plan } // Tag Partial aggregates that must not be converted to Comet because a diff --git a/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql index 51135545d44..a7490a7ce66 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql @@ -43,9 +43,8 @@ query SELECT max(q), min(q) FROM (SELECT 1.0D / (-d) AS q FROM nan_divisor UNION ALL SELECT 0.0D AS q FROM nan_divisor) --- percentile_approx orders a NaN with the sign bit set below every other value --- (https://github.com/apache/datafusion-comet/issues/6519). Until that is fixed, Comet keeps --- normalizing the divisor, which makes the quotient a canonical NaN. +-- The quotient keeps the sign bit of its NaN, and percentile_approx orders that NaN above every +-- other value, as Spark does. query SELECT percentile_approx(q, 1.0D), percentile_approx(q, 0.0D) FROM (SELECT /*+ COALESCE(1) */ q diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q83/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q83/extended.txt index a90327716a8..d7249c9e3e3 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q83/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q83/extended.txt @@ -101,4 +101,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 52 out of 52 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 52 out of 52 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q83.ansi/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q83.ansi/extended.txt index a90327716a8..d7249c9e3e3 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q83.ansi/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark4_0/q83.ansi/extended.txt @@ -101,4 +101,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 52 out of 52 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 52 out of 52 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q21/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q21/extended.txt index 11afabb84f4..a2a5110f431 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q21/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q21/extended.txt @@ -27,4 +27,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 24 out of 24 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 13 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 24 out of 24 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q34/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q34/extended.txt index da568f3647f..d3653090937 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q34/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q34/extended.txt @@ -37,4 +37,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.customer -Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 14 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39a/extended.txt index 5a1aaca450a..1ca6d90b337 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39a/extended.txt @@ -60,4 +60,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 46 out of 46 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 46 out of 46 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 10 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39b/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39b/extended.txt index 5a1aaca450a..1ca6d90b337 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39b/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q39b/extended.txt @@ -60,4 +60,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 46 out of 46 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 46 out of 46 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 10 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q73/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q73/extended.txt index da568f3647f..d3653090937 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q73/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q73/extended.txt @@ -37,4 +37,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.customer -Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 14 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q78/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q78/extended.txt index aca1629a7a3..311042490d0 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q78/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q78/extended.txt @@ -1,6 +1,6 @@ CometColumnarToRow +- CometTakeOrderedAndProject [COMET-INFO: JVM codegen dispatcher: add, cast, coalesce, divide, round] - +- CometProject [COMET-INFO: JVM codegen dispatcher: add, cast, coalesce, divide, knownfloatingpointnormalized, normalizenanandzero, round] + +- CometProject [COMET-INFO: JVM codegen dispatcher: add, cast, coalesce, divide, round] +- CometSortMergeJoin :- CometProject : +- CometSortMergeJoin @@ -76,4 +76,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 65 out of 65 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 7 codegen dispatch. \ No newline at end of file +Comet accelerated 65 out of 65 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 5 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q83/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q83/extended.txt index d6a2f432728..9dbc72b2944 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q83/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q83/extended.txt @@ -114,4 +114,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 62 out of 62 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 62 out of 62 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 11 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q34/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q34/extended.txt index da568f3647f..d3653090937 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q34/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q34/extended.txt @@ -37,4 +37,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.customer -Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 15 native, 0 codegen dispatch. \ No newline at end of file +Comet accelerated 33 out of 33 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 14 native, 0 codegen dispatch. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q78/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q78/extended.txt index 93e05678ac6..7e9d7e88e6a 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q78/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q78/extended.txt @@ -1,6 +1,6 @@ CometColumnarToRow +- CometTakeOrderedAndProject - +- CometProject [COMET-INFO: JVM codegen dispatcher: add, cast, coalesce, divide, knownfloatingpointnormalized, normalizenanandzero, round] + +- CometProject [COMET-INFO: JVM codegen dispatcher: add, cast, coalesce, divide, round] +- CometSortMergeJoin :- CometProject : +- CometSortMergeJoin @@ -76,4 +76,4 @@ CometColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 65 out of 65 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 7 codegen dispatch. \ No newline at end of file +Comet accelerated 65 out of 65 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. Accelerated expressions: 12 native, 5 codegen dispatch. \ No newline at end of file From 17c1413c8e879f17c8acea5c1a7987817063d648 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:29:25 +0000 Subject: [PATCH 78/89] test: migrate scalar hash coverage to SQL fixtures (#6881) Part of #6634. Preserve scalar input types, hash child order and dictionary configurations while replacing 18 Scala tests with SQL fixtures. --- .../sql-tests/expressions/hash/crc32.sql | 2 + .../sql-tests/expressions/hash/hash.sql | 65 +++++++++ .../sql-tests/expressions/hash/hash_types.sql | 97 +++++++++++++ .../apache/comet/CometExpressionSuite.scala | 94 ------------- .../comet/CometHashExpressionSuite.scala | 130 ------------------ 5 files changed, 164 insertions(+), 224 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/hash/hash_types.sql diff --git a/spark/src/test/resources/sql-tests/expressions/hash/crc32.sql b/spark/src/test/resources/sql-tests/expressions/hash/crc32.sql index 7b922df0cca..150cb11fe9a 100644 --- a/spark/src/test/resources/sql-tests/expressions/hash/crc32.sql +++ b/spark/src/test/resources/sql-tests/expressions/hash/crc32.sql @@ -15,6 +15,8 @@ -- specific language governing permissions and limitations -- under the License. +-- ConfigMatrix: parquet.enable.dictionary=false,true + -- crc32 function statement CREATE TABLE test(col string, a int, b float) USING parquet diff --git a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql index cbf90b18459..ca2065de384 100644 --- a/spark/src/test/resources/sql-tests/expressions/hash/hash.sql +++ b/spark/src/test/resources/sql-tests/expressions/hash/hash.sql @@ -16,6 +16,7 @@ -- under the License. -- Config: spark.sql.legacy.allowHashOnMapType=true +-- ConfigMatrix: parquet.enable.dictionary=false,true -- hash functions statement @@ -27,6 +28,67 @@ INSERT INTO test VALUES ('Spark SQL ', 10, 1.2), (NULL, NULL, NULL), ('', 0, 0. query SELECT md5(col), md5(cast(a as string)), md5(cast(b as string)), hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col), xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col), sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128), sha2(col, -1), sha1(col), sha1(cast(a as string)), sha1(cast(b as string)) FROM test +-- Single-column strings, including the repeated inputs from the dictionary and md5 tests. +statement +CREATE TABLE test_hash_strings(c STRING) USING parquet + +statement +INSERT INTO test_hash_strings VALUES + ('hello'), (''), ('Spark SQL'), ('苹果手机'), (NULL), + ('a'), ('b'), ('a'), ('b'), ('a'), ('c'), + ('test1'), ('test1'), ('test2'), ('test2') + +query expect_native(hash, xxhash64, md5) +SELECT c, hash(c), xxhash64(c), md5(c) FROM test_hash_strings + +-- Keep DOUBLE distinct from the FLOAT column above and exercise both child orders. +statement +CREATE TABLE test_hash_columns(a INT, b STRING, c DOUBLE) USING parquet + +statement +INSERT INTO test_hash_columns VALUES + (1, 'hello', 3.14), (2, '', 0.0), (NULL, NULL, NULL), (-1, 'test', -1.5), (3, 'test', 1.0) + +query expect_native(hash, xxhash64) +SELECT hash(a, b, c), xxhash64(a, b, c), hash(c, b, a), xxhash64(c, b, a), + hash(a), xxhash64(a), hash(b), xxhash64(b), hash(c), xxhash64(c) +FROM test_hash_columns + +-- Extra integers are hashed as children, not used as seeds. +query expect_native(hash, xxhash64) +SELECT hash(a, 0), xxhash64(a, 0), hash(a, 42), xxhash64(a, 42), + hash(a, -1), xxhash64(a, -1) FROM test_hash_columns + +-- A deterministic replacement for the generated corpus: 1000 rows with repeated strings, +-- integer/float boundaries and timestamp-like strings containing control characters. +statement +CREATE TABLE test_hash_boundaries(col STRING, a INT, b FLOAT) USING parquet + +statement +INSERT INTO test_hash_boundaries +SELECT /*+ COALESCE(1) */ col, a, b FROM VALUES + ('012345', 2147483647, float('3.4028235E38')), + ('6789/:', -2147483648, float('-3.4028235E38')), + ('T \t\r\n', 0, float('1.4E-45')), + (' \t\r\n/:T', -1, float('-1.4E-45')), + ('', 1, float('-0.0')), + (NULL, NULL, CAST(NULL AS FLOAT)), + ('12:34T', NULL, float('0.0')), + (NULL, -42, float('NaN')), + ('012345', 10, CAST(NULL AS FLOAT)), + ('6789/:', -10, float('-Infinity')) +AS v(col, a, b) CROSS JOIN range(100) + +-- FLOAT-to-string formatting has separate compatibility limits, so hash the float directly. +query expect_native(md5, hash, xxhash64, crc32, sha2, sha1) +SELECT md5(col), md5(CAST(a AS STRING)), + hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col), + xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col), + crc32(col), crc32(CAST(a AS STRING)), + sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), + sha2(col, 128), sha2(col, -1), sha1(col), sha1(CAST(a AS STRING)) +FROM test_hash_boundaries + -- literal arguments -- The SQL file test suite disables ConstantFolding, so these literal arguments reach Comet's -- native engine as scalar values rather than being folded away by Spark's optimizer. @@ -43,6 +105,9 @@ CREATE TABLE test_nan(d double, f float) USING parquet statement INSERT INTO test_nan VALUES (double('NaN'), float('NaN')), (0.0, 0.0), (double('-0.0'), float('-0.0')), (1.5, 1.5), (NULL, NULL), (double('Infinity'), float('Infinity')), (double('-Infinity'), float('-Infinity')) +statement +INSERT INTO test_nan VALUES (1.0, 1.0), (-1.0, -1.0), (3.14159265358979, 3.14159) + query SELECT hash(d), hash(-d), xxhash64(d), xxhash64(-d), hash(f), hash(-f), xxhash64(f), xxhash64(-f), hash(-d, -f), xxhash64(-d, -f) FROM test_nan diff --git a/spark/src/test/resources/sql-tests/expressions/hash/hash_types.sql b/spark/src/test/resources/sql-tests/expressions/hash/hash_types.sql new file mode 100644 index 00000000000..2922be456f3 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/hash/hash_types.sql @@ -0,0 +1,97 @@ +-- 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. + +-- Read typed columns from Parquet so hashing cannot be evaluated by the optimizer. +statement +CREATE TABLE test_hash_boolean(c BOOLEAN) USING parquet + +statement +INSERT INTO test_hash_boolean VALUES (true), (false), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_boolean + +statement +CREATE TABLE test_hash_tinyint(c TINYINT) USING parquet + +statement +INSERT INTO test_hash_tinyint VALUES (1), (0), (-1), (127), (-128), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_tinyint + +statement +CREATE TABLE test_hash_smallint(c SMALLINT) USING parquet + +statement +INSERT INTO test_hash_smallint VALUES (1), (0), (-1), (32767), (-32768), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_smallint + +statement +CREATE TABLE test_hash_int(c INT) USING parquet + +statement +INSERT INTO test_hash_int VALUES (1), (0), (-1), (2147483647), (-2147483648), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_int + +statement +CREATE TABLE test_hash_bigint(c BIGINT) USING parquet + +statement +INSERT INTO test_hash_bigint VALUES + (1), (0), (-1), (9223372036854775807), (-9223372036854775808), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_bigint + +statement +CREATE TABLE test_hash_binary(c BINARY) USING parquet + +statement +INSERT INTO test_hash_binary VALUES (X''), (X'00'), (X'0102030405'), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_binary + +statement +CREATE TABLE test_hash_date(c DATE) USING parquet + +statement +INSERT INTO test_hash_date VALUES + (DATE '2023-01-01'), (DATE '1970-01-01'), (DATE '2000-12-31'), + (DATE '1969-12-31'), (NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c) FROM test_hash_date + +statement +CREATE TABLE test_hash_timestamp(c TIMESTAMP, ntz TIMESTAMP_NTZ) USING parquet + +statement +INSERT INTO test_hash_timestamp VALUES + (TIMESTAMP '2023-01-01 12:00:00', TIMESTAMP_NTZ '2023-01-01 12:00:00'), + (TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00'), + (TIMESTAMP '2000-12-31 23:59:59', TIMESTAMP_NTZ '2000-12-31 23:59:59'), + (TIMESTAMP '1969-12-31 23:59:59.999999', TIMESTAMP_NTZ '1969-12-31 23:59:59.999999'), + (NULL, NULL) + +query expect_native(hash, xxhash64) +SELECT c, hash(c), xxhash64(c), ntz, hash(ntz), xxhash64(ntz) FROM test_hash_timestamp diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 6e1d77b3619..f79c6415567 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -2060,20 +2060,6 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("md5") { - Seq(false, true).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "test" - withTable(table) { - sql(s"create table $table(col String) using parquet") - sql( - s"insert into $table values ('test1'), ('test1'), ('test2'), ('test2'), (NULL), ('')") - checkSparkAnswerAndOperator(s"select md5(col) FROM $table") - } - } - } - } - test("unhex") { val table = "unhex_table" Seq(false, true).foreach { dictionaryEnabled => @@ -2617,48 +2603,6 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("hash functions") { - Seq(true, false).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "test" - withTable(table) { - sql(s"create table $table(col string, a int, b float) using parquet") - sql(s""" - |insert into $table values - |('Spark SQL ', 10, 1.2), (NULL, NULL, NULL), ('', 0, 0.0), ('苹果手机', NULL, 3.999999) - |, ('Spark SQL ', 10, 1.2), (NULL, NULL, NULL), ('', 0, 0.0), ('苹果手机', NULL, 3.999999) - |""".stripMargin) - checkSparkAnswerAndOperator(""" - |select - |md5(col), md5(cast(a as string)), md5(cast(b as string)), - |hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col), - |xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col), - |crc32(col), crc32(cast(a as string)), crc32(cast(b as string)), - |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128), sha2(col, -1), - |sha1(col), sha1(cast(a as string)), sha1(cast(b as string)) - |from test - |""".stripMargin) - } - } - } - } - - test("sha2 with all-literal arguments and ConstantFolding disabled") { - // https://github.com/apache/datafusion-comet/issues/3340 - // When ConstantFolding is disabled, an all-literal sha2() call reaches the native - // engine as scalar arguments. Verify it computes the correct result rather than crashing. - withSQLConf( - "spark.sql.optimizer.excludedRules" -> - "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") { - checkSparkAnswerAndOperator(""" - |select - |sha2('test', 0), sha2('test', 256), sha2('test', 224), - |sha2('test', 384), sha2('test', 512), sha2('test', 128), sha2('test', -1), - |sha2(cast(null as string), 256) - |""".stripMargin) - } - } - test("remainder function") { def withAnsiMode(enabled: Boolean)(f: => Unit): Unit = { withSQLConf( @@ -2738,44 +2682,6 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("hash functions with random input") { - val dataGen = DataGenerator.DEFAULT - // sufficient number of rows to create dictionary encoded ArrowArray. - val randomNumRows = 1000 - - val whitespaceChars = " \t\r\n" - val timestampPattern = "0123456789/:T" + whitespaceChars - Seq(true, false).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "test" - withTable(table) { - sql(s"create table $table(col string, a int, b float) using parquet") - val tableSchema = spark.table(table).schema - val rows = dataGen.generateRows( - randomNumRows, - tableSchema, - Some(() => dataGen.generateString(timestampPattern, 6))) - val data = spark.createDataFrame(spark.sparkContext.parallelize(rows), tableSchema) - data.write - .mode("append") - .insertInto(table) - // with random generated data - // disable cast(b as string) for now, as the cast from float to string may produce incompatible result - checkSparkAnswerAndOperator(""" - |select - |md5(col), md5(cast(a as string)), --md5(cast(b as string)), - |hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col), - |xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col), - |crc32(col), crc32(cast(a as string)), - |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128), sha2(col, -1), - |sha1(col), sha1(cast(a as string)) - |from test - |""".stripMargin) - } - } - } - } - test("hash function with decimal input") { val testPrecisionScales: Seq[(Int, Int)] = Seq( (1, 0), diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index c20e5899b37..dc9ae920ef6 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -34,100 +34,6 @@ import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGener */ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { - test("hash - boolean") { - withTable("t") { - sql("CREATE TABLE t(c BOOLEAN) USING parquet") - sql("INSERT INTO t VALUES (true), (false), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - byte/tinyint") { - withTable("t") { - sql("CREATE TABLE t(c TINYINT) USING parquet") - sql("INSERT INTO t VALUES (1), (0), (-1), (127), (-128), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - short/smallint") { - withTable("t") { - sql("CREATE TABLE t(c SMALLINT) USING parquet") - sql("INSERT INTO t VALUES (1), (0), (-1), (32767), (-32768), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - int") { - withTable("t") { - sql("CREATE TABLE t(c INT) USING parquet") - sql("INSERT INTO t VALUES (1), (0), (-1), (2147483647), (-2147483648), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - long/bigint") { - withTable("t") { - sql("CREATE TABLE t(c BIGINT) USING parquet") - sql( - "INSERT INTO t VALUES (1), (0), (-1), (9223372036854775807), (-9223372036854775808), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - float") { - withTable("t") { - sql("CREATE TABLE t(c FLOAT) USING parquet") - sql("INSERT INTO t VALUES (1.0), (0.0), (-0.0), (-1.0), (3.14159), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - double") { - withTable("t") { - sql("CREATE TABLE t(c DOUBLE) USING parquet") - sql("INSERT INTO t VALUES (1.0), (0.0), (-0.0), (-1.0), (3.14159265358979), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - string") { - withTable("t") { - sql("CREATE TABLE t(c STRING) USING parquet") - sql("INSERT INTO t VALUES ('hello'), (''), ('Spark SQL'), ('苹果手机'), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - binary") { - withTable("t") { - sql("CREATE TABLE t(c BINARY) USING parquet") - sql("INSERT INTO t VALUES (X''), (X'00'), (X'0102030405'), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - date") { - withTable("t") { - sql("CREATE TABLE t(c DATE) USING parquet") - sql( - "INSERT INTO t VALUES (DATE '2023-01-01'), (DATE '1970-01-01'), (DATE '2000-12-31'), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - - test("hash - timestamp") { - withTable("t") { - sql("CREATE TABLE t(c TIMESTAMP) USING parquet") - sql("""INSERT INTO t VALUES - (TIMESTAMP '2023-01-01 12:00:00'), - (TIMESTAMP '1970-01-01 00:00:00'), - (TIMESTAMP '2000-12-31 23:59:59'), - (null)""") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - test("hash - decimal (precision <= 18)") { Seq((10, 2), (18, 0), (18, 10)).foreach { case (precision, scale) => withTable("t") { @@ -312,19 +218,6 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } - test("hash - multiple primitive columns") { - withTable("t") { - sql("CREATE TABLE t(a INT, b STRING, c DOUBLE) USING parquet") - sql("""INSERT INTO t VALUES - (1, 'hello', 3.14), - (2, '', 0.0), - (null, null, null), - (-1, 'test', -1.5)""") - checkSparkAnswerAndOperator( - "SELECT hash(a, b, c), xxhash64(a, b, c), hash(c, b, a), xxhash64(c, b, a), hash(a), xxhash64(a), hash(b), xxhash64(b), hash(c), xxhash64(c) FROM t") - } - } - test("hash - multiple columns with arrays") { withTable("t") { sql("CREATE TABLE t(a INT, b ARRAY, c STRING) USING parquet") @@ -369,16 +262,6 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } - test("hash - with custom seed") { - withTable("t") { - sql("CREATE TABLE t(c INT) USING parquet") - sql("INSERT INTO t VALUES (1), (2), (3), (null)") - // Extra integer arguments are hashed as additional children, not as the Spark seed. - checkSparkAnswerAndOperator( - "SELECT hash(c), xxhash64(c), hash(c, 0), xxhash64(c, 0), hash(c, 42), xxhash64(c, 42), hash(c, -1), xxhash64(c, -1) FROM t ORDER BY c") - } - } - test("hash - large array") { withTable("t") { sql("CREATE TABLE t(c ARRAY) USING parquet") @@ -406,19 +289,6 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } - test("hash - with dictionary encoding") { - Seq(true, false).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - withTable("t") { - sql("CREATE TABLE t(c STRING) USING parquet") - // Repeated values to trigger dictionary encoding - sql("INSERT INTO t VALUES ('a'), ('b'), ('a'), ('b'), ('a'), ('c'), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } - } - } - test("hash - array with dictionary encoding") { Seq(true, false).foreach { dictionary => withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { From a17a76f4638ba2929b4607af396f530493d27165 Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:30:03 +0000 Subject: [PATCH 79/89] test: migrate Java RLIKE semantics and batch coverage to SQL fixtures (#6863) --- .../expressions/string/rlike_java.sql | 59 ++++++++- .../expressions/string/rlike_java_batches.sql | 42 +++++++ .../apache/comet/CometRegExpJvmSuite.scala | 117 +----------------- 3 files changed, 100 insertions(+), 118 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/string/rlike_java_batches.sql diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql index 7ea90ba90dc..1c7088f7ff2 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql @@ -17,12 +17,17 @@ -- Test RLIKE dispatcher patterns and the native empty-pattern case -- Raw literals keep Spark from dropping regex backslashes such as '\d'. +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.RLike.allowIncompatible=false statement CREATE TABLE test_rlike_java(s string) USING parquet statement -INSERT INTO test_rlike_java VALUES ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab') +INSERT INTO test_rlike_java VALUES + ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab'), + ('no digits'), ('mixed_42_data'), ('xyzzy'), ('foobar'), ('foobaz'), ('barfoo'), + ('barbar'), ('foofoo'), ('FOO'), ('foo'), ('fOO'), ('bar'), ('a1'), ('9z'), ('abc'), ('x') query expect_dispatch(rlike) SELECT s RLIKE r'^\d+$' FROM test_rlike_java @@ -31,11 +36,11 @@ query expect_dispatch(rlike) SELECT s RLIKE '^[a-z]+$' FROM test_rlike_java query expect_native(rlike) -SELECT s RLIKE '' FROM test_rlike_java +SELECT s, s RLIKE '' FROM test_rlike_java -- backreference (Java-only) query expect_dispatch(rlike) -SELECT s RLIKE r'^(\w)\1$' FROM test_rlike_java +SELECT s, s RLIKE r'^(\w)\1$' FROM test_rlike_java -- lookahead (Java-only) query expect_dispatch(rlike) @@ -48,3 +53,51 @@ SELECT s RLIKE '(?i)hello' FROM test_rlike_java -- Literal arguments mix native (empty pattern) and dispatched RLIKE, so only check parity. query SELECT 'hello' RLIKE '^[a-z]+$', '12345' RLIKE r'^\d+$', '' RLIKE '', NULL RLIKE 'a' + +-- Substring matches differ from the anchored digit pattern above. +query expect_dispatch(rlike) +SELECT s, s RLIKE r'\d+' AS m FROM test_rlike_java + +query expect_dispatch(rlike) +SELECT s FROM test_rlike_java WHERE s RLIKE r'\d+' + +query expect_dispatch(rlike) +SELECT s FROM test_rlike_java WHERE s RLIKE r'^(\w)\1$' + +-- Preserve the original lookaround, flag and named-group patterns and subjects. +query expect_dispatch(rlike) +SELECT s, s RLIKE r'foo(?=bar)' FROM test_rlike_java + +query expect_dispatch(rlike) +SELECT s FROM test_rlike_java WHERE s RLIKE r'foo(?=bar)' + +query expect_dispatch(rlike) +SELECT s, s RLIKE r'(?<=foo)bar' FROM test_rlike_java + +query expect_dispatch(rlike) +SELECT s, s RLIKE r'(?i)foo' FROM test_rlike_java + +query expect_dispatch(rlike) +SELECT s, s RLIKE r'(?\d)' FROM test_rlike_java + +-- Anchors use Java semantics, including the empty-subject case. +query expect_dispatch(rlike) +SELECT s, s RLIKE '^$' FROM test_rlike_java + +statement +CREATE TABLE test_rlike_java_nulls(s string) USING parquet + +statement +INSERT INTO test_rlike_java_nulls VALUES (NULL), (NULL), (NULL) + +query expect_dispatch(rlike) +SELECT s RLIKE r'\d+' FROM test_rlike_java_nulls + +statement +CREATE TABLE test_rlike_java_grouped(s string, k int) USING parquet + +statement +INSERT INTO test_rlike_java_grouped VALUES ('aa', 1), ('ab', 1), ('aa', 2), ('xyzzy', 2), ('aa', 3), (NULL, 3) + +query expect_dispatch(rlike) +SELECT k, COUNT(*) AS c FROM test_rlike_java_grouped WHERE s RLIKE r'^(\w)\1$' GROUP BY k ORDER BY k diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_java_batches.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_java_batches.sql new file mode 100644 index 00000000000..0b4c1d8e7db --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_java_batches.sql @@ -0,0 +1,42 @@ +-- 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. + +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.RLike.allowIncompatible=false +-- Config: spark.comet.batchSize=64 + +statement +CREATE TABLE test_rlike_java_batches(id int, s string) USING parquet + +-- One input partition ensures the 5,000 rows span multiple 64-row Arrow batches. +-- Matching, nonmatching and NULL subjects recur across the batch boundaries. +statement +INSERT INTO test_rlike_java_batches +SELECT CAST(id AS INT), CASE + WHEN id % 7 = 0 THEN NULL + WHEN id % 4 = 0 THEN 'aa' + WHEN id % 4 = 1 THEN 'bb' + WHEN id % 4 = 2 THEN 'ab' + ELSE concat('row_', CAST(id AS STRING), '_aa') +END +FROM range(0, 5000, 1, 1) + +query expect_dispatch(rlike) +SELECT id, s, s RLIKE r'^(\w)\1$' FROM test_rlike_java_batches + +query expect_dispatch(rlike) +SELECT id, s FROM test_rlike_java_batches WHERE s RLIKE r'^(\w)\1$' diff --git a/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala b/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala index d90a299bed9..46caadc31f8 100644 --- a/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRegExpJvmSuite.scala @@ -20,24 +20,16 @@ package org.apache.comet import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec} -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.col import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} // Regex expressions other than in-subset `rlike` run through the codegen dispatcher by default // (Spark's own code, enabled by default) rather than the native rust path. -class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { +class CometRegExpJvmSuite extends CometTestBase { - // Patterns that the Rust regex crate cannot handle. Using one of these proves - // the JVM path was taken: if the pattern reached native, native would have - // rejected it and the operator would not be Comet. - private val backreference = "^(\\\\w)\\\\1$" + // A Java-only pattern used to check the opt-in behavior when Rust rejects it. private val lookahead = "foo(?=bar)" - private val lookbehind = "(?<=foo)bar" - private val embeddedFlags = "(?i)foo" - private val namedGroup = "(?\\\\d)" // Same structure as Sunchao's 4096-branch / 131,072-capture reproducer // (PR #5415). CometRegexSuite asserts the full pattern on the analyzer @@ -62,85 +54,6 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { // ========== rlike tests ========== - test("rlike: projection produces Java regex semantics with null handling") { - withSubjects("abc123", "no digits", null, "mixed_42_data") { - val df = sql("SELECT s, s rlike '\\\\d+' AS m FROM t") - checkSparkAnswerAndOperator(df) - } - } - - test("rlike: predicate filters rows using Java regex semantics") { - withSubjects("abc123", "no digits", null, "mixed_42_data") { - val df = sql("SELECT s FROM t WHERE s rlike '\\\\d+'") - checkSparkAnswerAndOperator(df) - } - } - - test("rlike: backreference in projection (Java-only construct)") { - withSubjects("aa", "ab", "xyzzy", null) { - val df = sql(s"SELECT s, s rlike '$backreference' FROM t") - checkSparkAnswerAndOperator(df) - val plan = df.queryExecution.executedPlan - assert( - collect(plan) { case p: CometProjectExec => p }.nonEmpty, - s"Expected CometProjectExec in:\n$plan") - } - } - - test("rlike: backreference in predicate (Java-only construct)") { - withSubjects("aa", "ab", "xyzzy", null) { - val df = sql(s"SELECT s FROM t WHERE s rlike '$backreference'") - checkSparkAnswerAndOperator(df) - val plan = df.queryExecution.executedPlan - assert( - collect(plan) { case f: CometFilterExec => f }.nonEmpty, - s"Expected CometFilterExec in:\n$plan") - } - } - - test("rlike: lookahead pattern (Java-only construct)") { - withSubjects("foobar", "foobaz", "barfoo", null) { - checkSparkAnswerAndOperator(sql(s"SELECT s, s rlike '$lookahead' FROM t")) - checkSparkAnswerAndOperator(sql(s"SELECT s FROM t WHERE s rlike '$lookahead'")) - } - } - - test("rlike: lookbehind pattern (Java-only construct)") { - withSubjects("foobar", "barbar", "foofoo", null) { - checkSparkAnswerAndOperator(sql(s"SELECT s, s rlike '$lookbehind' FROM t")) - } - } - - test("rlike: embedded case-insensitive flag (Java-only construct)") { - withSubjects("FOO", "foo", "fOO", "bar") { - checkSparkAnswerAndOperator(sql(s"SELECT s, s rlike '$embeddedFlags' FROM t")) - } - } - - test("rlike: named groups (Java-only construct)") { - withSubjects("a1", "ab", "9z", null) { - checkSparkAnswerAndOperator(sql(s"SELECT s, s rlike '$namedGroup' FROM t")) - } - } - - test("rlike: empty pattern matches every non-null row") { - withSubjects("abc", "", null) { - checkSparkAnswerAndOperator(sql("SELECT s, s rlike '' FROM t")) - } - } - - test("rlike: empty subject string is handled correctly") { - withSubjects("", "x", null) { - checkSparkAnswerAndOperator(sql("SELECT s, s rlike '^$' FROM t")) - } - } - - test("rlike: all-null subject column produces all-null result") { - withSubjects(null, null, null) { - checkSparkAnswerAndOperator(sql("SELECT s rlike '\\\\d+' FROM t")) - } - } - test("rlike: invalid pattern falls back to Spark") { withSubjects("a") { val ex = intercept[Throwable](sql("SELECT s rlike '[' FROM t").collect()) @@ -152,32 +65,6 @@ class CometRegExpJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("rlike: combines with filter, projection, and aggregate") { - withTable("t") { - sql("CREATE TABLE t (s STRING, k INT) USING parquet") - sql("""INSERT INTO t VALUES - | ('aa', 1), ('ab', 1), ('aa', 2), ('xyzzy', 2), ('aa', 3), (NULL, 3)""".stripMargin) - val df = sql(s"""SELECT k, COUNT(*) AS c - |FROM t - |WHERE s rlike '$backreference' - |GROUP BY k - |ORDER BY k""".stripMargin) - checkSparkAnswerAndOperator(df) - } - } - - test("rlike: many rows spanning multiple batches") { - withTable("t") { - sql("CREATE TABLE t (s STRING) USING parquet") - val values = (0 until 5000) - .map(i => if (i % 7 == 0) "(NULL)" else s"('row_${i}_aa')") - .mkString(", ") - sql(s"INSERT INTO t VALUES $values") - checkSparkAnswerAndOperator(sql(s"SELECT s, s rlike '$backreference' FROM t")) - checkSparkAnswerAndOperator(sql(s"SELECT s FROM t WHERE s rlike '$backreference'")) - } - } - private def withRLikeExplain(f: => Unit): Unit = { withSQLConf( CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", From 033cb37fe6841adf879a0d340206280b9e77a1fb Mon Sep 17 00:00:00 2001 From: KUAN-HAO HUANG <101171023+rich7420@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:31:36 +0000 Subject: [PATCH 80/89] test: add SQL error-class parity checks (#6650) * test: add SQL error-class parity checks * docs: clarify SQL fixture locations --- .../contributor-guide/sql-file-tests.md | 69 ++++++++++-- .../framework/expect_error_class.sql | 34 ++++++ .../apache/comet/CometSqlFileTestSuite.scala | 104 ++++++++++++++++-- .../org/apache/comet/SqlFileTestParser.scala | 28 +++-- .../apache/comet/SqlFileTestParserSuite.scala | 60 ++++++++++ 5 files changed, 263 insertions(+), 32 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/framework/expect_error_class.sql diff --git a/docs/source/contributor-guide/sql-file-tests.md b/docs/source/contributor-guide/sql-file-tests.md index 62bff931bb1..78bb8806c95 100644 --- a/docs/source/contributor-guide/sql-file-tests.md +++ b/docs/source/contributor-guide/sql-file-tests.md @@ -38,17 +38,19 @@ Run a single test file by adding the file name (without `.sql` extension) after ``` This uses ScalaTest's substring matching, so the argument must match part of the test name. -Test names follow the pattern `sql-file: expressions//.sql []`. +Test names use paths relative to `sql-tests/`, for example +`sql-file: expressions//.sql []` or +`sql-file: framework/.sql []`. ## Test file location SQL test files live under: ``` -spark/src/test/resources/sql-tests/expressions/ +spark/src/test/resources/sql-tests/ ``` -Files are organized into category subdirectories: +Expression fixtures are organized into category subdirectories: ``` expressions/ @@ -67,8 +69,11 @@ expressions/ struct/ -- create_named_struct, get_struct_field, ... ``` -The test suite recursively discovers all `.sql` files in these directories. Each file becomes -one or more ScalaTest test cases. +SQL test framework regression fixtures live in `framework/`, such as +`framework/expect_error_class.sql`. + +The test suite recursively discovers all `.sql` files under `sql-tests/`. Each file becomes one +or more ScalaTest test cases. ## File format @@ -213,9 +218,10 @@ between matrix combinations. Keys changed by `SET` need not also appear in `Conf `SET spark.sql.optimizer.excludedRules=...` replaces the whole exclusion list, including the default ConstantFolding exclusion. -Enabling the codegen dispatcher through `SET` still requires a non-error sentinel when -the file contains `expect_error`, just as enabling it through `Config` does. Run that -sentinel under the same settings as the error query. +Enabling the codegen dispatcher through `SET` still requires an operator-checking sentinel +when the file contains `expect_error`, just as enabling it through `Config` does. A plain +`query`, `expect_native`, `expect_dispatch` or `expect_error_class` can serve as the sentinel. +Run it under the same settings as the error query. This isolation covers SQL configs, not `USE`, temporary views or functions, or SQL session variables. Manage those separately. @@ -316,8 +322,12 @@ SELECT space(n) FROM test_space WHERE n < 0 #### `query expect_error()` -Asserts that both Spark and Comet throw an exception containing the given pattern. Use this -for ANSI mode tests where invalid operations should throw errors. +Asserts that both Spark and Comet throw an exception containing the given pattern. This mode +permits fallback to Spark and does not compare structured error classes or SQLSTATE. Use it +for unstructured exceptions or errors on an intentional fallback path. For structured errors +from Comet execution, prefer `expect_error_class` below. + +The literal-only examples below check error messages; they do not establish Comet coverage. ```sql -- Config: spark.sql.ansi.enabled=true @@ -335,11 +345,46 @@ query expect_error(INVALID_ARRAY_INDEX) SELECT array(1, 2, 3)[10] ``` +#### `query expect_error_class()` + +Checks the planned Comet operators and requires Spark and Comet to throw the same exception +class, exact error class and SQLSTATE. Neither exception's cause chain may contain a +`CometNativeException`. Spark 4 also calls the error class a "condition". + +The argument is a single uppercase error class, including an optional subclass separated by +a dot. It is compared exactly, so `ARITHMETIC_OVERFLOW` does not match +`BINARY_ARITHMETIC_OVERFLOW`. + +```sql +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE division_input(n int, d int) USING parquet + +statement +INSERT INTO division_input SELECT 1, 0 + +query expect_error_class(DIVIDE_BY_ZERO) +SELECT n / d FROM division_input +``` + +Read the expression inputs from Parquet columns so the error occurs during execution. +Errors raised while analyzing or planning the query fail the test. The operator check runs +before execution, using the initial plan when AQE is enabled, like the other coverage modes. +Query-context assertions remain in Scala; this mode uses `checkSparkError`'s exception-class, +error-class and SQLSTATE checks. + +Fixtures that enable the codegen dispatcher and use legacy `expect_error` need a sentinel +query that checks Comet operators. A plain `query`, `expect_native`, `expect_dispatch` or +`expect_error_class` satisfies that requirement. The strict error-class mode checks its own +operators and needs no additional successful query. + ## Adding a new test -1. Create a `.sql` file under the appropriate subdirectory in +1. Create an expression fixture under the appropriate subdirectory in `spark/src/test/resources/sql-tests/expressions/`. Create a new subdirectory if no - existing category fits. + existing category fits. Put SQL test framework regression fixtures in + `spark/src/test/resources/sql-tests/framework/`. 2. Add the Apache license header as a SQL comment. diff --git a/spark/src/test/resources/sql-tests/framework/expect_error_class.sql b/spark/src/test/resources/sql-tests/framework/expect_error_class.sql new file mode 100644 index 00000000000..b58d58d3580 --- /dev/null +++ b/spark/src/test/resources/sql-tests/framework/expect_error_class.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. + +-- Config: spark.sql.ansi.enabled=true + +-- Column inputs ensure these errors occur during execution inside Comet. +statement +CREATE TABLE error_class_input(n int, d int, s string, arr array, idx int) USING parquet + +statement +INSERT INTO error_class_input SELECT 1, 0, 'invalid', array(1, 2), 3 + +query expect_error_class(DIVIDE_BY_ZERO) +SELECT n / d FROM error_class_input + +query expect_error_class(CAST_INVALID_INPUT) +SELECT CAST(s AS INT) FROM error_class_input + +query expect_error_class(INVALID_ARRAY_INDEX) +SELECT arr[idx] FROM error_class_input diff --git a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala index 72ffb23ebaf..ba6e804d664 100644 --- a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala @@ -23,6 +23,8 @@ import java.io.File import scala.util.{Failure, Success, Try} +import org.scalatest.exceptions.TestFailedException + import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.catalyst.expressions.{Abs, CreateArray, Literal} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -157,23 +159,24 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { case _ => false } if (!hasExpectError) return - // `ExpectDispatch` and `ExpectNative` both run `checkSparkAnswerAndImpl`, which performs the - // same answer and operator checks as a plain `query` before asserting anything extra, so - // either is a strictly stronger sentinel than `CheckCoverageAndAnswer`. Not accepting them - // meant upgrading a file's last positive query to one of the new modes made preflight reject - // the file for having no sentinel, forcing a redundant plain query alongside it. + // The implementation modes check answers and operators. ExpectErrorClass also checks + // operators, so each can serve as a sentinel without a redundant successful query. val hasSentinel = file.records.exists { - case SqlQuery(_, CheckCoverageAndAnswer | _: ExpectDispatch | _: ExpectNative, _) => true + case SqlQuery( + _, + CheckCoverageAndAnswer | _: ExpectDispatch | _: ExpectNative | _: ExpectErrorClass, + _) => + true case _ => false } assert( hasSentinel, s"SQL fixture $relativePath combines `expect_error` with " + - s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true but is missing a non-error " + - "sentinel query. Without one, a silent dispatcher fallback to Spark would let the " + - "`expect_error` queries pass vacuously (Spark raises the same error on the fallback " + - "path). Add at least one `query`, `expect_dispatch` or `expect_native` over valid " + - "input so the operator check fails if the expression did not execute natively.") + s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true but is missing a " + + "sentinel query that checks Comet operators. Without one, a silent fallback to Spark " + + "would let the `expect_error` queries pass vacuously (Spark raises the same error " + + "on the fallback path). Add a `query`, `expect_dispatch`, `expect_native` or " + + "`expect_error_class` so the operator check fails if the query falls back to Spark.") } private def runTestFile(relativePath: String, file: SqlTestFile): Unit = { @@ -225,6 +228,8 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { assert( cometError.get.getMessage.contains(pattern), s"Comet error '${cometError.get.getMessage}' does not contain '$pattern'") + case ExpectErrorClass(errorClass) => + checkSparkError(spark.sql(sql), errorClass) } } @@ -369,7 +374,7 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { "SELECT CAST('invalid' AS INT)")) val original = spark.sessionState.conf.getAllConfs val error = intercept[Exception](runTestFile("set-dispatcher-sentinel", file)) - assert(error.getMessage.contains("missing a non-error")) + assert(error.getMessage.contains("missing a sentinel query")) assert(spark.sessionState.conf.getAllConfs == original) } @@ -388,6 +393,81 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { assert(spark.sessionState.conf.getAllConfs == original) } + test("expect_error_class rejects Spark fallback that a message-only check accepts") { + withParquetTable(Seq(Tuple1(0)), "error_class_runner_input") { + def fixture(mode: String): SqlTestFile = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_EXEC_ENABLED.key}=false", + s"query $mode", + "SELECT 1 / _1 FROM error_class_runner_input")) + + runTestFile("fallback.sql", fixture("expect_error(DIVIDE_BY_ZERO)")) + val error = intercept[RuntimeException] { + runTestFile("fallback.sql", fixture("expect_error_class(DIVIDE_BY_ZERO)")) + } + assert(error.getCause.isInstanceOf[TestFailedException]) + assert(error.getCause.getMessage.contains("fallback.sql:4")) + assert(error.getCause.getMessage.contains("Expected only Comet native operators")) + } + } + + test("expect_error_class rejects substring-only error classes") { + withParquetTable(Seq(Tuple1(0)), "error_class_runner_input") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + "query expect_error_class(ZERO)", + "SELECT 1 / _1 FROM error_class_runner_input")) + val error = intercept[RuntimeException] { + runTestFile("substring.sql", file) + } + assert(error.getCause.isInstanceOf[TestFailedException]) + assert(error.getCause.getMessage.contains("substring.sql:3")) + assert(error.getCause.getMessage.contains("did not equal")) + // ScalaTest brackets the differing parts of strings in assertion messages. + assert( + error.getCause.getMessage.replace("[", "").replace("]", "").contains("DIVIDE_BY_ZERO")) + } + } + + test("expect_error_class checks operators without a successful codegen sentinel") { + withParquetTable(Seq(Tuple1(0)), "error_class_runner_input") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true", + "query expect_error(DIVIDE_BY_ZERO)", + "SELECT 1 / _1 FROM error_class_runner_input", + "", + "query expect_error_class(DIVIDE_BY_ZERO)", + "SELECT 1 / _1 FROM error_class_runner_input")) + val error = intercept[RuntimeException] { + runTestFile("sentinel.sql", file.copy(records = file.records.take(1))) + } + assert(error.getCause.isInstanceOf[TestFailedException]) + assert(error.getCause.getMessage.contains("sentinel.sql:4")) + assert(error.getCause.getMessage.contains("is missing a sentinel query")) + runTestFile("sentinel.sql", file) + } + } + + test("expect_error_class rejects successful queries") { + withParquetTable(Seq(Tuple1(1)), "error_class_runner_input") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + "query expect_error_class(DIVIDE_BY_ZERO)", + "SELECT 1 / _1 FROM error_class_runner_input")) + val error = intercept[RuntimeException] { + runTestFile("successful.sql", file) + } + assert(error.getCause.isInstanceOf[TestFailedException]) + assert(error.getCause.getMessage.contains("successful.sql:3")) + assert(error.getCause.getMessage.contains("Spark did not fail with DIVIDE_BY_ZERO")) + } + } + // Discover and register all .sql test files discoverTestFiles(testResourceDir).foreach { file => val relativePath = testResourceDir.toURI.relativize(file.toURI).getPath diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala index f7e653c2e13..8510bc499ba 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala @@ -77,15 +77,18 @@ case class ExpectNative(names: Seq[String]) extends QueryAssertionMode * Asserts that both Spark and Comet raise an error whose message contains `pattern`. * * Fixtures that combine `ExpectError` with `spark.comet.exec.scalaUDF.codegen.enabled=true` must - * also include at least one [[CheckCoverageAndAnswer]] sentinel query over valid input. If the - * dispatcher silently rejects the expression at plan time, the operator falls back to Spark and - * Spark itself raises the same error, so the `ExpectError` queries would pass vacuously. The - * sentinel query uses `checkSparkAnswerAndOperator`, which fails when the expression does not run - * inside Comet. `CometSqlFileTestSuite.requireSentinelForCodegenExpectError` enforces this shape - * at test-run time. + * also include a query that checks Comet operators, such as [[CheckCoverageAndAnswer]] over valid + * input or [[ExpectErrorClass]]. If the dispatcher silently rejects the expression at plan time, + * the operator falls back to Spark and Spark itself raises the same error, so the `ExpectError` + * queries would pass vacuously. The sentinel query checks Comet operators and fails when the + * query falls back to Spark. `CometSqlFileTestSuite.requireSentinelForCodegenExpectError` + * enforces this shape at test-run time. */ case class ExpectError(pattern: String) extends QueryAssertionMode +/** Checks Comet operators and Spark exception type, exact error class and SQLSTATE parity. */ +case class ExpectErrorClass(errorClass: String) extends QueryAssertionMode + /** * Parsed representation of a .sql test file. * @@ -209,7 +212,7 @@ object SqlFileTestParser { lineIdx = nextIdx case s if s.startsWith("query") => - val mode = parseQueryAssertionMode(s) + val mode = parseQueryAssertionMode(s, lineIdx + 1) lineIdx += 1 val startLine = lineIdx + 1 val (sql, nextIdx) = collectSql(lines, lineIdx) @@ -236,10 +239,13 @@ object SqlFileTestParser { private val FallbackPattern = """query\s+expect_fallback\((.+)\)""".r private val IgnorePattern = """query\s+ignore\((.+)\)""".r private val ErrorPattern = """query\s+expect_error\((.+)\)""".r + private val ErrorClassPattern = + """query\s+expect_error_class\(\s*([A-Z0-9_]+(?:\.[A-Z0-9_]+)*)\s*\)""".r + private val InvalidErrorClassPattern = """query\s+expect_error_class\b.*""".r private val DispatchPattern = """query\s+expect_dispatch\((.+)\)""".r private val NativePattern = """query\s+expect_native\((.+)\)""".r - private def parseQueryAssertionMode(directive: String): QueryAssertionMode = { + private def parseQueryAssertionMode(directive: String, line: Int): QueryAssertionMode = { directive match { case FallbackPattern(reason) => ExpectFallback(reason.trim) @@ -247,6 +253,12 @@ object SqlFileTestParser { Ignore(reason.trim) case ErrorPattern(pattern) => ExpectError(pattern.trim) + case ErrorClassPattern(errorClass) => + ExpectErrorClass(errorClass) + case InvalidErrorClassPattern() => + throw new IllegalArgumentException( + s"line $line: invalid directive '$directive'; " + + "expected query expect_error_class(ERROR_CLASS)") case DispatchPattern(names) => ExpectDispatch(splitNames(names)) case NativePattern(names) => diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala index 0d107395ba7..827b12ec1f1 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala @@ -218,9 +218,55 @@ class SqlFileTestParserSuite extends AnyFunSuite { assert(modeOf("query expect_dispatch(lower,,)") === ExpectDispatch(Seq("lower"))) } + test("expect_error_class parses an exact class, including subclasses and whitespace") { + assert( + modeOf("query expect_error_class( DIVIDE_BY_ZERO )") === + ExpectErrorClass("DIVIDE_BY_ZERO")) + assert( + modeOf("query expect_error_class(DATATYPE_MISMATCH.INVALID_MAP_KEY_TYPE)") === + ExpectErrorClass("DATATYPE_MISMATCH.INVALID_MAP_KEY_TYPE")) + } + + test("malformed expect_error_class directives fail instead of becoming plain queries") { + Seq( + "query expect_error_class", + "query expect_error_class()", + "query expect_error_class( )", + "query expect_error_class(DIVIDE_BY_ZERO", + "query expect_error_class(DIVIDE_BY_ZERO, OTHER)", + "query expect_error_class(DIVIDE BY ZERO)", + "query expect_error_class(divide_by_zero)", + "query expect_error_class(DIVIDE_BY_ZERO) trailing").foreach { directive => + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(Seq("-- fixture", directive, "SELECT 1")) + } + assert(error.getMessage.contains("line 2")) + assert(error.getMessage.contains(directive)) + } + } + + test("invalid error-class directives identify the source file") { + val tempDir = Files.createDirectories(new File(System.getProperty("java.io.tmpdir")).toPath) + val file = Files.createTempFile(tempDir, "invalid-error-class", ".sql") + try { + val _ = Files.write( + file, + "-- fixture\nquery expect_error_class()\nSELECT 1\n".getBytes(StandardCharsets.UTF_8)) + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(file.toFile) + } + assert(error.getMessage.contains(file.toString)) + assert(error.getMessage.contains("line 2")) + } finally { + val _ = Files.deleteIfExists(file) + } + } + test("the new modes do not shadow the existing ones") { assert(modeOf("query expect_fallback(some reason)") === ExpectFallback("some reason")) assert(modeOf("query expect_error(DIVIDE_BY_ZERO)") === ExpectError("DIVIDE_BY_ZERO")) + assert( + modeOf("query expect_error_class(DIVIDE_BY_ZERO)") === ExpectErrorClass("DIVIDE_BY_ZERO")) assert(modeOf("query spark_answer_only") === SparkAnswerOnly) assert(modeOf("query tolerance=0.001") === WithTolerance(0.001)) assert( @@ -238,6 +284,20 @@ class SqlFileTestParserSuite extends AnyFunSuite { assert(queries.map(_.sql) === Seq("SELECT abs(a) FROM t", "SELECT hypot(a, b) FROM t")) } + test("error-class queries retain their SQL text and source line") { + val queries = parseQueries( + "-- fixture", + "query expect_error_class(DIVIDE_BY_ZERO)", + "SELECT a / b FROM t", + "", + "query expect_error(BY_ZERO)", + "SELECT a / b FROM t") + assert( + queries.map(_.mode) === Seq(ExpectErrorClass("DIVIDE_BY_ZERO"), ExpectError("BY_ZERO"))) + assert(queries.map(_.sql) === Seq("SELECT a / b FROM t", "SELECT a / b FROM t")) + assert(queries.map(_.line) === Seq(3, 6)) + } + // #5702: Spark releases before SPARK-54918 keep -0.0 distinct in array_distinct and array_union, // and native distinct/union run by default only on Spark 4.2.0, the one release that normalizes // their arguments in the plan. The signed-zero cases therefore live in version-gated fixtures: From 0f4811dc4817efebaf5ff6e6f243ab1e098564de Mon Sep 17 00:00:00 2001 From: YuLun Mao <93830314+mizulun@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:35:12 +0000 Subject: [PATCH 81/89] fix: restore localReads assertion in AdaptiveQueryExecSuite diffs (#6382) * fix: restore localReads assertion in AdaptiveQueryExecSuite diffs * fix: restore localReads assertion in Spark 3.4.3 diff --- dev/diffs/3.4.3.diff | 70 ++++++++++++++++-------------------- dev/diffs/3.5.9.diff | 74 +++++++++++++++++--------------------- dev/diffs/4.0.4.diff | 84 +++++++++++++++++++------------------------- dev/diffs/4.1.3.diff | 84 +++++++++++++++++++------------------------- 4 files changed, 136 insertions(+), 176 deletions(-) diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index af31251b32e..1caa75a571a 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -1751,7 +1751,7 @@ index ac710c32296..2854b433dd3 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 593bd7bb4ba..47285d87416 100644 +index 593bd7bb4ba..7111252fc44 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._ @@ -1862,17 +1862,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", -@@ -273,7 +289,8 @@ class AdaptiveQueryExecSuite - val localReads = collect(adaptivePlan) { - case read: AQEShuffleReadExec if read.isLocalRead => read - } -- assert(localReads.length == 2) -+ // Comet shuffle changes shuffle metrics -+ assert(localReads.length == 1) - val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD] - val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD] - // the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2 -@@ -298,7 +315,9 @@ class AdaptiveQueryExecSuite +@@ -298,7 +314,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -1883,7 +1873,7 @@ index 593bd7bb4ba..47285d87416 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -312,7 +331,9 @@ class AdaptiveQueryExecSuite +@@ -312,7 +330,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -1894,7 +1884,7 @@ index 593bd7bb4ba..47285d87416 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -322,7 +343,7 @@ class AdaptiveQueryExecSuite +@@ -322,7 +342,7 @@ class AdaptiveQueryExecSuite } } @@ -1903,7 +1893,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -337,7 +358,7 @@ class AdaptiveQueryExecSuite +@@ -337,7 +357,7 @@ class AdaptiveQueryExecSuite } } @@ -1912,7 +1902,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -353,7 +374,7 @@ class AdaptiveQueryExecSuite +@@ -353,7 +373,7 @@ class AdaptiveQueryExecSuite } } @@ -1921,7 +1911,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -398,7 +419,7 @@ class AdaptiveQueryExecSuite +@@ -398,7 +418,7 @@ class AdaptiveQueryExecSuite } } @@ -1930,7 +1920,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -443,7 +464,7 @@ class AdaptiveQueryExecSuite +@@ -443,7 +463,7 @@ class AdaptiveQueryExecSuite } } @@ -1939,7 +1929,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { -@@ -489,7 +510,7 @@ class AdaptiveQueryExecSuite +@@ -489,7 +509,7 @@ class AdaptiveQueryExecSuite } } @@ -1948,7 +1938,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -508,7 +529,7 @@ class AdaptiveQueryExecSuite +@@ -508,7 +528,7 @@ class AdaptiveQueryExecSuite } } @@ -1957,7 +1947,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -539,7 +560,9 @@ class AdaptiveQueryExecSuite +@@ -539,7 +559,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -1968,7 +1958,7 @@ index 593bd7bb4ba..47285d87416 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.nonEmpty) -@@ -560,7 +583,9 @@ class AdaptiveQueryExecSuite +@@ -560,7 +582,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -1979,7 +1969,7 @@ index 593bd7bb4ba..47285d87416 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.isEmpty) -@@ -569,7 +594,8 @@ class AdaptiveQueryExecSuite +@@ -569,7 +593,8 @@ class AdaptiveQueryExecSuite } } @@ -1989,7 +1979,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000", -@@ -664,7 +690,8 @@ class AdaptiveQueryExecSuite +@@ -664,7 +689,8 @@ class AdaptiveQueryExecSuite val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) // There is still a SMJ, and its two shuffles can't apply local read. @@ -1999,7 +1989,7 @@ index 593bd7bb4ba..47285d87416 100644 } } -@@ -786,7 +813,8 @@ class AdaptiveQueryExecSuite +@@ -786,7 +812,8 @@ class AdaptiveQueryExecSuite } } @@ -2009,7 +1999,7 @@ index 593bd7bb4ba..47285d87416 100644 Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint => def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") { findTopLevelSortMergeJoin(plan) -@@ -1004,7 +1032,8 @@ class AdaptiveQueryExecSuite +@@ -1004,7 +1031,8 @@ class AdaptiveQueryExecSuite } } @@ -2019,7 +2009,7 @@ index 593bd7bb4ba..47285d87416 100644 withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT key FROM testData GROUP BY key") -@@ -1599,7 +1628,7 @@ class AdaptiveQueryExecSuite +@@ -1599,7 +1627,7 @@ class AdaptiveQueryExecSuite val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id") assert(collect(adaptivePlan) { @@ -2028,7 +2018,7 @@ index 593bd7bb4ba..47285d87416 100644 }.length == 1) } } -@@ -1679,7 +1708,8 @@ class AdaptiveQueryExecSuite +@@ -1679,7 +1707,8 @@ class AdaptiveQueryExecSuite } } @@ -2038,7 +2028,7 @@ index 593bd7bb4ba..47285d87416 100644 def hasRepartitionShuffle(plan: SparkPlan): Boolean = { find(plan) { case s: ShuffleExchangeLike => -@@ -1864,6 +1894,9 @@ class AdaptiveQueryExecSuite +@@ -1864,6 +1893,9 @@ class AdaptiveQueryExecSuite def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = { assert(collect(ds.queryExecution.executedPlan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -2048,7 +2038,7 @@ index 593bd7bb4ba..47285d87416 100644 }.size == 1) ds.collect() val plan = ds.queryExecution.executedPlan -@@ -1872,6 +1905,9 @@ class AdaptiveQueryExecSuite +@@ -1872,6 +1904,9 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(plan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -2058,7 +2048,7 @@ index 593bd7bb4ba..47285d87416 100644 }.size == 1) checkAnswer(ds, testData) } -@@ -2028,7 +2064,8 @@ class AdaptiveQueryExecSuite +@@ -2028,7 +2063,8 @@ class AdaptiveQueryExecSuite } } @@ -2068,7 +2058,7 @@ index 593bd7bb4ba..47285d87416 100644 withTempView("t1", "t2") { def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = { Seq("100", "100000").foreach { size => -@@ -2114,7 +2151,8 @@ class AdaptiveQueryExecSuite +@@ -2114,7 +2150,8 @@ class AdaptiveQueryExecSuite } } @@ -2078,7 +2068,7 @@ index 593bd7bb4ba..47285d87416 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", -@@ -2213,7 +2251,7 @@ class AdaptiveQueryExecSuite +@@ -2213,7 +2250,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2 GROUP BY key1") val shuffles1 = collect(adaptive1) { @@ -2087,7 +2077,7 @@ index 593bd7bb4ba..47285d87416 100644 } assert(shuffles1.size == 3) // shuffles1.head is the top-level shuffle under the Aggregate operator -@@ -2226,7 +2264,7 @@ class AdaptiveQueryExecSuite +@@ -2226,7 +2263,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2") val shuffles2 = collect(adaptive2) { @@ -2096,7 +2086,7 @@ index 593bd7bb4ba..47285d87416 100644 } if (hasRequiredDistribution) { assert(shuffles2.size == 3) -@@ -2260,7 +2298,8 @@ class AdaptiveQueryExecSuite +@@ -2260,7 +2297,8 @@ class AdaptiveQueryExecSuite } } @@ -2106,7 +2096,7 @@ index 593bd7bb4ba..47285d87416 100644 CostEvaluator.instantiate( classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf) intercept[IllegalArgumentException] { -@@ -2404,6 +2443,7 @@ class AdaptiveQueryExecSuite +@@ -2404,6 +2442,7 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult(query) assert(adaptive.collect { case sort: SortExec => sort @@ -2114,7 +2104,7 @@ index 593bd7bb4ba..47285d87416 100644 }.size == 1) val read = collect(adaptive) { case read: AQEShuffleReadExec => read -@@ -2421,7 +2461,8 @@ class AdaptiveQueryExecSuite +@@ -2421,7 +2460,8 @@ class AdaptiveQueryExecSuite } } @@ -2124,7 +2114,7 @@ index 593bd7bb4ba..47285d87416 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true", -@@ -2533,7 +2574,7 @@ class AdaptiveQueryExecSuite +@@ -2533,7 +2573,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value2 = value3") val shuffles1 = collect(adaptive1) { @@ -2133,7 +2123,7 @@ index 593bd7bb4ba..47285d87416 100644 } assert(shuffles1.size == 4) val smj1 = findTopLevelSortMergeJoin(adaptive1) -@@ -2544,7 +2585,7 @@ class AdaptiveQueryExecSuite +@@ -2544,7 +2584,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value1 = value3") val shuffles2 = collect(adaptive2) { @@ -2142,7 +2132,7 @@ index 593bd7bb4ba..47285d87416 100644 } assert(shuffles2.size == 4) val smj2 = findTopLevelSortMergeJoin(adaptive2) -@@ -2703,7 +2744,8 @@ class AdaptiveQueryExecSuite +@@ -2703,7 +2743,8 @@ class AdaptiveQueryExecSuite } } diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index 7549996767a..67e59a4a323 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -1947,7 +1947,7 @@ index 5a413c77754..207b66e1d7b 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 2f8e401e743..fc4ce35be4b 100644 +index 2f8e401e743..117bd9ffdde 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._ @@ -2065,17 +2065,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", -@@ -299,7 +316,8 @@ class AdaptiveQueryExecSuite - val localReads = collect(adaptivePlan) { - case read: AQEShuffleReadExec if read.isLocalRead => read - } -- assert(localReads.length == 2) -+ // Comet shuffle changes shuffle metrics -+ assert(localReads.length == 1) - val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD] - val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD] - // the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2 -@@ -324,7 +342,9 @@ class AdaptiveQueryExecSuite +@@ -324,7 +341,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2086,7 +2076,7 @@ index 2f8e401e743..fc4ce35be4b 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -338,7 +358,9 @@ class AdaptiveQueryExecSuite +@@ -338,7 +357,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2097,7 +2087,7 @@ index 2f8e401e743..fc4ce35be4b 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -348,7 +370,7 @@ class AdaptiveQueryExecSuite +@@ -348,7 +369,7 @@ class AdaptiveQueryExecSuite } } @@ -2106,7 +2096,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -363,7 +385,7 @@ class AdaptiveQueryExecSuite +@@ -363,7 +384,7 @@ class AdaptiveQueryExecSuite } } @@ -2115,7 +2105,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -379,7 +401,7 @@ class AdaptiveQueryExecSuite +@@ -379,7 +400,7 @@ class AdaptiveQueryExecSuite } } @@ -2124,7 +2114,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -424,7 +446,7 @@ class AdaptiveQueryExecSuite +@@ -424,7 +445,7 @@ class AdaptiveQueryExecSuite } } @@ -2133,7 +2123,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -469,7 +491,7 @@ class AdaptiveQueryExecSuite +@@ -469,7 +490,7 @@ class AdaptiveQueryExecSuite } } @@ -2142,7 +2132,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { -@@ -515,7 +537,7 @@ class AdaptiveQueryExecSuite +@@ -515,7 +536,7 @@ class AdaptiveQueryExecSuite } } @@ -2151,7 +2141,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -534,7 +556,7 @@ class AdaptiveQueryExecSuite +@@ -534,7 +555,7 @@ class AdaptiveQueryExecSuite } } @@ -2160,7 +2150,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -565,7 +587,9 @@ class AdaptiveQueryExecSuite +@@ -565,7 +586,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -2171,7 +2161,7 @@ index 2f8e401e743..fc4ce35be4b 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.nonEmpty) -@@ -586,7 +610,9 @@ class AdaptiveQueryExecSuite +@@ -586,7 +609,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -2182,7 +2172,7 @@ index 2f8e401e743..fc4ce35be4b 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.isEmpty) -@@ -595,7 +621,8 @@ class AdaptiveQueryExecSuite +@@ -595,7 +620,8 @@ class AdaptiveQueryExecSuite } } @@ -2192,7 +2182,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000", -@@ -690,7 +717,8 @@ class AdaptiveQueryExecSuite +@@ -690,7 +716,8 @@ class AdaptiveQueryExecSuite val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) // There is still a SMJ, and its two shuffles can't apply local read. @@ -2202,7 +2192,7 @@ index 2f8e401e743..fc4ce35be4b 100644 } } -@@ -812,7 +840,8 @@ class AdaptiveQueryExecSuite +@@ -812,7 +839,8 @@ class AdaptiveQueryExecSuite } } @@ -2212,7 +2202,7 @@ index 2f8e401e743..fc4ce35be4b 100644 Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint => def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") { findTopLevelSortMergeJoin(plan) -@@ -1030,7 +1059,8 @@ class AdaptiveQueryExecSuite +@@ -1030,7 +1058,8 @@ class AdaptiveQueryExecSuite } } @@ -2222,7 +2212,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT key FROM testData GROUP BY key") -@@ -1625,7 +1655,7 @@ class AdaptiveQueryExecSuite +@@ -1625,7 +1654,7 @@ class AdaptiveQueryExecSuite val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id") assert(collect(adaptivePlan) { @@ -2231,7 +2221,7 @@ index 2f8e401e743..fc4ce35be4b 100644 }.length == 1) } } -@@ -1705,7 +1735,8 @@ class AdaptiveQueryExecSuite +@@ -1705,7 +1734,8 @@ class AdaptiveQueryExecSuite } } @@ -2241,7 +2231,7 @@ index 2f8e401e743..fc4ce35be4b 100644 def hasRepartitionShuffle(plan: SparkPlan): Boolean = { find(plan) { case s: ShuffleExchangeLike => -@@ -1890,6 +1921,9 @@ class AdaptiveQueryExecSuite +@@ -1890,6 +1920,9 @@ class AdaptiveQueryExecSuite def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = { assert(collect(ds.queryExecution.executedPlan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -2251,7 +2241,7 @@ index 2f8e401e743..fc4ce35be4b 100644 }.size == 1) ds.collect() val plan = ds.queryExecution.executedPlan -@@ -1898,6 +1932,9 @@ class AdaptiveQueryExecSuite +@@ -1898,6 +1931,9 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(plan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -2261,7 +2251,7 @@ index 2f8e401e743..fc4ce35be4b 100644 }.size == 1) checkAnswer(ds, testData) } -@@ -2054,7 +2091,8 @@ class AdaptiveQueryExecSuite +@@ -2054,7 +2090,8 @@ class AdaptiveQueryExecSuite } } @@ -2271,7 +2261,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withTempView("t1", "t2") { def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = { Seq("100", "100000").foreach { size => -@@ -2140,7 +2178,8 @@ class AdaptiveQueryExecSuite +@@ -2140,7 +2177,8 @@ class AdaptiveQueryExecSuite } } @@ -2281,7 +2271,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", -@@ -2239,7 +2278,7 @@ class AdaptiveQueryExecSuite +@@ -2239,7 +2277,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2 GROUP BY key1") val shuffles1 = collect(adaptive1) { @@ -2290,7 +2280,7 @@ index 2f8e401e743..fc4ce35be4b 100644 } assert(shuffles1.size == 3) // shuffles1.head is the top-level shuffle under the Aggregate operator -@@ -2252,7 +2291,7 @@ class AdaptiveQueryExecSuite +@@ -2252,7 +2290,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2") val shuffles2 = collect(adaptive2) { @@ -2299,7 +2289,7 @@ index 2f8e401e743..fc4ce35be4b 100644 } if (hasRequiredDistribution) { assert(shuffles2.size == 3) -@@ -2286,7 +2325,8 @@ class AdaptiveQueryExecSuite +@@ -2286,7 +2324,8 @@ class AdaptiveQueryExecSuite } } @@ -2309,7 +2299,7 @@ index 2f8e401e743..fc4ce35be4b 100644 CostEvaluator.instantiate( classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf) intercept[IllegalArgumentException] { -@@ -2452,6 +2492,7 @@ class AdaptiveQueryExecSuite +@@ -2452,6 +2491,7 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult(query) assert(adaptive.collect { case sort: SortExec => sort @@ -2317,7 +2307,7 @@ index 2f8e401e743..fc4ce35be4b 100644 }.size == 1) val read = collect(adaptive) { case read: AQEShuffleReadExec => read -@@ -2469,7 +2510,8 @@ class AdaptiveQueryExecSuite +@@ -2469,7 +2509,8 @@ class AdaptiveQueryExecSuite } } @@ -2327,7 +2317,7 @@ index 2f8e401e743..fc4ce35be4b 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true", -@@ -2581,7 +2623,7 @@ class AdaptiveQueryExecSuite +@@ -2581,7 +2622,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value2 = value3") val shuffles1 = collect(adaptive1) { @@ -2336,7 +2326,7 @@ index 2f8e401e743..fc4ce35be4b 100644 } assert(shuffles1.size == 4) val smj1 = findTopLevelSortMergeJoin(adaptive1) -@@ -2592,7 +2634,7 @@ class AdaptiveQueryExecSuite +@@ -2592,7 +2633,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value1 = value3") val shuffles2 = collect(adaptive2) { @@ -2345,7 +2335,7 @@ index 2f8e401e743..fc4ce35be4b 100644 } assert(shuffles2.size == 4) val smj2 = findTopLevelSortMergeJoin(adaptive2) -@@ -2850,6 +2892,7 @@ class AdaptiveQueryExecSuite +@@ -2850,6 +2891,7 @@ class AdaptiveQueryExecSuite }.size == (if (firstAccess) 1 else 0)) assert(collect(initialExecutedPlan) { case s: SortExec => s @@ -2353,7 +2343,7 @@ index 2f8e401e743..fc4ce35be4b 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2932,13 +2975,18 @@ class AdaptiveQueryExecSuite +@@ -2932,13 +2974,18 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult( "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") @@ -2374,7 +2364,7 @@ index 2f8e401e743..fc4ce35be4b 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -2980,7 +3028,9 @@ class AdaptiveQueryExecSuite +@@ -2980,7 +3027,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 2133589aa92..83a19b27250 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -2636,7 +2636,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index fb8fab6a80f..cfcff43d557 100644 +index fb8fab6a80f..bbe8dcafea8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -28,12 +28,14 @@ import org.apache.spark.SparkException @@ -2766,17 +2766,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", -@@ -301,7 +327,8 @@ class AdaptiveQueryExecSuite - val localReads = collect(adaptivePlan) { - case read: AQEShuffleReadExec if read.isLocalRead => read - } -- assert(localReads.length == 2) -+ // Comet shuffle changes shuffle metrics -+ assert(localReads.length == 1) - val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD] - val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD] - // the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2 -@@ -326,7 +353,9 @@ class AdaptiveQueryExecSuite +@@ -326,7 +352,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2787,7 +2777,7 @@ index fb8fab6a80f..cfcff43d557 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -340,7 +369,9 @@ class AdaptiveQueryExecSuite +@@ -340,7 +368,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2798,7 +2788,7 @@ index fb8fab6a80f..cfcff43d557 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -350,7 +381,7 @@ class AdaptiveQueryExecSuite +@@ -350,7 +380,7 @@ class AdaptiveQueryExecSuite } } @@ -2807,7 +2797,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -365,7 +396,7 @@ class AdaptiveQueryExecSuite +@@ -365,7 +395,7 @@ class AdaptiveQueryExecSuite } } @@ -2816,7 +2806,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -381,7 +412,7 @@ class AdaptiveQueryExecSuite +@@ -381,7 +411,7 @@ class AdaptiveQueryExecSuite } } @@ -2825,7 +2815,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -426,7 +457,7 @@ class AdaptiveQueryExecSuite +@@ -426,7 +456,7 @@ class AdaptiveQueryExecSuite } } @@ -2834,7 +2824,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -471,7 +502,7 @@ class AdaptiveQueryExecSuite +@@ -471,7 +501,7 @@ class AdaptiveQueryExecSuite } } @@ -2843,7 +2833,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { -@@ -517,7 +548,7 @@ class AdaptiveQueryExecSuite +@@ -517,7 +547,7 @@ class AdaptiveQueryExecSuite } } @@ -2852,7 +2842,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -536,7 +567,7 @@ class AdaptiveQueryExecSuite +@@ -536,7 +566,7 @@ class AdaptiveQueryExecSuite } } @@ -2861,7 +2851,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -567,7 +598,9 @@ class AdaptiveQueryExecSuite +@@ -567,7 +597,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -2872,7 +2862,7 @@ index fb8fab6a80f..cfcff43d557 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.nonEmpty) -@@ -588,7 +621,9 @@ class AdaptiveQueryExecSuite +@@ -588,7 +620,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -2883,7 +2873,7 @@ index fb8fab6a80f..cfcff43d557 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.isEmpty) -@@ -597,7 +632,8 @@ class AdaptiveQueryExecSuite +@@ -597,7 +631,8 @@ class AdaptiveQueryExecSuite } } @@ -2893,7 +2883,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000", -@@ -692,7 +728,8 @@ class AdaptiveQueryExecSuite +@@ -692,7 +727,8 @@ class AdaptiveQueryExecSuite val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) // There is still a SMJ, and its two shuffles can't apply local read. @@ -2903,7 +2893,7 @@ index fb8fab6a80f..cfcff43d557 100644 } } -@@ -814,7 +851,8 @@ class AdaptiveQueryExecSuite +@@ -814,7 +850,8 @@ class AdaptiveQueryExecSuite } } @@ -2913,7 +2903,7 @@ index fb8fab6a80f..cfcff43d557 100644 Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint => def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") { findTopLevelSortMergeJoin(plan) -@@ -1096,7 +1134,8 @@ class AdaptiveQueryExecSuite +@@ -1096,7 +1133,8 @@ class AdaptiveQueryExecSuite } } @@ -2923,7 +2913,7 @@ index fb8fab6a80f..cfcff43d557 100644 withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT key FROM testData GROUP BY key") -@@ -1662,13 +1701,13 @@ class AdaptiveQueryExecSuite +@@ -1662,13 +1700,13 @@ class AdaptiveQueryExecSuite val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult( "SELECT key FROM testData WHERE key = 0 ORDER BY key, value") assert(findTopLevelSort(plan1).size == 1) @@ -2939,7 +2929,7 @@ index fb8fab6a80f..cfcff43d557 100644 } } -@@ -1676,18 +1715,18 @@ class AdaptiveQueryExecSuite +@@ -1676,18 +1714,18 @@ class AdaptiveQueryExecSuite withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult( "SELECT key, count(*) FROM testData WHERE value = 'no_match' GROUP BY key") @@ -2964,7 +2954,7 @@ index fb8fab6a80f..cfcff43d557 100644 } } -@@ -1708,7 +1747,7 @@ class AdaptiveQueryExecSuite +@@ -1708,7 +1746,7 @@ class AdaptiveQueryExecSuite |""".stripMargin) checkNumUnion(plan1, 1) checkNumUnion(adaptivePlan1, 0) @@ -2973,7 +2963,7 @@ index fb8fab6a80f..cfcff43d557 100644 val (plan2, adaptivePlan2) = runAdaptiveAndVerifyResult( """ -@@ -1718,7 +1757,7 @@ class AdaptiveQueryExecSuite +@@ -1718,7 +1756,7 @@ class AdaptiveQueryExecSuite |""".stripMargin) checkNumUnion(plan2, 1) checkNumUnion(adaptivePlan2, 0) @@ -2982,7 +2972,7 @@ index fb8fab6a80f..cfcff43d557 100644 } } -@@ -1730,7 +1769,7 @@ class AdaptiveQueryExecSuite +@@ -1730,7 +1768,7 @@ class AdaptiveQueryExecSuite val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id") assert(collect(adaptivePlan) { @@ -2991,7 +2981,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.length == 1) } } -@@ -1810,7 +1849,8 @@ class AdaptiveQueryExecSuite +@@ -1810,7 +1848,8 @@ class AdaptiveQueryExecSuite } } @@ -3001,7 +2991,7 @@ index fb8fab6a80f..cfcff43d557 100644 def hasRepartitionShuffle(plan: SparkPlan): Boolean = { find(plan) { case s: ShuffleExchangeLike => -@@ -1995,6 +2035,9 @@ class AdaptiveQueryExecSuite +@@ -1995,6 +2034,9 @@ class AdaptiveQueryExecSuite def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = { assert(collect(ds.queryExecution.executedPlan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -3011,7 +3001,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.size == 1) ds.collect() val plan = ds.queryExecution.executedPlan -@@ -2003,6 +2046,9 @@ class AdaptiveQueryExecSuite +@@ -2003,6 +2045,9 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(plan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -3021,7 +3011,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.size == 1) checkAnswer(ds, testData) } -@@ -2159,7 +2205,8 @@ class AdaptiveQueryExecSuite +@@ -2159,7 +2204,8 @@ class AdaptiveQueryExecSuite } } @@ -3031,7 +3021,7 @@ index fb8fab6a80f..cfcff43d557 100644 withTempView("t1", "t2") { def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = { Seq("100", "100000").foreach { size => -@@ -2245,7 +2292,8 @@ class AdaptiveQueryExecSuite +@@ -2245,7 +2291,8 @@ class AdaptiveQueryExecSuite } } @@ -3041,7 +3031,7 @@ index fb8fab6a80f..cfcff43d557 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", -@@ -2344,7 +2392,7 @@ class AdaptiveQueryExecSuite +@@ -2344,7 +2391,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2 GROUP BY key1") val shuffles1 = collect(adaptive1) { @@ -3050,7 +3040,7 @@ index fb8fab6a80f..cfcff43d557 100644 } assert(shuffles1.size == 3) // shuffles1.head is the top-level shuffle under the Aggregate operator -@@ -2357,7 +2405,7 @@ class AdaptiveQueryExecSuite +@@ -2357,7 +2404,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2") val shuffles2 = collect(adaptive2) { @@ -3059,7 +3049,7 @@ index fb8fab6a80f..cfcff43d557 100644 } if (hasRequiredDistribution) { assert(shuffles2.size == 3) -@@ -2391,7 +2439,8 @@ class AdaptiveQueryExecSuite +@@ -2391,7 +2438,8 @@ class AdaptiveQueryExecSuite } } @@ -3069,7 +3059,7 @@ index fb8fab6a80f..cfcff43d557 100644 CostEvaluator.instantiate( classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf) intercept[IllegalArgumentException] { -@@ -2557,6 +2606,7 @@ class AdaptiveQueryExecSuite +@@ -2557,6 +2605,7 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult(query) assert(adaptive.collect { case sort: SortExec => sort @@ -3077,7 +3067,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.size == 1) val read = collect(adaptive) { case read: AQEShuffleReadExec => read -@@ -2574,7 +2624,8 @@ class AdaptiveQueryExecSuite +@@ -2574,7 +2623,8 @@ class AdaptiveQueryExecSuite } } @@ -3087,7 +3077,7 @@ index fb8fab6a80f..cfcff43d557 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true", -@@ -2686,7 +2737,7 @@ class AdaptiveQueryExecSuite +@@ -2686,7 +2736,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value2 = value3") val shuffles1 = collect(adaptive1) { @@ -3096,7 +3086,7 @@ index fb8fab6a80f..cfcff43d557 100644 } assert(shuffles1.size == 4) val smj1 = findTopLevelSortMergeJoin(adaptive1) -@@ -2697,7 +2748,7 @@ class AdaptiveQueryExecSuite +@@ -2697,7 +2747,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value1 = value3") val shuffles2 = collect(adaptive2) { @@ -3105,7 +3095,7 @@ index fb8fab6a80f..cfcff43d557 100644 } assert(shuffles2.size == 4) val smj2 = findTopLevelSortMergeJoin(adaptive2) -@@ -2955,6 +3006,7 @@ class AdaptiveQueryExecSuite +@@ -2955,6 +3005,7 @@ class AdaptiveQueryExecSuite }.size == (if (firstAccess) 1 else 0)) assert(collect(initialExecutedPlan) { case s: SortExec => s @@ -3113,7 +3103,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2967,6 +3019,7 @@ class AdaptiveQueryExecSuite +@@ -2967,6 +3018,7 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(finalExecutedPlan) { case s: SortExec => s @@ -3121,7 +3111,7 @@ index fb8fab6a80f..cfcff43d557 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3042,13 +3095,18 @@ class AdaptiveQueryExecSuite +@@ -3042,13 +3094,18 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult( "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") @@ -3142,7 +3132,7 @@ index fb8fab6a80f..cfcff43d557 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3138,7 +3196,8 @@ class AdaptiveQueryExecSuite +@@ -3138,7 +3195,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index 929fbfd09da..c1714a7357f 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -2807,7 +2807,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 188a28ff1c0..d1d8714cd32 100644 +index 188a28ff1c0..5be29a7377d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,12 +27,14 @@ import org.apache.spark.SparkException @@ -2937,17 +2937,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", -@@ -303,7 +329,8 @@ class AdaptiveQueryExecSuite - val localReads = collect(adaptivePlan) { - case read: AQEShuffleReadExec if read.isLocalRead => read - } -- assert(localReads.length == 2) -+ // Comet shuffle changes shuffle metrics -+ assert(localReads.length == 1) - val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD] - val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD] - // the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2 -@@ -328,7 +355,9 @@ class AdaptiveQueryExecSuite +@@ -328,7 +354,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2958,7 +2948,7 @@ index 188a28ff1c0..d1d8714cd32 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -342,7 +371,9 @@ class AdaptiveQueryExecSuite +@@ -342,7 +370,9 @@ class AdaptiveQueryExecSuite .groupBy($"a").count() checkAnswer(testDf, Seq()) val plan = testDf.queryExecution.executedPlan @@ -2969,7 +2959,7 @@ index 188a28ff1c0..d1d8714cd32 100644 val coalescedReads = collect(plan) { case r: AQEShuffleReadExec => r } -@@ -352,7 +383,7 @@ class AdaptiveQueryExecSuite +@@ -352,7 +382,7 @@ class AdaptiveQueryExecSuite } } @@ -2978,7 +2968,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -367,7 +398,7 @@ class AdaptiveQueryExecSuite +@@ -367,7 +397,7 @@ class AdaptiveQueryExecSuite } } @@ -2987,7 +2977,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -383,7 +414,7 @@ class AdaptiveQueryExecSuite +@@ -383,7 +413,7 @@ class AdaptiveQueryExecSuite } } @@ -2996,7 +2986,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -428,7 +459,7 @@ class AdaptiveQueryExecSuite +@@ -428,7 +458,7 @@ class AdaptiveQueryExecSuite } } @@ -3005,7 +2995,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -473,7 +504,7 @@ class AdaptiveQueryExecSuite +@@ -473,7 +503,7 @@ class AdaptiveQueryExecSuite } } @@ -3014,7 +3004,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { -@@ -519,7 +550,7 @@ class AdaptiveQueryExecSuite +@@ -519,7 +549,7 @@ class AdaptiveQueryExecSuite } } @@ -3023,7 +3013,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -538,7 +569,7 @@ class AdaptiveQueryExecSuite +@@ -538,7 +568,7 @@ class AdaptiveQueryExecSuite } } @@ -3032,7 +3022,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") { -@@ -569,7 +600,9 @@ class AdaptiveQueryExecSuite +@@ -569,7 +599,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -3043,7 +3033,7 @@ index 188a28ff1c0..d1d8714cd32 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.nonEmpty) -@@ -590,7 +623,9 @@ class AdaptiveQueryExecSuite +@@ -590,7 +622,9 @@ class AdaptiveQueryExecSuite assert(smj.size == 1) val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) @@ -3054,7 +3044,7 @@ index 188a28ff1c0..d1d8714cd32 100644 // Even with local shuffle read, the query stage reuse can also work. val ex = findReusedExchange(adaptivePlan) assert(ex.isEmpty) -@@ -599,7 +634,8 @@ class AdaptiveQueryExecSuite +@@ -599,7 +633,8 @@ class AdaptiveQueryExecSuite } } @@ -3064,7 +3054,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000", -@@ -694,7 +730,8 @@ class AdaptiveQueryExecSuite +@@ -694,7 +729,8 @@ class AdaptiveQueryExecSuite val bhj = findTopLevelBroadcastHashJoin(adaptivePlan) assert(bhj.size == 1) // There is still a SMJ, and its two shuffles can't apply local read. @@ -3074,7 +3064,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } } -@@ -822,7 +859,8 @@ class AdaptiveQueryExecSuite +@@ -822,7 +858,8 @@ class AdaptiveQueryExecSuite } } @@ -3084,7 +3074,7 @@ index 188a28ff1c0..d1d8714cd32 100644 Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint => def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") { findTopLevelSortMergeJoin(plan) -@@ -1266,7 +1304,8 @@ class AdaptiveQueryExecSuite +@@ -1266,7 +1303,8 @@ class AdaptiveQueryExecSuite } } @@ -3094,7 +3084,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT key FROM testData GROUP BY key") -@@ -1832,13 +1871,13 @@ class AdaptiveQueryExecSuite +@@ -1832,13 +1870,13 @@ class AdaptiveQueryExecSuite val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult( "SELECT key FROM testData WHERE key = 0 ORDER BY key, value") assert(findTopLevelSort(plan1).size == 1) @@ -3110,7 +3100,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } } -@@ -1846,18 +1885,18 @@ class AdaptiveQueryExecSuite +@@ -1846,18 +1884,18 @@ class AdaptiveQueryExecSuite withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult( "SELECT key, count(*) FROM testData WHERE value = 'no_match' GROUP BY key") @@ -3135,7 +3125,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } } -@@ -1878,7 +1917,7 @@ class AdaptiveQueryExecSuite +@@ -1878,7 +1916,7 @@ class AdaptiveQueryExecSuite |""".stripMargin) checkNumUnion(plan1, 1) checkNumUnion(adaptivePlan1, 0) @@ -3144,7 +3134,7 @@ index 188a28ff1c0..d1d8714cd32 100644 val (plan2, adaptivePlan2) = runAdaptiveAndVerifyResult( """ -@@ -1888,7 +1927,7 @@ class AdaptiveQueryExecSuite +@@ -1888,7 +1926,7 @@ class AdaptiveQueryExecSuite |""".stripMargin) checkNumUnion(plan2, 1) checkNumUnion(adaptivePlan2, 0) @@ -3153,7 +3143,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } } -@@ -1900,7 +1939,7 @@ class AdaptiveQueryExecSuite +@@ -1900,7 +1938,7 @@ class AdaptiveQueryExecSuite val (_, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id") assert(collect(adaptivePlan) { @@ -3162,7 +3152,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.length == 1) } } -@@ -1981,7 +2020,8 @@ class AdaptiveQueryExecSuite +@@ -1981,7 +2019,8 @@ class AdaptiveQueryExecSuite } } @@ -3172,7 +3162,7 @@ index 188a28ff1c0..d1d8714cd32 100644 def hasRepartitionShuffle(plan: SparkPlan): Boolean = { find(plan) { case s: ShuffleExchangeLike => -@@ -2168,6 +2208,9 @@ class AdaptiveQueryExecSuite +@@ -2168,6 +2207,9 @@ class AdaptiveQueryExecSuite def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = { assert(collect(ds.queryExecution.executedPlan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -3182,7 +3172,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.size == 1) ds.collect() val plan = ds.queryExecution.executedPlan -@@ -2176,6 +2219,9 @@ class AdaptiveQueryExecSuite +@@ -2176,6 +2218,9 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(plan) { case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s @@ -3192,7 +3182,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.size == 1) checkAnswer(ds, testData) } -@@ -2340,7 +2386,8 @@ class AdaptiveQueryExecSuite +@@ -2340,7 +2385,8 @@ class AdaptiveQueryExecSuite } } @@ -3202,7 +3192,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withTempView("t1", "t2") { def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = { Seq("100", "100000").foreach { size => -@@ -2426,7 +2473,8 @@ class AdaptiveQueryExecSuite +@@ -2426,7 +2472,8 @@ class AdaptiveQueryExecSuite } } @@ -3212,7 +3202,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", -@@ -2525,7 +2573,7 @@ class AdaptiveQueryExecSuite +@@ -2525,7 +2572,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2 GROUP BY key1") val shuffles1 = collect(adaptive1) { @@ -3221,7 +3211,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } assert(shuffles1.size == 3) // shuffles1.head is the top-level shuffle under the Aggregate operator -@@ -2538,7 +2586,7 @@ class AdaptiveQueryExecSuite +@@ -2538,7 +2585,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " + s"JOIN skewData2 ON key1 = key2") val shuffles2 = collect(adaptive2) { @@ -3230,7 +3220,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } if (hasRequiredDistribution) { assert(shuffles2.size == 3) -@@ -2572,7 +2620,8 @@ class AdaptiveQueryExecSuite +@@ -2572,7 +2619,8 @@ class AdaptiveQueryExecSuite } } @@ -3240,7 +3230,7 @@ index 188a28ff1c0..d1d8714cd32 100644 CostEvaluator.instantiate( classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf) intercept[IllegalArgumentException] { -@@ -2738,6 +2787,7 @@ class AdaptiveQueryExecSuite +@@ -2738,6 +2786,7 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult(query) assert(adaptive.collect { case sort: SortExec => sort @@ -3248,7 +3238,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.size == 1) val read = collect(adaptive) { case read: AQEShuffleReadExec => read -@@ -2755,7 +2805,8 @@ class AdaptiveQueryExecSuite +@@ -2755,7 +2804,8 @@ class AdaptiveQueryExecSuite } } @@ -3258,7 +3248,7 @@ index 188a28ff1c0..d1d8714cd32 100644 withTempView("v") { withSQLConf( SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true", -@@ -2867,7 +2918,7 @@ class AdaptiveQueryExecSuite +@@ -2867,7 +2917,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value2 = value3") val shuffles1 = collect(adaptive1) { @@ -3267,7 +3257,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } assert(shuffles1.size == 4) val smj1 = findTopLevelSortMergeJoin(adaptive1) -@@ -2878,7 +2929,7 @@ class AdaptiveQueryExecSuite +@@ -2878,7 +2928,7 @@ class AdaptiveQueryExecSuite runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " + "JOIN skewData3 ON value1 = value3") val shuffles2 = collect(adaptive2) { @@ -3276,7 +3266,7 @@ index 188a28ff1c0..d1d8714cd32 100644 } assert(shuffles2.size == 4) val smj2 = findTopLevelSortMergeJoin(adaptive2) -@@ -3136,6 +3187,7 @@ class AdaptiveQueryExecSuite +@@ -3136,6 +3186,7 @@ class AdaptiveQueryExecSuite }.size == (if (firstAccess) 1 else 0)) assert(collect(initialExecutedPlan) { case s: SortExec => s @@ -3284,7 +3274,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3148,6 +3200,7 @@ class AdaptiveQueryExecSuite +@@ -3148,6 +3199,7 @@ class AdaptiveQueryExecSuite }.isEmpty) assert(collect(finalExecutedPlan) { case s: SortExec => s @@ -3292,7 +3282,7 @@ index 188a28ff1c0..d1d8714cd32 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3223,13 +3276,18 @@ class AdaptiveQueryExecSuite +@@ -3223,13 +3275,18 @@ class AdaptiveQueryExecSuite val (_, adaptive) = runAdaptiveAndVerifyResult( "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") @@ -3313,7 +3303,7 @@ index 188a28ff1c0..d1d8714cd32 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3319,7 +3377,8 @@ class AdaptiveQueryExecSuite +@@ -3319,7 +3376,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) From 695ca3ab7e6edefffffd6e073069561e66086c50 Mon Sep 17 00:00:00 2001 From: dustin Date: Mon, 12 Oct 2026 14:08:02 +0000 Subject: [PATCH 82/89] fix: do not resolve DPP subqueries when AQE reads a native scan's partitioning (#6268) * fix: do not resolve DPP subqueries when AQE reads a native scan's partitioning AQE validates a stage by reading outputPartitioning after CoalesceShufflePartitions changes it, which is before Comet converts the DPP placeholders. CometNativeScanExec answered with its post-pruning split count, which ran the SubqueryAdaptiveBroadcastExec placeholder and threw. A non-bucketed scan now reports UnknownPartitioning(0), as FileSourceScanExec does, and CometNativeExec sizes the native RDD from perPartitionData at execution. Part of #6133. * fix: do not resolve DPP subqueries when AQE reads an Iceberg scan's partitioning CometIcebergNativeScanExec counted its tasks in outputPartitioning, which runs BatchScanExec.inputRDD and translateRuntimeFilterV2 on a DPP subquery that has no result yet while AQE validates a stage. It now reports UnknownPartitioning(0), as BatchScanExec does on Spark 3.5 and later, and CometNativeExec sizes the native RDD from the scan's tasks at execution. * docs: say a plan-data scan's outputPartitioning must not run its DPP subqueries AQE reads outputPartitioning while it optimizes a stage, before the adaptive DPP placeholders are converted, so a scan that computes it from perPartitionData runs a placeholder subquery. Say so in the CometScanWithPlanData doc, and point a scan that cannot report a real partitioning otherwise at UnknownPartitioning(0), as CometNativeScanExec does for a non-bucketed scan. * test: count a native scan's partitions through perPartitionData A non-bucketed CometNativeScanExec now reports UnknownPartitioning(0), as FileSourceScanExec does, so the coalesced-scan metrics test cannot read the split count from outputPartitioning. perPartitionData holds the count the native RDD is sized from, which is what the precondition meant to check. --- .../comet/CometIcebergNativeScanExec.scala | 9 +- .../spark/sql/comet/CometNativeScanExec.scala | 30 ++-- .../apache/spark/sql/comet/operators.scala | 18 +++ .../comet/CometIcebergNativeSuite.scala | 98 +++++++++++- .../apache/comet/exec/CometExecSuite.scala | 147 ++++++++++++++++++ .../sql/comet/CometTaskMetricsSuite.scala | 2 +- 6 files changed, 280 insertions(+), 24 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index 0cb7b885b6e..31e518cfcf8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -128,7 +128,14 @@ case class CometIcebergNativeScanExec( // Only accessed during execution, not planning def numPartitions: Int = perPartitionData.length - override lazy val outputPartitioning: Partitioning = UnknownPartitioning(numPartitions) + // Like BatchScanExec on Spark 3.5 and later without a key-grouped partitioning, the scan does + // not know its partition count while planning, because dynamic partition pruning decides which + // tasks survive and its subqueries resolve at execution. AQE reads outputPartitioning while it + // optimizes a stage (ValidateRequirements after CoalesceShufflePartitions), which is before + // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so serializing the + // tasks here would run a placeholder subquery. At execution, CometNativeExec sizes the native + // RDD from perPartitionData instead. + override lazy val outputPartitioning: Partitioning = UnknownPartitioning(0) override lazy val outputOrdering: Seq[SortOrder] = Nil diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index bff854a2eee..296fc4e6ccf 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -120,27 +120,15 @@ case class CometNativeScanExec( // exposed for testing lazy val bucketedScan: Boolean = originalPlan.bucketedScan && !disableBucketedScan - override lazy val outputPartitioning: Partitioning = { - if (bucketedScan) { - originalPlan.outputPartitioning - } else { - // Use perPartitionData.length instead of originalPlan.inputRDD.getNumPartitions. - // - // originalPlan.inputRDD triggers FileSourceScanExec's full scan pipeline including - // codegen on partition filter expressions. With DPP, this calls - // InSubqueryExec.doGenCode which requires the subquery to have finished - but - // outputPartitioning can be accessed before prepare() runs (e.g., by - // ValidateRequirements during plan validation). - // - // perPartitionData goes through serializedPartitionData, which explicitly resolves - // DPP subqueries (via updateResult()) before accessing file partitions. This is the - // same pattern CometIcebergNativeScanExec uses. - // - // This is also more correct: perPartitionData.length reflects the post-DPP partition - // count, matching what CometExecRDD actually uses in doExecuteColumnar(). - UnknownPartitioning(perPartitionData.length) - } - } + // Like FileSourceScanExec, only a bucketed scan knows its partitioning while planning. The + // file-split count of any other scan depends on dynamic partition pruning, whose subqueries + // resolve at execution. AQE reads outputPartitioning while it optimizes a stage + // (ValidateRequirements after CoalesceShufflePartitions), which is before + // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so computing the file + // partitions here would run a placeholder subquery. At execution, CometNativeExec sizes the + // native RDD from perPartitionData instead. + override lazy val outputPartitioning: Partitioning = + if (bucketedScan) originalPlan.outputPartitioning else UnknownPartitioning(0) override lazy val outputOrdering: Seq[SortOrder] = originalPlan.outputOrdering diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index a712f44796a..de354502c59 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1110,6 +1110,13 @@ abstract class CometNativeExec extends CometExec { // broadcast plan. val (firstNonBroadcastPlanRDD, firstNonBroadcastPlanNumPartitions) = firstNonBroadcastPlan.get._1 match { + // Plan-data scans, Iceberg included, are sized from perPartitionData, which + // findAllPlanData above has resolved. A bucketed scan has one entry per bucket, after any + // bucket coalescing, so the count still matches its HashPartitioning. + case scan: CometScanWithPlanData => + (null.asInstanceOf[RDD[Any]], scan.perPartitionData.length) + case scan: CometIcebergNativeScanExec => + (null.asInstanceOf[RDD[Any]], scan.numPartitions) case plan: CometNativeExec => (null.asInstanceOf[RDD[Any]], plan.outputPartitioning.numPartitions) case plan => @@ -1349,6 +1356,17 @@ abstract class CometLeafExec extends CometNativeExec with LeafExecNode { * parent's native execution receives an empty input. (`CometIcebergNativeScanExec` does NOT use * this trait; it has a dedicated `findAllPlanData` case.) * + * `perPartitionData.length` is the partition count the native block runs with: + * `CometNativeExec.buildNativeContext` sizes the native RDD from it at execution, and the count + * in `outputPartitioning` is not used for that. + * + * `outputPartitioning` must not read `perPartitionData`, or anything else that evaluates the + * scan's DPP subqueries. AQE calls it while optimizing a stage, before + * `CometPlanAdaptiveDynamicPruningFilters` has converted the adaptive DPP placeholders, so + * computing it from those runs a placeholder subquery. A scan that cannot report a real + * partitioning without them should report `UnknownPartitioning(0)`, as `CometNativeScanExec` does + * for a non-bucketed scan. + * * Each implementation also resolves its own DPP subqueries via `ensureSubqueriesResolved` before * `commonData`/`perPartitionData` are read. That method lives on [[CometLeafExec]], so the `self: * CometLeafExec` self-type below makes "is a leaf scan" a compile-time requirement: an diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 02d2640d10d..4180e3789e6 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{InSubqueryExec, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, BroadcastQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, BroadcastQueryStageExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.functions.col @@ -5635,6 +5635,102 @@ class CometIcebergNativeSuite } } + // AQE runs CoalesceShufflePartitions before CometPlanAdaptiveDynamicPruningFilters. When that + // rule changes a stage, AQE validates the stage by reading outputPartitioning on every node in + // it. With Spark's UnionExec, the aggregate's shuffle read shares the final stage with the DPP + // scan and gets coalesced, so the scan is asked for its partitioning while its DPP subquery is + // still the adaptive placeholder. Answering must neither run that subquery nor lose the pruning. + test("AQE DPP - shuffle coalescing in the scan's stage keeps DPP working") { + assume(icebergAvailable, "Iceberg not available") + withTempIcebergDir { warehouseDir => + val dimDir = new File(warehouseDir, "dim_parquet") + val otherDir = new File(warehouseDir, "other_parquet") + withSQLConf( + "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.aqe_cat.type" -> "hadoop", + "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1KB", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_PARALLELISM_FIRST.key -> "false", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true", + CometConf.COMET_EXEC_UNION_ENABLED.key -> "false") { + + spark.sql(""" + CREATE TABLE aqe_cat.db.coalesce_fact (amount INT, store_id INT) + USING iceberg PARTITIONED BY (store_id) + """) + // One data file per store, so an unpruned scan reads 10 files and a pruned one reads 1. + spark + .range(100) + .selectExpr("cast(id as int) as amount", "cast(id % 10 as int) as store_id") + .coalesce(1) + .sortWithinPartitions("store_id") + .writeTo("aqe_cat.db.coalesce_fact") + .append() + + spark + .range(10) + .selectExpr("cast(id as int) as store_id", "cast(id as string) as country") + .write + .parquet(dimDir.getAbsolutePath) + spark.read.parquet(dimDir.getAbsolutePath).createOrReplaceTempView("coalesce_dim") + spark + .range(100) + .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v") + .write + .parquet(otherDir.getAbsolutePath) + spark.read.parquet(otherDir.getAbsolutePath).createOrReplaceTempView("coalesce_other") + + val query = + """SELECT /*+ BROADCAST(d) */ f.amount AS a, f.store_id AS b + |FROM aqe_cat.db.coalesce_fact f JOIN coalesce_dim d ON f.store_id = d.store_id + |WHERE d.country = '3' + |UNION ALL + |SELECT k AS a, CAST(count(*) AS INT) AS b + |FROM coalesce_other GROUP BY k""".stripMargin + val (_, cometPlan) = checkSparkAnswer(query) + assertIcebergNativeScanPresent(cometPlan) + assertNoLeftoverCSAB(cometPlan) + + // On 3.4, Spark's own DPP rule plans the subquery before coalescing, so check results only. + if (isSpark35Plus) { + val dppScans = collectIcebergNativeScans(cometPlan).filter( + _.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression])) + assert(dppScans.size == 1, s"Expected one DPP Iceberg scan:\n${cometPlan.treeString}") + val scan = dppScans.head + // The test only covers the bug if coalescing changed the stage holding the DPP scan. + // That is the final stage, and TreeNode.collect stops at its query stage leaves. + val finalStage = stripAQEPlan(cometPlan) + assert( + finalStage.find(_ eq scan).nonEmpty && + finalStage.collect { + case r: AQEShuffleReadExec if r.isCoalescedRead => r + }.nonEmpty, + "Expected the DPP scan and a coalesced shuffle read in the final stage:\n" + + cometPlan.treeString) + assertCsbBroadcastReuse(collectIcebergDPPSubqueries(cometPlan), cometPlan) + // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 data files. + // Iceberg filters its already planned tasks at runtime, so count the serialized tasks + // and the splits read rather than Iceberg's planning metrics. + val plannedTasks = scan.perPartitionData + .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount) + .sum + val numSplits = scan.metrics("num_splits").value + assert( + plannedTasks == 1 && numSplits == 1, + s"Expected DPP to prune to 1 of 10 files, planned $plannedTasks tasks and read " + + s"$numSplits splits:\n${cometPlan.treeString}") + } + + spark.sql("DROP TABLE aqe_cat.db.coalesce_fact") + } + } + } + test("AQE DPP - cross-stage scalar subquery with broadcast reuse") { assume(icebergAvailable, "Iceberg not available") withTempIcebergDir { warehouseDir => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index a59eff0e02a..f40fd42d1bd 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -1285,6 +1285,153 @@ class CometExecSuite extends CometTestBase { } } + // AQE runs CoalesceShufflePartitions before CometPlanAdaptiveDynamicPruningFilters. When that + // rule changes a stage, AQE validates the stage by reading outputPartitioning on every node in + // it. With Spark's UnionExec, the aggregate's shuffle read shares the final stage with the DPP + // scan and gets coalesced, so the scan is asked for its partitioning while its DPP subquery is + // still the adaptive placeholder. Answering must not run that subquery. + test("AQE DPP: shuffle coalescing in the scan's stage keeps DPP working") { + withTempDir { dir => + val path = dir.getAbsolutePath + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + spark + .range(100) + .selectExpr("cast(id % 10 as int) as store_id", "cast(id as int) as amount") + .write + .partitionBy("store_id") + .parquet(s"$path/fact") + spark + .range(10) + .selectExpr("cast(id as int) as store_id", "cast(id as string) as country") + .write + .parquet(s"$path/dim") + spark + .range(100) + .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v") + .write + .parquet(s"$path/other") + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_PARALLELISM_FIRST.key -> "false", + CometConf.COMET_EXEC_UNION_ENABLED.key -> "false") { + spark.read.parquet(s"$path/fact").createOrReplaceTempView("dpp_coalesce_fact") + spark.read.parquet(s"$path/dim").createOrReplaceTempView("dpp_coalesce_dim") + spark.read.parquet(s"$path/other").createOrReplaceTempView("dpp_coalesce_other") + + val df = spark.sql("""SELECT f.amount AS a, f.store_id AS b + |FROM dpp_coalesce_fact f JOIN dpp_coalesce_dim d ON f.store_id = d.store_id + |WHERE d.country = '3' + |UNION ALL + |SELECT k AS a, CAST(count(*) AS INT) AS b + |FROM dpp_coalesce_other GROUP BY k""".stripMargin) + val (_, cometPlan) = checkSparkAnswer(df) + + if (isSpark35Plus) { + val dppScans = collect(cometPlan) { + case s: CometNativeScanExec + if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => + s + } + assert(dppScans.size == 1, s"Expected one DPP native scan:\n${cometPlan.treeString}") + // The test only covers the bug if coalescing changed the stage holding the DPP scan. + // That is the final stage, and TreeNode.collect stops at its query stage leaves. + val finalStage = stripAQEPlan(cometPlan) + assert( + finalStage.find(_ eq dppScans.head).nonEmpty && + finalStage.collect { + case r: AQEShuffleReadExec if r.isCoalescedRead => r + }.nonEmpty, + "Expected the DPP scan and a coalesced shuffle read in the final stage:\n" + + cometPlan.treeString) + assertAqeDppShape(cometPlan, expectedCometSubqueryBroadcasts = Some(1)) + // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 partitions. + assert( + dppScans.head.metrics("numPartitions").value == 1, + s"Expected DPP to prune to 1 partition:\n${cometPlan.treeString}") + } + } + } + } + + // A bucketed scan reports its bucket layout, which DPP does not change. The sort-merge join on + // the bucket column keeps relying on it, so neither join side is shuffled, while DPP from the + // broadcast dimension join still prunes the fact table's partitions. + test("AQE DPP: bucketed scan keeps its partitioning for a sort-merge join") { + withTable("dpp_bucket_fact", "dpp_bucket_dim", "dpp_bucket_other") { + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + spark + .range(100) + .selectExpr( + "cast(id % 10 as int) as store_id", + "cast(id % 7 as int) as k", + "cast(id as int) as amount") + .write + .partitionBy("store_id") + .bucketBy(4, "k") + .format("parquet") + .saveAsTable("dpp_bucket_fact") + spark + .range(10) + .selectExpr("cast(id as int) as store_id", "cast(id as string) as country") + .write + .format("parquet") + .saveAsTable("dpp_bucket_dim") + spark + .range(50) + .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v") + .write + .bucketBy(4, "k") + .format("parquet") + .saveAsTable("dpp_bucket_other") + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.BUCKETING_ENABLED.key -> "true") { + val df = sql("""SELECT /*+ BROADCAST(d), MERGE(o) */ f.amount, f.store_id, o.v + |FROM dpp_bucket_fact f + |JOIN dpp_bucket_dim d ON f.store_id = d.store_id + |JOIN dpp_bucket_other o ON f.k = o.k + |WHERE d.country = '3'""".stripMargin) + val (_, cometPlan) = checkSparkAnswer(df) + + val smjs = collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j + } + assert(smjs.size == 1, s"Expected one sort-merge join:\n${cometPlan.treeString}") + val shuffles = collect(cometPlan) { case s: ShuffleExchangeLike => s } + assert( + shuffles.isEmpty, + s"Expected no shuffle around the bucketed join:\n${cometPlan.treeString}") + + if (isSpark35Plus) { + val dppScans = collect(cometPlan) { + case s: CometNativeScanExec + if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => + s + } + assert(dppScans.size == 1, s"Expected one DPP native scan:\n${cometPlan.treeString}") + assert( + dppScans.head.bucketedScan, + s"Expected a bucketed scan:\n${cometPlan.treeString}") + assertAqeDppShape(cometPlan, expectedCometSubqueryBroadcasts = Some(1)) + // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 partitions. + assert( + dppScans.head.metrics("numPartitions").value == 1, + s"Expected DPP to prune to 1 partition:\n${cometPlan.treeString}") + } + } + } + } + // DPP subquery reuses the join's broadcast via AQE stageCache. test("AQE DPP: broadcast exchange reuse") { withDppTables { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index 318a416ce47..9bf5e955954 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -1489,7 +1489,7 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { run(plainQuery, confs) val plainScan = collectFirst(plainPlan) { case s: CometNativeScanExec => s }.get assert( - plainScan.outputPartitioning.numPartitions == 4, + plainScan.perPartitionData.length == 4, s"Expected one scan partition per file:\n${plainPlan.treeString}") assert(plainTasks == 4, s"Expected 4 tasks at interval $interval, got $plainTasks") assert(plainRows == totalRows, s"output_rows at interval $interval: $plainRows") From 38f394e2707814020cbf057510e25e85f07e47a8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 15:00:48 +0000 Subject: [PATCH 83/89] docs: prefer behavior-named parameters over Spark versions (#6776) * docs: prefer behavior-named parameters over Spark versions When behavior varies between Spark versions, resolve the version in Scala and pass native code a parameter that describes the behavior. Add the guidance to the contributor guide and the expression skills. * docs: address review on behavior-named parameters Reconcile the review skill's shim paragraph with the new rule: the version check may live in a shim or a named predicate in Scala, and protobuf and native code see only the behavior. Extend the rule to native function names, point to RegrSparkVersions for fixes that land in patch releases, ask comments to cite the Spark JIRA, state the fork benefit as it is today, and replace the removed isSpark33Plus helper. The two existing version-named names are tracked in #6891. --- .ai/skills/implement-comet-expression/SKILL.md | 3 ++- .ai/skills/review-comet-expression-pr/SKILL.md | 4 +++- .ai/skills/review-comet-pr/SKILL.md | 17 +++++++++++++---- .ai/skills/wire-datafusion-function/SKILL.md | 2 ++ .../adding_a_new_expression.md | 16 +++++++++++++++- 5 files changed, 35 insertions(+), 7 deletions(-) diff --git a/.ai/skills/implement-comet-expression/SKILL.md b/.ai/skills/implement-comet-expression/SKILL.md index 529293f3784..1ee9eff091f 100644 --- a/.ai/skills/implement-comet-expression/SKILL.md +++ b/.ai/skills/implement-comet-expression/SKILL.md @@ -77,7 +77,8 @@ Follow `adding_a_new_expression.md`: 2. Register it in the matching map in `QueryPlanSerde.scala`. 3. If the function name collides with a DataFusion built-in that has a different signature, use `scalarFunctionExprToProtoWithReturnType` (see "When to set the return type explicitly"). 4. For a new scalar function, add a match case in `native/spark-expr/src/comet_scalar_funcs.rs::create_comet_physical_fun`. If step 2 found an upstream implementation, wire that in. Otherwise implement the function under `native/spark-expr/src/`. -5. Add at least one Comet SQL Test at `spark/src/test/resources/sql-tests/expressions//$ARGUMENTS.sql` exercising column references, literals, and `NULL`. +5. If Spark's behavior differs between versions, resolve the version in the serde or a shim and pass native code a parameter named for the behavior (`wrap_second_millisecond_overflow`), not the version (`spark_420_plus`). See "Name the behavior, not the Spark version" in `adding_a_new_expression.md`. +6. Add at least one Comet SQL Test at `spark/src/test/resources/sql-tests/expressions//$ARGUMENTS.sql` exercising column references, literals, and `NULL`. Build and smoke-test: diff --git a/.ai/skills/review-comet-expression-pr/SKILL.md b/.ai/skills/review-comet-expression-pr/SKILL.md index 17f8e2f04ce..c1a80773961 100644 --- a/.ai/skills/review-comet-expression-pr/SKILL.md +++ b/.ai/skills/review-comet-expression-pr/SKILL.md @@ -266,7 +266,9 @@ reference in that doc is the only place they are documented. 3. **Wrong return type**, it must match Spark exactly 4. **Tests in the wrong framework**, Scala tests where a SQL file test would do 5. **Missing `getSupportLevel`**, divergences left undeclared rather than marked `Incompatible` -6. **Version-specific Spark behavior implemented once**, with no shim +6. **Version-specific Spark behavior implemented once**, with no shim or version check, or with a + protobuf field, native argument, or native function named for the Spark version + (`spark_420_plus`) instead of the behavior (`wrap_second_millisecond_overflow`) 7. **Name collides with a DataFusion built-in** and no explicit return type 8. **Timestamp result mislabelled**, with the session timezone or no timezone instead of `"UTC"`. It only shows once the result is compared or fed to another expression. diff --git a/.ai/skills/review-comet-pr/SKILL.md b/.ai/skills/review-comet-pr/SKILL.md index e93361b69f9..ea1f57a9d1f 100644 --- a/.ai/skills/review-comet-pr/SKILL.md +++ b/.ai/skills/review-comet-pr/SKILL.md @@ -171,10 +171,19 @@ Report the comparison in the review even when nothing changed, so the reviewer k ### Spark version coverage -Comet supports several Spark versions. Version-specific behavior belongs in the shims under -`spark/src/main/spark-{3.4,3.5,3.x,4.0,4.1,4.2}/org/apache/comet/shims/`, not in branches on a -version string in shared code, and not in native Rust. If the PR adds a shim for one 4.x version, -check that the sibling 4.x source sets got it too. +Comet supports several Spark versions. The check on the Spark version belongs in Scala, either in a +shim under `spark/src/main/spark-*/` or in one named predicate such as `isSpark42Plus` or +`RegrSparkVersions.slopeFiltersVarByPairNulls`. It does not belong in an ad hoc comparison on the +version string in shared code, or in native Rust. If the PR adds a shim for one 4.x version, check +that the sibling 4.x source sets got it too. + +Protobuf fields, native function arguments, and native function names only ever see the behavior, +such as `wrap_second_millisecond_overflow`, never the version, such as `spark_420_plus` (#6740). A +fork that backports a Spark fix then changes the one Scala check that sets the flag and leaves the +protobuf and Rust alone. Flag a version-named field, argument, or function, and comments that say +"Spark 4.2" for behavior later releases inherit or that leave out the Spark JIRA. +`Cast.is_spark4_plus` and `get_json_object_spark34` predate the rule and #6891 tracks renaming them, +so flag them only when a PR adds behavior behind them. When Spark changed the behavior in a patch release, such as SPARK-55969 or SPARK-54918, a check on the minor version is wrong for every earlier patch. CI builds only the newest patch of each line, so diff --git a/.ai/skills/wire-datafusion-function/SKILL.md b/.ai/skills/wire-datafusion-function/SKILL.md index 881d24291bf..4688f076c15 100644 --- a/.ai/skills/wire-datafusion-function/SKILL.md +++ b/.ai/skills/wire-datafusion-function/SKILL.md @@ -102,6 +102,8 @@ Helpers from `QueryPlanSerde`: - `optExprWithInfo(optExpr, expr, children*)` — wrap final result; propagates "why we couldn't convert" tags. - `withInfo(expr, "reason")` — tag a fallback when returning `None`. +If the Spark behavior you are matching varies by Spark version, resolve the version in Scala and pass the native side a parameter, or pick a function name, that describes the behavior rather than the version (see "Name the behavior, not the Spark version" in `adding_a_new_expression.md`). + `getSupportLevel` returning `Incompatible(Some("…"))` gates behind `spark.comet.expr..allowIncompatible=true`. `Unsupported(…)` always falls back. ### 4. Register the UDF (Pattern B only) diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index a246ea70bb6..20bcae06e69 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -582,7 +582,21 @@ Spark. If the expression you're adding has different behavior across different Spark versions, you'll need to account for that in your implementation. There are two tools at your disposal to help with this: 1. Shims that exist in `spark/src/main/spark-$SPARK_VERSION/org/apache/comet/shims/CometExprShim.scala` for each Spark version. These shims are used to provide compatibility between different Spark versions. -2. Variables that correspond to the Spark version, such as `isSpark33Plus`, which can be used to conditionally execute code based on the Spark version. +2. Variables that correspond to the Spark version, such as `isSpark42Plus`, which can be used to conditionally execute code based on the Spark version. + +#### Name the behavior, not the Spark version + +When native code has to match a behavior that differs between Spark versions, decide which behavior applies in Scala and pass the result to native code under a name that describes the behavior, not the Spark version that introduced it. This covers protobuf fields, native function arguments, and native function names. For example, `TruncTimestamp` carries a `wrap_second_millisecond_overflow` flag, which the serde sets from `!isSpark42Plus`. It is not called `spark_420_plus`. The `filter_var_by_pair_nulls` field of `Regr` is another example. + +Set the flag in one place on the Scala side, in a shim or in the serde, so that a fork overrides only that. A helper such as `isSpark42Plus` is enough when the behavior changed in a minor release. When Spark changed it in patch releases, compare the full version in a named predicate, because a check on the minor version is wrong for every earlier patch. `RegrSparkVersions.r2DegenerateCasesSwapped` in `serde/aggregates.scala` is the model. It resolves SPARK-55969, which shipped in 3.5.9, 4.0.3, 4.1.2, and 4.2.0. + +Reasons to prefer behavior names: + +- Some deployments run a Spark fork that backports fixes from newer open-source releases. Such a fork changes the one Scala check that sets the flag and leaves the protobuf and native code alone. A `spark_420_plus` flag cannot describe a fork that has one 4.2 fix but not the others. +- Spark occasionally changes behavior in a patch release or between minor releases, so a version number is a poor proxy for the behavior. +- The native code stays free of Spark version logic, and a reader of the Rust or proto definition can tell what the flag does without looking up a release. + +Write comments the same way. Cite the Spark JIRA that changed the behavior, as the `TruncTimestamp` serde does with SPARK-56663, because forks track backports by JIRA rather than by release. Say "Spark 4.2 and later" for a behavior that future releases inherit, rather than "Spark 4.2". If a later release changes the behavior again, add a new parameter or shim for that change. ## Shimming to Support Different Spark Versions From 907c0e074c4e34812b3ae808736db28ecb173d6e Mon Sep 17 00:00:00 2001 From: dustin Date: Mon, 12 Oct 2026 17:44:34 +0000 Subject: [PATCH 84/89] test: run the S3 suites against an available MinIO image (#6887) * test: run the S3 suites against an available MinIO image Docker Hub no longer serves minio/minio, so every suite built on CometS3TestBase aborted in beforeAll. Use cgr.dev/chainguard/minio, declared as a substitute for minio/minio so MinIOContainer accepts it. The suites also relied on an AWS region from the environment. Give the bucket setup client an explicit region, and add client.region to the REST catalog vended properties, since Iceberg's JVM S3FileIO reads that key rather than s3.region. * test: pin the MinIO test image by digest Chainguard rebuilds the latest tag from its MinIO fork, and newer releases tighten request checks, so a rebuild could break the S3 suites with no change on our side. Pin the multi-arch index digest of RELEASE.2026-09-22T19-25-18Z so every run uses the same server. --- .../test/scala/org/apache/comet/CometS3TestBase.scala | 10 +++++++++- .../org/apache/comet/IcebergReadFromS3Suite.scala | 2 ++ .../comet/cloud/s3/CometS3CredentialBridgeSuite.scala | 1 + 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala b/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala index d51f0ccedc1..a15e3f75be0 100644 --- a/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus import software.amazon.awssdk.auth.credentials.{AwsBasicCredentials, StaticCredentialsProvider} +import software.amazon.awssdk.regions.Region import software.amazon.awssdk.services.s3.S3Client import software.amazon.awssdk.services.s3.model.{CreateBucketRequest, HeadBucketRequest} @@ -45,8 +46,14 @@ trait CometS3TestBase extends CometTestBase { protected def testBucketName: String + // Docker Hub no longer serves minio/minio. This is Chainguard's build of MinIO + // RELEASE.2026-09-22T19-25-18Z, pinned by its multi-arch index digest. + private val minioImage = "cgr.dev/chainguard/minio@sha256:" + + "dc455cdd4db6e390b1b26514f15ff37681a777bf8dab41aac445a0259f751e56" + override def beforeAll(): Unit = { - minioContainer = new MinIOContainer(DockerImageName.parse("minio/minio:latest")) + minioContainer = new MinIOContainer( + DockerImageName.parse(minioImage).asCompatibleSubstituteFor("minio/minio")) .withUserName(userName) .withPassword(password) minioContainer.start() @@ -120,6 +127,7 @@ trait CometS3TestBase extends CometTestBase { val s3Client = S3Client .builder() .endpointOverride(URI.create(minioContainer.getS3URL)) + .region(Region.US_EAST_1) .credentialsProvider(StaticCredentialsProvider.create(credentials)) .forcePathStyle(true) .build() diff --git a/spark/src/test/scala/org/apache/comet/IcebergReadFromS3Suite.scala b/spark/src/test/scala/org/apache/comet/IcebergReadFromS3Suite.scala index 1271fa841e8..168bd0b49f9 100644 --- a/spark/src/test/scala/org/apache/comet/IcebergReadFromS3Suite.scala +++ b/spark/src/test/scala/org/apache/comet/IcebergReadFromS3Suite.scala @@ -247,6 +247,7 @@ class IcebergReadFromS3Suite extends CometS3TestBase with RESTCatalogHelper { "s3.secret-access-key" -> "WRONG_SECRET_KEY", "s3.endpoint" -> minioContainer.getS3URL, "s3.region" -> "us-east-1", + "client.region" -> "us-east-1", "s3.path-style-access" -> "true") val warehouse = s"s3a://$testBucketName/warehouse-bad-creds" @@ -280,6 +281,7 @@ class IcebergReadFromS3Suite extends CometS3TestBase with RESTCatalogHelper { "s3.secret-access-key" -> password, "s3.endpoint" -> minioContainer.getS3URL, "s3.region" -> "us-east-1", + "client.region" -> "us-east-1", "s3.path-style-access" -> "true") val warehouse = s"s3a://$testBucketName/warehouse-vending" diff --git a/spark/src/test/scala/org/apache/comet/cloud/s3/CometS3CredentialBridgeSuite.scala b/spark/src/test/scala/org/apache/comet/cloud/s3/CometS3CredentialBridgeSuite.scala index 8df9b5bb90d..f3bc317c295 100644 --- a/spark/src/test/scala/org/apache/comet/cloud/s3/CometS3CredentialBridgeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/cloud/s3/CometS3CredentialBridgeSuite.scala @@ -151,6 +151,7 @@ class CometS3CredentialBridgeSuite "s3.secret-access-key" -> password, "s3.endpoint" -> minioContainer.getS3URL, "s3.region" -> "us-east-1", + "client.region" -> "us-east-1", "s3.path-style-access" -> "true", "comet.test-tenant-id" -> "tenant-A") val warehouse = s"s3a://$testBucketName/warehouse-rest-spi" From 1768216b8c3d730355894f56338da3104dea1eaf Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 18:00:12 +0000 Subject: [PATCH 85/89] build: stop publishing comet-common, the parent pom and comet-spark-integration (#6888) Only comet-spark is useful to a consumer. Its published pom is flattened and shade-reduced, so it has no and no comet-common dependency, and comet-common's six classes are inside the shaded jar. The parent pom is installed as the raw root pom.xml with unresolved ${...} placeholders in its artifactId, and comet-spark-integration only checks the jar. The nightly snapshot workflow passed -Dmaven.deploy.skip=false, so it uploaded all four modules for each of the five variants. Release staging uploads whatever `install` leaves in its temporary repository, which includes comet-common and the parent pom. Skip deploying every module except comet-spark, drop the override from the snapshot workflow, and have publish-to-maven.sh upload only comet-spark-spark* directories. Installing is unchanged: builds run from inside spark/ (make benchmark-*, the TPC-H and TPC-DS data generation in pr_build_linux.yml) resolve comet-common from the local repository. --- .github/workflows/publish_snapshot.yml | 9 +++++---- common/pom.xml | 9 ++++----- dev/release/publish-to-maven.sh | 5 +++++ docs/source/contributor-guide/release_process.md | 5 +++++ pom.xml | 7 ++++++- spark-integration/pom.xml | 6 ++---- spark/pom.xml | 2 +- 7 files changed, 28 insertions(+), 15 deletions(-) diff --git a/.github/workflows/publish_snapshot.yml b/.github/workflows/publish_snapshot.yml index ae98bc855ee..e5d8eeea1d8 100644 --- a/.github/workflows/publish_snapshot.yml +++ b/.github/workflows/publish_snapshot.yml @@ -279,11 +279,12 @@ jobs: # target directory the check ran against, and the parent pom pins # project.build.outputTimestamp, so it re-creates the same jar. # - # The root pom skips deploying the parent pom, but consumers need it - # to resolve the child poms, and the release publishes it. The user - # property overrides every module's setting. + # Only comet-spark opts in to deploy (see the root pom), so this + # publishes its shaded jar, sources jars and pom and nothing else. + # Leave -Dmaven.deploy.skip=false off: as a user property it would + # override every module's setting. # shellcheck disable=SC2086 - ./mvnw -B "$GOAL" -DskipTests -Dmaven.deploy.skip=false $profiles + ./mvnw -B "$GOAL" -DskipTests $profiles echo "::endgroup::" done ls -l jars diff --git a/common/pom.xml b/common/pom.xml index fa72c9c1d27..24e28838f5a 100644 --- a/common/pom.xml +++ b/common/pom.xml @@ -39,12 +39,11 @@ under the License. Kept small on purpose so `mvn compile -pl common` is a cheap prerequisite for `cargo test` in the Rust workspace. All other JVM code lives in comet-spark. - --> - - - false - + Not deployed: comet-spark bundles these classes into its shaded jar, so the deploy skip + inherited from the parent pom stays in force. It is still installed locally, which is where + builds run from inside spark/ look for it. + --> diff --git a/dev/release/publish-to-maven.sh b/dev/release/publish-to-maven.sh index 314b0855cfd..1e01fbeb3fd 100755 --- a/dev/release/publish-to-maven.sh +++ b/dev/release/publish-to-maven.sh @@ -120,6 +120,11 @@ echo "Deploying artifacts from $LOCAL_REPO" pushd $LOCAL_REPO/org/apache/datafusion +# Only comet-spark is published. comet-common is shaded into the comet-spark jar and the comet-spark +# pom is flattened, so nothing published refers to comet-common or the parent pom. The build installs +# both locally, and comet-spark-integration only checks the jar. +find . -mindepth 1 -maxdepth 1 ! -name 'comet-spark-spark*' -exec rm -rf {} + + # Remove any extra files generated during install find . -type f |grep -v \.jar |grep -v \.pom | xargs rm diff --git a/docs/source/contributor-guide/release_process.md b/docs/source/contributor-guide/release_process.md index cd3ee4b6d84..ef6d26bcc50 100644 --- a/docs/source/contributor-guide/release_process.md +++ b/docs/source/contributor-guide/release_process.md @@ -499,6 +499,11 @@ Creating Nexus staging repository In the Nexus repository UI (https://repository.apache.org/) locate and verify the artifacts in staging (https://central.sonatype.org/publish/release/#locate-and-examine-your-staging-repository). +The staging repository should hold only the `comet-spark-spark_` artifacts, one for +each Spark and Scala combination that `build-release-comet.sh` builds. The build installs `comet-common` and the +parent pom too, but `publish-to-maven.sh` leaves them out of the upload: `comet-common` is shaded into the +`comet-spark` jar and the published `comet-spark` pom has no parent, so nothing refers to either of them. + The script closes the staging repository but does not release it. Releasing to Maven Central is a manual step performed only after the vote passes (see [Publishing Maven Artifacts](#publishing-maven-artifacts) below). diff --git a/pom.xml b/pom.xml index 0704b7a86c9..084c086e447 100644 --- a/pom.xml +++ b/pom.xml @@ -117,7 +117,12 @@ under the License. ${session.executionRootDirectory} - + true diff --git a/spark-integration/pom.xml b/spark-integration/pom.xml index f656bdd8583..643447e2d76 100644 --- a/spark-integration/pom.xml +++ b/spark-integration/pom.xml @@ -34,10 +34,8 @@ under the License. comet-spark-integration pom - - - false - + diff --git a/spark/pom.xml b/spark/pom.xml index 0dfd7bb259d..6b71b199eb8 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -34,7 +34,7 @@ under the License. comet-spark - + false From 23a9a7cfb9e0775b4f3876c19dd2e80b4242aa0d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 19:12:55 +0000 Subject: [PATCH 86/89] feat: close fanout partitions early when the memory pool refuses the native Iceberg writer (#6773) * feat: close fanout partitions early when the memory pool refuses the native Iceberg writer A native fanout write keeps a file open for every partition a task writes to, and its buffers count against the task's memory pool, so a task over many partitions could fail where iceberg-java's on-heap writer succeeded. When the pool refuses the writer's reservation, it now writes out and closes the partitions holding the most memory, in their open file or in rows not yet handed to it, until what is left fits. A closed partition's next rows open a new file with the same writer properties, so the write finishes with more, smaller files instead of failing. A files_closed_early metric counts them. Writes that keep one file open still fail when that file outgrows the pool. * docs: describe how a native fanout write closes partitions early under memory pressure * test: compare the early-closing fanout write with its source rows The control write is native too, so it could lose the same rows. Also fix a stale FanoutWriter mention in the user guide and a missing noun in the FanoutFiles doc comment. * refactor: keep each fanout partition's feed, writer and memory in one record FanoutPartitions holds every partition in first-seen order, each with its feed, its open writer, the OpenFileMemory child its files report to, and the properties its files are written with. That replaces FanoutFiles and FanoutFeeds and their join by key, the seen counter, PartitionProperties::get and the partition field threaded through the metered builder. InnerWriter::reserve now owns the close-until-it-fits loop, ranking partitions by index without cloning keys and tracking pending bytes without rescanning. WriteMetrics gets a constructor, the Rust memory tests share their setup, and the Scala test captures the write with captureWrite. * docs: say the pool cannot fail a fanout write, and fix a stale comment Everything a fanout write reserves belongs to one of its partitions, so once every partition holding memory is closed the resize is down to nothing, which the pool cannot refuse. The contributor guide said a fanout write could still fail with nothing left to close; only an unpartitioned or clustered write can. FanoutPartition's key is used to open every file of the partition, not only to write held rows out at close. * test: check that a fanout write short of memory closes the partitions holding the most No test pinned the order in which a write the pool refuses closes its partitions: closing every partition at once, or closing them smallest first or in the order they were first seen, passed every test. The new test has 16 or 64 partitions take turns through 32 batches in an eighth of the memory the write needs, and asserts fewer than six files per partition. Closing the largest first writes between four and five, closing all at once seven, and the other two orders seventeen. The user guide now says how many files to expect, and that Iceberg's rewrite_data_files compacts them. * test: give the close-order test's count for 64 partitions too Closing partitions smallest first or in first-seen order writes 17 files per partition with 16 partitions and 20 with 64; the comment only gave the first. * refactor: drop the expect from a fanout partition's write and check its counters in debug builds - `FanoutPartition::write` takes the writer it opens from `Option::insert`, so it no longer needs an `expect`. - `OpenFileMemory::child` asserts that it is called on the task's total, since a share counts only one level up. - `FanoutPartitions::write` asserts that `held_bytes` is still the sum of what the feeds hold back, and the field's doc names the places that keep it so. - `by_memory` says why it ranks partitions by index. - `payload_rows_from` takes its regions by value instead of copying them. * test: check the files_closed_early count in the fanout early-close tests The Rust tests passed `WriteMetrics::default()` and never read it back, so a miscount in `reserve`'s return went unnoticed. `write_reserving_from` now takes the metrics, and `fanout_write_in_part_of_the_pool` gives its tight write a real `WriteMetrics` and checks `closed_early <= files <= closed_early + partitions`: no file rolls at the target size, so every file was closed either early or when the task ended, at most one per partition. Counting one file fewer or one more per refusal fails `a_fanout_write_short_of_memory_closes_the_partitions_holding_the_most`. * fix: pass on a resize error that is not the pool refusing the fanout writer `reserve` treated any error from `try_resize` as the pool refusing, and closed partitions until the resize fit. A pool backed by Spark can also fail a resize with an exception Spark threw, such as a spill of another consumer that failed or the `InterruptedException` of a task killed while it waited for memory. That comes back as `External` (a Java throwable) or `Execution` (any other Comet error), while every pool reports a refusal as `ResourcesExhausted`. Once enough partitions had closed, the resize was a shrink that never called Spark, so the write carried on and the exception was lost. Both `try_resize` calls now retry only on `ResourcesExhausted` and pass anything else back, which fails the task. The new test's pool refuses twice, so two partitions close early, then fails the third resize with an `External` or an `Execution` error. It asserts that the write returns that error and deletes every file it created, the two closed early included. The contributor guide now says a refusal from the pool cannot fail a fanout write, rather than the pool, and the comment on `reserve` no longer says an open file holds at most a row group, which is wrong on S3 and GCS. * docs: say a fanout partition stays in FanoutPartitions, not open, until the task ends The comment on `held_bytes` still said every partition stays open, which stopped being true once partitions could close early. What makes the held-back rows share one limit is that every partition, and its feed, stays in `partitions`. --- .../review-comet-iceberg-write-pr/SKILL.md | 4 +- .../contributor-guide/iceberg-writes.md | 44 +- .../contributor-guide/memory_management.md | 18 +- .../user-guide/latest/iceberg-writes.md | 37 +- .../src/execution/operators/iceberg_write.rs | 828 ++++++++++++++---- .../sql/comet/CometIcebergWriteExec.scala | 9 +- .../comet/CometIcebergWriteActionSuite.scala | 90 +- 7 files changed, 828 insertions(+), 202 deletions(-) diff --git a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md index a835f45441a..0344f0fa267 100644 --- a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md +++ b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md @@ -156,7 +156,9 @@ Read the ownership table in the contributor guide before reviewing any change ne `MeteredParquetWriterBuilder`, and rows the writer holds outside iceberg-rust (the `PartitionFeed`s, or anything a new feed holds back) must be added to what `run_write_task` reserves. A builder that constructs `ParquetWriterBuilder` directly leaves its files out of the task's - memory reservation, so a wide fanout write grows past the pool instead of failing its task. + memory reservation, so a wide fanout write grows past the pool instead of closing + partitions early or failing its task. A fanout file must also report to its partition's + `OpenFileMemory`, or the partitions holding the most are not the ones closed. A storage scheme newly supported for writes needs its entry in `StorageWrites::for_location`: a file reports its flushed row groups until its storage has written them out, which a local file does at once and an object store only part by part. diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 2cdfb9038b5..6f930733b17 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -223,7 +223,7 @@ per task, and `PartitionWriterBuilder` builds the rest of the stack once per par partition's Parquet writer properties: ```text -UnpartitionedWriter | FanoutWriter | ClusteredWriter +UnpartitionedWriter | FanoutPartitions | ClusteredWriter -> PartitionWriterBuilder, per partition: ParquetWriterBuilder -> RollingFileWriterBuilder -> DataFileWriterBuilder ``` @@ -270,8 +270,9 @@ Points where Comet adapts iceberg-rust to match iceberg-java: - **Field ids and casting.** `decorate_batch_with_field_ids` casts each batch to the field-id-annotated Arrow schema derived from the Iceberg schema, with `safe: false`, so a type mismatch fails the task instead of writing NULLs. -- **Deterministic output order.** iceberg-rust's `FanoutWriter` closes its writers out of a - `HashMap`, so the fanout path sorts its `DataFile`s by path before returning them +- **Deterministic output order.** `FanoutPartitions`, Comet's version of iceberg-rust's + `FanoutWriter`, closes its files in an order that depends on how the rows arrived and on which + partitions it closed early, so the fanout path sorts its `DataFile`s by path before returning them ([#5776](https://github.com/apache/datafusion-comet/issues/5776)). Manifest order becomes the row order of an unordered read, so any map iteration that reaches the output needs the same care. @@ -281,8 +282,21 @@ hold in memory. It hands `ParquetWriter` each file's `OutputFile` behind a `Coun writer counts the bytes that leave memory on their way to storage, and a file reports what it has written less those. When they leave depends on the storage (`StorageWrites`). After every batch `run_write_task` resizes the task's reservation to what the open files report plus the rows each -`PartitionFeed` holds, for the dictionary choice or for pacing, and a resize the pool refuses fails -the task. What that figure covers, and what it misses, is described under +`PartitionFeed` holds, for the dictionary choice or for pacing (`InnerWriter::reserve`). When the +pool refuses a fanout write, `reserve` writes out and closes the partitions holding the most, in +their open file and their feed, until the resize succeeds, and `run_write_task` counts them in the +`files_closed_early` metric. Taking the largest first frees the most memory per file closed, so +each refusal closes as few files as it can. Each partition's files report to a child of the task's +`OpenFileMemory`, which is how `reserve` finds them. That is why the fanout path uses +`FanoutPartitions` rather than iceberg-rust's `FanoutWriter`, which cannot close one partition's +writer. A closed partition's next rows open a new file with the properties its first file used, +which the partition keeps. A refusal from the pool cannot fail a fanout write: everything the write +reserves belongs to one of its partitions, so once all of them are closed it reserves nothing, and a +resize down always succeeds. Only a refusal (`ResourcesExhausted`) makes `reserve` close +partitions. Any other error from a resize fails the task, such as an exception Spark threw because +a spill failed or because the task was killed while it waited for memory. An unpartitioned or +clustered write has no file it can close early, so a resize the pool refuses fails its task. What +the reserved figure covers, and what it misses, is described under [Native writers](memory_management.md#native-writers). `FileIO` comes from `load_file_io` in `iceberg_common.rs`, shared with the native scan. It picks @@ -416,16 +430,16 @@ the writer: run the write suites and the Iceberg Spark tests. The pin policy is ## Testing -| Suite | What it covers | -| ---------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, a fanout write outgrowing the memory pool, AQE re-planning. | -| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | -| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | -| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | -| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | -| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, memory reservation, partition path rendering, partition values past `chrono`'s calendar. | -| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | -| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | +| Suite | What it covers | +| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, writes outgrowing the memory pool, AQE re-planning. | +| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | +| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | +| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | +| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | +| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, memory reservation, partition path rendering, partition values past `chrono`'s calendar. | +| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | +| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | The Comet suites run against the Iceberg version each Spark profile pins in `spark/pom.xml`: 1.5.2 for Spark 3.4, 1.8.1 for 3.5, 1.10.0 for 4.0 and 4.2, and 1.11.0 for 4.1. Only the default profile diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index af84e4a2786..9c31d3fecb2 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -396,8 +396,11 @@ An operator that never calls `try_grow` is invisible to the pool no matter how m ### Native writers Both native writers reserve what they hold between batches through a single consumer per task, -`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch. Neither can spill, so -when the pool refuses a resize the task fails with a `CometNativeException` whose message starts +`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch. Neither can spill. A +fanout Iceberg write gives memory back instead, by closing partitions early when the pool refuses a +resize, as described below, so a refusal never fails it. An exception Spark throws during a resize, +such as a failed spill or a killed task, still fails the task. When the pool refuses any other +write's resize, the task fails with a `CometNativeException` whose message starts `Additional allocation failed for` and names the consumer. That is a task failure Spark can retry. Unreserved, the same memory would count only toward the container limit, where exceeding it kills the executor. @@ -418,10 +421,13 @@ the executor. it has flushed since its last part was uploaded. At the default row-group size (`write.parquet.row-group-size-bytes`, 128 MiB) that is the last row group it flushed. The reservation also covers the rows each partition holds back, first for its dictionary choice and - then until they fill the 1000-row unit the rolling writer is fed in. It does not cover what - parquet-rs holds beyond the encoded size: dictionary hash tables, unencoded dictionary indices - and buffer capacity. Native writes decline Bloom filters today, and neither figure would include - them. + then until they fill the 1000-row unit the rolling writer is fed in. When the pool refuses a + fanout write's resize, the writer writes out and closes the partitions holding the most, file + and held rows together, until the resize succeeds, so the write ends with more files rather than + failing. Each file also reports to its partition's total, which is how the writer finds them. + The reservation does not cover what parquet-rs holds beyond the encoded size: dictionary hash + tables, unencoded dictionary indices and buffer capacity. Native writes decline Bloom filters + today, and neither figure would include them. The writers register one consumer per task rather than one per open file because every consumer registered with `fair_unified` lowers the share of every other consumer in the task. A consumer per diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index 2d84caeaca0..cea6ca42a92 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -315,13 +315,24 @@ The native writer's buffers are charged to Comet's memory pool, the off-heap bud native operators draw on, where iceberg-java's buffers sit on the JVM heap. A fanout write keeps a data file open for every partition a task writes to. Each open file holds the row group it is writing in memory, up to `write.parquet.row-group-size-bytes`, and on S3 or GCS also the last row -group it flushed, which is uploaded once the next one is complete or the file closes. So a task -writing to many partitions needs memory in proportion to them. When the pool cannot grant it, the -task fails with a `CometNativeException` reading `Additional allocation failed for IcebergWriteExec` -instead of exceeding the executor's memory, and Spark retries it like any other task failure. Such a -write fits in less memory with the fanout writer disabled (`write.spark.fanout.enabled=false`): -Spark then sorts each task's rows by partition, and the task keeps one file open at a time. A -smaller row-group size also helps. Otherwise the write needs a larger `spark.memory.offHeap.size`. +group it flushed, which is uploaded once the next one is complete or the file closes. Each +partition also holds the rows that have not reached its file yet. So a task writing to many +partitions needs memory in proportion to them. When the pool cannot grant it, the task writes out +and closes the partitions holding the most memory until what is left fits, and a closed +partition's next rows open a new file. The write then finishes with more, smaller files than +iceberg-java's would. When every partition keeps getting rows, a task given an eighth of the memory +its partitions need writes about four files per partition instead of one, however many partitions +it writes to. The `files closed early to free memory` metric of its `CometIcebergWrite` operator +counts the files it closed early. Disabling the fanout writer (`write.spark.fanout.enabled=false`) +avoids the extra files: Spark then sorts each task's rows by partition, and the task keeps one file +open at a time. A larger `spark.memory.offHeap.size` also helps, and Iceberg's `rewrite_data_files` +procedure compacts small files once they are written. + +A write that keeps one file open, unpartitioned or clustered, has no partition to close. When that +file outgrows the pool, the task fails with a `CometNativeException` reading +`Additional allocation failed for IcebergWriteExec` instead of exceeding the executor's memory, and +Spark retries it like any other task failure. A smaller `write.parquet.row-group-size-bytes` or a +larger `spark.memory.offHeap.size` lets such a write fit. Partial results are never committed. The commit set is exactly the commit messages returned by successful tasks — a failed task contributes none — and if the job fails, the driver-side @@ -416,6 +427,11 @@ a data file but not what any reader computes from it: they may cross the target several grid steps apart, and the resulting files can differ in row count by an arbitrary number of 1000-row blocks. Do not rely on file-layout parity between the two writers; rely only on each file rolling on its own 1000-row boundary. +- A fanout write that the memory pool cannot hold closes partitions before the task ends (see + [Failure handling](#failure-handling)), where iceberg-java's fanout writer keeps every file open + until then. So a partition can have more, smaller files than iceberg-java writes, and a file + closed early ends off the 1000-row grid, as the last file of a task does. A partition closed + before its rows filled a first page makes its dictionary choice from the rows it has. - A fanout write lists a task's data files in file-path order, where iceberg-java lists them in its own `StructLikeMap` iteration order. Both are stable across runs, and neither is a documented ordering, but the manifest entry order becomes the scan-task order and so the row @@ -423,10 +439,9 @@ a data file but not what any reader computes from it: commit gives each file's rows, so the same rows can get different `_row_id` values from the two writers. The ids are unique either way, and across tasks iceberg-java's own assignment already depends on the order in which the tasks finish. Only the sorted order is reproducible on the - native path: iceberg-rust's `FanoutWriter` closes its per-partition writers out of a `HashMap`, - which under Rust's per-process `RandomState` would otherwise give a different order on every - run. Clustered and unpartitioned writes append in creation order on both paths and are - unaffected. + native path: the order the native writer closes its files in depends on the order rows arrive in + and on which partitions it closes early to free memory. Clustered and unpartitioned writes append + in creation order on both paths and are unaffected. - Compressed page bytes are implementation-defined: the codec and any explicit level are translated, but parquet-rs and parquet-mr embed different encoder implementations and defaults (zstd default levels, LZ4 framing), so byte-identical output is not achievable even diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index ba9b08bd9da..a2996119ed5 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -25,6 +25,7 @@ //! `ManifestWriter` against an in-memory `FileIO`. The JVM decodes the bytes with //! `ManifestFiles.read(...)` to recover the `DataFile`s for commit. +use std::cmp::Reverse; use std::collections::HashMap; use std::fmt; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -44,7 +45,7 @@ use datafusion::execution::TaskContext; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::metrics::{ - ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, + Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, }; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -69,9 +70,9 @@ use iceberg::writer::file_writer::{ FileWriter, FileWriterBuilder, ParquetWriter, ParquetWriterBuilder, }; use iceberg::writer::partitioning::clustered_writer::ClusteredWriter; -use iceberg::writer::partitioning::fanout_writer::FanoutWriter; use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter; -use iceberg::writer::{CurrentFileStatus, IcebergWriterBuilder}; +use iceberg::writer::partitioning::PartitioningWriter; +use iceberg::writer::{CurrentFileStatus, IcebergWriter, IcebergWriterBuilder}; use iceberg::ErrorKind; #[cfg(test)] use parquet::arrow::PARQUET_FIELD_ID_META_KEY; @@ -123,17 +124,18 @@ struct PartitionWriterBuilder { storage: StorageWrites, } -#[async_trait::async_trait] -impl IcebergWriterBuilder for PartitionWriterBuilder { - type R = PartitionDataFileWriter; - - async fn build(&self, partition_key: Option) -> iceberg::Result { - let properties = self - .properties - .take(partition_key.as_ref().map(PartitionKey::data)); +impl PartitionWriterBuilder { + /// The data file writer for `partition_key`, writing with `properties`, whose files report + /// what they hold to `open_files`. + async fn build_with( + &self, + partition_key: Option, + properties: WriterProperties, + open_files: OpenFileMemory, + ) -> iceberg::Result { let parquet_builder = MeteredParquetWriterBuilder { inner: ParquetWriterBuilder::new(properties, Arc::clone(&self.schema)), - open_files: self.open_files.clone(), + open_files, storage: self.storage, }; let rolling_builder = RollingFileWriterBuilder::new( @@ -149,13 +151,28 @@ impl IcebergWriterBuilder for PartitionWriterBuilder { } } +/// How the unpartitioned and clustered writers open a partition's writer: once, with the +/// properties chosen for it. +#[async_trait::async_trait] +impl IcebergWriterBuilder for PartitionWriterBuilder { + type R = PartitionDataFileWriter; + + async fn build(&self, partition_key: Option) -> iceberg::Result { + let properties = self + .properties + .take(partition_key.as_ref().map(PartitionKey::data)); + self.build_with(partition_key, properties, self.open_files.clone()) + .await + } +} + /// The Parquet writer properties each partition's files are written with: the table's, with the /// dictionary choice [`DictionaryChooser`] made from the partition's first rows. /// /// A partition's [`PartitionFeed`] records the choice before it hands the partition's first rows -/// to the writer, and [`PartitionWriterBuilder`] takes it back out when that first write opens the -/// partition's writer. The partitioning writers open a partition's writer exactly once per task, -/// on its first write. +/// to the writer, and the partition takes it back out when that first write opens the partition's +/// writer, once per task. A fanout partition keeps what it took for the files it opens after +/// closing one early (see [`FanoutPartition`]). struct PartitionProperties { chooser: DictionaryChooser, /// Keyed by partition value; `None` is the whole task of an unpartitioned write. @@ -196,16 +213,35 @@ impl PartitionProperties { /// What a task's open data files hold in memory between them. /// -/// iceberg-rust keeps each open file writer private inside its rolling and partitioning writers, -/// so the files report their own shares here through [`MeteredParquetWriter`], and `run_write_task` -/// reserves the total. A fanout write keeps one file open per partition, so this is what grows -/// with the partition count. +/// iceberg-rust keeps each open file writer private inside its rolling writer, so the files report +/// their own shares here through [`MeteredParquetWriter`], and `run_write_task` reserves the +/// total. A fanout write keeps one file open per partition, so this is what grows with the +/// partition count. Each fanout partition's files report to a [`child`](Self::child) of it, which +/// is how the write finds the partitions to close when the pool refuses the total. #[derive(Clone, Debug, Default)] -struct OpenFileMemory(Arc); +struct OpenFileMemory { + bytes: Arc, + /// The total this one is part of, for a fanout partition's. + parent: Option>, +} impl OpenFileMemory { fn bytes(&self) -> usize { - self.0.load(Ordering::Relaxed) + self.bytes.load(Ordering::Relaxed) + } + + /// A counter for one fanout partition's files, which also counts towards this one. Only the + /// task's total has children: a share counts towards its own counter and that counter's + /// parent, so a child's child would never reach the total `run_write_task` reserves. + fn child(&self) -> OpenFileMemory { + debug_assert!( + self.parent.is_none(), + "a child's child would not count towards the task's total" + ); + OpenFileMemory { + bytes: Arc::default(), + parent: Some(Arc::clone(&self.bytes)), + } } fn share(&self) -> OpenFileShare { @@ -226,14 +262,12 @@ struct OpenFileShare { impl OpenFileShare { fn set(&mut self, bytes: usize) { - if bytes > self.bytes { - self.total - .0 - .fetch_add(bytes - self.bytes, Ordering::Relaxed); - } else { - self.total - .0 - .fetch_sub(self.bytes - bytes, Ordering::Relaxed); + for total in std::iter::once(&self.total.bytes).chain(&self.total.parent) { + if bytes > self.bytes { + total.fetch_add(bytes - self.bytes, Ordering::Relaxed); + } else { + total.fetch_sub(self.bytes - bytes, Ordering::Relaxed); + } } self.bytes = bytes; } @@ -245,7 +279,7 @@ impl Drop for OpenFileShare { } } -/// [`ParquetWriterBuilder`] whose files report what they hold in memory to the task's +/// [`ParquetWriterBuilder`] whose files report what they hold in memory to an /// [`OpenFileMemory`]. #[derive(Clone, Debug)] struct MeteredParquetWriterBuilder { @@ -816,9 +850,7 @@ impl ExecutionPlan for IcebergWriteExec { partition: usize, context: Arc, ) -> DFResult { - // Time spent inside the iceberg-rust writer stack (write + close), excluding time spent - // waiting on the upstream input stream. Surfaced on the JVM exec's SQL metrics by name. - let write_time = MetricBuilder::new(&self.metrics).subset_time("write_time", partition); + let metrics = WriteMetrics::new(&self.metrics, partition); // One consumer for the whole task, however many files it opens: a consumer per file // would shrink every other consumer's share of the fair pool as a fanout write widened. let reservation = MemoryConsumer::new(format!("IcebergWriteExec[{partition}]")) @@ -843,7 +875,7 @@ impl ExecutionPlan for IcebergWriteExec { writer_properties.as_ref().clone(), partition_id, task_attempt_id, - write_time, + metrics, reservation, ) .await?; @@ -909,14 +941,35 @@ impl DisplayAs for IcebergWriteExec { } } +/// What a write task reports, surfaced on the JVM exec's SQL metrics by name. +#[derive(Clone, Debug, Default)] +struct WriteMetrics { + /// Time spent inside the iceberg-rust writer stack (write + close), excluding time spent + /// waiting on the upstream input stream. + write_time: Time, + /// Data files a fanout write closed before the task ended, to give back the memory they held. + files_closed_early: Count, +} + +impl WriteMetrics { + fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + write_time: MetricBuilder::new(metrics).subset_time("write_time", partition), + files_closed_early: MetricBuilder::new(metrics) + .counter("files_closed_early", partition), + } + } +} + /// One-shot per-task write coroutine. Builds the iceberg-rust writer stack, decorates each input /// batch with `PARQUET_FIELD_ID_META_KEY` metadata so iceberg-rust can match Arrow columns to -/// Iceberg field IDs, and routes through `UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter` -/// depending on `writer_mode`. +/// Iceberg field IDs, and routes through `UnpartitionedWriter`/[`FanoutPartitions`]/ +/// `ClusteredWriter` depending on `writer_mode`. /// -/// After every batch, `reservation` is resized to what the task's writers hold between batches: -/// the open files' shares (see [`MeteredParquetWriter`]) and the rows waiting in their -/// [`PartitionFeed`]s. The writer cannot spill, so a reservation the pool refuses fails the task. +/// After every batch, `reservation` is resized to what the task's writers hold between batches +/// (see [`InnerWriter::reserve`]), and a fanout write the pool refuses closes partitions early +/// until what is left fits. The writer cannot spill, so a reservation the pool still refuses +/// fails the task. /// /// On success the still-armed [`AbortOnDrop`] is returned along with the data files: the caller /// owns cleanup until the JVM acknowledges the output after recording its locations. @@ -930,7 +983,7 @@ async fn run_write_task( writer_properties: WriterProperties, partition_id: Option, task_attempt_id: Option, - write_time: Time, + metrics: WriteMetrics, reservation: MemoryReservation, ) -> DFResult<(Vec, AbortOnDrop)> { // The JVM exec wrapper stamps both ids per task; a missing id means the plan template was @@ -1011,11 +1064,9 @@ async fn run_write_task( UnpartitionedWriter::new(data_file_builder), PartitionFeed::new(None, slicer), ), - (false, ProtoIcebergWriterMode::IcebergWriterFanout) => InnerWriter::Fanout( - FanoutWriter::new(data_file_builder), - splitter()?, - FanoutFeeds::default(), - ), + (false, ProtoIcebergWriterMode::IcebergWriterFanout) => { + InnerWriter::Fanout(FanoutPartitions::new(data_file_builder), splitter()?) + } (false, ProtoIcebergWriterMode::IcebergWriterClustered) => { InnerWriter::Clustered(ClusteredWriter::new(data_file_builder), splitter()?, None) } @@ -1033,12 +1084,15 @@ async fn run_write_task( if nests_floats { decorated = drop_unwritten_values(decorated)?; } - let timer = write_time.timer(); + let timer = metrics.write_time.timer(); writer.write(decorated, &properties).await?; + let closed = writer + .reserve(&reservation, &open_files, &properties) + .await?; + metrics.files_closed_early.add(closed); timer.done(); - reservation.try_resize(open_files.bytes() + writer.pending_bytes())?; } - let _timer = write_time.timer(); + let _timer = metrics.write_time.timer(); writer.close(&properties).await } .await; @@ -1057,7 +1111,218 @@ async fn run_write_task( } } -/// Enum-based dispatch over the three iceberg-rust partitioning writers, each paired with the +/// A fanout write: like iceberg-rust's `FanoutWriter`, a data file writer open for every +/// partition the task has written to, except that a partition can be closed before the task ends, +/// to give back the memory it holds. Its next rows then open a new file, with the same properties. +/// `FanoutWriter` keeps its writers private and closes them only all at once, so the fanout path +/// keeps its own. +/// +/// iceberg-java's fanout writer keeps every file open until the task ends, its buffers growing on +/// the JVM heap. The native writer's buffers count against the task's memory pool instead, so when +/// the pool refuses them, closing partitions early lets the write finish with more, smaller files +/// rather than fail (see [`InnerWriter::reserve`]). +struct FanoutPartitions { + builder: PartitionWriterBuilder, + /// Every partition the task has seen, in the order it first saw them. + partitions: Vec, + /// Where each partition value sits in `partitions`. + index: HashMap, + /// Bytes the feeds hold back for their dictionary choice, across all partitions. Every + /// partition stays in `partitions` until the task ends, so the rows they hold back share one + /// limit: once they reach it, every partition still holding makes its choice from the rows it + /// has. A task fanning out to many partitions that each get less than a page would otherwise + /// hold all of its rows back, uncompressed, until it closed. + /// + /// Always the sum of the feeds' [`PartitionFeed::held_bytes`], kept in step wherever a feed + /// changes what it holds back: a push, the release once the limit is reached, and an early + /// close. `write` checks this in debug builds. + held_bytes: usize, + /// The data files of the partitions closed early. + closed: Vec, +} + +/// One partition of a fanout write. +struct FanoutPartition { + /// What each of the partition's files is opened with: its first, and any it opens after + /// closing one early. + key: PartitionKey, + feed: PartitionFeed, + /// The writer of the partition's open file, if it has one open. + writer: Option, + /// What the open file holds, as part of the task's [`OpenFileMemory`]. + memory: OpenFileMemory, + /// The properties the partition's files are written with, taken when its first file opens and + /// kept for the files it opens after closing one early. + properties: Option, +} + +impl FanoutPartition { + /// What the partition holds: its open file and the rows its feed has not handed over. + fn bytes(&self) -> usize { + self.memory.bytes() + self.feed.pending_bytes() + } + + /// Writes `unit` to the partition's open file, opening one first if it has none. + async fn write( + &mut self, + builder: &PartitionWriterBuilder, + unit: RecordBatch, + ) -> iceberg::Result<()> { + let writer = match self.writer.as_mut() { + Some(writer) => writer, + None => { + let properties = self + .properties + .get_or_insert_with(|| builder.properties.take(Some(self.key.data()))) + .clone(); + self.writer.insert( + builder + .build_with(Some(self.key.clone()), properties, self.memory.clone()) + .await?, + ) + } + }; + writer.write(unit).await + } + + /// Writes out every row the partition still holds back and closes its file, returning that + /// file's data files. A file closed at a partial unit ends off iceberg-java's row grid, as the + /// last file of a task does. + async fn close( + &mut self, + builder: &PartitionWriterBuilder, + properties: &PartitionProperties, + ) -> DFResult> { + for unit in self.feed.finish(properties)? { + self.write(builder, unit).await.map_err(iceberg_err)?; + } + match self.writer.take() { + Some(mut writer) => writer.close().await.map_err(iceberg_err), + None => Ok(Vec::new()), + } + } +} + +impl FanoutPartitions { + fn new(builder: PartitionWriterBuilder) -> Self { + Self { + builder, + partitions: Vec::new(), + index: HashMap::new(), + held_bytes: 0, + closed: Vec::new(), + } + } + + /// Memory held by the rows waiting in the feeds, for every partition the task has seen. + fn pending_bytes(&self) -> usize { + self.partitions + .iter() + .map(|partition| partition.feed.pending_bytes()) + .sum() + } + + /// Writes `batch`, feeding each of its partitions separately. + async fn write( + &mut self, + splitter: &PartitionSplitter, + batch: &RecordBatch, + properties: &PartitionProperties, + ) -> DFResult<()> { + for (key, part) in splitter.split_groups(batch)? { + let next = self.partitions.len(); + let i = *self.index.entry(key.data().clone()).or_insert(next); + if i == next { + self.partitions.push(FanoutPartition { + feed: PartitionFeed::new(Some(key.data().clone()), splitter.slicer), + key, + writer: None, + memory: self.builder.open_files.child(), + properties: None, + }); + } + let partition = &mut self.partitions[i]; + let held_before = partition.feed.held_bytes(); + let units = partition.feed.push(part, properties)?; + self.held_bytes = self.held_bytes - held_before + partition.feed.held_bytes(); + for unit in units { + partition + .write(&self.builder, unit) + .await + .map_err(iceberg_err)?; + } + } + if !properties.chooser.may_hold(self.held_bytes) { + for partition in &mut self.partitions { + for unit in partition.feed.release(properties)? { + partition + .write(&self.builder, unit) + .await + .map_err(iceberg_err)?; + } + } + self.held_bytes = 0; + } + debug_assert_eq!( + self.held_bytes, + self.partitions + .iter() + .map(|partition| partition.feed.held_bytes()) + .sum::(), + "the held-back total drifted from what the feeds hold back" + ); + Ok(()) + } + + /// The partitions holding memory, as indexes into `partitions`, the most first. Of two holding + /// as much, the one seen first comes first, so which files a write closes early does not + /// depend on `HashMap` order. Indexes rather than references, because `reserve` closes each + /// partition it picks while it walks the ranking, and a close borrows the whole + /// `FanoutPartitions` mutably. + fn by_memory(&self) -> Vec { + let mut holding: Vec<(usize, usize)> = self + .partitions + .iter() + .enumerate() + .map(|(i, partition)| (partition.bytes(), i)) + .filter(|&(bytes, _)| bytes > 0) + .collect(); + holding.sort_unstable_by_key(|&(bytes, i)| (Reverse(bytes), i)); + holding.into_iter().map(|(_, i)| i).collect() + } + + /// Closes the partition at `i` before the task ends, keeping its data files for the output. + async fn close_early(&mut self, i: usize, properties: &PartitionProperties) -> DFResult<()> { + let partition = &mut self.partitions[i]; + self.held_bytes -= partition.feed.held_bytes(); + let data_files = partition.close(&self.builder, properties).await?; + self.closed.extend(data_files); + Ok(()) + } + + /// Closes every partition and returns the task's data files in file-path order. + async fn close(mut self, properties: &PartitionProperties) -> DFResult> { + let mut data_files = std::mem::take(&mut self.closed); + for partition in &mut self.partitions { + data_files.extend(partition.close(&self.builder, properties).await?); + } + // The files come back in the order the partitions closed, which depends on the order the + // rows arrived in and on which partitions closed early. That order becomes the manifest + // entry order, then the scan-task order, then the row order of an unordered `SELECT *`, + // which Iceberg's own + // `TestMetadataTablesWithPartitionEvolution.testPartitionColumnNamedPartition` compares + // positionally against what iceberg-java wrote (apache/datafusion-comet#5776). + // + // Sorting by path is what makes the task's output reproducible. Path order sorts by + // partition directory, then by the file name counter within a partition, and needs nothing + // from the write order. The clustered and unpartitioned writers append in creation order + // and are already deterministic, so only the fanout path sorts. + data_files.sort_unstable_by(|a, b| a.file_path().cmp(b.file_path())); + Ok(data_files) + } +} + +/// Enum-based dispatch over the three partitioning writers, each paired with the /// [`PartitionFeed`] state that holds a partition's first rows back for its dictionary choice and /// keeps its rolling writer on iceberg-java's row grid, and the two partitioned ones with the /// [`PartitionSplitter`] that cuts their input. Each variant takes the same builder so we can keep @@ -1066,11 +1331,7 @@ async fn run_write_task( enum InnerWriter { Unpartitioned(UnpartitionedWriter, PartitionFeed), /// The fanout writer keeps one file open per partition, so every partition is fed separately. - Fanout( - FanoutWriter, - PartitionSplitter, - FanoutFeeds, - ), + Fanout(FanoutPartitions, PartitionSplitter), /// The clustered writer closes a partition's file as soon as the next key arrives, so only the /// current key's held rows are live; they are written out before the switch. Clustered( @@ -1087,17 +1348,50 @@ impl InnerWriter { fn pending_bytes(&self) -> usize { match self { InnerWriter::Unpartitioned(_, feed) => feed.pending_bytes(), - InnerWriter::Fanout(_, _, fanout) => fanout - .feeds - .values() - .map(|(_, feed)| feed.pending_bytes()) - .sum(), + InnerWriter::Fanout(fanout, _) => fanout.pending_bytes(), InnerWriter::Clustered(_, _, live) => { live.as_ref().map_or(0, |(_, feed)| feed.pending_bytes()) } } } + /// Resizes `reservation` to what the task's writers hold between batches: the open files' + /// shares (`open_files`, see [`MeteredParquetWriter`]) and the rows waiting in their + /// [`PartitionFeed`]s. When the pool refuses, a fanout write writes out and closes its + /// partitions, the one holding the most first, until what is left fits, and this returns how + /// many it closed. The unpartitioned and clustered writers hold one file open and cannot close + /// it early, so the pool's refusal stands. Any other error from a resize fails the task (see + /// [`is_refusal`]). + async fn reserve( + &mut self, + reservation: &MemoryReservation, + open_files: &OpenFileMemory, + properties: &PartitionProperties, + ) -> DFResult { + let mut pending = self.pending_bytes(); + let mut refused = match reservation.try_resize(open_files.bytes() + pending) { + Ok(()) => return Ok(0), + Err(e) if is_refusal(&e) => e, + Err(e) => return Err(e), + }; + let InnerWriter::Fanout(fanout, _) = self else { + return Err(refused); + }; + for (closed, i) in fanout.by_memory().into_iter().enumerate() { + pending -= fanout.partitions[i].feed.pending_bytes(); + fanout.close_early(i, properties).await?; + match reservation.try_resize(open_files.bytes() + pending) { + Ok(()) => return Ok(closed + 1), + Err(e) if is_refusal(&e) => refused = e, + Err(e) => return Err(e), + } + } + // Not reached: everything a fanout write reserves belongs to one of its partitions, so + // once every partition holding memory is closed the resize is down to nothing, which the + // pool cannot refuse. + Err(refused) + } + /// Writes `batch` through the writer this task built, in the [`ROWS_DIVISOR`]-row units the /// rolling writer needs to re-check the target file size on iceberg-java's cadence. Rows are /// fed after partition splitting, so each partition's dictionary choice is made from its own @@ -1108,7 +1402,6 @@ impl InnerWriter { batch: RecordBatch, properties: &PartitionProperties, ) -> DFResult<()> { - use iceberg::writer::partitioning::PartitioningWriter; match self { InnerWriter::Unpartitioned(w, feed) => { for unit in feed.push(batch, properties)? { @@ -1116,29 +1409,8 @@ impl InnerWriter { } Ok(()) } - InnerWriter::Fanout(w, splitter, fanout) => { - for (key, part) in splitter.split_groups(&batch)? { - let (key, feed) = fanout.feeds.entry(key.data().clone()).or_insert_with(|| { - let partition = Some(key.data().clone()); - (key, PartitionFeed::new(partition, splitter.slicer)) - }); - let held_before = feed.held_bytes(); - let units = feed.push(part, properties)?; - fanout.held_bytes = fanout.held_bytes - held_before + feed.held_bytes(); - for unit in units { - w.write(key.clone(), unit).await.map_err(iceberg_err)?; - } - } - // Every partition stays open, so the rows they hold back share one limit. - if !properties.chooser.may_hold(fanout.held_bytes) { - for (key, feed) in fanout.feeds.values_mut() { - for unit in feed.release(properties)? { - w.write(key.clone(), unit).await.map_err(iceberg_err)?; - } - } - fanout.held_bytes = 0; - } - Ok(()) + InnerWriter::Fanout(fanout, splitter) => { + fanout.write(splitter, &batch, properties).await } InnerWriter::Clustered(w, splitter, live) => { for (key, part) in splitter.split_runs(&batch)? { @@ -1175,7 +1447,6 @@ impl InnerWriter { /// same at close: the leftovers land in the file that is open at that point, unless the /// pending target-size check rolls first -- exactly as they would have on the JVM path. async fn close(self, properties: &PartitionProperties) -> DFResult> { - use iceberg::writer::partitioning::PartitioningWriter; match self { InnerWriter::Unpartitioned(mut w, mut feed) => { for unit in feed.finish(properties)? { @@ -1183,29 +1454,7 @@ impl InnerWriter { } w.close().await.map_err(iceberg_err) } - InnerWriter::Fanout(mut w, _, fanout) => { - for (_, (key, mut feed)) in fanout.feeds { - for unit in feed.finish(properties)? { - w.write(key.clone(), unit).await.map_err(iceberg_err)?; - } - } - let mut data_files = w.close().await.map_err(iceberg_err)?; - // `FanoutWriter` holds its per-partition writers in a `HashMap` and `close` - // iterates it directly, so the order it returns follows Rust's per-process - // `RandomState` and differs on every run. That order becomes the manifest entry - // order, then the scan-task order, then the row order of an unordered - // `SELECT *` -- which Iceberg's own - // `TestMetadataTablesWithPartitionEvolution.testPartitionColumnNamedPartition` - // compares positionally against what iceberg-java wrote - // (apache/datafusion-comet#5776). - // - // Sorting by path is what makes the task's output reproducible. Path order sorts - // by partition directory, then by the file name counter within a partition, and - // needs nothing from the write order. The clustered and unpartitioned writers - // append in creation order and are already deterministic, so only this arm sorts. - data_files.sort_unstable_by(|a, b| a.file_path().cmp(b.file_path())); - Ok(data_files) - } + InnerWriter::Fanout(fanout, _) => fanout.close(properties).await, InnerWriter::Clustered(mut w, splitter, live) => { if let Some((key, mut feed)) = live { // Feeding defers a partition's rows to here, so the unclustered-input @@ -1225,6 +1474,16 @@ impl InnerWriter { // --- helpers ------------------------------------------------------------- +/// Whether `e`, from resizing a reservation, is the pool refusing the memory, which every pool +/// reports as `ResourcesExhausted` and a fanout write answers by closing partitions early. A pool +/// backed by Spark can also fail a resize with an exception Spark threw: a spill of another +/// consumer that failed, or the `InterruptedException` of a task killed while it waited for +/// memory. That comes back as another kind of error and has to fail the task. Closing partitions +/// would swallow it, because once the resize is down to a shrink it no longer calls Spark. +fn is_refusal(e: &DataFusionError) -> bool { + matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) +} + fn parse_iceberg_schema(json: &str) -> DFResult { let schema: IcebergSchema = serde_json::from_str(json).map_err(|e| { DataFusionError::Internal(format!("Failed to parse iceberg schema JSON: {e}")) @@ -1728,21 +1987,6 @@ impl PartitionFeed { } } -/// The feeds of a fanout write, one per partition, all open at once. -/// -/// Between them they hold back no more than a single partition may: once the rows held across all -/// partitions reach that, every partition still holding makes its choice from the rows it has. A -/// task fanning out to many partitions that each get less than a page would otherwise hold all of -/// its rows back, uncompressed, until it closed. -#[derive(Default)] -struct FanoutFeeds { - /// The `PartitionKey` is kept alongside each feed so held rows can still be written out at - /// close. - feeds: HashMap, - /// Bytes held back across all feeds. - held_bytes: usize, -} - /// `true` when `data_type` puts a float or double under a list or map, where a slice's offset /// window is invisible to iceberg-rust's NaN-count visitor. fn float_under_list_or_map(data_type: &DataType) -> bool { @@ -2832,7 +3076,7 @@ mod tests { writer_properties, Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await?; @@ -2932,7 +3176,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -2980,7 +3224,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -3047,11 +3291,10 @@ mod tests { /// A fanout task's data files must come back in a deterministic order. /// - /// iceberg-rust's `FanoutWriter` keeps its per-partition writers in a `HashMap` and - /// `close` iterates it directly, so the `DataFile` order it returns follows Rust's - /// per-process `RandomState` -- a different order on every run. That order becomes the - /// manifest entry order, which becomes the scan-task order, which becomes the row order - /// of an unordered `SELECT *`. Iceberg's own + /// The order a fanout write closes its files in depends on the order its rows arrived in + /// and on which partitions it closed early. That order would become the manifest entry + /// order, which becomes the scan-task order, which becomes the row order of an unordered + /// `SELECT *`. Iceberg's own /// `TestMetadataTablesWithPartitionEvolution.testPartitionColumnNamedPartition` compares /// such a `SELECT *` positionally against what iceberg-java wrote, and fails when the /// two disagree. See https://github.com/apache/datafusion-comet/issues/5776. @@ -4257,7 +4500,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -4386,7 +4629,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -4454,7 +4697,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -4547,7 +4790,7 @@ mod tests { WriterProperties::builder().build(), Some(0), Some(0), - Time::default(), + WriteMetrics::default(), unbounded_reservation(), ) .await @@ -4621,6 +4864,31 @@ mod tests { rows: usize, payload_bytes: usize, ) -> RecordBatch { + payload_rows_from(first, vec![region; rows], payload_bytes) + } + + /// [`round_robin_batch_from`]'s rows, each with its own `payload_bytes`-byte payload. + fn round_robin_payload_batch_from( + first: usize, + rows: usize, + partitions: usize, + payload_bytes: usize, + ) -> RecordBatch { + let regions: Vec = (0..rows) + .map(|row| format!("r{}", row % partitions)) + .collect(); + let regions: Vec<&str> = regions.iter().map(String::as_str).collect(); + payload_rows_from(first, regions, payload_bytes) + } + + /// One row for each of `regions`, numbered from `first`, each with its own + /// `payload_bytes`-byte payload. + fn payload_rows_from( + first: usize, + regions: Vec<&str>, + payload_bytes: usize, + ) -> RecordBatch { + let rows = regions.len(); let ids: Vec = (first..first + rows).map(|id| id as i32).collect(); let payloads: Vec = (first..first + rows) .map(|id| format!("{id:0>payload_bytes$}")) @@ -4633,7 +4901,7 @@ mod tests { ])), vec![ Arc::new(Int32Array::from(ids)), - Arc::new(StringArray::from(vec![region; rows])), + Arc::new(StringArray::from(regions)), Arc::new(StringArray::from(payloads)), ], ) @@ -4642,10 +4910,11 @@ mod tests { /// Writes `batches` into a fresh table of `schema`, partitioned by `region` unless the /// mode is unpartitioned, reserving from `pool` the way `IcebergWriteExec::execute` - /// reserves from the task's pool. A successful write's files are left for the caller to - /// inspect. - async fn write_reserving_from( - pool: &Arc, + /// reserves from the task's pool, and counting into `metrics`. A successful write's files + /// are left for the caller to inspect. + async fn write_reserving_from( + pool: &Arc

, + metrics: WriteMetrics, schema: Schema, writer_mode: ProtoIcebergWriterMode, batches: Vec, @@ -4670,7 +4939,7 @@ mod tests { ), target_file_size_bytes, ); - let pool: Arc = Arc::::clone(pool); + let pool: Arc = Arc::

::clone(pool); let reservation = MemoryConsumer::new("IcebergWriteExec[0]").register(&pool); let written = run_write_task( input_stream(batches), @@ -4681,7 +4950,7 @@ mod tests { writer_properties, Some(0), Some(0), - Time::default(), + metrics, reservation, ) .await @@ -4707,6 +4976,7 @@ mod tests { .build(); let (_dir, written) = write_reserving_from( &pool, + WriteMetrics::default(), iceberg_user_schema(), writer_mode, vec![batch], @@ -4747,6 +5017,7 @@ mod tests { let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); let (_dir, written) = write_reserving_from( &pool, + WriteMetrics::default(), iceberg_user_schema(), ProtoIcebergWriterMode::IcebergWriterFanout, vec![round_robin_batch( @@ -4829,6 +5100,7 @@ mod tests { let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); let (_dir, written) = write_reserving_from( &pool, + WriteMetrics::default(), iceberg_user_schema(), ProtoIcebergWriterMode::IcebergWriterUnpartitioned, batches, @@ -4871,6 +5143,7 @@ mod tests { let pool = Arc::new(PeakMemoryPool::new(8 * row_group_bytes)); let (_dir, written) = write_reserving_from( &pool, + WriteMetrics::default(), iceberg_payload_schema(), ProtoIcebergWriterMode::IcebergWriterFanout, batches, @@ -4894,32 +5167,182 @@ mod tests { assert_eq!(pool.reserved(), 0, "the write kept a reservation"); } - /// A write whose open files outgrow what the pool grants fails with the pool's error, - /// rather than holding memory nothing accounts for, and deletes the files it had opened. + /// Rows per partition in each data file of `data_files`, keyed by partition value. + fn rows_per_partition(data_files: &[DataFile]) -> HashMap { + let mut rows = HashMap::new(); + for file in data_files { + *rows.entry(file.partition().clone()).or_default() += file.record_count(); + } + rows + } + + /// Writes `batches` into a fanout table of `schema` twice: with a pool to spare, then with + /// a pool of the most that first write reserved divided by `divisor`, which it can only + /// fit by closing partitions early. Returns both writes' data files and the second write's + /// directory, after checking that the second fit, gave its reservation back and counted + /// the files it closed early. + async fn fanout_write_in_part_of_the_pool( + divisor: usize, + schema: Schema, + batches: impl Fn() -> Vec, + properties: impl Fn() -> WriterProperties, + ) -> (Vec, Vec, TempDir) { + let roomy = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (_dir, written) = write_reserving_from( + &roomy, + WriteMetrics::default(), + schema.clone(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + let roomy_files = written.unwrap(); + let tight = Arc::new(PeakMemoryPool::new(roomy.peak() / divisor)); + let metrics = WriteMetrics::default(); + let (dir, written) = write_reserving_from( + &tight, + metrics.clone(), + schema, + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + let tight_files = written.expect("closing partitions early lets the write fit"); + assert_eq!(tight.reserved(), 0, "the write kept a reservation"); + // No file rolls at `TARGET_FILE_SIZE`, so each file was closed either early, and + // counted, or when the task ended, which closes at most one file per partition. + let closed_early = metrics.files_closed_early.value(); + let partitions = rows_per_partition(&roomy_files).len(); + assert!( + closed_early > 0 + && closed_early <= tight_files.len() + && tight_files.len() <= closed_early + partitions, + "{closed_early} files counted as closed early, of {} for {partitions} partitions", + tight_files.len() + ); + (roomy_files, tight_files, dir) + } + + /// A fanout write whose partitions outgrow what the pool grants writes out and closes the + /// ones holding the most until what is left fits, and carries on: a closed partition's + /// next rows open a new file. Every row still lands, in more files than partitions. + #[tokio::test] + async fn a_fanout_write_the_pool_cannot_hold_closes_partitions_until_it_fits() { + // Four batches, each a unit for every one of 16 partitions, so a partition closed + // early gets more rows after it. One unit is one page, so every unit goes straight to + // its partition's file. + let (roomy, tight, dir) = fanout_write_in_part_of_the_pool( + 2, + iceberg_user_schema(), + || { + (0..4) + .map(|batch| { + round_robin_batch_from(batch * 16 * ROWS_DIVISOR, 16 * ROWS_DIVISOR, 16) + }) + .collect() + }, + || { + WriterProperties::builder() + .set_data_page_row_count_limit(ROWS_DIVISOR) + .build() + }, + ) + .await; + assert_eq!(roomy.len(), 16); + assert!(tight.len() > 16, "{} files for 16 partitions", tight.len()); + assert_eq!(rows_per_partition(&tight), rows_per_partition(&roomy)); + assert_eq!(parquet_files_under(dir.path()).len(), tight.len()); + } + + /// A partition's memory counts the rows its feed still holds back, and closing it writes + /// them out first. Here every partition's rows wait for their dictionary choice, so no + /// file is open when the pool refuses: the write only fits by writing partitions out. + #[tokio::test] + async fn closing_a_partition_early_writes_out_the_rows_it_holds_back() { + let (roomy, tight, _dir) = fanout_write_in_part_of_the_pool( + 2, + iceberg_user_schema(), + || vec![round_robin_batch(16 * ROWS_DIVISOR, 16)], + || WriterProperties::builder().build(), + ) + .await; + assert_eq!(record_counts(&tight), record_counts(&roomy)); + } + + /// Closing the partitions holding the most frees the most memory per file closed, so a + /// write short of memory closes few files. Here 16 or 64 partitions take turns through 32 + /// batches, so every partition keeps getting rows, and the pool grants an eighth of what + /// the write needs. The write ends with between four and five files per partition, + /// however many partitions it has. Closing every partition at once would leave seven, and + /// closing them in the order they were first seen, or smallest first, seventeen to twenty. + #[tokio::test] + async fn a_fanout_write_short_of_memory_closes_the_partitions_holding_the_most() { + for partitions in [16, 64] { + let (roomy, tight, _dir) = fanout_write_in_part_of_the_pool( + 8, + iceberg_payload_schema(), + || { + (0..32) + .map(|batch| { + round_robin_payload_batch_from(batch * 1024, 1024, partitions, 100) + }) + .collect() + }, + || WriterProperties::builder().build(), + ) + .await; + assert_eq!(roomy.len(), partitions); + assert_eq!(rows_per_partition(&tight), rows_per_partition(&roomy)); + assert!( + tight.len() < 6 * partitions, + "{} files for {partitions} partitions", + tight.len() + ); + } + } + + /// A write whose open file outgrows what the pool grants, with no partition it can close + /// early, fails with the pool's error rather than holding memory nothing accounts for, + /// and deletes the files it had opened. An unpartitioned write keeps its one file open. #[tokio::test] async fn a_write_the_pool_cannot_hold_fails_and_deletes_its_files() { - let batches = || vec![round_robin_batch(16 * ROWS_DIVISOR, 16)]; - let properties = || WriterProperties::builder().build(); + // One unit a batch, and one unit is one page, so the open file grows with every batch. + let batches = || { + (0..4) + .map(|batch| round_robin_batch_from(batch * ROWS_DIVISOR, ROWS_DIVISOR, 16)) + .collect::>() + }; + let properties = || { + WriterProperties::builder() + .set_data_page_row_count_limit(ROWS_DIVISOR) + .build() + }; - // With room to spare, every partition's unit reaches its own file. + // With room to spare, the rows reach one file. let roomy = Arc::new(PeakMemoryPool::new(usize::MAX)); let (dir, written) = write_reserving_from( &roomy, + WriteMetrics::default(), iceberg_user_schema(), - ProtoIcebergWriterMode::IcebergWriterFanout, + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, batches(), properties(), TARGET_FILE_SIZE, ) .await; - assert_eq!(written.unwrap().len(), 16); - assert_eq!(parquet_files_under(dir.path()).len(), 16); + assert_eq!(written.unwrap().len(), 1); + assert_eq!(parquet_files_under(dir.path()).len(), 1); let tight = Arc::new(PeakMemoryPool::new(roomy.peak() / 2)); let (dir, written) = write_reserving_from( &tight, + WriteMetrics::default(), iceberg_user_schema(), - ProtoIcebergWriterMode::IcebergWriterFanout, + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, batches(), properties(), TARGET_FILE_SIZE, @@ -4939,6 +5362,111 @@ mod tests { assert_eq!(tight.reserved(), 0, "the failed write kept a reservation"); } + /// A pool with room to spare that refuses the first `refusals` grows it is asked for, then + /// fails the next with `error`, the way a pool backed by Spark passes on an exception + /// Spark threw, and grants every grow after that. + #[derive(Debug)] + struct FailingPool { + inner: UnboundedMemoryPool, + refusals: AtomicUsize, + error: Mutex>, + } + + impl FailingPool { + fn new(refusals: usize, error: DataFusionError) -> Self { + Self { + inner: UnboundedMemoryPool::default(), + refusals: AtomicUsize::new(refusals), + error: Mutex::new(Some(error)), + } + } + } + + impl fmt::Display for FailingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "FailingPool({})", self.inner) + } + } + + impl MemoryPool for FailingPool { + fn name(&self) -> &str { + "FailingPool" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> DFResult<()> { + if self + .refusals + .try_update(Ordering::Relaxed, Ordering::Relaxed, |left| { + left.checked_sub(1) + }) + .is_ok() + { + return Err(DataFusionError::ResourcesExhausted(format!( + "refused {additional} bytes" + ))); + } + if let Some(error) = self.error.lock().unwrap().take() { + return Err(error); + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + } + + /// Only the pool's refusal makes a fanout write close partitions early. Any other error + /// from a resize, such as an exception Spark threw because a spill failed or the task was + /// killed, fails the write, which deletes every file it created, the ones it closed early + /// included. + #[tokio::test] + async fn a_fanout_write_fails_on_a_resize_error_that_is_not_a_refusal() { + for error in [ + DataFusionError::External("an exception Spark threw".into()), + DataFusionError::Execution("a JNI call that failed".into()), + ] { + let message = error.to_string(); + // Two refusals close two partitions early, then the third resize fails. + let pool = Arc::new(FailingPool::new(2, error)); + let (dir, written) = write_reserving_from( + &pool, + WriteMetrics::default(), + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + // A unit for each of 16 partitions, and one unit is one page, so every + // partition has a file open when the pool refuses. + vec![round_robin_batch(16 * ROWS_DIVISOR, 16)], + WriterProperties::builder() + .set_data_page_row_count_limit(ROWS_DIVISOR) + .build(), + TARGET_FILE_SIZE, + ) + .await; + let error = + written.expect_err("an error that is not a refusal must fail the write"); + assert_eq!(error.to_string(), message); + assert_eq!( + pool.refusals.load(Ordering::Relaxed), + 0, + "the write did not close partitions early before the error" + ); + assert_eq!( + parquet_files_under(dir.path()), + Vec::::new(), + "the failed write left files behind" + ); + } + } + /// A task that fails mid-write drops its open files without closing them, so a file's /// share has to come back either way. #[tokio::test] diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 83149ced77c..fedad6b3f92 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -126,15 +126,18 @@ case class CometIcebergWriteExec( // Names mirror Spark's stock `BatchWriteHelper` metrics (`numFiles` / `numOutputRows` / // `numOutputBytes`) so the Spark SQL UI shows the same row as it would for a non-Comet - // Iceberg write. `write_time` is pushed from the native operator by name (see - // `iceberg_write.rs`). + // Iceberg write. `write_time` and `files_closed_early` are pushed from the native operator by + // name (see `iceberg_write.rs`). override lazy val metrics: Map[String, SQLMetric] = Map( "numFiles" -> SQLMetrics.createMetric(sparkContext, "number of files written"), "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), "numOutputBytes" -> SQLMetrics.createSizeMetric(sparkContext, "written output"), "write_time" -> SQLMetrics.createNanoTimingMetric( sparkContext, - "time in native Iceberg writer")) + "time in native Iceberg writer"), + "files_closed_early" -> SQLMetrics.createMetric( + sparkContext, + "files closed early to free memory")) override def doExecute(): RDD[InternalRow] = { val columnarRdd = doExecuteColumnar() diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 5f19ac20430..489a5db6689 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -3051,21 +3051,21 @@ class CometIcebergWriteActionSuite } } - // The native writer reserves what its open files hold from the task's memory pool. A fanout - // write keeps a file open for every partition, so over enough partitions it outgrows the pool - // and fails its task, which Spark can retry, instead of growing in memory that no budget - // accounts for until the executor is killed. - test("native acceleration: a fanout write that outgrows the memory pool fails its task") { + // The native writer reserves what its partitions hold from the task's memory pool. A fanout + // write keeps every partition it has seen open, so over enough partitions it outgrows the pool. + // It then writes out and closes the partitions holding the most, whose next rows open new + // files, and finishes with more, smaller files where it would otherwise fail its task. + test("native acceleration: a fanout write that outgrows the memory pool closes partitions") { assumeNativeAcceleration() withIcebergCatalog { _ => val session = spark import session.implicits._ - // One task writes 64 partitions of 1000 random 500-byte strings, a run of rows per - // partition. Each open file holds about 500 KB, so the pool of about 4 MiB (0.002 of the - // suite's 2 GiB off-heap size) is outgrown long before all 64 files are open. + // One task writes 64 partitions of 1000 random 500-byte strings, the partitions taking + // turns, so every partition keeps getting rows. That is about 32 MB, and the pool of about + // 4 MiB (0.002 of the suite's 2 GiB off-heap size) holds an eighth of it. val rows = 64000 (0 until rows) - .map(i => (i, s"r${i / 1000}", new scala.util.Random(i).alphanumeric.take(500).mkString)) + .map(i => (i, s"r${i % 64}", new scala.util.Random(i).alphanumeric.take(500).mkString)) .toDF("id", "region", "payload") .coalesce(1) .createOrReplaceTempView("fanout_oom_src") @@ -3082,10 +3082,69 @@ class CometIcebergWriteActionSuite } withNativeEnabled { + // withSQLConf returns Unit before Spark 4.0, so the assertions run inside it. withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { - val (failedPlans, error) = captureFailedPlans(spark) { + val snapshot = captureWrite("fanout_oom") { spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom SELECT * FROM fanout_oom_src") } + assert( + snapshot.snapshotDelta == 1L, + s"expected 1 commit, got ${snapshot.snapshotDelta}") + val writers = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case w: CometIcebergWriteExec => w } + } + assert( + writers.nonEmpty, + s"the write did not run natively:\n${snapshot.plans.mkString("\n--\n")}") + assert( + writers.map(_.metrics("files_closed_early").value).sum > 0, + "the write closed no partition early") + } + + // With the whole pool, every partition keeps its one file open. + spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom_control SELECT * FROM fanout_oom_src") + assert(parquetFiles(dataDir("fanout_oom_control")).size == 64) + assert( + parquetFiles(dataDir("fanout_oom")).size > 64, + s"${parquetFiles(dataDir("fanout_oom")).size} files for 64 partitions") + // Compared with the source rather than the control, which is native too and could lose + // the same rows. + val written = spark.sql(s"SELECT id, region, payload FROM $catalog.$ns.fanout_oom") + val source = spark.table("fanout_oom_src") + assert(written.count() == rows) + assert(written.exceptAll(source).isEmpty && source.exceptAll(written).isEmpty) + } + } + } + + // A write that keeps one file open has no partition to close early, so when the file outgrows + // the pool the task fails, which Spark can retry, instead of growing in memory that no budget + // accounts for until the executor is killed. + test("native acceleration: a write whose open file outgrows the memory pool fails its task") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + val session = spark + import session.implicits._ + // One task writes about 32 MB of random 500-byte strings into one file, whose row group + // outgrows the pool of about 4 MiB (0.002 of the suite's 2 GiB off-heap size). + val rows = 64000 + (0 until rows) + .map(i => (i, "r", new scala.util.Random(i).alphanumeric.take(500).mkString)) + .toDF("id", "region", "payload") + .coalesce(1) + .createOrReplaceTempView("file_oom_src") + Seq("file_oom", "file_oom_control").foreach { table => + spark.sql(s""" + CREATE TABLE $catalog.$ns.$table (id INT, region STRING, payload STRING) + USING iceberg + """) + } + + withNativeEnabled { + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { + val (failedPlans, error) = captureFailedPlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.file_oom SELECT * FROM file_oom_src") + } assert( error.toSeq .flatMap(exceptionChain) @@ -3098,23 +3157,22 @@ class CometIcebergWriteActionSuite collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), s"the failed write did not run natively:\n${failedPlans.mkString("\n--\n")}") } - assert(countSnapshots("fanout_oom") == 0L, "the failed write must not commit") + assert(countSnapshots("file_oom") == 0L, "the failed write must not commit") assert( - parquetFiles(dataDir("fanout_oom")).isEmpty, - s"the failed task left data files behind: ${parquetFiles(dataDir("fanout_oom"))}") + parquetFiles(dataDir("file_oom")).isEmpty, + s"the failed task left data files behind: ${parquetFiles(dataDir("file_oom"))}") // The executor outlived the failure, and with the whole pool the same write succeeds. val controlPlans = capturePlans(spark) { - spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom_control SELECT * FROM fanout_oom_src") + spark.sql(s"INSERT INTO $catalog.$ns.file_oom_control SELECT * FROM file_oom_src") } assert( controlPlans.exists(p => collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), "the control write did not run natively") assert( - spark.sql(s"SELECT count(*) FROM $catalog.$ns.fanout_oom_control").head().getLong(0) + spark.sql(s"SELECT count(*) FROM $catalog.$ns.file_oom_control").head().getLong(0) == rows) - assert(parquetFiles(dataDir("fanout_oom_control")).size == 64) } } } From daf33bd61b44d80e3be7f81f55367ed8c4d23645 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 19:22:55 +0000 Subject: [PATCH 87/89] fix: skip an array or map function's later arguments where the array or map is NULL (#6716) * fix: skip an array or map function's later arguments where the array or map is NULL Spark's array_contains, array_position, array_remove, arrays_overlap, array_union, slice, m[k] and a[i] are null-intolerant BinaryExpressions or TernaryExpressions, which return NULL for a NULL array or map without evaluating the other arguments. The native functions evaluated every argument over the whole batch, so an argument that raises, such as an ANSI cast of a malformed string, failed a query that Spark completes, and a nondeterministic argument saw rows Spark never gives it. Add a native NullShortCircuit wrapper that evaluates the first argument over the whole batch and each later one only for the rows where no earlier argument is NULL, as ArrayInsert already does for its own arguments. A null_short_circuit flag on ScalarFunc and ListExtract enables it, set through QueryPlanSerde.withNullShortCircuit. The planner leaves a function unwrapped when every later argument can be evaluated for any row. Closes #6613. * Simplify NullShortCircuit and skip masking for infallible arguments Track the skipped rows with NullBuffer::union and collect an argument's columns with collect_columns. Keep each argument's batch field on the argument itself instead of in a parallel list. Evaluate a later argument that can be evaluated for any row over the whole batch, as literals already were, and build a projection only for the arguments that are masked. Leave a column whose index does not change unrewritten. Point the guide and the proto comment at withNullShortCircuit instead of repeating its conditions. * Evaluate a shared subexpression for every row, as Spark's generated code does Spark's subexpression elimination evaluates a subexpression that an operator's expressions share for every row, before the expressions that use it, in projections and aggregations. A shared ANSI cast therefore raises in Spark even on rows whose array is NULL. Before serializing an operator, ask Spark's EquivalentExpressions which subexpressions its expressions share, and leave unmasked any call whose later argument holds one, so it raises there as well. * Evaluate a masked argument through one copy, so shuffle keeps its state A masked argument was evaluated through its rebuilt copy, which reads its columns from a batch of just those columns, only for a batch with a NULL in an earlier argument, and through the original for every other batch. Rebuilding an argument whose column indices change gives it new state, so array_union(a, shuffle(b)) restarted shuffle's random sequence each time evaluation moved between the two, and drew different permutations from Spark's. The rebuilt copy now evaluates a masked argument for every batch. A unit test evaluates array_union(a, shuffle(b)) over a batch without a NULL array, one with, and one without again, against the unwrapped function, which draws the same permutations as Spark when b is NULL wherever a is. A Spark 4.0+ SQL fixture reads 8200 rows from one file, so that only the scan's second batch has NULL arrays. Both fail with the previous evaluation. * fix: count LEGACY and TRY string-to-integer casts as infallible, and find shared subexpressions lazily Cast::is_infallible now accepts a cast from Utf8 or LargeUtf8 to an integer in LEGACY or TRY mode, which returns NULL for a string it cannot parse. Comet's parsers for those modes never raise. With ANSI off, the native planner then leaves array_contains(ids, CAST(s AS INT)) unwrapped, and CASE keeps its eager path for such a cast. withSharedSubexpressions now runs EquivalentExpressions on the operator's first canShortCircuitNulls call instead of for every converted operator. canShortCircuitNulls skips it entirely when every later argument is a leaf, which subexpression elimination never shares. array_null_short_circuit_non_ansi.sql runs the array shapes with ANSI off. The casts return NULL there, while a slice start of 0 and a rand argument still depend on the short-circuit. The shuffle fixture pins spark.comet.batchSize, and the map fixture covers map_from_arrays, which its CASE WHEN guards already protect. * perf: keep CASE lazy for string-parsing casts, and give NullShortCircuit its own check Counting a LEGACY or TRY string-to-integer cast as infallible also put CASE and IF on their eager path, which parses every row for each branch that some row chooses. With a selective predicate, that made IF about 8x slower and a three-branch CASE about 3x slower. Cast::is_infallible keeps its original set: casts that cannot fail and cost little more than reading their input. Cast::cannot_fail adds the LEGACY and TRY string-to-integer casts. case_when.rs exposes both checks. CASE uses is_infallible, so it chooses the path it chose before. NullShortCircuit uses cannot_fail, since it masks an argument only to keep it from failing on rows that Spark skips. The conditional bench gains IF and CASE shapes whose branches parse strings. * fix: find shared subexpressions where shuffle keys and TakeOrderedAndProject are serialized The native shuffle writer serializes its partitioning keys, and CometTakeOrderedAndProjectExec its sort keys and projection, when the plan runs, outside the operator conversion that sets the shared subexpressions. canShortCircuitNulls took the missing set as an empty one, so a cast that two partitioning keys share was skipped on the rows whose array is NULL, where Spark's projection evaluates it for every row and raises. canShortCircuitNulls now skips no argument outside withSharedSubexpressions. The shuffle writer calls it with the partitioning keys, which Spark evaluates together in one projection, and CometTakeOrderedAndProjectExec with its own expressions. * fix: scope a TakeOrderedAndProject's sort keys and projection apart, as Spark evaluates them 3090d859a gave CometTakeOrderedAndProjectExec's sort keys and projection one shared set. Spark's ordering evaluates each sort key on its own, without subexpression elimination, and its projection runs separately over the rows the sort keeps. So a cast in two sort keys, or in a sort key and the projection, raised in Comet where Spark returns the rows. The top-K plans now use an empty set, and the projection its projection list. --- .../adding_a_new_expression.md | 4 + native/core/src/execution/jni_api.rs | 1 + native/core/src/execution/planner.rs | 20 +- native/proto/src/proto/expr.proto | 6 + native/spark-expr/benches/conditional.rs | 60 +- .../src/conditional_funcs/case_when.rs | 68 +- .../spark-expr/src/conditional_funcs/mod.rs | 2 + .../conditional_funcs/null_short_circuit.rs | 733 ++++++++++++++++++ .../spark-expr/src/conversion_funcs/cast.rs | 19 +- .../apache/comet/rules/CometExecRule.scala | 33 +- .../apache/comet/serde/QueryPlanSerde.scala | 67 ++ .../scala/org/apache/comet/serde/arrays.scala | 24 +- .../scala/org/apache/comet/serde/maps.scala | 4 +- .../CometTakeOrderedAndProjectExec.scala | 21 +- .../shuffle/CometNativeShuffleWriter.scala | 17 +- .../array/array_null_short_circuit.sql | 125 +++ .../array_null_short_circuit_non_ansi.sql | 66 ++ .../array_null_short_circuit_shuffle.sql | 38 + .../map/map_null_short_circuit.sql | 63 ++ .../comet/serde/QueryPlanSerdeSuite.scala | 23 +- 20 files changed, 1329 insertions(+), 65 deletions(-) create mode 100644 native/spark-expr/src/conditional_funcs/null_short_circuit.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index 20bcae06e69..3ce1b4e4732 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -152,6 +152,10 @@ object CometLevenshtein extends CometExpressionSerde[Levenshtein] { When the return type is set on the proto, the native planner skips the registry lookup entirely and routes straight to the Comet UDF registered in `create_comet_physical_fun_with_eval_mode`. +#### Arguments that Spark skips after a NULL + +A null-intolerant `BinaryExpression` or `TernaryExpression` returns NULL as soon as one of its arguments is NULL, without evaluating the arguments after it. Native execution evaluates every argument of a `ScalarFunc` over the whole batch, so an argument that can fail, such as an ANSI cast of a malformed string, would fail on rows that Spark skips. Wrap the serialized call of such an expression in `withNullShortCircuit`, as `CometArrayContains` and `CometSlice` do, after checking the conditions in its Scaladoc. + #### Registering the Expression Handler Once you've created your `CometExpressionSerde` implementation, register it in `QueryPlanSerde.scala` by adding it to the appropriate expression map (e.g., `mathExpressions`, `stringExpressions`, `predicateExpressions`, etc.): diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index f8e9e497f6f..262d8b8fd41 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -2784,6 +2784,7 @@ mod tests { }], return_type: None, fail_on_error: false, + null_short_circuit: false, })), query_context: None, expr_id: None, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ceca9e657af..15a5c48acb3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -146,8 +146,8 @@ use datafusion_comet_spark_expr::{ ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, - NormalizeNestedFloats, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal, - ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, + NormalizeNestedFloats, NullShortCircuit, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, + SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -968,7 +968,7 @@ impl PhysicalPlanner { .as_ref() .map(|e| self.create_expr(e, Arc::clone(&input_schema))) .transpose()?; - Ok(Arc::new(ListExtract::new( + let list_extract: Arc = Arc::new(ListExtract::new( child, ordinal, default_value, @@ -976,7 +976,12 @@ impl PhysicalPlanner { expr.fail_on_error, spark_expr.expr_id, Arc::clone(&self.query_context_registry), - ))) + )); + if expr.null_short_circuit { + Ok(NullShortCircuit::wrap(list_extract, &input_schema)?) + } else { + Ok(list_extract) + } } ExprStruct::GetArrayStructFields(expr) => { let child = @@ -4055,6 +4060,11 @@ impl PhysicalPlanner { Arc::new(Field::new(fun_name, data_type.clone(), true)), Arc::new(ConfigOptions::default()), )); + let scalar_expr = if expr.null_short_circuit { + NullShortCircuit::wrap(scalar_expr, &input_schema)? + } else { + scalar_expr + }; // DF53 changed some UDFs (e.g. md5) to return StringViewArray at execution // time (apache/datafusion#20045). Comet does not yet support view types, so @@ -6866,6 +6876,7 @@ mod tests { args: vec![array_col, array_col_1], return_type: None, fail_on_error: false, + null_short_circuit: false, })), query_context: None, expr_id: None, @@ -6992,6 +7003,7 @@ mod tests { args: vec![array_col, array_col_1], return_type: None, fail_on_error: false, + null_short_circuit: false, })), query_context: None, expr_id: None, diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index b85f08f6531..50d47160125 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -558,6 +558,10 @@ message ScalarFunc { repeated Expr args = 2; DataType return_type = 3; bool fail_on_error = 4; + // Evaluate each argument after the first only for the rows where no earlier argument is NULL, + // as Spark does for a null-intolerant BinaryExpression or TernaryExpression. See + // QueryPlanSerde.withNullShortCircuit. + bool null_short_circuit = 5; } message CaseWhen { @@ -631,6 +635,8 @@ message ListExtract { Expr default_value = 3; bool one_based = 4; bool fail_on_error = 5; + // Evaluate the ordinal only for the rows where the child is not NULL, as in ScalarFunc + bool null_short_circuit = 6; } message GetArrayStructFields { diff --git a/native/spark-expr/benches/conditional.rs b/native/spark-expr/benches/conditional.rs index 2febcddee99..8ded4762409 100644 --- a/native/spark-expr/benches/conditional.rs +++ b/native/spark-expr/benches/conditional.rs @@ -45,9 +45,10 @@ type Expr = Arc; type Shape<'a> = (&'a str, fn() -> Expr, &'a RecordBatch); /// Columns shaped like the ones in `CometConditionalExpressionBenchmark`: -/// `c1` a random long, `c2` an int in `0..100`, `c3` another random long, and `c4` / `c5` short -/// and `c6` / `c7` long strings. `null_density` applies to every column. With `sorted`, `c1` and -/// `c2` ascend through the batch, so every predicate over them selects one contiguous run of rows. +/// `c1` a random long, `c2` an int in `0..100`, `c3` another random long, `c4` / `c5` short +/// and `c6` / `c7` long strings, and `c8` / `c9` integers written as strings. `null_density` +/// applies to every column. With `sorted`, `c1` and `c2` ascend through the batch, so every +/// predicate over them selects one contiguous run of rows. fn make_batch(null_density: f32, sorted: bool) -> RecordBatch { let mut rng = StdRng::seed_from_u64(42); let mut c1: Vec = (0..NUM_ROWS).map(|_| rng.random::()).collect(); @@ -57,7 +58,7 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch { c2.sort_unstable(); } let c3: Vec = (0..NUM_ROWS).map(|_| rng.random::()).collect(); - let valid: Vec> = (0..7) + let valid: Vec> = (0..9) .map(|_| { (0..NUM_ROWS) .map(|_| rng.random::() >= null_density) @@ -84,6 +85,11 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch { .map(|i| valid[col][i].then(|| format!("{tag}{i}-").repeat(12))) .collect() }; + let numeric = |col: usize, seed: usize| -> StringArray { + (0..NUM_ROWS) + .map(|i| valid[col][i].then(|| format!("{}", (i * 7919 + seed) % 2_000_000))) + .collect() + }; let columns: Vec = vec![ Arc::new(c1), Arc::new(c2), @@ -92,6 +98,8 @@ fn make_batch(null_density: f32, sorted: bool) -> RecordBatch { Arc::new(short(4, "t")), Arc::new(long(5, "long value ")), Arc::new(long(6, "other value ")), + Arc::new(numeric(7, 17)), + Arc::new(numeric(8, 4242)), ]; RecordBatch::try_new(Arc::new(schema()), columns).unwrap() } @@ -105,6 +113,8 @@ fn schema() -> Schema { Field::new("c5", DataType::Utf8, true), Field::new("c6", DataType::Utf8, true), Field::new("c7", DataType::Utf8, true), + Field::new("c8", DataType::Utf8, true), + Field::new("c9", DataType::Utf8, true), ]) } @@ -309,13 +319,45 @@ fn case_fallible_branch() -> Expr { ) } +/// A string-to-integer cast with ANSI off, which returns NULL for a string it cannot parse. +fn parse_int(column: &str) -> Expr { + spark_cast(col(column), DataType::Int32) +} + +/// A selective predicate whose branch parses a string: 5% of the rows choose it. +fn if_selective_parse() -> Expr { + if_expr(c2_below(5), parse_int("c8"), lit(ScalarValue::Int32(None))) +} + +/// A predicate that most rows match, with a branch that parses a string. +fn if_mostly_parse() -> Expr { + if_expr(c2_below(95), parse_int("c8"), lit(ScalarValue::Int32(None))) +} + +/// Both branches parse a string, and half of the rows choose each. +fn if_parse_either() -> Expr { + if_expr(c2_below(50), parse_int("c8"), parse_int("c9")) +} + +/// Three selective branches, each parsing a string for 2% of the rows. +fn case_selective_parse_3_branches() -> Expr { + case_when( + vec![ + (c2_below(2), parse_int("c8")), + (c2_below(4), parse_int("c9")), + (c2_below(6), parse_int("c8")), + ], + None, + ) +} + fn criterion_benchmark(c: &mut Criterion) { let no_nulls = make_batch(0.0, false); let sparse_nulls = make_batch(0.1, false); let dense_nulls = make_batch(0.9, false); let sorted = make_batch(0.0, true); - let shapes: [Shape; 26] = [ + let shapes: [Shape; 30] = [ ( "case literal 3 branches", case_literal_3_branches, @@ -374,6 +416,14 @@ fn criterion_benchmark(c: &mut Criterion) { case_column_10_branches, &sorted, ), + ("if selective parse", if_selective_parse, &no_nulls), + ("if mostly parse", if_mostly_parse, &no_nulls), + ("if parse either", if_parse_either, &no_nulls), + ( + "case selective parse 3 branches", + case_selective_parse_3_branches, + &no_nulls, + ), ]; let mut group = c.benchmark_group("conditional"); group.throughput(Throughput::Elements(NUM_ROWS as u64)); diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index 192e6046f61..c5478bb1d23 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -389,12 +389,27 @@ impl PhysicalExpr for CaseWhenExpr { } } -/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for. +/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at little cost. /// -/// That holds when it can neither fail nor return something different for seeing more rows: a -/// column, a literal, and comparisons, boolean logic, null checks, widening casts and wrapping -/// arithmetic over them. Anything else, including every function, is assumed to be able to fail. -fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { +/// That holds when it can neither fail nor return something different for seeing more rows, and is +/// cheap: a column, a literal, and comparisons, boolean logic, null checks, widening casts and +/// wrapping arithmetic over them. Anything else, including every function and a cast that parses a +/// string, is treated as able to fail. `CASE` uses this to decide whether to evaluate every branch +/// over the whole batch, which pays off only when that costs little more than selecting rows. +pub(super) fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { + infallible(expr, input_schema, false) +} + +/// Whether `expr` can be evaluated for rows that Spark would not evaluate it for, at any cost: what +/// [`is_infallible`] accepts, and casts that [`Cast::cannot_fail`] accepts beyond it, such as a +/// LEGACY or TRY cast from a string to an integer, which returns NULL for a string it cannot parse. +/// `NullShortCircuit` uses this, since it masks an argument only to keep it from failing on rows +/// that Spark skips. +pub(super) fn cannot_fail(expr: &Arc, input_schema: &Schema) -> bool { + infallible(expr, input_schema, true) +} + +fn infallible(expr: &Arc, input_schema: &Schema, parsing: bool) -> bool { if expr.is::() || expr.is::() { return true; } @@ -403,7 +418,11 @@ fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { } else if let Some(comparison) = expr.downcast_ref::() { comparison.is_infallible(input_schema) } else if let Some(cast) = expr.downcast_ref::() { - cast.is_infallible(input_schema) + if parsing { + cast.cannot_fail(input_schema) + } else { + cast.is_infallible(input_schema) + } } else if let Some(normalize) = expr.downcast_ref::() { normalize.is_infallible() } else if let Some(not) = expr.downcast_ref::() { @@ -421,7 +440,7 @@ fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { && expr .children() .into_iter() - .all(|child| is_infallible(child, input_schema)) + .all(|child| infallible(child, input_schema, parsing)) } fn binary_is_infallible(binary: &BinaryExpr, input_schema: &Schema) -> bool { @@ -1151,6 +1170,7 @@ mod tests { Field::new("d", DataType::Float64, true), Field::new("e", DataType::Float64, true), Field::new("s", DataType::Utf8, true), + Field::new("ls", DataType::LargeUtf8, true), Field::new("dec", DataType::Decimal128(10, 2), true), Field::new("l", list_of(DataType::Float64), true), Field::new("m", list_of(DataType::Float64), true), @@ -1164,15 +1184,17 @@ mod tests { ), ]); let c = |name: &str| col(name, &schema).unwrap(); - let cast = |e: Arc, to: DataType| -> Arc { - Arc::new(Cast::new( + let cast_in = |mode: EvalMode, e: Arc, to: DataType| { + let cast: Arc = Arc::new(Cast::new( e, to, - SparkCastOptions::new_without_timezone(EvalMode::Ansi), + SparkCastOptions::new_without_timezone(mode), None, None, - )) + )); + cast }; + let cast = |e: Arc, to: DataType| cast_in(EvalMode::Ansi, e, to); let infallible = |e: Arc| is_infallible(&e, &schema); // A comparison as the planner builds it, which follows Spark's ordering for floats let compare = |left: &str, op: Operator, right: Arc| { @@ -1229,9 +1251,31 @@ mod tests { assert!(!infallible(checked)); // Decimal arithmetic can overflow assert!(!infallible(binary(c("dec"), Operator::Plus, c("dec")))); - // Narrowing and parsing casts can fail + // Narrowing and parsing casts can fail under ANSI assert!(!infallible(cast(c("i"), DataType::Int32))); assert!(!infallible(cast(c("s"), DataType::Int32))); + // A string that LEGACY or TRY mode cannot parse as an integer becomes NULL, so the cast + // cannot fail. Parsing costs too much for CASE to do it for rows it would skip, though. + for mode in [EvalMode::Legacy, EvalMode::Try] { + for to in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + ] { + for column in ["s", "ls"] { + let parse = cast_in(mode, c(column), to.clone()); + assert!(!infallible(Arc::clone(&parse)), "{mode:?} {column} {to}"); + assert!(cannot_fail(&parse, &schema), "{mode:?} {column} {to}"); + } + } + // Other string casts are not claimed + assert!( + !cannot_fail(&cast_in(mode, c("s"), DataType::Float64), &schema), + "{mode:?}" + ); + } + assert!(!cannot_fail(&cast(c("s"), DataType::Int32), &schema)); // So can anything built from a part that can fail assert!(!infallible(binary( binary(c("i"), Operator::Divide, c("i")), diff --git a/native/spark-expr/src/conditional_funcs/mod.rs b/native/spark-expr/src/conditional_funcs/mod.rs index dc13347b3a1..36d854aecca 100644 --- a/native/spark-expr/src/conditional_funcs/mod.rs +++ b/native/spark-expr/src/conditional_funcs/mod.rs @@ -17,6 +17,8 @@ mod case_when; mod if_expr; +mod null_short_circuit; pub use case_when::{create_case_when, create_if_expr, CaseWhenExpr}; pub use if_expr::IfExpr; +pub use null_short_circuit::NullShortCircuit; diff --git a/native/spark-expr/src/conditional_funcs/null_short_circuit.rs b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs new file mode 100644 index 00000000000..232b9725cf6 --- /dev/null +++ b/native/spark-expr/src/conditional_funcs/null_short_circuit.rs @@ -0,0 +1,733 @@ +// 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 super::case_when::cannot_fail; +use arrow::array::{new_null_array, Array, BooleanArray}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{internal_datafusion_err, internal_err, Result}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::{Column, Literal}; +use datafusion::physical_expr::utils::collect_columns; +use datafusion::physical_expr::PhysicalExpr; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +/// Evaluates a function's arguments the way Spark evaluates those of a null-intolerant +/// `BinaryExpression` or `TernaryExpression`: the first for every row, and each later one only for +/// the rows where no earlier argument is NULL. +/// +/// Spark returns NULL as soon as an argument is NULL, without evaluating the ones after it. +/// DataFusion evaluates every argument over the whole batch before it calls a function, so an +/// argument that can fail, such as an ANSI cast of a malformed string, would fail on a row that +/// Spark answers with NULL, and a nondeterministic argument would see rows that it does not see in +/// Spark. An argument that can be evaluated for any row, such as a column, is evaluated for all of +/// them. +/// +/// The function still sees every row, with NULL for each argument skipped there, so it has to +/// return NULL wherever an argument other than the last is NULL. Its arguments have to be its +/// children, in the order Spark evaluates them, and it has to compute its result from them alone. +/// `ArrayInsert` evaluates its own arguments the same way. +#[derive(Debug)] +pub struct NullShortCircuit { + /// The function over its arguments + inner: Arc, + /// The function over the evaluated arguments, reading each one that is not a literal from a + /// column of the batch that holds them + body: Arc, + args: Vec, +} + +impl NullShortCircuit { + /// Wraps `inner` if evaluating one of its arguments after the first for a row that Spark skips + /// could change the outcome, and returns it unchanged otherwise. + pub fn wrap( + inner: Arc, + input_schema: &Schema, + ) -> Result> { + let masked = inner + .children() + .iter() + .enumerate() + .map(|(i, arg)| i > 0 && !cannot_fail(arg, input_schema)) + .collect::>(); + if !masked.contains(&true) { + return Ok(inner); + } + let plans = inner + .children() + .into_iter() + .zip(masked) + .map(|(arg, masked)| { + if arg.is::() { + return Ok((None, masked)); + } + // A masked argument is NULL on the rows it is skipped for + let field = arg.return_field(input_schema)?; + Ok(( + Some(Arc::new(field.as_ref().clone().with_nullable(true))), + masked, + )) + }) + .collect::>>()?; + Ok(Arc::new(Self::try_new(inner, plans)?)) + } + + /// `plans` holds each argument's [`Arg::field`] and whether it is masked. + fn try_new(inner: Arc, plans: Vec<(Option, bool)>) -> Result { + let children = inner.children(); + if children.len() != plans.len() { + return internal_err!( + "NullShortCircuit has {} plans for {} arguments", + plans.len(), + children.len() + ); + } + let args = children + .into_iter() + .zip(plans) + .map(|(expr, (field, masked))| Arg::try_new(expr, field, masked)) + .collect::>>()?; + let mut body_children = Vec::with_capacity(args.len()); + let mut column = 0; + for arg in &args { + if let Some(field) = &arg.field { + body_children.push(Arc::new(Column::new(field.name(), column)) as _); + column += 1; + } else { + body_children.push(Arc::clone(&arg.expr)); + } + } + let body = Arc::clone(&inner).with_new_children(body_children)?; + Ok(Self { inner, body, args }) + } +} + +impl Hash for NullShortCircuit { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl PartialEq for NullShortCircuit { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +impl Eq for NullShortCircuit {} + +impl Display for NullShortCircuit { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "NullShortCircuit({})", self.inner) + } +} + +impl PhysicalExpr for NullShortCircuit { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.inner.nullable(input_schema) + } + + fn return_field(&self, input_schema: &Schema) -> Result { + self.inner.return_field(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let num_rows = batch.num_rows(); + // The rows where an argument evaluated so far is NULL, or `None` while there are none + let mut nulls: Option = None; + let mut fields = vec![]; + let mut columns = vec![]; + let last = self.args.len() - 1; + for (i, arg) in self.args.iter().enumerate() { + let value = arg.evaluate(batch, nulls.as_ref())?; + if i < last { + let value_nulls = match &value { + ColumnarValue::Array(array) if array.len() != num_rows => { + return internal_err!( + "{} returned {} rows for a batch of {num_rows}", + arg.expr, + array.len() + ); + } + ColumnarValue::Array(array) => array.logical_nulls(), + ColumnarValue::Scalar(scalar) if scalar.is_null() => { + Some(NullBuffer::new_null(num_rows)) + } + ColumnarValue::Scalar(_) => None, + }; + nulls = NullBuffer::union(nulls.as_ref(), value_nulls.as_ref()) + .filter(|n| n.null_count() > 0); + if nulls.as_ref().is_some_and(|n| n.null_count() == num_rows) { + // Every row is NULL, and no later argument is evaluated for any of them + let data_type = self.inner.data_type(batch.schema_ref())?; + return Ok(ColumnarValue::Array(new_null_array(&data_type, num_rows))); + } + } + if let Some(field) = &arg.field { + let column = value.into_array(num_rows)?; + // The batch is checked against the types the arguments have at run time, which + // can differ from their planned types in the names and nullability of nested fields + fields.push(if field.data_type() == column.data_type() { + Arc::clone(field) + } else { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(column.data_type().clone()), + ) + }); + columns.push(column); + } + } + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + let args = + RecordBatch::try_new_with_options(Arc::new(Schema::new(fields)), columns, &options)?; + self.body.evaluate(&args) + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let inner = Arc::clone(&self.inner).with_new_children(children)?; + let plans = self + .args + .iter() + .map(|arg| (arg.field.clone(), arg.masked.is_some())) + .collect(); + Ok(Arc::new(Self::try_new(inner, plans)?)) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.inner.fmt_sql(f) + } + + fn is_volatile_node(&self) -> bool { + self.inner.is_volatile_node() + } +} + +/// An argument, and how it is evaluated. +#[derive(Debug)] +struct Arg { + expr: Arc, + /// The field of the column that `body` reads the argument from, or `None` for a literal, + /// which `body` keeps + field: Option, + /// For an argument evaluated for only the rows where no earlier argument is NULL, the input + /// columns it reads, and the argument reading them from a batch of just those columns, so + /// that skipping rows filters no other column. It is evaluated in place of `expr` for every + /// batch, even one with no NULL, because rebuilding the argument gave it its own copy of any + /// state kept across batches, such as the random generator of `shuffle` + masked: Option<(Vec, Arc)>, +} + +impl Arg { + fn try_new( + expr: &Arc, + field: Option, + masked: bool, + ) -> Result { + let masked = masked.then(|| project(expr)).transpose()?; + Ok(Self { + expr: Arc::clone(expr), + field, + masked, + }) + } + + /// Evaluates the argument, for just the rows outside `nulls` if it is masked, returning a + /// value for every row of the batch. + fn evaluate(&self, batch: &RecordBatch, nulls: Option<&NullBuffer>) -> Result { + let Some((projection, projected)) = &self.masked else { + return self.expr.evaluate(batch); + }; + let batch = batch.project(projection)?; + match nulls { + Some(nulls) => { + let rows = BooleanArray::new(nulls.inner().clone(), None); + projected.evaluate_selection(&batch, &rows) + } + None => projected.evaluate(&batch), + } + } +} + +/// The input columns that `expr` reads, and `expr` reading them from a batch of just those columns. +fn project(expr: &Arc) -> Result<(Vec, Arc)> { + let mut projection = collect_columns(expr) + .iter() + .map(Column::index) + .collect::>(); + projection.sort_unstable(); + projection.dedup(); + let projected = Arc::clone(expr) + .transform_down(|e| { + let Some(column) = e.downcast_ref::() else { + return Ok(Transformed::no(e)); + }; + let index = projection + .binary_search(&column.index()) + .map_err(|_| internal_datafusion_err!("{column} is not projected"))?; + if index == column.index() { + return Ok(Transformed::no(e)); + } + Ok(Transformed::yes( + Arc::new(Column::new(column.name(), index)) as Arc, + )) + })? + .data; + Ok((projection, projected)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + create_query_context_map, Cast, EvalMode, ListExtract, ShuffleExpr, SparkArrayPositionFunc, + SparkArraySlice, SparkCastOptions, + }; + use arrow::array::{ArrayRef, Int32Array, Int64Array, ListArray, StringArray}; + use arrow::datatypes::{Field, Int32Type}; + use datafusion::common::config::ConfigOptions; + use datafusion::common::ScalarValue; + use datafusion::functions_nested::set_ops::array_union_udf; + use datafusion::logical_expr::ScalarUDF; + use datafusion::physical_expr::expressions::{col, lit, CaseExpr, IsNotNullExpr}; + use datafusion::physical_expr::ScalarFunctionExpr; + use std::sync::Mutex; + + /// An array per row, a string that is not an integer on most rows whose array is NULL, and + /// the integer it is elsewhere. + fn test_batch() -> RecordBatch { + let arrays = ListArray::from_iter_primitive::(vec![ + None, + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3)]), + None, + Some(vec![Some(5), None, Some(6)]), + None, + ]); + let ints = Int32Array::from(vec![7, 2, 1, 7, 6, 7]); + let strings = StringArray::from(vec!["bad", "2", "3", "worse", "3", "x"]); + let lengths = Int64Array::from(vec![Some(1), Some(1), None, Some(1), Some(2), Some(1)]); + let schema = Schema::new(vec![ + Field::new("a", arrays.data_type().clone(), true), + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + Field::new("len", DataType::Int64, true), + ]); + RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(arrays), + Arc::new(ints), + Arc::new(strings), + Arc::new(lengths), + ], + ) + .unwrap() + } + + fn ansi_cast(expr: Arc, to: DataType) -> Arc { + Arc::new(Cast::new( + expr, + to, + SparkCastOptions::new_without_timezone(EvalMode::Ansi), + None, + None, + )) + } + + fn function( + udf: ScalarUDF, + args: Vec>, + return_type: DataType, + ) -> Arc { + let name = udf.name().to_string(); + Arc::new(ScalarFunctionExpr::new( + &name, + Arc::new(udf), + args, + Arc::new(Field::new(&name, return_type, true)), + Arc::new(ConfigOptions::default()), + )) + } + + fn array_position( + array: Arc, + element: Arc, + ) -> Arc { + let udf = ScalarUDF::new_from_impl(SparkArrayPositionFunc::new()); + function(udf, vec![array, element], DataType::Int64) + } + + fn slice(args: Vec>, schema: &Schema) -> Arc { + let udf = ScalarUDF::new_from_impl(SparkArraySlice::new()); + function(udf, args, schema.field(0).data_type().clone()) + } + + fn evaluate(expr: &Arc, batch: &RecordBatch) -> Result { + expr.evaluate(batch)?.into_array(batch.num_rows()) + } + + /// `array_position(a, CAST(s AS INT))` over `test_batch` + fn expected_positions() -> ArrayRef { + Arc::new(Int64Array::from(vec![ + None, + Some(2), + Some(1), + None, + Some(0), + None, + ])) + } + + /// `slice(a, CAST(s AS INT), len)` over `test_batch` + fn expected_slices() -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::(vec![ + None, + Some(vec![Some(2)]), + None, + None, + Some(vec![Some(6)]), + None, + ])) + } + + #[test] + fn later_argument_is_not_evaluated_where_the_array_is_null() { + let batch = test_batch(); + let schema = batch.schema(); + let array = Arc::new(RowRecorder::new(col("a", &schema).unwrap())); + let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + let unguarded = array_position(Arc::clone(&array) as _, element); + // The cast fails on "bad", whose array is NULL + let error = evaluate(&unguarded, &batch).unwrap_err().to_string(); + assert!(error.contains("CAST_INVALID_INPUT"), "{error}"); + array.seen.lock().unwrap().clear(); + + let guarded = NullShortCircuit::wrap(unguarded, &schema).unwrap(); + assert!(guarded.is::()); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected_positions().as_ref()); + // The array is evaluated once, over the whole batch + assert_eq!(*array.seen.lock().unwrap(), vec![(6, 4)]); + // An offset into the batch reaches the arrays' buffers and the filter + let actual = evaluate(&guarded, &batch.slice(1, 4)).unwrap(); + assert_eq!(actual.as_ref(), expected_positions().slice(1, 4).as_ref()); + } + + #[test] + fn later_argument_still_fails_where_the_array_is_not_null() { + let batch = test_batch(); + let schema = batch.schema(); + let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + let guarded = + NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema) + .unwrap(); + let arrays = + ListArray::from_iter_primitive::(vec![None, Some(vec![Some(1)])]); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(arrays), + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(StringArray::from(vec!["bad", "x"])), + Arc::new(Int64Array::from(vec![1, 1])), + ], + ) + .unwrap(); + // "bad" is skipped, but the array of "x" is not NULL + let error = evaluate(&guarded, &batch).unwrap_err().to_string(); + assert!(error.contains("'x'"), "{error}"); + } + + #[test] + fn infallible_arguments_are_not_wrapped() { + let batch = test_batch(); + let schema = batch.schema(); + let column = array_position(col("a", &schema).unwrap(), col("i", &schema).unwrap()); + let wrapped = NullShortCircuit::wrap(Arc::clone(&column), &schema).unwrap(); + assert!(Arc::ptr_eq(&column, &wrapped)); + let widened = slice( + vec![ + col("a", &schema).unwrap(), + ansi_cast(col("i", &schema).unwrap(), DataType::Int64), + lit(1i64), + ], + &schema, + ); + let wrapped = NullShortCircuit::wrap(Arc::clone(&widened), &schema).unwrap(); + assert!(Arc::ptr_eq(&widened, &wrapped)); + // With ANSI off, a string that the cast cannot parse becomes NULL, so it cannot fail either + let parsed = array_position( + col("a", &schema).unwrap(), + Arc::new(Cast::new( + col("s", &schema).unwrap(), + DataType::Int32, + SparkCastOptions::new_without_timezone(EvalMode::Legacy), + None, + None, + )), + ); + let wrapped = NullShortCircuit::wrap(Arc::clone(&parsed), &schema).unwrap(); + assert!(Arc::ptr_eq(&parsed, &wrapped)); + } + + #[test] + fn each_argument_waits_for_every_earlier_one() { + let batch = test_batch(); + let schema = batch.schema(); + // The length is only evaluated where neither the array nor the start is NULL + let start = ansi_cast( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + DataType::Int64, + ); + let length = Arc::new(RowRecorder::new(col("len", &schema).unwrap())); + let args = vec![col("a", &schema).unwrap(), start, Arc::clone(&length) as _]; + let guarded = NullShortCircuit::wrap(slice(args, &schema), &schema).unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected_slices().as_ref()); + // The length saw the three rows that have an array, and only the column it reads + assert_eq!(*length.seen.lock().unwrap(), vec![(3, 1)]); + } + + #[test] + fn infallible_later_argument_is_evaluated_for_every_row() { + let batch = test_batch(); + let schema = batch.schema(); + let start = ansi_cast( + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + DataType::Int64, + ); + let args = vec![ + col("a", &schema).unwrap(), + start, + col("len", &schema).unwrap(), + ]; + let guarded = NullShortCircuit::wrap(slice(args, &schema), &schema).unwrap(); + // Only the start, which can fail, skips the rows whose array is NULL + let masked = guarded + .downcast_ref::() + .unwrap() + .args + .iter() + .map(|arg| arg.masked.is_some()) + .collect::>(); + assert_eq!(masked, vec![false, true, false]); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected_slices().as_ref()); + } + + #[test] + fn null_literal_array_evaluates_no_later_argument() { + let batch = test_batch(); + let schema = batch.schema(); + let null_array = lit(ScalarValue::try_new_null(schema.field(0).data_type()).unwrap()); + let element = Arc::new(RowRecorder::new(ansi_cast( + col("s", &schema).unwrap(), + DataType::Int32, + ))); + let guarded = NullShortCircuit::wrap( + array_position(null_array, Arc::clone(&element) as _), + &schema, + ) + .unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.data_type(), &DataType::Int64); + assert_eq!(actual.null_count(), batch.num_rows()); + assert!(element.seen.lock().unwrap().is_empty()); + } + + #[test] + fn list_extract_is_guarded() { + let batch = test_batch(); + let schema = batch.schema(); + // element_at(a, CAST(s AS INT)), with a NULL for an index past the end + let extract = Arc::new(ListExtract::new( + col("a", &schema).unwrap(), + ansi_cast(col("s", &schema).unwrap(), DataType::Int32), + None, + true, + false, + None, + create_query_context_map(), + )); + let guarded = NullShortCircuit::wrap(extract, &schema).unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + let expected: ArrayRef = Arc::new(Int32Array::from(vec![ + None, + Some(2), + None, + None, + Some(6), + None, + ])); + assert_eq!(actual.as_ref(), expected.as_ref()); + } + + #[test] + fn column_indices_can_be_rewritten() { + let batch = test_batch(); + let schema = batch.schema(); + let element = ansi_cast(col("s", &schema).unwrap(), DataType::Int32); + let guarded = + NullShortCircuit::wrap(array_position(col("a", &schema).unwrap(), element), &schema) + .unwrap(); + // DataFusion's CASE projects the batch down to the columns that it reads, and rewrites + // their indices in its branches through `with_new_children` + let when = Arc::new(IsNotNullExpr::new(col("len", &schema).unwrap())); + let case: Arc = + Arc::new(CaseExpr::try_new(None, vec![(when, guarded)], None).unwrap()); + let actual = evaluate(&case, &batch).unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![ + None, + Some(2), + None, + None, + Some(0), + None, + ])); + assert_eq!(actual.as_ref(), expected.as_ref()); + } + + #[test] + fn stateful_argument_keeps_its_state_across_batches() { + let list = DataType::new_list(DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list.clone(), true), + Field::new("b", list.clone(), true), + ])); + // array_union(a, shuffle(b)). The masked shuffle reads `b` from the first column of a + // batch of just that column, so it is rebuilt, with a generator of its own. + let union = || { + let shuffle = Arc::new(ShuffleExpr::new(col("b", &schema).unwrap(), 42)); + let udf = array_union_udf().as_ref().clone(); + function(udf, vec![col("a", &schema).unwrap(), shuffle], list.clone()) + }; + // An empty array and [1, 2, 3, 4] on each row that is not NULL. A shuffle draws nothing + // for a NULL row, so the unwrapped function, which evaluates it for the rows whose array + // is NULL as well, draws the same permutations as Spark. + let batch = |valid: &[bool]| { + let four = || (1..=4).map(Some).collect::>(); + let a = valid.iter().map(|&v| v.then(Vec::>::new)); + let b = valid.iter().map(|&v| v.then(four)); + let columns: Vec = vec![ + Arc::new(ListArray::from_iter_primitive::(a)), + Arc::new(ListArray::from_iter_primitive::(b)), + ]; + RecordBatch::try_new(Arc::clone(&schema), columns).unwrap() + }; + let unguarded = union(); + let guarded = NullShortCircuit::wrap(union(), &schema).unwrap(); + assert!(guarded.is::()); + // The permutations carry on from one batch to the next, whether or not it has a NULL + for valid in [vec![true, true], vec![false, true, true], vec![true, true]] { + let batch = batch(&valid); + let expected = evaluate(&unguarded, &batch).unwrap(); + let actual = evaluate(&guarded, &batch).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + } + } + + /// Evaluates `child`, recording the rows and columns of each batch that it is evaluated for. + #[derive(Debug)] + struct RowRecorder { + child: Arc, + /// Shared with the copies that `with_new_children` makes + seen: Arc>>, + } + + impl RowRecorder { + fn new(child: Arc) -> Self { + Self { + child, + seen: Arc::new(Mutex::new(vec![])), + } + } + } + + impl Hash for RowRecorder { + fn hash(&self, state: &mut H) { + self.child.hash(state); + } + } + + impl PartialEq for RowRecorder { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) + } + } + + impl Eq for RowRecorder {} + + impl Display for RowRecorder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RowRecorder({})", self.child) + } + } + + impl PhysicalExpr for RowRecorder { + fn data_type(&self, input_schema: &Schema) -> Result { + self.child.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.child.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + self.seen + .lock() + .unwrap() + .push((batch.num_rows(), batch.num_columns())); + self.child.evaluate(batch) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + child: Arc::clone(&children[0]), + seen: Arc::clone(&self.seen), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.child.fmt_sql(f) + } + } +} diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 30b3fe997a7..884573d01ef 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -119,8 +119,9 @@ impl Cast { } } - /// Whether the cast cannot fail for any value, in any eval mode: a cast to the child's own - /// type, from NULL, or to a wider number. `cast_array` rejects a dictionary target first. + /// Whether the cast cannot fail for any value, in any eval mode, and costs little more than + /// reading its input: a cast to the child's own type, from NULL, or to a wider number. + /// `cast_array` rejects a dictionary target first. pub(crate) fn is_infallible(&self, input_schema: &Schema) -> bool { use DataType::*; let Ok(from) = self.child.data_type(input_schema) else { @@ -138,6 +139,20 @@ impl Cast { | (Float32, Float64) )) } + + /// Whether the cast cannot fail for any value: what [`Self::is_infallible`] accepts, and a cast + /// from a string to an integer in LEGACY or TRY mode, which returns NULL for a string it cannot + /// parse. Parsing costs far more per row than those casts. + pub(crate) fn cannot_fail(&self, input_schema: &Schema) -> bool { + use DataType::*; + self.is_infallible(input_schema) + || (self + .child + .data_type(input_schema) + .is_ok_and(|from| matches!(from, Utf8 | LargeUtf8)) + && matches!(self.data_type, Int8 | Int16 | Int32 | Int64) + && self.cast_options.eval_mode != EvalMode::Ansi) + } } /// Spark cast options diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 32faea60e65..5a6e761d44b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1134,24 +1134,21 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } val builder = OperatorOuterClass.Operator.newBuilder().setPlanId(op.id) - if (op.children.nonEmpty && op.children.forall(_.isInstanceOf[CometNativeExec])) { - val childOp = op.children.map(_.asInstanceOf[CometNativeExec].nativeOp) - childOp.foreach(builder.addChildren) - return serde - .convert(op, builder, childOp: _*) - .map { nativeOp => - val exec = serde.createExec(nativeOp, op) - rollUpInfoMessages(op, exec) - exec - } - } else { - return serde - .convert(op, builder) - .map { nativeOp => - val exec = serde.createExec(nativeOp, op) - rollUpInfoMessages(op, exec) - exec - } + // Which subexpressions Spark evaluates for every row decides where native execution may + // skip a function's arguments (see `QueryPlanSerde.canShortCircuitNulls`) + val nativeOp = QueryPlanSerde.withSharedSubexpressions(op.expressions) { + if (op.children.nonEmpty && op.children.forall(_.isInstanceOf[CometNativeExec])) { + val childOp = op.children.map(_.asInstanceOf[CometNativeExec].nativeOp) + childOp.foreach(builder.addChildren) + serde.convert(op, builder, childOp: _*) + } else { + serde.convert(op, builder) + } + } + return nativeOp.map { nativeOp => + val exec = serde.createExec(nativeOp, op) + rollUpInfoMessages(op, exec) + exec } } None diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 43960ad5604..56cb421ae84 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicLong import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ +import scala.util.DynamicVariable import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions._ @@ -1297,6 +1298,72 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { Some(ExprOuterClass.Expr.newBuilder().setScalarFunc(builder).build()) } + /** + * The subexpressions that occur more than once in the expressions of the operator being + * serialized, as Spark's subexpression elimination finds them. Spark's generated code for a + * projection or an aggregation evaluates each of them for every row, before the expressions + * that use them. They are treated that way for every operator, and whatever + * `spark.sql.subexpressionElimination.enabled` says, which only keeps them evaluated for every + * row where Spark might not. They are found on the first `canShortCircuitNulls` call for the + * operator, so an operator without an argument to skip never runs `EquivalentExpressions`. + * `None` outside `withSharedSubexpressions`, where they are unknown. + */ + private val sharedSubexpressions = new DynamicVariable[Option[SharedSubexpressions]](None) + + private final class SharedSubexpressions(expressions: Seq[Expression]) { + lazy val common: Seq[Expression] = { + val equivalence = new EquivalentExpressions + expressions.foreach(equivalence.addExprTree(_)) + equivalence.getCommonSubexpressions + } + } + + /** + * Runs `f`, which serializes an operator whose expressions are `expressions`. Code that + * serializes expressions outside an operator's conversion, such as when the plan runs, calls it + * with the expressions Spark evaluates together, or no argument is skipped. + */ + def withSharedSubexpressions[T](expressions: Seq[Expression])(f: => T): T = + sharedSubexpressions.withValue(Some(new SharedSubexpressions(expressions)))(f) + + /** + * Whether native execution may skip the arguments of `expr` after the first on the rows where + * an earlier argument is NULL, as Spark's evaluation of a null-intolerant `BinaryExpression` or + * `TernaryExpression` does. Not if one of them holds a subexpression that the operator's + * expressions share: Spark evaluates that for every row, and raises if it fails on a row whose + * array is NULL. Nor outside `withSharedSubexpressions`, where what they share is unknown. + */ + def canShortCircuitNulls(expr: Expression): Boolean = { + val later = expr.children.tail + // Subexpression elimination never shares a leaf, such as a column or a literal + later.forall(_.isInstanceOf[LeafExpression]) || sharedSubexpressions.value.exists { s => + !later.exists(_.exists(e => s.common.exists(_.semanticEquals(e)))) + } + } + + /** + * Has native execution evaluate the arguments of `scalarFunc`, the serialized `expr`, the way + * Spark evaluates those of a null-intolerant `BinaryExpression` or `TernaryExpression`: each + * argument after the first only for the rows where no earlier argument is NULL, unless + * `canShortCircuitNulls` rules it out. Otherwise every argument is evaluated over the whole + * batch, so one that can fail, such as an ANSI cast of a malformed string, fails on a row that + * Spark returns NULL for without evaluating it (#6613). The arguments must be in the order + * Spark evaluates them, and the function must return NULL wherever an argument other than the + * last is NULL. The native planner drops the guard when every later argument can be evaluated + * for any row, such as a column or a literal. + */ + def withNullShortCircuit(expr: Expression, scalarFunc: Option[Expr]): Option[Expr] = + scalarFunc.map { proto => + require(proto.hasScalarFunc, s"not a scalar function: $proto") + if (canShortCircuitNulls(expr)) { + proto.toBuilder + .setScalarFunc(proto.getScalarFunc.toBuilder.setNullShortCircuit(true)) + .build() + } else { + proto + } + } + /** * If `handler` is a `CodegenDispatchFallback`, run `expr` through the JVM codegen dispatcher * and return `Some((handler, proto))` on success; otherwise return `None`. Shared by the diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index a16d42c94da..aa17b6ef53b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -56,7 +56,7 @@ object CometArrayRemove } else { "array_remove_all" } - scalarFunctionExprToProto(function, arrayExprProto, keyExprProto) + withNullShortCircuit(expr, scalarFunctionExprToProto(function, arrayExprProto, keyExprProto)) } } @@ -184,7 +184,7 @@ object CometArrayContains } else { "array_contains" } - scalarFunctionExprToProto(function, arrayExprProto, keyExprProto) + withNullShortCircuit(expr, scalarFunctionExprToProto(function, arrayExprProto, keyExprProto)) } } @@ -376,7 +376,7 @@ object CometArraysOverlap false, leftArrayExprProto, rightArrayExprProto) - arraysOverlapScalarExpr + withNullShortCircuit(expr, arraysOverlapScalarExpr) } } @@ -583,11 +583,13 @@ object CometSlice extends CometExpressionSerde[Slice] { // output, so only `return_field_from_args` is guaranteed to match. Spark's `expr.dataType` is // not: `CometCreateArray` may have widened the input to a deeply-nullable element type, and // DataFusion list elements are nullable where Spark's `containsNull` says otherwise. - scalarFunctionExprToProto( - "spark_array_slice", - arrayExprProto, - startExprProto, - lengthExprProto) + withNullShortCircuit( + expr, + scalarFunctionExprToProto( + "spark_array_slice", + arrayExprProto, + startExprProto, + lengthExprProto)) } } @@ -674,7 +676,7 @@ object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { val arraysUnionScalarExpr = scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto) - arraysUnionScalarExpr + withNullShortCircuit(expr, arraysUnionScalarExpr) } } @@ -734,6 +736,8 @@ object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] { .setOrdinal(ordinalExpr.get) .setOneBased(false) .setFailOnError(expr.failOnError) + // Spark evaluates the ordinal only where the array is not NULL, see withNullShortCircuit + .setNullShortCircuit(canShortCircuitNulls(expr)) Some( ExprOuterClass.Expr @@ -1007,7 +1011,7 @@ object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with Array // (matching Spark's behavior) val optExpr = scalarFunctionExprToProto("spark_array_position", arrayExprProto, elementExprProto) - optExpr + withNullShortCircuit(expr, optExpr) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/maps.scala b/spark/src/main/scala/org/apache/comet/serde/maps.scala index f9d3e2dbc00..ae32794d7fb 100644 --- a/spark/src/main/scala/org/apache/comet/serde/maps.scala +++ b/spark/src/main/scala/org/apache/comet/serde/maps.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, hasNonDefaultStringCollation, scalarFunctionExprToProto, withNullShortCircuit} import org.apache.comet.shims.CometTypeShim /** @@ -128,7 +128,7 @@ object CometMapExtract extends CometExpressionSerde[GetMapValue] { val mapExpr = exprToProtoInternal(expr.child, inputs, binding) val keyExpr = exprToProtoInternal(expr.key, inputs, binding) val mapExtractExpr = scalarFunctionExprToProto("map_extract", mapExpr, keyExpr) - mapExtractExpr + withNullShortCircuit(expr, mapExtractExpr) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala index 88a797d123d..af78e16e825 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.isCometShuffleEnabled import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported} -import org.apache.comet.serde.QueryPlanSerde.{exprToProto, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, supportedSortType, withSharedSubexpressions} import org.apache.comet.serde.operator.CometSink object CometTakeOrderedAndProjectExec extends CometSink[TakeOrderedAndProjectExec] { @@ -135,11 +135,21 @@ case class CometTakeOrderedAndProjectExec( // physical limit includes the offset; the native limit applies that offset once. CometExecUtils.getLimitNativePlan(child.output, limit, offset) } else { - CometExecUtils.getTopKNativePlan(child.output, sortOrder, child, limit, offset) + topKNativePlan(offset) + } + // Spark's projection runs on its own over the rows the sort keeps, so only its list can share + // a subexpression (see QueryPlanSerde.canShortCircuitNulls) + withSharedSubexpressions(projectList) { + selection.flatMap(CometExecUtils.getProjectionNativePlan(projectList, child.output, _)) } - selection.flatMap(CometExecUtils.getProjectionNativePlan(projectList, child.output, _)) } + // Spark's ordering evaluates each sort key on its own, without subexpression elimination + private def topKNativePlan(offset: Int): Option[OperatorOuterClass.Operator] = + withSharedSubexpressions(Seq.empty) { + CometExecUtils.getTopKNativePlan(child.output, sortOrder, child, limit, offset) + } + protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { val childRDD = child.executeColumnar() if (childRDD.getNumPartitions == 0 || limit == 0) { @@ -157,10 +167,7 @@ case class CometTakeOrderedAndProjectExec( } else { val numParts = childRDD.getNumPartitions // Serialize the plan once before mapping to avoid repeated serialization per partition - val topK = - CometExecUtils - .getTopKNativePlan(child.output, sortOrder, child, limit) - .get + val topK = topKNativePlan(offset = 0).get val serializedTopK = CometExec.serializeNativePlan(topK) val numOutputCols = child.output.length val inputSchema = CometUtils.fromAttributes(child.output) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index fcd85b6f6a1..b7264d07047 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -343,8 +343,13 @@ class CometNativeShuffleWriter[K, V]( val partitioning = PartitioningOuterClass.HashPartition.newBuilder() partitioning.setNumPartitions(effectivePartitionCount) - val partitionExprs = hashPartitioning.expressions - .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + // Spark evaluates the keys in one projection, whose subexpression elimination sees what + // they share (see QueryPlanSerde.canShortCircuitNulls) + val partitionExprs = + QueryPlanSerde.withSharedSubexpressions(hashPartitioning.expressions) { + hashPartitioning.expressions + .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + } if (partitionExprs.length != hashPartitioning.expressions.length) { throw new UnsupportedOperationException( @@ -377,8 +382,12 @@ class CometNativeShuffleWriter[K, V]( } { - val orderingExprs = rangePartitioning.ordering - .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + // As for hash partitioning, Spark evaluates the keys in one projection + val orderingExprs = + QueryPlanSerde.withSharedSubexpressions(rangePartitioning.ordering.map(_.child)) { + rangePartitioning.ordering + .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + } if (orderingExprs.length != rangePartitioning.ordering.length) { throw new UnsupportedOperationException( s"Partitioning $rangePartitioning is not supported.") diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql new file mode 100644 index 00000000000..559574fc38f --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit.sql @@ -0,0 +1,125 @@ +-- 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. + +-- These array functions are null-intolerant BinaryExpressions or TernaryExpressions in Spark, +-- which return NULL for a NULL array without evaluating the other arguments. Comet evaluates each +-- later argument only for the rows where every earlier one is not NULL, so an argument that raises, +-- here an ANSI cast of a malformed string, cannot raise on a row Spark skips. Each query names the +-- expression with expect_native, because the codegen dispatcher would pass by running Spark's code. +-- https://github.com/apache/datafusion-comet/issues/6613 + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE array_null_short_circuit(ai array, s string, i string) USING parquet + +-- `s` and `i` are malformed only where the array is NULL. COALESCE(1) writes a single file, so the +-- NULL arrays share a batch with the others. +statement +INSERT INTO array_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES + (NULL, 'bad', 'bad'), + (array(1, 2), '2', '1'), + (array(3), '7', '0'), + (NULL, 'worse', 'worse'), + (array(4, NULL, 5), '5', '2') + AS v(ai, s, i) + +query expect_native(array_contains) +SELECT array_contains(ai, CAST(s AS INT)) FROM array_null_short_circuit + +query expect_native(array_position) +SELECT array_position(ai, CAST(s AS INT)) FROM array_null_short_circuit + +query expect_native(array_remove) +SELECT array_remove(ai, CAST(s AS INT)) FROM array_null_short_circuit + +query expect_native(arrays_overlap) +SELECT arrays_overlap(ai, array(CAST(s AS INT))) FROM array_null_short_circuit + +query expect_native(array_union) +SELECT array_union(ai, array(CAST(s AS INT))) FROM array_null_short_circuit + +-- slice is a TernaryExpression: the length is evaluated only where the start is not NULL either +query expect_native(slice) +SELECT slice(ai, CAST(s AS INT), 1), slice(ai, 1, CAST(i AS INT)) FROM array_null_short_circuit + +query expect_native(getarrayitem) +SELECT ai[CAST(i AS INT)] FROM array_null_short_circuit + +-- A start of 0 makes slice raise in any mode, and only rows whose array is NULL have one +query expect_native(array_union) +SELECT array_union(ai, slice(array(1), IF(ai IS NULL, 0, 1), 1)) FROM array_null_short_circuit + +-- A nondeterministic argument sees the same rows as in Spark, so it returns the same values +query expect_native(array_position) +SELECT array_position(ai, CAST(rand(7L) * 3 AS INT)) FROM array_null_short_circuit + +-- Spark's subexpression elimination evaluates a subexpression that an operator's expressions share +-- for every row, before them, so this cast raises on the rows whose array is NULL as well, in a +-- projection and in an aggregation. Comet does not skip an argument that holds one. +query expect_error(CAST_INVALID_INPUT) +SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)) +FROM array_null_short_circuit + +query expect_error(CAST_INVALID_INPUT) +SELECT sum(array_position(ai, CAST(s AS INT))), max(array_contains(ai, CAST(s AS INT))) +FROM array_null_short_circuit + +-- Here the whole call is shared, and Spark evaluates its cast only where the array is not NULL +query expect_native(array_contains) +SELECT array_contains(ai, CAST(s AS INT)), NOT array_contains(ai, CAST(s AS INT)) +FROM array_null_short_circuit + +-- A shuffle evaluates its partitioning keys together in one projection, and a TakeOrderedAndProject +-- its projection list, so a cast that two of them share raises on the rows whose array is NULL too. +-- Both are serialized when the plan runs, after the operator's conversion. +query expect_error(CAST_INVALID_INPUT) +SELECT ai, s FROM array_null_short_circuit +DISTRIBUTE BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)) + +query expect_error(CAST_INVALID_INPUT) +SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)) +FROM (SELECT * FROM array_null_short_circuit ORDER BY i LIMIT 10) + +-- A cast that only one of them holds is skipped, for hash and range partitioning keys and for a +-- TakeOrderedAndProject +query +SELECT ai, s FROM array_null_short_circuit DISTRIBUTE BY array_contains(ai, CAST(s AS INT)) + +query +SELECT ai, s FROM array_null_short_circuit ORDER BY array_contains(ai, CAST(s AS INT)), s + +query expect_native(array_contains) +SELECT array_contains(ai, CAST(s AS INT)) +FROM (SELECT * FROM array_null_short_circuit ORDER BY i LIMIT 10) + +-- A TakeOrderedAndProject compares rows one sort key at a time and evaluates its projection apart +-- from the keys, so a cast in two sort keys, or in a sort key and the projection, is skipped too +query expect_native(array_contains, array_position) +SELECT i FROM (SELECT * FROM array_null_short_circuit + ORDER BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)), i LIMIT 10) + +query expect_native(array_contains, array_position) +SELECT array_contains(ai, CAST(s AS INT)) +FROM (SELECT * FROM array_null_short_circuit ORDER BY array_position(ai, CAST(s AS INT)), i LIMIT 10) + +-- Where the array is not NULL, the argument is evaluated and raises as in Spark +query expect_error(CAST_INVALID_INPUT) +SELECT array_contains(ai, CAST(s || 'x' AS INT)) FROM array_null_short_circuit + +query expect_error(CAST_INVALID_INPUT) +SELECT ai[CAST(i || 'x' AS INT)] FROM array_null_short_circuit diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql new file mode 100644 index 00000000000..1f557326d4d --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_non_ansi.sql @@ -0,0 +1,66 @@ +-- 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. + +-- The queries of array_null_short_circuit.sql with ANSI off. A cast of a malformed string then +-- returns NULL instead of raising, so the cast queries match Spark however the arguments are +-- evaluated, and the native planner drops the guard for them. A start of 0 still makes slice raise, +-- and a nondeterministic argument still advances on every row it sees, so those two shapes depend on +-- skipping the rows whose array is NULL in every mode. +-- https://github.com/apache/datafusion-comet/issues/6613 + +-- Config: spark.sql.ansi.enabled=false + +statement +CREATE TABLE array_null_short_circuit_non_ansi(ai array, s string, i string) USING parquet + +statement +INSERT INTO array_null_short_circuit_non_ansi SELECT /*+ COALESCE(1) */ * FROM VALUES + (NULL, 'bad', 'bad'), + (array(1, 2), '2', '1'), + (array(3), '7', '0'), + (NULL, 'worse', 'worse'), + (array(4, NULL, 5), '5', '2') + AS v(ai, s, i) + +query expect_native(array_contains) +SELECT array_contains(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi + +query expect_native(array_position) +SELECT array_position(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi + +query expect_native(array_remove) +SELECT array_remove(ai, CAST(s AS INT)) FROM array_null_short_circuit_non_ansi + +query expect_native(arrays_overlap) +SELECT arrays_overlap(ai, array(CAST(s AS INT))) FROM array_null_short_circuit_non_ansi + +query expect_native(array_union) +SELECT array_union(ai, array(CAST(s AS INT))) FROM array_null_short_circuit_non_ansi + +query expect_native(slice) +SELECT slice(ai, CAST(s AS INT), 1), slice(ai, 1, CAST(i AS INT)) FROM array_null_short_circuit_non_ansi + +query expect_native(getarrayitem) +SELECT ai[CAST(i AS INT)] FROM array_null_short_circuit_non_ansi + +-- A start of 0 makes slice raise in any mode, and only rows whose array is NULL have one +query expect_native(array_union) +SELECT array_union(ai, slice(array(1), IF(ai IS NULL, 0, 1), 1)) FROM array_null_short_circuit_non_ansi + +-- A nondeterministic argument sees the same rows as in Spark, so it returns the same values +query expect_native(array_position) +SELECT array_position(ai, CAST(rand(7L) * 3 AS INT)) FROM array_null_short_circuit_non_ansi diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql new file mode 100644 index 00000000000..e5fe87017fb --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_shuffle.sql @@ -0,0 +1,38 @@ +-- 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. + +-- array_union evaluates shuffle(b) only for the rows whose array is not NULL, as in +-- array_null_short_circuit.sql. shuffle keeps one random generator for every batch of a partition, +-- so its permutations have to carry on from a batch with no NULL array to one with some. The two +-- argument shuffle(array, seed) form only exists in Spark 4.0+. +-- https://github.com/apache/datafusion-comet/issues/6613 + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.batchSize=8192 + +statement +CREATE TABLE array_null_short_circuit_shuffle(id bigint, a array, b array) USING parquet + +-- A single file, which the scan reads in batches of 8192 rows (the batch size is pinned above). Only +-- the second batch has NULL arrays. +statement +INSERT INTO array_null_short_circuit_shuffle +SELECT id, IF(id >= 8192 AND id % 2 = 0, NULL, array(0)), array(1, 2, 3, 4) +FROM range(0, 8200, 1, 1) + +query expect_native(array_union, shuffle) +SELECT id, array_union(a, shuffle(b, 42)) FROM array_null_short_circuit_shuffle diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql b/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql new file mode 100644 index 00000000000..6f295f275a0 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit.sql @@ -0,0 +1,63 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Spark returns NULL for a NULL map without evaluating the key, so a key that raises, here an ANSI +-- cast of a malformed string, cannot raise on a row whose map is NULL. map_contains_key is +-- array_contains(map_keys(m), k) after analysis. The array counterparts, and why each query names +-- its expression with expect_native, are in array/array_null_short_circuit.sql. +-- https://github.com/apache/datafusion-comet/issues/6613 + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE map_null_short_circuit(m map, s string) USING parquet + +-- `s` is malformed only where the map is NULL. COALESCE(1) writes a single file, so the NULL maps +-- share a batch with the others. +statement +INSERT INTO map_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES + (NULL, 'bad'), + (map(1, 10, 2, 20), '2'), + (map(3, 30), '4'), + (NULL, 'worse') + AS v(m, s) + +query expect_native(array_contains) +SELECT map_contains_key(m, CAST(s AS INT)) FROM map_null_short_circuit + +query expect_native(getmapvalue) +SELECT m[CAST(s AS INT)] FROM map_null_short_circuit + +-- Where the map is not NULL, the key is evaluated and raises as in Spark +query expect_error(CAST_INVALID_INPUT) +SELECT m[CAST(s || 'x' AS INT)] FROM map_null_short_circuit + +-- map_from_arrays builds a map rather than reading one, but it is null-intolerant too: Spark skips +-- the values where the keys are NULL. Comet guards it with CASE WHEN keys IS NOT NULL instead. +statement +CREATE TABLE map_from_arrays_null_short_circuit(k array, s string) USING parquet + +statement +INSERT INTO map_from_arrays_null_short_circuit SELECT /*+ COALESCE(1) */ * FROM VALUES + (NULL, 'bad'), + (array(1), '10'), + (NULL, 'worse'), + (array(2), '20') + AS v(k, s) + +query +SELECT map_from_arrays(k, array(CAST(s AS INT))) FROM map_from_arrays_null_short_circuit diff --git a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala index 19476f8e13a..8d248669b0a 100644 --- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala @@ -21,10 +21,11 @@ package org.apache.comet.serde import org.scalatest.funsuite.AnyFunSuite +import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayPosition, AttributeReference, Cast} import org.apache.spark.sql.types._ import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} -import org.apache.comet.serde.QueryPlanSerde.supportedDataType +import org.apache.comet.serde.QueryPlanSerde.{canShortCircuitNulls, supportedDataType, withSharedSubexpressions} class QueryPlanSerdeSuite extends AnyFunSuite { @@ -100,4 +101,24 @@ class QueryPlanSerdeSuite extends AnyFunSuite { rejected.foreach(dt => assert(!supports(dt), s"$name should reject $dt")) } } + + test("canShortCircuitNulls skips an argument only where the shared subexpressions are known") { + val array = AttributeReference("ai", ArrayType(IntegerType))() + val value = AttributeReference("i", IntegerType)() + val cast = Cast(AttributeReference("s", StringType)(), IntegerType) + val contains = ArrayContains(array, cast) + val position = ArrayPosition(array, cast) + + // A column is never shared, so it can be skipped anywhere + assert(canShortCircuitNulls(ArrayContains(array, value))) + // Something serialized outside an operator's conversion might share the cast + assert(!canShortCircuitNulls(contains)) + withSharedSubexpressions(Seq(contains)) { + assert(canShortCircuitNulls(contains)) + } + withSharedSubexpressions(Seq(contains, position)) { + assert(!canShortCircuitNulls(contains)) + assert(!canShortCircuitNulls(position)) + } + } } From c937df09c2cbd302772310a85f14ca52f32455c7 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 12 Oct 2026 19:25:01 +0000 Subject: [PATCH 88/89] fix: share the missing task entry retry and stop retrying it after the task has ended (#6813) * refactor: share the missing task entry retry between the shuffle allocator and the native pool #6403 added a retry for Spark's missing task entry (SPARK-59444) to CometUnifiedShuffleMemoryAllocator, and #6310 added its own copy to CometTaskMemoryManager.acquireMemory. Move both into MissingTaskEntryRetry under org.apache.comet. CometTaskMemoryManager passes its task id, which keeps #6310's exact "key not found: " match; the shuffle allocator keeps #6403's prefix match. After the retries run out the native pool still gets a zero grant and the shuffle allocator still throws SparkOutOfMemoryError, and the warning now carries the last exception for both. The fixtures and the waiting-request harness of both suites move to TaskMemoryTestUtils, and CometTaskMemoryManagerSuite keeps one race test next to the refusal and rethrow tests. memory_management.md describes the retry once for both call sites, says what a zero grant does to try_grow and grow, and drops the stale #6304 reference. Co-authored-by: Isaac * fix: don't retry a missing task entry once the task has ended releasePlan drops a plan's stream when the task ends, while a Tokio worker may still wait in Spark's acquireMemory for the memory that stream holds. When the release empties the task's balance, the worker loses the task's entry. Before #6310 it failed at once; with the retry it registered the ended task again and, below its minimum share, waited for other tasks to release memory while holding the TaskMemoryManager monitor that the executor's cleanUpAllAllocatedMemory needs. MissingTaskEntryRetry now refuses without retrying when the task has completed or been killed: CometTaskMemoryManager returns a zero grant and the shuffle allocator throws SparkOutOfMemoryError. Both callers pass the TaskContext they were created in, since native threads have none of their own. Co-authored-by: Isaac --------- Co-authored-by: Isaac --- .../contributor-guide/memory_management.md | 35 ++- .../apache/comet/MissingTaskEntryRetry.java | 162 +++++++++++ .../apache/spark/CometTaskMemoryManager.java | 68 ++--- .../CometUnifiedShuffleMemoryAllocator.java | 66 ++--- .../spark/CometTaskMemoryManagerSuite.scala | 258 ++++++------------ .../apache/spark/TaskMemoryTestUtils.scala | 219 +++++++++++++++ ...etUnifiedShuffleMemoryAllocatorSuite.scala | 218 ++++----------- 7 files changed, 562 insertions(+), 464 deletions(-) create mode 100644 spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java create mode 100644 spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 9c31d3fecb2..184bf0de275 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -137,15 +137,10 @@ allocator before reading anything into the split. returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the same task like any other allocation. A page or pointer array allocation that waits in Spark for -memory can wake up to find the task's entry gone from Spark's execution pool, because another -consumer of the task released its last bytes in the meantime, and Spark then throws a -`NoSuchElementException` ("key not found") instead of a grant -([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). The allocator retries that -case, since the failed call was granted nothing, and after three attempts throws -`SparkOutOfMemoryError`, which the shuffle writers handle like any other refused page. Only the -shuffle allocator's callers are guarded this way. Spark's operators in the same task, such as its -sorters and aggregates, can still hit the exception until Spark re-registers a waiting task in -`ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). +memory can wake up to find the task's entry gone from Spark's execution pool; see +[Constraints on a Comet memory consumer](#constraints-on-a-comet-memory-consumer). The allocator +retries it, and after three attempts throws `SparkOutOfMemoryError`, which the shuffle writers +handle like any other refused page. Which allocator each call site uses, and who ends up charged for the bytes: @@ -216,12 +211,22 @@ that grant, and the only figure available for doing so is the task-wide one abov **A parked acquire can wake up to a missing task entry.** Spark removes the task's `memoryForTask` entry when its balance reaches zero, and an acquire that was waiting in `lock.wait()` and wakes -afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant. -`CometTaskMemoryManager` retries that one case. A missing entry means the task held nothing from -Spark at that moment, so the failed call has no partial grant to reconcile, and the retry registers -the task again and waits for its share as the first call would have. After a few attempts it returns -a zero grant, which the native side treats as a refusal and spills. JVM consumers that call -`allocatePage` directly are not covered; that gap is tracked in #6304. +afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant +([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). Releases take no task monitor, +so another native thread or plan, or a JVM consumer of the same task, can empty the balance while an +acquire waits. Both of Comet's consumers that ask Spark directly, `CometTaskMemoryManager` and the +JVM shuffle allocator, retry that one case through `MissingTaskEntryRetry`. A missing entry means +the task held nothing from Spark at that moment, so the failed call has no partial grant to +reconcile, and the retry registers the task again and waits for its share as the first call would +have. After three attempts `CometTaskMemoryManager` returns a zero grant. `try_grow` treats that as +a refusal and the operator spills, while `grow` records the request as overcommit. The shuffle +allocator throws `SparkOutOfMemoryError` instead. Both refuse without retrying once the task has +completed or been killed: releasing a native plan at the end of the task can empty the balance under +a Tokio worker still waiting in Spark, and a retry would register the ended task again and could wait +for other tasks while holding the `TaskMemoryManager` monitor that Spark's final cleanup needs. +Spark's own operators in the same task, such as +its sorters and aggregates, are not guarded and can still hit the exception until Spark re-registers +a waiting task in `ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). **A consumer whose `spill` returns zero takes budget it can never give back.** `NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim diff --git a/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java new file mode 100644 index 00000000000..ae35a4be923 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java @@ -0,0 +1,162 @@ +/* + * 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. + */ + +package org.apache.comet; + +import java.util.NoSuchElementException; +import java.util.OptionalLong; +import java.util.function.Function; +import java.util.function.Supplier; + +import org.slf4j.Logger; + +import org.apache.spark.TaskContext; + +/** + * Retries a request for execution memory that Spark failed because the task's entry was gone from + * its execution pool. + * + *

Spark registers a task's entry in {@code ExecutionMemoryPool} once, before its wait loop, and + * removes it when the task's balance reaches zero. A request that waits for memory below the task's + * minimum share then fails when it wakes up with a NoSuchElementException ("key not found: " and + * the task id). Releases take no task monitor, so any other consumer of the same task, JVM or + * native, can empty the balance while a request waits. The failed call was granted nothing, because + * Spark only drops the entry at a zero balance, which cannot hold while the call has a partial + * grant. So trying again is safe: it registers the task again and waits for its share as Spark + * would have. Nothing here holds a lock across attempts, so a retry parks in Spark exactly as the + * first attempt did. The Spark issue is SPARK-59444. + * + *

A request of a task that has completed or been killed is refused instead of retried. Releasing + * a native plan when the task ends gives back the plan's memory while a Tokio worker may still wait + * in Spark for it, and when that empties the task's balance the worker loses the entry. Retried, it + * would register the ended task again and could wait for other tasks to release memory, holding the + * TaskMemoryManager monitor that Spark's final cleanup of the task needs. + */ +public final class MissingTaskEntryRetry { + + /** Attempts at a request whose task entry Spark keeps losing. */ + public static final int MAX_ATTEMPTS = 3; + + private MissingTaskEntryRetry() {} + + /** + * Runs {@code request}, a call that asks Spark for {@code required} bytes for task {@code + * taskAttemptId}, and retries it when it lost that task's entry. See {@link #retry(Logger, + * TaskContext, long, Supplier, Function)}. + */ + public static T retry( + Logger logger, + TaskContext context, + long taskAttemptId, + long required, + Supplier request, + Function refuse) { + return retry(logger, context, OptionalLong.of(taskAttemptId), required, request, refuse); + } + + /** + * Runs {@code request}, a call that asks Spark for {@code required} bytes, and retries it when it + * lost the task's entry. Without the task id, any task's missing entry counts. Other exceptions + * are rethrown. After {@link #MAX_ATTEMPTS} attempts that all lost the entry, returns what {@code + * refuse} makes of the last attempt's exception, or throws what it throws. It does the same + * without retrying once {@code context}, the task's context or null when there is none, says the + * task has completed or been killed. Both the retries and the refusal are logged to {@code + * logger}. + */ + public static T retry( + Logger logger, + TaskContext context, + long required, + Supplier request, + Function refuse) { + return retry(logger, context, OptionalLong.empty(), required, request, refuse); + } + + private static T retry( + Logger logger, + TaskContext context, + OptionalLong taskAttemptId, + long required, + Supplier request, + Function refuse) { + for (int attempt = 1; ; attempt++) { + try { + return request.get(); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e, taskAttemptId)) { + throw e; + } + String task = taskAttemptId.isPresent() ? "Task " + taskAttemptId.getAsLong() : "The task"; + if (hasEnded(context)) { + logger.info( + "{} lost its execution memory entry in Spark while waiting to acquire {} bytes after " + + "it ended, refusing the request", + task, + required); + return refuse.apply(e); + } + if (attempt >= MAX_ATTEMPTS) { + logger.warn( + "{} lost its execution memory entry in Spark on {} attempts to acquire {} bytes, " + + "refusing the request", + task, + attempt, + required, + e); + return refuse.apply(e); + } + logger.info( + "{} lost its execution memory entry in Spark while waiting to acquire {} bytes, " + + "trying again", + task, + required); + } + } + } + + /** + * Whether the task has completed or been killed. TaskContextImpl.isCompleted takes the context's + * monitor, but Spark runs the task's completion listeners, which release native plans, without + * holding it. + */ + private static boolean hasEnded(TaskContext context) { + return context != null && (context.isInterrupted() || context.isCompleted()); + } + + /** + * Whether {@code e} is Spark's execution pool failing to find the task's entry: the message names + * the task, or any task when its id is unknown, and the exception comes from the pool. + */ + static boolean isMissingTaskEntry(NoSuchElementException e, OptionalLong taskAttemptId) { + String message = e.getMessage(); + boolean matches = + taskAttemptId.isPresent() + ? ("key not found: " + taskAttemptId.getAsLong()).equals(message) + : message != null && message.startsWith("key not found: "); + if (!matches) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; + } +} diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index a468446bae7..1c76627c5b5 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -20,7 +20,6 @@ package org.apache.spark; import java.io.IOException; -import java.util.NoSuchElementException; import java.util.concurrent.atomic.AtomicLong; import org.slf4j.Logger; @@ -30,6 +29,8 @@ import org.apache.spark.memory.MemoryMode; import org.apache.spark.memory.TaskMemoryManager; +import org.apache.comet.MissingTaskEntryRetry; + /** * A adapter class that is used by Comet native to acquire & release memory through Spark's unified * memory manager. This assumes Spark's off-heap memory mode is enabled. All the native plans in a @@ -39,14 +40,14 @@ public class CometTaskMemoryManager { private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class); - /** Attempts at an acquire whose task entry Spark keeps losing, see acquireFromSpark. */ - private static final int MAX_ACQUIRE_ATTEMPTS = 3; - /** The id of the native plan this memory manager was created for, the first in its task. */ private final long id; private final long taskAttemptId; + /** The task's context, read from native threads, which have none of their own. */ + private final TaskContext taskContext; + public final TaskMemoryManager internal; private final NativeMemoryConsumer nativeMemoryConsumer; private final AtomicLong used = new AtomicLong(); @@ -54,7 +55,8 @@ public class CometTaskMemoryManager { public CometTaskMemoryManager(long id, long taskAttemptId) { this.id = id; this.taskAttemptId = taskAttemptId; - this.internal = TaskContext$.MODULE$.get().taskMemoryManager(); + this.taskContext = TaskContext$.MODULE$.get(); + this.internal = taskContext.taskMemoryManager(); this.nativeMemoryConsumer = new NativeMemoryConsumer(); } @@ -99,54 +101,18 @@ public long acquireMemory(long size) { } /** - * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry. Spark - * removes a task's entry from its execution pool when the task's balance reaches zero, and an - * acquire that was waiting for memory then fails when it wakes up with a NoSuchElementException - * ("key not found: " and the task id). Releases take no task monitor, so another native thread or - * a JVM consumer of the same task can empty the balance while an acquire waits. The failed call - * was granted nothing, because Spark only drops the entry at a zero balance, which cannot hold - * while this call has a partial grant. So trying again is safe: it registers the task again and - * waits for its share as Spark would have. After a few failed attempts the grant is refused. The - * helper holds no lock of its own across attempts, so a retry parks in Spark exactly as the first - * attempt did. + * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry; see + * {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the grant is + * refused, which the native pools treat like any other refusal. */ private long acquireFromSpark(long size) { - for (int attempt = 1; ; attempt++) { - try { - return internal.acquireExecutionMemory(size, nativeMemoryConsumer); - } catch (NoSuchElementException e) { - if (!isMissingTaskEntry(e)) { - throw e; - } - if (attempt >= MAX_ACQUIRE_ATTEMPTS) { - logger.warn( - "Task {} lost its execution memory entry in Spark on {} attempts to acquire {} " - + "bytes, refusing the request", - taskAttemptId, - attempt, - size); - return 0; - } - logger.info( - "Task {} lost its execution memory entry in Spark while waiting to acquire {} bytes, " - + "trying again", - taskAttemptId, - size); - } - } - } - - /** Whether {@code e} is Spark's execution pool failing to find this task's entry. */ - private boolean isMissingTaskEntry(NoSuchElementException e) { - if (!("key not found: " + taskAttemptId).equals(e.getMessage())) { - return false; - } - for (StackTraceElement frame : e.getStackTrace()) { - if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { - return true; - } - } - return false; + return MissingTaskEntryRetry.retry( + logger, + taskContext, + taskAttemptId, + size, + () -> internal.acquireExecutionMemory(size, nativeMemoryConsumer), + e -> 0L); } // Called by Comet native through JNI diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index e3a2269cacc..f7bc13ea152 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -21,12 +21,12 @@ import java.io.IOException; import java.util.Map; -import java.util.NoSuchElementException; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.spark.TaskContext; import org.apache.spark.memory.MemoryConsumer; import org.apache.spark.memory.MemoryMode; import org.apache.spark.memory.SparkOutOfMemoryError; @@ -34,6 +34,8 @@ import org.apache.spark.unsafe.array.LongArray; import org.apache.spark.unsafe.memory.MemoryBlock; +import org.apache.comet.MissingTaskEntryRetry; + /** * A simple memory allocator used by `CometShuffleExternalSorter` to allocate memory blocks which * store serialized rows. This class is simply an implementation of `MemoryConsumer` that delegates @@ -49,11 +51,12 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory private static final Logger logger = LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class); - /** Attempts at an allocation whose task entry Spark keeps losing, see retryMissingTaskEntry. */ - private static final int MAX_ALLOCATE_ATTEMPTS = 3; + /** The task's context, or null outside a task; see retryMissingTaskEntry. */ + private final TaskContext taskContext; CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); + this.taskContext = TaskContext.get(); if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) { throw new IllegalArgumentException( "CometUnifiedShuffleMemoryAllocator should be used with off-heap " @@ -83,31 +86,17 @@ public LongArray allocateArray(long size) { /** * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's - * entry. Spark removes a task's entry from its execution pool when the task's balance reaches - * zero, and an allocation that was waiting for memory then fails when it wakes up with a - * NoSuchElementException ("key not found: " and the task id). Releases take no task monitor, so a - * native consumer or another JVM consumer of the same task can empty the balance while an - * allocation waits. The failed call was granted nothing and took no page, because Spark only - * drops the entry at a zero balance, so trying again is safe: it registers the task again and - * waits for its share as Spark would have. After a few failed attempts the allocation fails with - * a SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused - * by the last attempt's exception. The Spark issue is SPARK-59444. + * entry; see {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the + * allocation fails with a SparkOutOfMemoryError, which the shuffle writers handle like any other + * refused page, caused by the last attempt's exception. */ private T retryMissingTaskEntry(long required, Supplier allocation) { - for (int attempt = 1; ; attempt++) { - try { - return allocation.get(); - } catch (NoSuchElementException e) { - if (!isMissingTaskEntry(e)) { - throw e; - } - if (attempt >= MAX_ALLOCATE_ATTEMPTS) { - logger.warn( - "Lost the task's execution memory entry in Spark on {} attempts to allocate {} " - + "bytes, refusing the allocation", - attempt, - required, - e); + return MissingTaskEntryRetry.retry( + logger, + taskContext, + required, + allocation, + e -> { SparkOutOfMemoryError error = new SparkOutOfMemoryError( "UNABLE_TO_ACQUIRE_MEMORY", @@ -116,30 +105,7 @@ private T retryMissingTaskEntry(long required, Supplier allocation) { "receivedBytes", String.valueOf(0))); error.initCause(e); throw error; - } - logger.info( - "Lost the task's execution memory entry in Spark while waiting to allocate {} bytes, " - + "trying again", - required); - } - } - } - - /** - * Whether {@code e} is Spark's execution pool failing to find the task's entry. The task id is - * not available here, so this matches the message prefix and the pool's stack frame. - */ - private static boolean isMissingTaskEntry(NoSuchElementException e) { - String message = e.getMessage(); - if (message == null || !message.startsWith("key not found: ")) { - return false; - } - for (StackTraceElement frame : e.getStackTrace()) { - if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { - return true; - } - } - return false; + }); } public synchronized long free(MemoryBlock block) { diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index efd899252ac..82c5731da7a 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -19,16 +19,14 @@ package org.apache.spark -import java.util.Properties import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference} +import java.util.concurrent.atomic.AtomicReference -import org.apache.logging.log4j.{Level, LogManager} -import org.apache.logging.log4j.core.{LogEvent, LoggerContext} -import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} +import org.apache.logging.log4j.Level +import org.apache.logging.log4j.core.LogEvent +import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} -class CometTaskMemoryManagerSuite extends SparkFunSuite { +class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { test("native memory usage is visible to Spark's memory consumer") { withTaskMemoryManager { taskMemoryManager => @@ -136,62 +134,80 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } } - /** The events that Comet's and Spark's task memory managers log at `level` or above. */ - private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { - val appender = new LogAppender("task memory manager") - appender.setThreshold(level) - val loggers = Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName) - val context = LogManager.getContext(false).asInstanceOf[LoggerContext] - val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) - try withLogAppender(appender, loggers, Some(level))(f) - finally { - // For a logger with no config of its own, withLogAppender adds one and never removes it. - // It copies the root config's additivity, which is off, so left in place it would keep the - // logger's events out of the test log for the rest of the run. - unconfigured.foreach(context.getConfiguration.removeLogger) - context.updateLoggers() - } - appender.loggingEvents.toSeq - } - test("an acquire waiting in Spark survives a release that empties the task's balance") { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + checkRequestWhileBalanceEmpties(loggers) { taskMemoryManager => val manager = new CometTaskMemoryManager(1L, 0L) - assert(otherTask.acquireMemory(90L) == 90L) - assert(manager.acquireMemory(10L) == 10L) - // Another native thread releases the task's last 10 bytes while the acquire waits. - acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { - manager.releaseMemory(10L) - } + WaitingRequest( + manager.acquireMemory(_), + manager.releaseMemory(_), + () => manager.acquireMemory(20L), + granted => { + assert(granted == 20L) + assert(manager.getUsed == 20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) + manager.releaseMemory(20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + }) } } - test("an acquire waiting in Spark survives a sibling consumer freeing the task's last bytes") { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + for ((state, end) <- Seq[(String, TaskContextImpl => Unit)]( + ("completed", _.markTaskCompleted(None)), + ("was killed", _.markInterrupted("killed")))) { + test(s"an acquire that loses its task entry after the task $state is not made again") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => - val manager = new CometTaskMemoryManager(1L, 0L) - val sibling = new OffHeapConsumer(taskMemoryManager) - assert(otherTask.acquireMemory(90L) == 90L) - assert(sibling.acquireMemory(10L) == 10L) - - // A JVM consumer of the same task frees with freeMemory, which takes no task monitor. - acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { - sibling.freeMemory(10L) + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(otherTask.acquireMemory(90L) == 90L) + assert(manager.acquireMemory(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val acquire = new Thread(() => + try granted.set(manager.acquireMemory(20L)) + catch { case t: Throwable => failure.set(t) }) + acquire.setDaemon(true) + // Spark's executor frees what the task still holds once the task has ended, under the + // task memory manager's monitor, which a parked acquire holds. + val cleanUp = new Thread(() => taskMemoryManager.cleanUpAllAllocatedMemory()) + cleanUp.setDaemon(true) + + try { + acquire.start() + awaitWaitingInSpark(acquire) + // The task ends and releasing its plan gives back the task's last bytes while a native + // thread still waits in Spark, as releasePlan does when it drops a plan's stream. + end(TaskContext.get().asInstanceOf[TaskContextImpl]) + manager.releaseMemory(10L) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + // Parked again, the acquire would wait for the other task while holding the monitor. + assert(!acquire.isAlive, s"the acquire is ${acquire.getState} after the task $state") + cleanUp.start() + cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + assert(!cleanUp.isAlive, s"the clean up is ${cleanUp.getState}") + } finally { + // Free the other task's memory so that neither thread outlives a failed test. + otherTask.freeMemory(otherTask.getUsed) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + + assert(failure.get == null, s"the acquire failed: ${failure.get}") + assert(granted.get == 0L) + assert(manager.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) } } } test("a NoSuchElementException other than Spark's missing task entry is rethrown") { - val wrongTask = missingEntryError("key not found: 7", fromExecutionMemoryPool = true) - val notFromPool = missingEntryError("key not found: 0", fromExecutionMemoryPool = false) + val wrongTask = noSuchElement("key not found: 7", fromExecutionMemoryPool = true) + val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false) for (error <- Seq(wrongTask, notFromPool)) { - val taskMemoryManager = new FailingTaskMemoryManager(error) + val taskMemoryManager = new FailingTaskMemoryManager(() => error, failures = Int.MaxValue) withTaskContext(taskMemoryManager) { _ => val manager = new CometTaskMemoryManager(1L, 0L) val thrown = intercept[NoSuchElementException](manager.acquireMemory(10L)) @@ -203,108 +219,30 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } test("an acquire that keeps losing its task entry gives up with a zero grant") { - val error = missingEntryError("key not found: 0", fromExecutionMemoryPool = true) - val taskMemoryManager = new FailingTaskMemoryManager(error) + // A new exception on every call, so the test can tell which one is logged. + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) withTaskContext(taskMemoryManager) { _ => val manager = new CometTaskMemoryManager(1L, 0L) - assert(manager.acquireMemory(10L) == 0L) - assert(taskMemoryManager.calls.get == maxAcquireAttempts) + val events = logEvents(Level.INFO) { + // A refusal rather than an exception, so native spills as for any refused reservation. + assert(manager.acquireMemory(10L) == 0L) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) + assert(events.last.getThrown eq taskMemoryManager.errors.last) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) assert(manager.getUsed == 0L) assert(nativeMemoryConsumer(manager).getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) } } - private val TimeoutSeconds = 10L - - /** - * The task holds 10 bytes of a 100 byte pool and another task holds 90. An acquire of 20 bytes - * waits in Spark below its minimum share of 25 while `release` takes the task's balance to - * zero, which removes the task's entry from Spark's pool. Once the other task frees its memory - * the acquire must be granted all 20 bytes. - */ - private def acquireWhileBalanceEmpties( - manager: CometTaskMemoryManager, - taskMemoryManager: TaskMemoryManager, - otherTask: OffHeapConsumer)(release: => Unit): Unit = { - val granted = new AtomicLong(-1L) - val failure = new AtomicReference[Throwable]() - val acquire = new Thread(() => - try granted.set(manager.acquireMemory(20L)) - catch { case t: Throwable => failure.set(t) }) - acquire.setDaemon(true) - - try { - acquire.start() - val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) - while (!waitingInSpark(acquire) && System.nanoTime() < deadline) Thread.sleep(10) - assert(waitingInSpark(acquire), s"the acquire is ${acquire.getState}") - release - } finally { - // Free the other task's memory so that the acquire does not outlive a failed test. - otherTask.freeMemory(otherTask.getUsed) - acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) - } - - assert(!acquire.isAlive, s"the acquire is ${acquire.getState}") - assert(failure.get == null, s"the acquire failed: ${failure.get}") - assert(granted.get == 20L) - assert(manager.getUsed == 20L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) - manager.releaseMemory(20L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) - } - - private def waitingInSpark(thread: Thread): Boolean = - thread.getState == Thread.State.WAITING && - thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") - - /** A 100 byte off-heap execution pool. */ - private def offHeapMemoryManager(): UnifiedMemoryManager = { - val conf = new SparkConf() - .set("spark.memory.offHeap.enabled", "true") - .set("spark.memory.offHeap.size", "100") - .set("spark.memory.storageFraction", "0") - new UnifiedMemoryManager(conf, 1000L, 500L, 1) - } - - /** An off-heap consumer that never spills. */ - private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) - extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { - override def spill(size: Long, trigger: MemoryConsumer): Long = 0L - } - - private val MaxFailingCalls = 10 - - /** A task memory manager whose every acquire throws `error`. */ - private class FailingTaskMemoryManager(error: Throwable) - extends TaskMemoryManager(new TestMemoryManager(new SparkConf()), 0L) { - val calls = new AtomicInteger() + private val loggers = + Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName) - override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { - // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. - if (calls.incrementAndGet() > MaxFailingCalls) { - throw new IllegalStateException( - s"acquireExecutionMemory called more than $MaxFailingCalls times") - } - throw error - } - } - - private def missingEntryError( - message: String, - fromExecutionMemoryPool: Boolean): NoSuchElementException = { - val error = new NoSuchElementException(message) - if (fromExecutionMemoryPool) { - val poolFrame = new StackTraceElement( - "org.apache.spark.memory.ExecutionMemoryPool", - "acquireMemory", - "ExecutionMemoryPool.scala", - 115) - error.setStackTrace(poolFrame +: error.getStackTrace) - } - error - } + /** The events that Comet's and Spark's task memory managers log at `level` or above. */ + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = logEvents(level, loggers)(f) private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = { val memoryManager = new TestMemoryManager(new SparkConf()) @@ -312,40 +250,6 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { withTaskContext(new TaskMemoryManager(memoryManager, 0L))(f) } - private def withTaskContext(taskMemoryManager: TaskMemoryManager)( - f: TaskMemoryManager => Unit): Unit = { - val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, - partitionId = 0, - numPartitions = 1, - taskAttemptId = 0L, - attemptNumber = 0, - taskMemoryManager = taskMemoryManager, - localProperties = new Properties, - metricsSystem = null, - taskMetrics = TaskMetrics.empty, - cpus = 1, - resources = Map.empty) - - TaskContext.setTaskContext(taskContext) - try { - f(taskMemoryManager) - } finally { - try { - taskMemoryManager.cleanUpAllAllocatedMemory() - } finally { - TaskContext.unset() - } - } - } - - private def maxAcquireAttempts: Int = { - val field = classOf[CometTaskMemoryManager].getDeclaredField("MAX_ACQUIRE_ATTEMPTS") - field.setAccessible(true) - field.getInt(null) - } - private def nativeMemoryConsumer(manager: CometTaskMemoryManager): MemoryConsumer = { val field = classOf[CometTaskMemoryManager].getDeclaredField("nativeMemoryConsumer") field.setAccessible(true) diff --git a/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala new file mode 100644 index 00000000000..a82e2377e82 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala @@ -0,0 +1,219 @@ +/* + * 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. + */ + +package org.apache.spark + +import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable.ArrayBuffer + +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.{LogEvent, LoggerContext} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, UnifiedMemoryManager} + +import org.apache.comet.MissingTaskEntryRetry + +/** + * Helpers for suites that drive Spark's task memory manager from Comet's memory consumers, among + * them a reproduction of a request that waits in Spark's execution pool and loses the task's + * entry there (SPARK-59444). + */ +trait TaskMemoryTestUtils extends SparkFunSuite { + + protected val TimeoutSeconds = 10L + protected val MaxAcquireAttempts: Int = MissingTaskEntryRetry.MAX_ATTEMPTS + + /** A 100 byte off-heap execution pool. */ + protected def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + protected class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** + * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the + * result of `newError`, recording each one in `errors`. + */ + protected class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) + extends TaskMemoryManager(offHeapMemoryManager(), 0L) { + val calls = new AtomicInteger() + val errors = new ArrayBuffer[Throwable]() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + val call = calls.incrementAndGet() + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (call > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + if (call <= failures) { + val error = newError() + errors += error + throw error + } + super.acquireExecutionMemory(required, consumer) + } + } + + protected def noSuchElement( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + + /** Whether `thread` is parked in Spark's execution pool waiting for memory. */ + protected def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** Waits until `thread` is parked in Spark's execution pool, failing after the timeout. */ + protected def awaitWaitingInSpark(thread: Thread): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(thread), s"the request is ${thread.getState}") + } + + /** + * A request for memory made while the task's balance empties; see + * [[checkRequestWhileBalanceEmpties]]. `hold` and `release` take and give back the task's bytes + * through a consumer other than the one `request` asks through, and `request` returns what it + * was granted. `check` runs on the test thread once the request has returned. + */ + protected case class WaitingRequest( + hold: Long => Long, + release: Long => Unit, + request: () => Long, + check: Long => Unit) + + /** + * The task holds 10 bytes of a 100 byte pool through `hold` and another task holds 90. The + * request is made on another thread and waits in Spark below the task's minimum share of 25 + * while `release` gives back the task's last 10 bytes, which removes the task's entry from + * Spark's pool. Once the other task frees its memory the request must not fail, and `check` + * sees what it was granted. The first attempt lost the task's entry, so `loggers` must have + * logged a single retry at INFO. + */ + protected def checkRequestWhileBalanceEmpties(loggers: Seq[String])( + setup: TaskMemoryManager => WaitingRequest): Unit = { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val waiting = setup(taskMemoryManager) + // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. + assert(otherTask.acquireMemory(90L) == 90L) + assert(waiting.hold(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val thread = new Thread(() => + try granted.set(waiting.request()) + catch { case t: Throwable => failure.set(t) }) + thread.setDaemon(true) + + val events = logEvents(Level.INFO, loggers) { + try { + thread.start() + awaitWaitingInSpark(thread) + waiting.release(10L) + } finally { + // Free the other task's memory so that the request does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + } + + assert(!thread.isAlive, s"the request is ${thread.getState}") + assert(failure.get == null, s"the request failed: ${failure.get}") + waiting.check(granted.get) + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + } + } + + /** The events that `loggers` log at `level` or above while `f` runs. */ + protected def logEvents(level: Level, loggers: Seq[String])(f: => Unit): Seq[LogEvent] = { + val appender = new LogAppender(loggers.mkString(", ")) + appender.setThreshold(level) + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) + try withLogAppender(appender, loggers, Some(level))(f) + finally { + // For a logger with no config of its own, withLogAppender adds one and never removes it. + // It copies the root config's additivity, which is off, so left in place it would keep the + // logger's events out of the test log for the rest of the run. + unconfigured.foreach(context.getConfiguration.removeLogger) + context.updateLoggers() + } + appender.loggingEvents.toSeq + } + + protected def messages(events: Seq[LogEvent]): String = + events.map(_.getMessage.getFormattedMessage).mkString("\n") + + /** Runs `f` as task 0 with `taskMemoryManager`, then frees whatever the task still holds. */ + protected def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { + val taskContext = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = 0L, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + + TaskContext.setTaskContext(taskContext) + try { + f(taskMemoryManager) + } finally { + try { + taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + TaskContext.unset() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala index b978c272802..ecf2e47be21 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala @@ -19,23 +19,14 @@ package org.apache.spark.shuffle.comet -import java.util.Properties -import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} +import org.apache.logging.log4j.Level +import org.apache.logging.log4j.core.LogEvent +import org.apache.spark.{CometTaskMemoryManager, TaskContext, TaskContextImpl, TaskMemoryTestUtils} +import org.apache.spark.memory.SparkOutOfMemoryError -import scala.collection.mutable.ArrayBuffer - -import org.apache.logging.log4j.{Level, LogManager} -import org.apache.logging.log4j.core.{LogEvent, LoggerContext} -import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkFunSuite, TaskContext, TaskContextImpl} -import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, UnifiedMemoryManager} - -class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { +class CometUnifiedShuffleMemoryAllocatorSuite extends TaskMemoryTestUtils { private val PageSize = 20L - private val TimeoutSeconds = 10L - private val MaxAllocateAttempts = 3 test("a page waiting in Spark survives a release that empties the task's balance") { checkWaitingAllocationSurvives(20L)(allocatePage) @@ -79,25 +70,45 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { } assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) assert(events.last.getThrown eq taskMemoryManager.errors.last) - assert(taskMemoryManager.calls.get == MaxAllocateAttempts) - assert(taskMemoryManager.errors.distinct.size == MaxAllocateAttempts) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) + assert(taskMemoryManager.errors.distinct.size == MaxAcquireAttempts) assert(allocator.getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) } } + test("an allocation that loses its task entry after the task ended is not made again") { + for ((allocation, requested) <- allocations.zip(Seq("20", "24"))) { + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) + withTaskContext(taskMemoryManager) { _ => + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None) + val events = logEvents(Level.INFO) { + val thrown = intercept[SparkOutOfMemoryError](allocation(allocator)) + assert(thrown.getMessageParameters.get("requestedBytes") == requested) + assert(thrown.getCause eq taskMemoryManager.errors.last) + } + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + assert(taskMemoryManager.calls.get == 1) + assert(allocator.getUsed == 0L) + } + } + } + test("an allocation that loses its task entry twice is granted on the last attempt") { val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true) for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) { val taskMemoryManager = - new FailingTaskMemoryManager(() => error, failures = MaxAllocateAttempts - 1) + new FailingTaskMemoryManager(() => error, failures = MaxAcquireAttempts - 1) val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) val events = logEvents(Level.INFO) { assert(allocation(allocator) == requested) } assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO), messages(events)) - assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) assert(allocator.getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) @@ -124,162 +135,27 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { Seq(allocatePage _, allocatePointerArray _) /** - * The task holds 10 bytes of a 100 byte pool through a native consumer and another task holds - * 90. `allocate` asks for `requested` bytes on another thread and waits in Spark below the - * task's minimum share of 25 while the native consumer releases the task's last 10 bytes, which - * removes the task's entry from Spark's pool. Once the other task frees its memory the - * allocation must be granted in full. + * An allocation of `requested` bytes waits in Spark while a native consumer of the task + * releases the task's last bytes, and must then be granted in full. */ private def checkWaitingAllocationSurvives(requested: Long)( - allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => - val native = new CometTaskMemoryManager(1L, 0L) - val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) - // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. - assert(otherTask.acquireMemory(90L) == 90L) - assert(native.acquireMemory(10L) == 10L) - - val granted = new AtomicReference[java.lang.Long]() - val failure = new AtomicReference[Throwable]() - val thread = new Thread(() => - try granted.set(allocate(allocator)) - catch { case t: Throwable => failure.set(t) }) - thread.setDaemon(true) - - val events = logEvents(Level.INFO) { - try { - thread.start() - val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) - while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) - assert(waitingInSpark(thread), s"the allocation is ${thread.getState}") - native.releaseMemory(10L) - } finally { - // Free the other task's memory so that the allocation does not outlive a failed test. - otherTask.freeMemory(otherTask.getUsed) - thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) - } - } - - assert(!thread.isAlive, s"the allocation is ${thread.getState}") - assert(failure.get == null, s"the allocation failed: ${failure.get}") - assert(granted.get == requested) - assert(allocator.getUsed == 0L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) - // The first attempt lost the task's entry and the allocator tried again. - assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = + checkRequestWhileBalanceEmpties(Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName)) { + taskMemoryManager => + val native = new CometTaskMemoryManager(1L, 0L) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + WaitingRequest( + native.acquireMemory(_), + native.releaseMemory(_), + () => allocate(allocator), + granted => { + assert(granted == requested) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + }) } - } - - private def waitingInSpark(thread: Thread): Boolean = - thread.getState == Thread.State.WAITING && - thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") - - /** A 100 byte off-heap execution pool. */ - private def offHeapMemoryManager(): UnifiedMemoryManager = { - val conf = new SparkConf() - .set("spark.memory.offHeap.enabled", "true") - .set("spark.memory.offHeap.size", "100") - .set("spark.memory.storageFraction", "0") - new UnifiedMemoryManager(conf, 1000L, 500L, 1) - } - - /** An off-heap consumer that never spills. */ - private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) - extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { - override def spill(size: Long, trigger: MemoryConsumer): Long = 0L - } - - private val MaxFailingCalls = 10 - - /** - * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the - * result of `newError`, recording each one in `errors`. - */ - private class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) - extends TaskMemoryManager(offHeapMemoryManager(), 0L) { - val calls = new AtomicInteger() - val errors = new ArrayBuffer[Throwable]() - - override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { - val call = calls.incrementAndGet() - // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. - if (call > MaxFailingCalls) { - throw new IllegalStateException( - s"acquireExecutionMemory called more than $MaxFailingCalls times") - } - if (call <= failures) { - val error = newError() - errors += error - throw error - } - super.acquireExecutionMemory(required, consumer) - } - } - - private def noSuchElement( - message: String, - fromExecutionMemoryPool: Boolean): NoSuchElementException = { - val error = new NoSuchElementException(message) - if (fromExecutionMemoryPool) { - val poolFrame = new StackTraceElement( - "org.apache.spark.memory.ExecutionMemoryPool", - "acquireMemory", - "ExecutionMemoryPool.scala", - 115) - error.setStackTrace(poolFrame +: error.getStackTrace) - } - error - } /** The events that the allocator logs at `level` or above. */ - private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { - val appender = new LogAppender("unified shuffle memory allocator") - appender.setThreshold(level) - val loggers = Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName) - val context = LogManager.getContext(false).asInstanceOf[LoggerContext] - val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) - try withLogAppender(appender, loggers, Some(level))(f) - finally { - // For a logger with no config of its own, withLogAppender adds one and never removes it. - // It copies the root config's additivity, which is off, so left in place it would keep the - // logger's events out of the test log for the rest of the run. - unconfigured.foreach(context.getConfiguration.removeLogger) - context.updateLoggers() - } - appender.loggingEvents.toSeq - } - - private def messages(events: Seq[LogEvent]): String = - events.map(_.getMessage.getFormattedMessage).mkString("\n") - - private def withTaskContext(taskMemoryManager: TaskMemoryManager)( - f: TaskMemoryManager => Unit): Unit = { - val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, - partitionId = 0, - numPartitions = 1, - taskAttemptId = 0L, - attemptNumber = 0, - taskMemoryManager = taskMemoryManager, - localProperties = new Properties, - metricsSystem = null, - taskMetrics = TaskMetrics.empty, - cpus = 1, - resources = Map.empty) - - TaskContext.setTaskContext(taskContext) - try { - f(taskMemoryManager) - } finally { - try { - taskMemoryManager.cleanUpAllAllocatedMemory() - } finally { - TaskContext.unset() - } - } - } + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = + logEvents(level, Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName))(f) } From 42db91b2e892a28965248c8ba84c941eeb6769e8 Mon Sep 17 00:00:00 2001 From: dustin Date: Mon, 12 Oct 2026 19:30:53 +0000 Subject: [PATCH 89/89] test: expect Spark 3.5's empty LocalTableScan in the single-split DPP join test (#6902) With AQE on, once DPP prunes one shuffled input, AQEPropagateEmptyRelation replaces the join with an empty relation. Spark 4.0 plans that as EmptyRelationExec, which Comet converts to CometEmptyRelationExec, but Spark 3.5 has no EmptyRelation node and plans an empty LocalTableScanExec. Check for the node each version plans. Closes #6901. --- .../apache/comet/CometIcebergNativeSuite.scala | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 4180e3789e6..b00f3353e46 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} -import org.apache.spark.sql.execution.{InSubqueryExec, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec} +import org.apache.spark.sql.execution.{InSubqueryExec, LocalTableScanExec, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, BroadcastQueryStageExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} @@ -5322,10 +5322,21 @@ class CometIcebergNativeSuite if (isSpark35Plus && aqeEnabled && scans.isEmpty) { // Newer Spark versions shuffle the join inputs, allowing AQE to replace // the join and its scans with an empty relation after DPP prunes one input. + // Spark 4.0+ (SPARK-48466) plans it as EmptyRelationExec, which Comet + // converts to CometEmptyRelationExec. Spark 3.5 has no EmptyRelation node + // and plans an empty LocalTableScanExec instead. // This checks result compatibility; the SinglePartition regression assertions // below target Spark 3.4, where both scans must remain in one native block. + val emptyRelations = if (isSpark40Plus) { + collect(plan) { case empty: CometEmptyRelationExec => empty } + } else { + collect(plan) { + case empty: LocalTableScanExec if empty.rows.isEmpty => empty + case empty: CometLocalTableScanExec if empty.rows.isEmpty => empty + } + } assert( - collect(plan) { case empty: CometEmptyRelationExec => empty }.nonEmpty, + emptyRelations.nonEmpty, s"Expected AQE to replace the pruned join with an empty relation:\n$plan") } else { assert(scans.size == 2, s"Expected two native Iceberg scans:\n$plan")