From fbc06a2f2c349bd16ac49f6d48f81a6134ede3b1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 15:27:31 -0600 Subject: [PATCH 01/18] feat: support Iceberg system functions natively Implement Iceberg's `bucket`, `truncate`, `years`, `months`, `days`, and `hours` system functions as native scalar functions and route the `StaticInvoke` calls Spark binds for them through `CometStaticInvoke`, keyed on the Iceberg implementation class names. With these native, the hash distribution and local sort that Iceberg requests in front of a partitioned write stay in Comet, so the native Iceberg writer no longer declines a partitioned table that uses the default `write.distribution-mode`. Filters, projections, and sort keys over hidden-partitioning expressions stay native as well. The kernels match Iceberg's Java implementations exactly: the spec's byte encodings hashed with standard 32-bit Murmur3, Java's wrapping integer arithmetic for truncate, code-point counting for strings, and UTC-only calendar math for the temporal transforms. A truncated decimal that no longer fits its precision becomes null, as it does in Spark. Also name the declaring class in the fallback reason for an unlisted static invoke, since every Iceberg system function is called `invoke`. Closes #5635 --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../user-guide/latest/iceberg-writes.md | 5 +- docs/source/user-guide/latest/iceberg.md | 26 ++ native/spark-expr/src/comet_scalar_funcs.rs | 17 +- native/spark-expr/src/iceberg_funcs/bucket.rs | 376 +++++++++++++++ native/spark-expr/src/iceberg_funcs/mod.rs | 135 ++++++ .../spark-expr/src/iceberg_funcs/temporal.rs | 332 +++++++++++++ .../spark-expr/src/iceberg_funcs/truncate.rs | 404 ++++++++++++++++ native/spark-expr/src/lib.rs | 2 + .../apache/comet/serde/icebergFunctions.scala | 230 +++++++++ .../org/apache/comet/serde/statics.scala | 18 +- .../CometIcebergSystemFunctionSuite.scala | 436 ++++++++++++++++++ 13 files changed, 1979 insertions(+), 4 deletions(-) create mode 100644 native/spark-expr/src/iceberg_funcs/bucket.rs create mode 100644 native/spark-expr/src/iceberg_funcs/mod.rs create mode 100644 native/spark-expr/src/iceberg_funcs/temporal.rs create mode 100644 native/spark-expr/src/iceberg_funcs/truncate.rs create mode 100644 spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala create mode 100644 spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 2426db65fe4..230b434961f 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -359,6 +359,7 @@ jobs: org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite org.apache.comet.CometIcebergWriteDetectionSuite + org.apache.comet.CometIcebergSystemFunctionSuite org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.serde.operator.IcebergWriteProtoTranslationSuite org.apache.comet.csv.CometCsvNativeReadSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 6e4b50a8617..53907e2aef1 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -132,6 +132,7 @@ jobs: org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite org.apache.comet.CometIcebergWriteDetectionSuite + org.apache.comet.CometIcebergSystemFunctionSuite org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.serde.operator.IcebergWriteProtoTranslationSuite org.apache.comet.csv.CometCsvNativeReadSuite diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index db4f82c8ed8..bfee40a1cff 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -128,7 +128,10 @@ per-task Parquet write is delegated to [iceberg-rust](https://github.com/apache/ The native writer must produce the same outcome as iceberg-java — the same Parquet features, statistics, and manifest metadata — so a write is only eligible when every table property it depends on is one the native path reproduces exactly, and additionally only when the plan -feeding the write is fully Comet-native. Ineligible writes run through iceberg-java unchanged, +feeding the write is fully Comet-native. For a partitioned table that plan includes the hash +distribution and local sort Iceberg requests on its partition transforms; those stay native +because the transforms themselves have native implementations (see +[Iceberg system functions](iceberg.md)). Ineligible writes run through iceberg-java unchanged, with the reason reported as a fall-back reason in Comet's extended EXPLAIN output. **Most Iceberg write settings are not supported.** Detection is an allowlist: a write is diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 125105b9a0b..3b57a7efd88 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -196,6 +196,32 @@ the project, exchange, and sort operators around them stay on the Comet path end Spark, which forces a columnar-to-row roundtrip and demotes the surrounding shuffle from `CometExchange` to `CometColumnarExchange`. +### Iceberg system functions + +Iceberg's system functions `bucket`, `truncate`, `years`, `months`, `days`, and `hours` (the SQL +form of its partition transforms, for example `SELECT system.bucket(16, id) FROM t`) run natively. +Spark binds them as static invocations of Iceberg's per-type implementations under +`org.apache.iceberg.spark.functions`, and Comet recognizes those classes wherever the expression +appears: in a projection, a filter, a sort key, or the hash partitioning of a shuffle. + +The native kernels reproduce Iceberg's Java semantics exactly rather than approximately: + +- `bucket` hashes the spec's byte encoding of each value (8-byte little-endian for integers, dates, + and timestamps; UTF-8 for strings; raw bytes for binary; the minimal big-endian two's complement + of the unscaled value for decimals) with 32-bit Murmur3 and masks the sign bit before taking the + modulus. +- `truncate` uses Java's wrapping integer arithmetic, keeps the decimal's precision and scale + (a negative decimal whose truncated value no longer fits the precision becomes null, as it + does in Spark), and counts code points (not bytes) for strings. +- `years`, `months`, `days`, and `hours` are evaluated in UTC regardless of the session timezone + and go negative before the epoch; `days` returns a date, the other three an int. + +This matters most for writes. A partitioned table with the default `write.distribution-mode` +(`hash`) is planned with a shuffle and a local sort keyed on the partition transforms, and with +these functions native the whole sub-plan feeding the [native Iceberg writer](iceberg-writes.md) +stays in Comet. A `numBuckets` or `width` argument that is not a positive integer literal makes +the expression fall back to Spark. + ### Task input metrics The native Iceberg reader populates Spark's task-level `inputMetrics.bytesRead` (visible in the Spark UI Stages tab) using the `bytes_read` counter from iceberg-rust's `ScanMetrics`. This counter includes bytes read from both data files and delete files. diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index c68b998e617..a2225df2d5c 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -28,7 +28,8 @@ use crate::{ spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, spark_round, spark_rpad, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, - SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkMakeDate, SparkMakeInterval, + SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket, + SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, }; use arrow::datatypes::DataType; @@ -293,6 +294,20 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkDateFromUnixDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkDateTrunc::default())), Arc::new(ScalarUDF::new_from_impl(SparkFlatten::default())), + Arc::new(ScalarUDF::new_from_impl(SparkIcebergBucket::default())), + Arc::new(ScalarUDF::new_from_impl(SparkIcebergTruncate::default())), + Arc::new(ScalarUDF::new_from_impl( + SparkIcebergTemporalTransform::years(), + )), + Arc::new(ScalarUDF::new_from_impl( + SparkIcebergTemporalTransform::months(), + )), + Arc::new(ScalarUDF::new_from_impl( + SparkIcebergTemporalTransform::days(), + )), + Arc::new(ScalarUDF::new_from_impl( + SparkIcebergTemporalTransform::hours(), + )), Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), Arc::new(ScalarUDF::new_from_impl(SparkNextDay::default())), diff --git a/native/spark-expr/src/iceberg_funcs/bucket.rs b/native/spark-expr/src/iceberg_funcs/bucket.rs new file mode 100644 index 00000000000..e6b332d03fc --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/bucket.rs @@ -0,0 +1,376 @@ +// 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. + +//! Iceberg's `bucket(numBuckets, value)` transform: +//! `(murmur3_32(bytes(value)) & Integer.MAX_VALUE) % numBuckets`, where `bytes(value)` is the +//! encoding from Appendix B of the Iceberg spec (8-byte little-endian for integers, dates, and +//! timestamps; UTF-8 for strings; raw bytes for binary; minimal big-endian two's complement of +//! the unscaled value for decimals). + +use super::{apply_unary, positive_int_param, unsupported_type}; +use arrow::array::{Array, ArrayRef, AsArray, Int32Array}; +use arrow::compute::cast; +use arrow::datatypes::{ + DataType, Date32Type, Decimal128Type, Int32Type, Int64Type, TimeUnit, TimestampMicrosecondType, +}; +use datafusion::common::{utils::take_function_args, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// 32-bit MurmurHash3 (x86 variant) with seed 0, matching Guava's `Hashing.murmur3_32_fixed()` +/// that Iceberg's `BucketUtil` hashes with. +/// +/// Comet's Spark-compatible murmur3 (`spark_compatible_murmur3_hash`) cannot be reused: Spark +/// mixes the trailing 1 to 3 bytes into the hash one byte at a time, whereas the reference +/// algorithm packs them into a single little-endian word first, so the two disagree on every +/// input whose length is not a multiple of four. +pub(crate) fn murmur3_32(data: &[u8]) -> i32 { + const C1: u32 = 0xcc9e_2d51; + const C2: u32 = 0x1b87_3593; + + #[inline] + fn mix_k1(k1: u32) -> u32 { + k1.wrapping_mul(C1).rotate_left(15).wrapping_mul(C2) + } + + let mut h1: u32 = 0; + let mut chunks = data.chunks_exact(4); + for chunk in &mut chunks { + let k1 = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); + h1 ^= mix_k1(k1); + h1 = h1.rotate_left(13).wrapping_mul(5).wrapping_add(0xe654_6b64); + } + let tail = chunks.remainder(); + if !tail.is_empty() { + let mut k1: u32 = 0; + for (i, byte) in tail.iter().enumerate() { + k1 |= (*byte as u32) << (8 * i); + } + h1 ^= mix_k1(k1); + } + h1 ^= data.len() as u32; + h1 ^= h1 >> 16; + h1 = h1.wrapping_mul(0x85eb_ca6b); + h1 ^= h1 >> 13; + h1 = h1.wrapping_mul(0xc2b2_ae35); + h1 ^= h1 >> 16; + h1 as i32 +} + +/// `BucketUtil.hash(long)`: ints, longs, dates, and timestamps all hash as 8 little-endian bytes. +#[inline] +fn hash_long(value: i64) -> i32 { + murmur3_32(&value.to_le_bytes()) +} + +/// `BucketUtil.hash(BigDecimal)`: hashes `unscaledValue().toByteArray()`, the shortest big-endian +/// two's complement encoding of the unscaled value (at least one byte). +#[inline] +fn hash_decimal(unscaled: i128) -> i32 { + let bytes = unscaled.to_be_bytes(); + let mut start = 0; + while start < bytes.len() - 1 { + let redundant_sign_byte = match bytes[start] { + 0x00 => bytes[start + 1] & 0x80 == 0, + 0xFF => bytes[start + 1] & 0x80 != 0, + _ => false, + }; + if !redundant_sign_byte { + break; + } + start += 1; + } + murmur3_32(&bytes[start..]) +} + +fn bucket_array(fn_name: &str, array: &ArrayRef, num_buckets: i32) -> Result { + let bucket = |hash: i32| (hash & i32::MAX) % num_buckets; + let result: Int32Array = match array.data_type() { + // Iceberg binds tinyint and smallint inputs to `BucketInt`, hashing them as ints. + DataType::Int8 | DataType::Int16 => { + let widened = cast(array, &DataType::Int32)?; + return bucket_array(fn_name, &widened, num_buckets); + } + DataType::Int32 => array + .as_primitive::() + .unary(|v| bucket(hash_long(v as i64))), + DataType::Date32 => array + .as_primitive::() + .unary(|v| bucket(hash_long(v as i64))), + DataType::Int64 => array + .as_primitive::() + .unary(|v| bucket(hash_long(v))), + DataType::Timestamp(TimeUnit::Microsecond, _) => array + .as_primitive::() + .unary(|v| bucket(hash_long(v))), + DataType::Decimal128(_, _) => array + .as_primitive::() + .unary(|v| bucket(hash_decimal(v))), + DataType::Utf8 => array + .as_string::() + .iter() + .map(|v| v.map(|s| bucket(murmur3_32(s.as_bytes())))) + .collect(), + DataType::LargeUtf8 => array + .as_string::() + .iter() + .map(|v| v.map(|s| bucket(murmur3_32(s.as_bytes())))) + .collect(), + DataType::Binary => array + .as_binary::() + .iter() + .map(|v| v.map(|b| bucket(murmur3_32(b)))) + .collect(), + DataType::LargeBinary => array + .as_binary::() + .iter() + .map(|v| v.map(|b| bucket(murmur3_32(b)))) + .collect(), + DataType::FixedSizeBinary(_) => array + .as_fixed_size_binary() + .iter() + .map(|v| v.map(|b| bucket(murmur3_32(b)))) + .collect(), + other => return Err(unsupported_type(fn_name, other)), + }; + Ok(Arc::new(result)) +} + +/// `iceberg_bucket(numBuckets, value)`; see the module docs for the semantics. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkIcebergBucket { + signature: Signature, +} + +impl SparkIcebergBucket { + pub fn new() -> Self { + Self { + signature: Signature::variadic_any(Volatility::Immutable), + } + } +} + +impl Default for SparkIcebergBucket { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for SparkIcebergBucket { + fn name(&self) -> &str { + "iceberg_bucket" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [num_buckets, value] = take_function_args(self.name(), &args.args)?; + let num_buckets = positive_int_param(self.name(), "numBuckets", num_buckets)?; + apply_unary(value, |array| bucket_array(self.name(), array, num_buckets)) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_util::invoke; + use super::*; + use arrow::array::{ + BinaryArray, Date32Array, Decimal128Array, DictionaryArray, Int32Array, Int64Array, + Int8Array, StringArray, TimestampMicrosecondArray, + }; + use arrow::datatypes::Int8Type; + use datafusion::common::ScalarValue; + + /// Hash values from Appendix B of the Iceberg table spec. + #[test] + fn hashes_match_iceberg_spec_vectors() { + assert_eq!(hash_long(34), 2_017_239_379); + assert_eq!(hash_decimal(1420), -500_754_589); // decimal 14.20 + assert_eq!(hash_long(17_486), -653_330_422); // date 2017-11-16 + assert_eq!(hash_long(81_068_000_000), -662_762_989); // time 22:31:08 + assert_eq!(hash_long(1_510_871_468_000_000), -2_047_944_441); // 2017-11-16T22:31:08 + assert_eq!(murmur3_32("iceberg".as_bytes()), 1_210_000_089); + assert_eq!(murmur3_32(&[0x00, 0x01, 0x02, 0x03]), -188_683_207); + assert_eq!( + murmur3_32(&0xf79c_3e09_677c_4bbd_a479_3f34_9cb7_85e7_u128.to_be_bytes()), + 1_488_055_340 + ); // uuid f79c3e09-677c-4bbd-a479-3f349cb785e7 + } + + /// `BigInteger.toByteArray()` keeps exactly one sign byte. + #[test] + fn decimal_hash_uses_minimal_two_complement_bytes() { + assert_eq!(hash_decimal(0), murmur3_32(&[0x00])); + assert_eq!(hash_decimal(-1), murmur3_32(&[0xFF])); + assert_eq!(hash_decimal(127), murmur3_32(&[0x7F])); + assert_eq!(hash_decimal(128), murmur3_32(&[0x00, 0x80])); + assert_eq!(hash_decimal(-128), murmur3_32(&[0x80])); + assert_eq!(hash_decimal(-129), murmur3_32(&[0xFF, 0x7F])); + assert_eq!( + hash_decimal(i128::MAX), + murmur3_32(&i128::MAX.to_be_bytes()) + ); + assert_eq!( + hash_decimal(i128::MIN), + murmur3_32(&i128::MIN.to_be_bytes()) + ); + } + + fn bucket(num_buckets: i32, value: ArrayRef) -> Int32Array { + let result = invoke( + &SparkIcebergBucket::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(num_buckets))), + ColumnarValue::Array(value), + ], + ) + .unwrap(); + result.as_primitive::().clone() + } + + #[test] + fn buckets_every_supported_type_and_keeps_nulls() { + // bucket(100, 34) -> 79 is the example in Iceberg's function description. + let ints = bucket(100, Arc::new(Int32Array::from(vec![Some(34), None]))); + assert_eq!(ints, Int32Array::from(vec![Some(79), None])); + let longs = bucket(100, Arc::new(Int64Array::from(vec![Some(34), None]))); + assert_eq!(longs, Int32Array::from(vec![Some(79), None])); + let small = bucket(100, Arc::new(Int8Array::from(vec![Some(34), None]))); + assert_eq!(small, Int32Array::from(vec![Some(79), None])); + + let expected = |hash: i32| (hash & i32::MAX) % 16; + let dates = bucket(16, Arc::new(Date32Array::from(vec![Some(17_486), None]))); + assert_eq!(dates.value(0), expected(-653_330_422)); + assert!(dates.is_null(1)); + let timestamps = bucket( + 16, + Arc::new( + TimestampMicrosecondArray::from(vec![Some(1_510_871_468_000_000), None]) + .with_timezone("America/Los_Angeles"), + ), + ); + assert_eq!(timestamps.value(0), expected(-2_047_944_441)); + let ntz = bucket( + 16, + Arc::new(TimestampMicrosecondArray::from(vec![Some( + 1_510_871_468_000_000, + )])), + ); + assert_eq!(ntz.value(0), expected(-2_047_944_441)); + let decimals = bucket( + 16, + Arc::new( + Decimal128Array::from(vec![Some(1420), None]) + .with_precision_and_scale(4, 2) + .unwrap(), + ), + ); + assert_eq!(decimals.value(0), expected(-500_754_589)); + let strings = bucket( + 16, + Arc::new(StringArray::from(vec![Some("iceberg"), None, Some("")])), + ); + assert_eq!(strings.value(0), expected(1_210_000_089)); + assert!(strings.is_null(1)); + assert_eq!(strings.value(2), expected(murmur3_32(&[]))); + let binary = bucket( + 16, + Arc::new(BinaryArray::from(vec![ + Some([0x00u8, 0x01, 0x02, 0x03].as_slice()), + None, + ])), + ); + assert_eq!(binary.value(0), expected(-188_683_207)); + } + + #[test] + fn negative_hashes_never_produce_negative_buckets() { + // hash_long(17_486) is negative; masking with Integer.MAX_VALUE keeps the result in range. + let dates = bucket(7, Arc::new(Date32Array::from(vec![17_486]))); + assert_eq!(dates.value(0), (-653_330_422_i32 & i32::MAX) % 7); + assert!(dates.value(0) >= 0); + } + + #[test] + fn dictionary_input_is_unpacked() { + let dict: DictionaryArray = vec![Some("iceberg"), None, Some("iceberg")] + .into_iter() + .collect(); + let result = bucket(16, Arc::new(dict)); + let expected = (1_210_000_089_i32 & i32::MAX) % 16; + assert_eq!( + result, + Int32Array::from(vec![Some(expected), None, Some(expected)]) + ); + } + + #[test] + fn scalar_input_returns_scalar() { + let result = SparkIcebergBucket::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(100))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(34))), + ], + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(arrow::datatypes::Field::new("b", DataType::Int32, true)), + config_options: Arc::new(datafusion::config::ConfigOptions::default()), + }) + .unwrap(); + match result { + ColumnarValue::Scalar(ScalarValue::Int32(Some(79))) => {} + other => panic!("expected scalar 79, got {other:?}"), + } + } + + #[test] + fn rejects_non_positive_or_non_literal_num_buckets() { + let value = ColumnarValue::Array(Arc::new(Int32Array::from(vec![1]))); + for bad in [ + ColumnarValue::Scalar(ScalarValue::Int32(Some(0))), + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ColumnarValue::Array(Arc::new(Int32Array::from(vec![4]))), + ] { + let err = invoke(&SparkIcebergBucket::new(), vec![bad, value.clone()]).unwrap_err(); + assert!(err + .to_string() + .contains("numBuckets must be a positive Int32 literal")); + } + } + + #[test] + fn rejects_unsupported_types() { + let value = ColumnarValue::Array(Arc::new(arrow::array::Float64Array::from(vec![1.0]))); + let err = invoke( + &SparkIcebergBucket::new(), + vec![ColumnarValue::Scalar(ScalarValue::Int32(Some(4))), value], + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("does not support input type Float64")); + } +} diff --git a/native/spark-expr/src/iceberg_funcs/mod.rs b/native/spark-expr/src/iceberg_funcs/mod.rs new file mode 100644 index 00000000000..241c8d1cbdd --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/mod.rs @@ -0,0 +1,135 @@ +// 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. + +//! Native implementations of Iceberg's Spark system functions: `bucket`, `truncate`, `years`, +//! `months`, `days`, and `hours`. +//! +//! Spark binds these as `StaticInvoke` calls on the classes under +//! `org.apache.iceberg.spark.functions`, and they show up wherever hidden partitioning does: +//! the hash distribution and local sort in front of a partitioned Iceberg write, and row-level +//! filters, projections, and sort orders that mention a partition transform. The kernels here +//! reproduce Iceberg's Java implementations exactly (see the partition transforms section of the +//! Iceberg table spec), so a row lands in the same bucket or day whether Comet or Iceberg computes +//! it. + +mod bucket; +mod temporal; +mod truncate; + +pub use bucket::SparkIcebergBucket; +pub use temporal::SparkIcebergTemporalTransform; +pub use truncate::SparkIcebergTruncate; + +use arrow::array::{Array, ArrayRef}; +use arrow::compute::cast; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, Result, ScalarValue}; +use datafusion::logical_expr::ColumnarValue; +use std::sync::Arc; + +/// Unpacks a dictionary-encoded array to its value type so that the kernels only ever see plain +/// arrays. Any other array is returned unchanged. +fn unpack_dictionary(array: ArrayRef) -> Result { + match array.data_type() { + DataType::Dictionary(_, value_type) => Ok(cast(&array, value_type)?), + _ => Ok(array), + } +} + +/// The type a kernel sees for an input of type `data_type`, after dictionary unpacking. +fn unpacked_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Dictionary(_, value_type) => value_type.as_ref().clone(), + other => other.clone(), + } +} + +/// Applies an array kernel to a `ColumnarValue`, round-tripping a scalar through a one-row array. +fn apply_unary( + value: &ColumnarValue, + kernel: impl Fn(&ArrayRef) -> Result, +) -> Result { + match value { + ColumnarValue::Array(array) => { + let array = unpack_dictionary(Arc::clone(array))?; + Ok(ColumnarValue::Array(kernel(&array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = unpack_dictionary(scalar.to_array()?)?; + let result = kernel(&array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } +} + +/// Reads the `numBuckets` / `width` parameter. The Comet serde only converts these functions when +/// the parameter is a positive integer literal, so anything else here is a wiring bug. +fn positive_int_param(fn_name: &str, param: &str, value: &ColumnarValue) -> Result { + match value { + ColumnarValue::Scalar(ScalarValue::Int32(Some(n))) if *n > 0 => Ok(*n), + other => Err(DataFusionError::Execution(format!( + "{fn_name}: {param} must be a positive Int32 literal, got {other:?}" + ))), + } +} + +fn unsupported_type(fn_name: &str, data_type: &DataType) -> DataFusionError { + DataFusionError::Execution(format!( + "{fn_name} does not support input type {data_type:?}" + )) +} + +#[cfg(test)] +mod test_util { + use arrow::array::ArrayRef; + use arrow::datatypes::{DataType, Field}; + use datafusion::common::Result; + use datafusion::config::ConfigOptions; + use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + use std::sync::Arc; + + /// Invokes `udf` on `args` and returns the resulting array (scalars are widened). + pub(super) fn invoke(udf: &dyn ScalarUDFImpl, args: Vec) -> Result { + let arg_fields = args + .iter() + .enumerate() + .map(|(i, a)| Arc::new(Field::new(format!("arg{i}"), a.data_type(), true))) + .collect::>(); + let arg_types = arg_fields + .iter() + .map(|f| f.data_type().clone()) + .collect::>(); + let return_type = udf.return_type(&arg_types)?; + let number_rows = args + .iter() + .find_map(|a| match a { + ColumnarValue::Array(array) => Some(array.len()), + ColumnarValue::Scalar(_) => None, + }) + .unwrap_or(1); + let result = udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Arc::new(Field::new(udf.name(), return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + })?; + result.to_array(number_rows) + } +} diff --git a/native/spark-expr/src/iceberg_funcs/temporal.rs b/native/spark-expr/src/iceberg_funcs/temporal.rs new file mode 100644 index 00000000000..fc25937ccd2 --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/temporal.rs @@ -0,0 +1,332 @@ +// 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. + +//! Iceberg's `years`, `months`, `days`, and `hours` transforms. +//! +//! Iceberg's `DateTimeUtil` evaluates all four in UTC regardless of the Spark session timezone +//! (`TimestampType` and `TimestampNTZType` are handled identically), and all four floor: a value +//! before the epoch maps to a negative period. `years` and `months` are calendar-aware, `days` and +//! `hours` are plain floor division of the epoch value. `days` returns a date (Iceberg's +//! `DaysFunction.resultType()` is `DateType`), the other three return an int. +//! +//! The kernels work on the raw epoch values rather than going through Arrow's timezone-aware +//! `date_part`, which would otherwise shift a `TimestampType` column by the session offset. + +use super::{apply_unary, unsupported_type}; +use arrow::array::{ArrayRef, AsArray, Date32Array, Int32Array}; +use arrow::datatypes::{DataType, Date32Type, Int32Type, TimeUnit, TimestampMicrosecondType}; +use chrono::Datelike; +use datafusion::common::{utils::take_function_args, DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use num::integer::div_floor; +use std::sync::Arc; + +const MICROS_PER_HOUR: i64 = 3_600_000_000; +const MICROS_PER_DAY: i64 = 86_400_000_000; +const UNIX_EPOCH_YEAR: i32 = 1970; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum TemporalUnit { + Years, + Months, + Days, + Hours, +} + +impl TemporalUnit { + fn fn_name(self) -> &'static str { + match self { + TemporalUnit::Years => "iceberg_years", + TemporalUnit::Months => "iceberg_months", + TemporalUnit::Days => "iceberg_days", + TemporalUnit::Hours => "iceberg_hours", + } + } + + fn return_type(self) -> DataType { + match self { + TemporalUnit::Days => DataType::Date32, + _ => DataType::Int32, + } + } +} + +/// `DateTimeUtil.microsToDays`: floor division, so `-1` micros is day `-1`. +#[inline] +fn micros_to_days(micros: i64) -> i32 { + div_floor(micros, MICROS_PER_DAY) as i32 +} + +/// `DateTimeUtil.microsToHours`. +#[inline] +fn micros_to_hours(micros: i64) -> i32 { + div_floor(micros, MICROS_PER_HOUR) as i32 +} + +/// `DateTimeUtil.daysToYears`: whole calendar years between the epoch and the day, floored. +fn days_to_years(days: i32) -> Result { + Ok(civil_date(days)?.year() - UNIX_EPOCH_YEAR) +} + +/// `DateTimeUtil.daysToMonths`: whole calendar months between the epoch and the day, floored. +fn days_to_months(days: i32) -> Result { + let date = civil_date(days)?; + Ok((date.year() - UNIX_EPOCH_YEAR) * 12 + date.month0() as i32) +} + +fn civil_date(days: i32) -> Result { + Date32Type::to_naive_date_opt(days).ok_or_else(|| { + DataFusionError::Execution(format!("day {days} is out of the supported date range")) + }) +} + +/// Reduces both supported input types to days since the epoch. +fn to_epoch_days(fn_name: &str, array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Date32 => Ok(array.as_primitive::().clone()), + DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(array + .as_primitive::() + .unary(micros_to_days)), + other => Err(unsupported_type(fn_name, other)), + } +} + +fn transform_array(unit: TemporalUnit, array: &ArrayRef) -> Result { + let fn_name = unit.fn_name(); + let result: ArrayRef = match unit { + TemporalUnit::Years => { + Arc::new(to_epoch_days(fn_name, array)?.try_unary::<_, Int32Type, _>(days_to_years)?) + } + TemporalUnit::Months => { + Arc::new(to_epoch_days(fn_name, array)?.try_unary::<_, Int32Type, _>(days_to_months)?) + } + TemporalUnit::Days => Arc::new(to_epoch_days(fn_name, array)?), + TemporalUnit::Hours => match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, _) => { + let hours: Int32Array = array + .as_primitive::() + .unary(micros_to_hours); + Arc::new(hours) + } + other => return Err(unsupported_type(fn_name, other)), + }, + }; + Ok(result) +} + +/// `iceberg_years(value)`, `iceberg_months(value)`, `iceberg_days(value)`, and +/// `iceberg_hours(value)`; see the module docs for the semantics. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkIcebergTemporalTransform { + unit: TemporalUnit, + signature: Signature, +} + +impl SparkIcebergTemporalTransform { + pub fn new(unit: TemporalUnit) -> Self { + Self { + unit, + signature: Signature::variadic_any(Volatility::Immutable), + } + } + + pub fn years() -> Self { + Self::new(TemporalUnit::Years) + } + + pub fn months() -> Self { + Self::new(TemporalUnit::Months) + } + + pub fn days() -> Self { + Self::new(TemporalUnit::Days) + } + + pub fn hours() -> Self { + Self::new(TemporalUnit::Hours) + } +} + +impl ScalarUDFImpl for SparkIcebergTemporalTransform { + fn name(&self) -> &str { + self.unit.fn_name() + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.unit.return_type()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [value] = take_function_args(self.name(), &args.args)?; + apply_unary(value, |array| transform_array(self.unit, array)) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_util::invoke; + use super::*; + use arrow::array::{Array, TimestampMicrosecondArray}; + + fn transform(unit: TemporalUnit, value: ArrayRef) -> ArrayRef { + invoke( + &SparkIcebergTemporalTransform::new(unit), + vec![ColumnarValue::Array(value)], + ) + .unwrap() + } + + // Boundaries around the epoch, as (epoch days, years, months). + const DAY_CASES: &[(i32, i32, i32)] = &[ + (17_486, 47, 574), // 2017-11-16, the Iceberg spec example + (0, 0, 0), // 1970-01-01 + (-1, -1, -1), // 1969-12-31 + (-365, -1, -12), // 1969-01-01 + (-366, -2, -13), // 1968-12-31 + (365, 1, 12), // 1971-01-01 + (364, 0, 11), // 1970-12-31 + (31, 0, 1), // 1970-02-01 + (30, 0, 0), // 1970-01-31 + ]; + + #[test] + fn dates_match_iceberg_date_time_util() { + let days: Vec> = DAY_CASES.iter().map(|c| Some(c.0)).chain([None]).collect(); + let input: ArrayRef = Arc::new(Date32Array::from(days.clone())); + + let years = transform(TemporalUnit::Years, Arc::clone(&input)); + let months = transform(TemporalUnit::Months, Arc::clone(&input)); + let dates = transform(TemporalUnit::Days, Arc::clone(&input)); + for (i, (_, y, m)) in DAY_CASES.iter().enumerate() { + assert_eq!( + years.as_primitive::().value(i), + *y, + "years of {:?}", + DAY_CASES[i] + ); + assert_eq!( + months.as_primitive::().value(i), + *m, + "months of {:?}", + DAY_CASES[i] + ); + } + assert_eq!(dates.data_type(), &DataType::Date32); + assert_eq!(dates.as_primitive::(), &Date32Array::from(days)); + let last = DAY_CASES.len(); + assert!(years.is_null(last) && months.is_null(last) && dates.is_null(last)); + + let err = invoke( + &SparkIcebergTemporalTransform::hours(), + vec![ColumnarValue::Array(input)], + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("does not support input type Date32")); + } + + #[test] + fn timestamps_match_iceberg_date_time_util_in_utc() { + // (micros, years, months, days, hours) + let cases: &[(i64, i32, i32, i32, i32)] = &[ + (1_510_871_468_000_000, 47, 574, 17_486, 419_686), // 2017-11-16T22:31:08 (spec) + (0, 0, 0, 0, 0), + (-1, -1, -1, -1, -1), // 1969-12-31T23:59:59.999999 + (-MICROS_PER_HOUR, -1, -1, -1, -1), // 1969-12-31T23:00:00 + (-MICROS_PER_HOUR - 1, -1, -1, -1, -2), // 1969-12-31T22:59:59.999999 + (-MICROS_PER_DAY, -1, -1, -1, -24), // 1969-12-31T00:00:00 + (-MICROS_PER_DAY - 1, -1, -1, -2, -25), // 1969-12-30T23:59:59.999999 + (365 * MICROS_PER_DAY, 1, 12, 365, 8_760), // 1971-01-01T00:00:00 + (365 * MICROS_PER_DAY - 1, 0, 11, 364, 8_759), + ]; + let micros: Vec> = cases.iter().map(|c| Some(c.0)).chain([None]).collect(); + // A non-UTC timezone tag must not change the result. + for tz in [ + None, + Some("UTC"), + Some("America/Los_Angeles"), + Some("Asia/Kathmandu"), + ] { + let mut array = TimestampMicrosecondArray::from(micros.clone()); + if let Some(tz) = tz { + array = array.with_timezone(tz); + } + let input: ArrayRef = Arc::new(array); + let years = transform(TemporalUnit::Years, Arc::clone(&input)); + let months = transform(TemporalUnit::Months, Arc::clone(&input)); + let days = transform(TemporalUnit::Days, Arc::clone(&input)); + let hours = transform(TemporalUnit::Hours, Arc::clone(&input)); + assert_eq!(days.data_type(), &DataType::Date32); + for (i, (_, y, m, d, h)) in cases.iter().enumerate() { + let case = cases[i]; + assert_eq!( + years.as_primitive::().value(i), + *y, + "years {case:?} {tz:?}" + ); + assert_eq!( + months.as_primitive::().value(i), + *m, + "months {case:?} {tz:?}" + ); + assert_eq!( + days.as_primitive::().value(i), + *d, + "days {case:?} {tz:?}" + ); + assert_eq!( + hours.as_primitive::().value(i), + *h, + "hours {case:?} {tz:?}" + ); + } + let last = cases.len(); + assert!( + years.is_null(last) + && months.is_null(last) + && days.is_null(last) + && hours.is_null(last) + ); + } + } + + #[test] + fn rejects_unsupported_types() { + for unit in [ + TemporalUnit::Years, + TemporalUnit::Months, + TemporalUnit::Days, + TemporalUnit::Hours, + ] { + let err = invoke( + &SparkIcebergTemporalTransform::new(unit), + vec![ColumnarValue::Array(Arc::new(Int32Array::from(vec![1])))], + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("does not support input type Int32")); + } + } +} diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs new file mode 100644 index 00000000000..579d8f94a63 --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -0,0 +1,404 @@ +// 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. + +//! Iceberg's `truncate(width, value)` transform: `v - ((v % W) + W) % W` for integers (with +//! Java's wrapping arithmetic), the same on the unscaled value for decimals, the first `W` code +//! points of a string, and the first `W` bytes of a binary value. + +use super::{apply_unary, positive_int_param, unpacked_type, unsupported_type}; +use arrow::array::{ArrayRef, AsArray, Decimal128Array}; +use arrow::compute::kernels::substring::{substring, substring_by_char}; +use arrow::datatypes::{DataType, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type}; +use datafusion::common::{utils::take_function_args, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// `TruncateUtil.truncateInt`. Java's `int` arithmetic wraps on overflow, which can happen both in +/// `(v % w) + w` (for widths above 2^30) and in the final subtraction (near `Integer.MIN_VALUE`). +/// `TruncateUtil.truncateByte` / `truncateShort` evaluate the same expression in `int` and then +/// narrow, so tinyint and smallint inputs go through this function and are cast afterwards. +#[inline] +fn truncate_i32(v: i32, w: i32) -> i32 { + v.wrapping_sub((v % w).wrapping_add(w) % w) +} + +/// `TruncateUtil.truncateLong`, with the width promoted to `long` as Java does. +#[inline] +fn truncate_i64(v: i64, w: i64) -> i64 { + v.wrapping_sub((v % w).wrapping_add(w) % w) +} + +/// `TruncateUtil.truncateDecimal` on the unscaled value; `BigInteger` never overflows and neither +/// does an `i128` holding a 38-digit unscaled value minus a 31-bit width. +#[inline] +fn truncate_i128(v: i128, w: i128) -> i128 { + v - ((v % w) + w) % w +} + +fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result { + let result: ArrayRef = match array.data_type() { + DataType::Int8 => Arc::new( + array + .as_primitive::() + .unary::<_, Int8Type>(|v| truncate_i32(v as i32, width) as i8), + ), + DataType::Int16 => Arc::new( + array + .as_primitive::() + .unary::<_, Int16Type>(|v| truncate_i32(v as i32, width) as i16), + ), + DataType::Int32 => Arc::new( + array + .as_primitive::() + .unary::<_, Int32Type>(|v| truncate_i32(v, width)), + ), + DataType::Int64 => Arc::new( + array + .as_primitive::() + .unary::<_, Int64Type>(|v| truncate_i64(v, width as i64)), + ), + DataType::Decimal128(precision, scale) => { + // Truncating a negative value grows its magnitude by up to `width - 1` units of the + // last digit, so the result can need one more digit than the column allows. Spark's + // `UnsafeRowWriter` writes such a `Decimal` as null (`changePrecision` fails), so + // match that rather than emit a value the column's precision cannot hold. + let max_unscaled = 10_i128.pow(*precision as u32) - 1; + let truncated: Decimal128Array = + array.as_primitive::().unary_opt(|v| { + let truncated = truncate_i128(v, width as i128); + (truncated.abs() <= max_unscaled).then_some(truncated) + }); + Arc::new(truncated.with_precision_and_scale(*precision, *scale)?) + } + // `UTF8String.substring(0, width)` counts code points, not bytes. A width that covers + // the whole values buffer cannot truncate anything, so the input is returned as is; that + // also keeps a large width away from Arrow's substring kernels, which add it to the byte + // offsets and overflow on `i32::MAX`. + DataType::Utf8 => { + let strings = array.as_string::(); + if width as usize >= strings.value_data().len() { + Arc::clone(array) + } else { + Arc::new(substring_by_char(strings, 0, Some(width as u64))?) + } + } + DataType::LargeUtf8 => { + let strings = array.as_string::(); + if width as usize >= strings.value_data().len() { + Arc::clone(array) + } else { + Arc::new(substring_by_char(strings, 0, Some(width as u64))?) + } + } + // `BinaryUtil.truncateBinaryUnsafe` keeps the first `width` bytes. + DataType::Binary => { + if width as usize >= array.as_binary::().value_data().len() { + Arc::clone(array) + } else { + substring(array.as_ref(), 0, Some(width as u64))? + } + } + DataType::LargeBinary => { + if width as usize >= array.as_binary::().value_data().len() { + Arc::clone(array) + } else { + substring(array.as_ref(), 0, Some(width as u64))? + } + } + other => return Err(unsupported_type(fn_name, other)), + }; + Ok(result) +} + +/// `iceberg_truncate(width, value)`; see the module docs for the semantics. The result has the +/// same type as `value`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkIcebergTruncate { + signature: Signature, +} + +impl SparkIcebergTruncate { + pub fn new() -> Self { + Self { + signature: Signature::variadic_any(Volatility::Immutable), + } + } +} + +impl Default for SparkIcebergTruncate { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for SparkIcebergTruncate { + fn name(&self) -> &str { + "iceberg_truncate" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + let [_width, value] = take_function_args(self.name(), arg_types)?; + Ok(unpacked_type(value)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [width, value] = take_function_args(self.name(), &args.args)?; + let width = positive_int_param(self.name(), "width", width)?; + apply_unary(value, |array| truncate_array(self.name(), array, width)) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_util::invoke; + use super::*; + use arrow::array::{ + Array, BinaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray, + }; + use datafusion::common::ScalarValue; + + fn truncate(width: i32, value: ArrayRef) -> ArrayRef { + invoke( + &SparkIcebergTruncate::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(width))), + ColumnarValue::Array(value), + ], + ) + .unwrap() + } + + /// Examples from the Iceberg spec: truncate(10, 1) = 0, truncate(10, -1) = -10, + /// truncate(50, 10.65) = 10.50, truncate(3, "iceberg") = "ice". + #[test] + fn matches_iceberg_spec_examples() { + let ints = truncate( + 10, + Arc::new(Int32Array::from(vec![Some(1), Some(-1), None])), + ); + assert_eq!( + ints.as_primitive::(), + &Int32Array::from(vec![Some(0), Some(-10), None]) + ); + let longs = truncate( + 10, + Arc::new(Int64Array::from(vec![Some(1), Some(-1), None])), + ); + assert_eq!( + longs.as_primitive::(), + &Int64Array::from(vec![Some(0), Some(-10), None]) + ); + let decimals = truncate( + 50, + Arc::new( + Decimal128Array::from(vec![Some(1065), Some(-1065), None]) + .with_precision_and_scale(4, 2) + .unwrap(), + ), + ); + assert_eq!(decimals.data_type(), &DataType::Decimal128(4, 2)); + assert_eq!( + decimals.as_primitive::().values().as_ref(), + &[1050, -1100, 0] + ); + assert!(decimals.is_null(2)); + // -99.99 truncated to a width of 10 is -100.00, which does not fit decimal(4, 2); Spark + // produces null for it. + let overflow = truncate( + 10, + Arc::new( + Decimal128Array::from(vec![Some(-9999), Some(9999), Some(-9990)]) + .with_precision_and_scale(4, 2) + .unwrap(), + ), + ); + assert!(overflow.is_null(0)); + assert_eq!(overflow.as_primitive::().value(1), 9990); + assert_eq!(overflow.as_primitive::().value(2), -9990); + let strings = truncate( + 3, + Arc::new(StringArray::from(vec![ + Some("iceberg"), + Some("ic"), + Some(""), + Some("日本語テキスト"), + Some("a😀b😀c"), + None, + ])), + ); + assert_eq!( + strings.as_string::(), + &StringArray::from(vec![ + Some("ice"), + Some("ic"), + Some(""), + Some("日本語"), + Some("a😀b"), + None + ]) + ); + let binary = truncate( + 3, + Arc::new(BinaryArray::from(vec![ + Some([1u8, 2, 3, 4, 5].as_slice()), + Some([1u8].as_slice()), + None, + ])), + ); + assert_eq!( + binary.as_binary::(), + &BinaryArray::from(vec![ + Some([1u8, 2, 3].as_slice()), + Some([1u8].as_slice()), + None + ]) + ); + } + + /// Java narrows the `int` result back to `byte` / `short` and wraps `int` / `long` overflow. + #[test] + fn matches_java_wrapping_arithmetic() { + let bytes = truncate( + 1000, + Arc::new(Int8Array::from(vec![ + Some(i8::MIN), + Some(i8::MAX), + Some(-1), + ])), + ); + assert_eq!( + bytes.as_primitive::(), + &Int8Array::from(vec![Some(24), Some(0), Some(-1000i32 as i8)]) + ); + let shorts = truncate( + 100_000, + Arc::new(Int16Array::from(vec![Some(i16::MIN), Some(i16::MAX)])), + ); + assert_eq!( + shorts.as_primitive::(), + &Int16Array::from(vec![Some(-100_000i32 as i16), Some(0)]) + ); + let ints = truncate(1000, Arc::new(Int32Array::from(vec![i32::MIN, i32::MAX]))); + assert_eq!( + ints.as_primitive::(), + &Int32Array::from(vec![i32::MIN.wrapping_sub(352), 2_147_483_000]) + ); + let wide = truncate(i32::MAX, Arc::new(Int32Array::from(vec![i32::MAX - 1, -2]))); + // Java: (v % w) + w overflows for v = MAX - 1, then wraps back through % w. + let w = i32::MAX; + let expected = |v: i32| v.wrapping_sub((v % w).wrapping_add(w) % w); + assert_eq!( + wide.as_primitive::(), + &Int32Array::from(vec![expected(i32::MAX - 1), expected(-2)]) + ); + let longs = truncate(1000, Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX]))); + assert_eq!( + longs.as_primitive::(), + &Int64Array::from(vec![ + i64::MIN.wrapping_sub(((i64::MIN % 1000) + 1000) % 1000), + 9_223_372_036_854_775_000 + ]) + ); + } + + /// A width larger than any value is a no-op for strings and binary, and must not trip + /// Arrow's offset arithmetic (`i32::MAX` plus a non-zero offset overflows there). + #[test] + fn huge_width_leaves_strings_and_binary_unchanged() { + let strings: ArrayRef = Arc::new(StringArray::from(vec![ + Some("iceberg"), + None, + Some("日本語"), + Some(""), + ])); + let binary: ArrayRef = Arc::new(BinaryArray::from(vec![ + Some([1u8, 2, 3].as_slice()), + None, + Some([4u8, 5].as_slice()), + Some([].as_slice()), + ])); + for width in [10, 1000, i32::MAX] { + assert_eq!( + truncate(width, Arc::clone(&strings)).as_ref(), + strings.as_ref() + ); + assert_eq!( + truncate(width, Arc::clone(&binary)).as_ref(), + binary.as_ref() + ); + } + // The same widths still truncate rows that are longer than the width. + let long = truncate( + 5, + Arc::new(StringArray::from(vec![Some("ab"), Some("abcdefgh"), None])), + ); + assert_eq!( + long.as_string::(), + &StringArray::from(vec![Some("ab"), Some("abcde"), None]) + ); + } + + #[test] + fn return_type_follows_value_type() { + let udf = SparkIcebergTruncate::new(); + assert_eq!( + udf.return_type(&[DataType::Int32, DataType::Decimal128(10, 2)]) + .unwrap(), + DataType::Decimal128(10, 2) + ); + assert_eq!( + udf.return_type(&[ + DataType::Int32, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + ]) + .unwrap(), + DataType::Utf8 + ); + } + + #[test] + fn rejects_non_positive_width_and_unsupported_types() { + let err = invoke( + &SparkIcebergTruncate::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(0))), + ColumnarValue::Array(Arc::new(Int32Array::from(vec![1]))), + ], + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("width must be a positive Int32 literal")); + let err = invoke( + &SparkIcebergTruncate::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Array(Arc::new(arrow::array::Date32Array::from(vec![1]))), + ], + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("does not support input type Date32")); + } +} diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 705c08a2d26..37c66db51be 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -47,7 +47,9 @@ pub mod hash_funcs; mod string_funcs; mod datetime_funcs; +mod iceberg_funcs; pub use agg_funcs::*; +pub use iceberg_funcs::{SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate}; pub use cast::{spark_cast, Cast, SparkCastOptions}; diff --git a/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala new file mode 100644 index 00000000000..6ff587c134f --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala @@ -0,0 +1,230 @@ +/* + * 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.serde + +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} +import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke +import org.apache.spark.sql.types._ + +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType} + +/** + * Native support for Iceberg's Spark system functions (`bucket`, `truncate`, `years`, `months`, + * `days`, `hours`). + * + * Iceberg exposes each of these through Spark's static magic method, so + * `V2ExpressionUtils.resolveScalarFunction` binds them as `StaticInvoke(cls, "invoke", args)` + * where `cls` is one of the per-type implementations under `org.apache.iceberg.spark.functions` + * (e.g. `BucketFunction$BucketInt`). The same expressions appear in the hash distribution and + * local sort that Iceberg requests in front of a partitioned write, and in predicates and + * projections that users write against hidden partitioning, so routing them through + * [[CometStaticInvoke]] covers shuffle, sort, filter, and projection at once. + * + * Iceberg is not on Comet's compile classpath, so the handlers are keyed on class name rather + * than `Class[_]`. The list of classes is Iceberg's; `IcebergVersionFunction` is a zero-argument + * constant and is deliberately left out. + */ +object CometIcebergSystemFunctions { + + private val FunctionsPackage = "org.apache.iceberg.spark.functions." + + private def implementations( + outer: String, + handler: CometExpressionSerde[StaticInvoke], + inner: String*): Seq[(String, CometExpressionSerde[StaticInvoke])] = + inner.map(name => s"$FunctionsPackage$outer$$$name" -> handler) + + /** + * Handlers keyed by the fully qualified name of the Iceberg class that `StaticInvoke` calls. + */ + val staticInvokeHandlers: Map[String, CometExpressionSerde[StaticInvoke]] = ( + implementations( + "BucketFunction", + CometIcebergBucket, + "BucketInt", + "BucketLong", + "BucketString", + "BucketBinary", + "BucketDecimal") ++ + implementations( + "TruncateFunction", + CometIcebergTruncate, + "TruncateTinyInt", + "TruncateSmallInt", + "TruncateInt", + "TruncateBigInt", + "TruncateString", + "TruncateBinary", + "TruncateDecimal") ++ + implementations( + "YearsFunction", + CometIcebergYears, + "DateToYearsFunction", + "TimestampToYearsFunction", + "TimestampNtzToYearsFunction") ++ + implementations( + "MonthsFunction", + CometIcebergMonths, + "DateToMonthsFunction", + "TimestampToMonthsFunction", + "TimestampNtzToMonthsFunction") ++ + implementations( + "DaysFunction", + CometIcebergDays, + "DateToDaysFunction", + "TimestampToDaysFunction", + "TimestampNtzToDaysFunction") ++ + implementations( + "HoursFunction", + CometIcebergHours, + "TimestampToHoursFunction", + "TimestampNtzToHoursFunction") + ).toMap + + /** + * The `numBuckets` / `width` argument as a positive int, if it is a literal. Iceberg declares + * the parameter as `IntegerType`, so a tinyint or smallint literal arrives already cast and + * folded; the narrower literal types are matched anyway in case folding did not run. + */ + private[serde] def positiveIntLiteral(expr: Expression): Option[Int] = expr match { + case Literal(v: Int, IntegerType) if v > 0 => Some(v) + case Literal(v: Short, ShortType) if v > 0 => Some(v.toInt) + case Literal(v: Byte, ByteType) if v > 0 => Some(v.toInt) + case _ => None + } +} + +/** + * Shared shape of `bucket(numBuckets, value)` and `truncate(width, value)`: a positive integer + * parameter followed by the value. The parameter has to be a literal because the native kernel + * takes it as a constant, and it has to be positive because Iceberg's Java implementation divides + * by it (zero throws, which the fallback preserves by leaving the expression to Spark). + */ +abstract class CometIcebergParameterizedTransform( + nativeName: String, + parameterName: String, + valueTypeSupported: DataType => Boolean) + extends CometExpressionSerde[StaticInvoke] { + + override def getSupportLevel(expr: StaticInvoke): SupportLevel = expr.arguments match { + case Seq(parameter, value) => + if (CometIcebergSystemFunctions.positiveIntLiteral(parameter).isEmpty) { + Unsupported(Some(s"$parameterName must be a positive integer literal, got $parameter")) + } else if (!valueTypeSupported(value.dataType)) { + Unsupported(Some(s"$nativeName does not support input type ${value.dataType}")) + } else { + Compatible() + } + case other => + Unsupported(Some(s"expected ($parameterName, value) arguments, got ${other.size}")) + } + + override def convert( + expr: StaticInvoke, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = expr.arguments match { + case Seq(parameter, value) => + // Normalize to an int literal so the native side always sees an Int32 scalar. + val parameterProto = CometIcebergSystemFunctions + .positiveIntLiteral(parameter) + .flatMap(n => exprToProtoInternal(Literal(n, IntegerType), inputs, binding)) + val valueProto = exprToProtoInternal(value, inputs, binding) + scalarFunctionExprToProtoWithReturnType( + nativeName, + expr.dataType, + failOnError = false, + parameterProto, + valueProto) + case _ => None + } +} + +/** `bucket(numBuckets, value)` over the types `BucketFunction.bind` accepts. */ +object CometIcebergBucket + extends CometIcebergParameterizedTransform( + "iceberg_bucket", + "numBuckets", + { + case ByteType | ShortType | IntegerType | LongType | DateType | + TimestampType | TimestampNTZType | StringType | BinaryType => + true + case _: DecimalType => true + case _ => false + }) + +/** `truncate(width, value)` over the types `TruncateFunction.bind` accepts. */ +object CometIcebergTruncate + extends CometIcebergParameterizedTransform( + "iceberg_truncate", + "width", + { + case ByteType | ShortType | IntegerType | LongType | StringType | BinaryType => true + case _: DecimalType => true + case _ => false + }) + +/** Shared shape of the single-argument `years`, `months`, `days`, and `hours` transforms. */ +abstract class CometIcebergTemporalTransform( + nativeName: String, + valueTypeSupported: DataType => Boolean) + extends CometExpressionSerde[StaticInvoke] { + + override def getSupportLevel(expr: StaticInvoke): SupportLevel = expr.arguments match { + case Seq(value) if valueTypeSupported(value.dataType) => Compatible() + case Seq(value) => + Unsupported(Some(s"$nativeName does not support input type ${value.dataType}")) + case other => Unsupported(Some(s"expected one argument, got ${other.size}")) + } + + override def convert( + expr: StaticInvoke, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + val valueProto = exprToProtoInternal(expr.arguments.head, inputs, binding) + scalarFunctionExprToProtoWithReturnType( + nativeName, + expr.dataType, + failOnError = false, + valueProto) + } +} + +private object IcebergTemporalTypes { + val dateOrTimestamp: DataType => Boolean = { + case DateType | TimestampType | TimestampNTZType => true + case _ => false + } + val timestampOnly: DataType => Boolean = { + case TimestampType | TimestampNTZType => true + case _ => false + } +} + +object CometIcebergYears + extends CometIcebergTemporalTransform("iceberg_years", IcebergTemporalTypes.dateOrTimestamp) + +object CometIcebergMonths + extends CometIcebergTemporalTransform("iceberg_months", IcebergTemporalTypes.dateOrTimestamp) + +object CometIcebergDays + extends CometIcebergTemporalTransform("iceberg_days", IcebergTemporalTypes.dateOrTimestamp) + +object CometIcebergHours + extends CometIcebergTemporalTransform("iceberg_hours", IcebergTemporalTypes.timestampOnly) diff --git a/spark/src/main/scala/org/apache/comet/serde/statics.scala b/spark/src/main/scala/org/apache/comet/serde/statics.scala index 6501a1d46cb..6219b335bd8 100644 --- a/spark/src/main/scala/org/apache/comet/serde/statics.scala +++ b/spark/src/main/scala/org/apache/comet/serde/statics.scala @@ -51,17 +51,31 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] { // node survives and is handled directly (see CometBase64). ("encode", classOf[Base64]) -> CometBase64StaticInvoke) + /** + * Iceberg's system functions (`bucket`, `truncate`, `years`, ...) are matched on class name + * because Iceberg is not on Comet's compile classpath; see [[CometIcebergSystemFunctions]]. + */ + private def handlerFor(expr: StaticInvoke): Option[CometExpressionSerde[StaticInvoke]] = + staticInvokeExpressions + .get((expr.functionName, expr.staticObject)) + .orElse(CometIcebergSystemFunctions.staticInvokeHandlers.get(expr.staticObject.getName)) + + override def getSupportLevel(expr: StaticInvoke): SupportLevel = + handlerFor(expr).map(_.getSupportLevel(expr)).getOrElse(Compatible()) + override def convert( expr: StaticInvoke, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - staticInvokeExpressions.get((expr.functionName, expr.staticObject)) match { + handlerFor(expr) match { case Some(handler) => handler.convert(expr, inputs, binding) case None => + // Every Iceberg system function is named `invoke`, so name the declaring class too. withFallbackReason( expr, - s"Static invoke expression: ${expr.functionName} is not supported") + s"Static invoke expression: ${expr.functionName} is not supported " + + s"(declared on ${expr.staticObject.getName})") None } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala new file mode 100644 index 00000000000..6521d9f7b3f --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -0,0 +1,436 @@ +/* + * 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.math.{BigDecimal => JBigDecimal, BigInteger} +import java.time.{Instant, LocalDate, LocalDateTime, ZoneOffset} + +import scala.collection.mutable +import scala.util.Random + +import org.scalactic.source.Position +import org.scalatest.Tag + +import org.apache.spark.{CometListenerBusUtils, SparkConf} +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} +import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSortExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.serde.{CometIcebergBucket, CometIcebergTruncate, CometStaticInvoke, Compatible, Unsupported} + +/** + * Native support for Iceberg's system functions (`bucket`, `truncate`, `years`, `months`, `days`, + * `hours`). + * + * Every comparison runs the same query with Comet on and off, so the reference values come from + * Iceberg's own JVM implementations (`BucketFunction`, `TruncateFunction`, ...) evaluated by + * Spark, over seeded random data plus the boundary values of each type. A native result that + * disagreed with Iceberg would only fail loudly on the write path (the clustered writer rejects + * out-of-order partitions); in a filter or projection it would be a silently wrong answer, which + * is why the coverage is per type rather than a few hand-picked rows. + */ +class CometIcebergSystemFunctionSuite + extends CometTestBase + with CometIcebergTestBase + with AdaptiveSparkPlanHelper { + + override protected def sparkConf: SparkConf = { + super.sparkConf + .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + .set(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key, "true") + } + + override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit + pos: Position): Unit = { + super.test(testName, testTags: _*) { + assume(icebergAvailable, "Iceberg not available in classpath") + testFun + } + } + + private val catalog = "ice" + private val source = "system_function_source" + private val bucketColumns = + Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin", "dt", "ts", "ts_ntz") + private val truncateColumns = Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin") + + test("bucket matches Iceberg for every supported type") { + withSourceTable { + for (column <- bucketColumns; numBuckets <- Seq(1, 7, 16, Int.MaxValue)) { + checkSparkAnswerAndOperator( + s"SELECT $column, $catalog.system.bucket($numBuckets, $column) FROM $source") + } + } + } + + test("truncate matches Iceberg for every supported type") { + withSourceTable { + for (column <- truncateColumns; width <- Seq(1, 3, 10, 1000, Int.MaxValue)) { + checkSparkAnswerAndOperator( + s"SELECT $column, $catalog.system.truncate($width, $column) FROM $source") + } + } + } + + test("years, months, days, and hours match Iceberg regardless of session timezone") { + withSourceTable { + // Iceberg evaluates the temporal transforms in UTC; a shifted session timezone must not + // leak into the native result either. + for (timezone <- Seq("UTC", "America/Los_Angeles", "Asia/Kathmandu")) { + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { + for (column <- Seq("dt", "ts", "ts_ntz"); function <- Seq("years", "months", "days")) { + checkSparkAnswerAndOperator( + s"SELECT $column, $catalog.system.$function($column) FROM $source") + } + for (column <- Seq("ts", "ts_ntz")) { + checkSparkAnswerAndOperator( + s"SELECT $column, $catalog.system.hours($column) FROM $source") + } + } + } + } + } + + test("system functions in filters stay native") { + withSourceTable { + checkSparkAnswerAndOperator( + s"SELECT i32 FROM $source WHERE $catalog.system.bucket(8, i32) IN (0, 3)") + checkSparkAnswerAndOperator( + s"SELECT str FROM $source WHERE $catalog.system.truncate(1, str) = 'a'") + checkSparkAnswerAndOperator( + s"SELECT ts FROM $source WHERE $catalog.system.days(ts) >= DATE '2000-01-01'") + checkSparkAnswerAndOperator(s"SELECT dt FROM $source WHERE $catalog.system.months(dt) < 0") + } + } + + test("sorting on system functions stays native") { + withSourceTable { + val df = sql( + s"SELECT i32, str FROM $source " + + s"ORDER BY $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), i32, str") + checkSparkAnswerAndOperator(df) + assert( + collect(stripAQEPlan(df.queryExecution.executedPlan)) { case s: CometSortExec => + s + }.nonEmpty, + "expected a native sort") + } + } + + test("hash partitioning on system functions uses the native shuffle") { + withSourceTable { + val df = sql( + s"SELECT i32, str, ts FROM $source " + + s"DISTRIBUTE BY $catalog.system.bucket(8, i32), $catalog.system.truncate(2, str), " + + s"$catalog.system.hours(ts)") + checkSparkAnswerAndOperator(df) + checkCometExchange(df, 1, native = true) + } + } + + test("partitioned Iceberg write with default distribution mode stays native end to end") { + withSourceTable { + val table = s"$catalog.db.hidden_partitioning" + // No `write.distribution-mode`: Iceberg picks hash distribution for a partitioned table, + // which plans a shuffle and a local sort on the partition transforms. + sql(s""" + CREATE TABLE $table (i32 INT, str STRING, ts TIMESTAMP, dt DATE) + USING iceberg + PARTITIONED BY (bucket(4, i32), truncate(2, str), days(ts), months(dt))""") + try { + val plans = capturePlans { + sql(s"INSERT INTO $table SELECT i32, str, ts, dt FROM $source") + } + val writePlans = plans.filter(plan => + collectWithSubqueries(plan) { case w: CometIcebergWriteExec => w }.nonEmpty) + assert( + writePlans.nonEmpty, + s"expected a native Iceberg write in the captured plans:\n${plans.mkString("\n--\n")}") + writePlans.foreach { plan => + val cometShuffles = collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => + s + } + assert(cometShuffles.nonEmpty, s"expected a Comet shuffle in $plan") + cometShuffles.foreach(s => assert(s.shuffleType == CometNativeShuffle, s"$s")) + assert( + collectWithSubqueries(plan) { case s: ShuffleExchangeExec => s }.isEmpty, + s"the distribution shuffle stayed on Spark:\n$plan") + } + + checkAnswer( + sql(s"SELECT i32, str, ts, dt FROM $table"), + sql(s"SELECT i32, str, ts, dt FROM $source").collect()) + + // Iceberg's own view of the partitions must match what the JVM transforms compute over + // the written rows: one partition per distinct transform tuple. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val expected = sql(s""" + SELECT COUNT(*) FROM ( + SELECT DISTINCT $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), + $catalog.system.days(ts), $catalog.system.months(dt) + FROM $table)""").collect().head.getLong(0) + val actual = sql(s"SELECT COUNT(*) FROM $table.partitions").collect().head.getLong(0) + assert(actual == expected, s"expected $expected Iceberg partitions, found $actual") + } + } finally { + sql(s"DROP TABLE IF EXISTS $table") + } + } + } + + test("non-literal or non-positive parameters fall back to Spark") { + withSourceTable { + checkSparkAnswerAndFallbackReason( + s"SELECT $catalog.system.bucket(pmod(i32, 100) + 1, i32) FROM $source " + + "WHERE i32 IS NOT NULL", + "numBuckets must be a positive integer literal") + // Iceberg's Java implementation divides by the width, so a zero width has to stay with + // Spark to raise the same error; only the planning decision can be checked here. + val plan = + sql(s"SELECT $catalog.system.truncate(0, str) FROM $source").queryExecution.executedPlan + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert( + reasons.exists(_.contains("width must be a positive integer literal")), + s"unexpected fallback reasons: $reasons") + } + } + + test("support levels follow Iceberg's bind rules") { + val value = AttributeReference("v", IntegerType)() + def invoke(cls: Class[_], args: Seq[org.apache.spark.sql.catalyst.expressions.Expression]) = + StaticInvoke(cls, IntegerType, "invoke", args, propagateNull = false) + + val bucketInt = Class.forName("org.apache.iceberg.spark.functions.BucketFunction$BucketInt") + assert( + CometIcebergBucket.getSupportLevel( + invoke(bucketInt, Seq(Literal(4), value))) == Compatible()) + assert( + CometIcebergBucket + .getSupportLevel(invoke(bucketInt, Seq(Literal(4.toShort), value))) == Compatible()) + assert( + CometIcebergBucket + .getSupportLevel(invoke(bucketInt, Seq(Literal(0), value))) + .isInstanceOf[Unsupported]) + assert( + CometIcebergBucket + .getSupportLevel(invoke(bucketInt, Seq(Literal(-4), value))) + .isInstanceOf[Unsupported]) + assert( + CometIcebergBucket + .getSupportLevel(invoke(bucketInt, Seq(value, value))) + .isInstanceOf[Unsupported]) + assert( + CometIcebergBucket + .getSupportLevel(invoke(bucketInt, Seq(Literal(4), AttributeReference("f", FloatType)()))) + .isInstanceOf[Unsupported]) + + val truncateInt = + Class.forName("org.apache.iceberg.spark.functions.TruncateFunction$TruncateInt") + assert( + CometIcebergTruncate + .getSupportLevel(invoke(truncateInt, Seq(Literal(10), value))) == Compatible()) + assert( + CometIcebergTruncate + .getSupportLevel( + invoke(truncateInt, Seq(Literal(10), AttributeReference("d", DateType)()))) + .isInstanceOf[Unsupported]) + + // The dispatch in CometStaticInvoke goes by class name, so the same expressions resolve to + // the Iceberg handlers there too. + assert( + CometStaticInvoke.getSupportLevel( + invoke(bucketInt, Seq(Literal(4), value))) == Compatible()) + assert( + CometStaticInvoke + .getSupportLevel(invoke(bucketInt, Seq(Literal(0), value))) + .isInstanceOf[Unsupported]) + } + + test("fallback reason for an unlisted static invoke names the declaring class") { + val expr = StaticInvoke( + classOf[java.lang.Math], + IntegerType, + "abs", + Seq(AttributeReference("v", IntegerType)()), + propagateNull = false) + assert(CometStaticInvoke.convert(expr, Seq.empty, binding = false).isEmpty) + val reasons = expr.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty) + assert( + reasons.exists(r => + r.contains("Static invoke expression: abs is not supported") && + r.contains("java.lang.Math")), + s"unexpected fallback reasons: $reasons") + } + + /** Runs `f` with the Iceberg catalog registered and the source parquet table in scope. */ + private def withSourceTable(f: => Unit): Unit = withTempIcebergDir { warehouseDir => + withSQLConf( + s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog", + s"spark.sql.catalog.$catalog.type" -> "hadoop", + s"spark.sql.catalog.$catalog.warehouse" -> warehouseDir.getAbsolutePath) { + withTempPath { dir => + sourceData().write.parquet(dir.getAbsolutePath) + withParquetTable(dir.getAbsolutePath, source)(f) + } + } + } + + private val sourceSchema = StructType( + Seq( + StructField("i8", ByteType), + StructField("i16", ShortType), + StructField("i32", IntegerType), + StructField("i64", LongType), + StructField("dec18", DecimalType(18, 4)), + StructField("dec38", DecimalType(38, 10)), + StructField("str", StringType), + StructField("bin", BinaryType), + StructField("dt", DateType), + StructField("ts", TimestampType), + StructField("ts_ntz", TimestampNTZType))) + + /** + * Seeded random rows with nulls in every column, followed by the boundary values of each type + * (numeric extremes, the epoch and the microsecond before it, empty and multi-byte strings). + */ + private def sourceData(): DataFrame = { + val random = new Random(42) + // Code points rather than chars so that a surrogate pair is never split. + val alphabet = Seq("a", "b", "c", " ", "é", "日", "本", "語", "😀") + def maybeNull(value: => Any): Any = if (random.nextInt(8) == 0) null else value + def randomString(): String = + Seq.fill(random.nextInt(12))(alphabet(random.nextInt(alphabet.size))).mkString + def randomBinary(): Array[Byte] = { + val bytes = new Array[Byte](random.nextInt(10)) + random.nextBytes(bytes) + bytes + } + def randomDecimal38(): JBigDecimal = { + val unscaled = new BigInteger(126, random.self) + new JBigDecimal(if (random.nextBoolean()) unscaled else unscaled.negate(), 10) + } + def randomMicros(): Long = random.nextLong() % 4000000000000000L + def instant(micros: Long): Instant = + Instant.ofEpochSecond( + Math.floorDiv(micros, 1000000L), + Math.floorMod(micros, 1000000L) * 1000) + def localDateTime(micros: Long): LocalDateTime = + LocalDateTime.ofEpochSecond( + Math.floorDiv(micros, 1000000L), + (Math.floorMod(micros, 1000000L) * 1000).toInt, + ZoneOffset.UTC) + + val randomRows = (0 until 400).map { _ => + Row( + maybeNull(random.nextInt().toByte), + maybeNull(random.nextInt().toShort), + maybeNull(random.nextInt()), + maybeNull(random.nextLong()), + maybeNull(JBigDecimal.valueOf(random.nextLong() % 100000000000000000L, 4)), + maybeNull(randomDecimal38()), + maybeNull(randomString()), + maybeNull(randomBinary()), + maybeNull(LocalDate.ofEpochDay(random.nextInt(40000) - 20000)), + maybeNull(instant(randomMicros())), + maybeNull(localDateTime(randomMicros()))) + } + val dec38Max = new JBigDecimal(BigInteger.TEN.pow(38).subtract(BigInteger.ONE), 10) + val boundaryRows = Seq( + Row( + Byte.MinValue, + Short.MinValue, + Int.MinValue, + Long.MinValue, + new JBigDecimal("-99999999999999.9999"), + dec38Max.negate(), + "", + Array.empty[Byte], + LocalDate.ofEpochDay(0), + Instant.EPOCH, + LocalDateTime.of(1970, 1, 1, 0, 0)), + Row( + Byte.MaxValue, + Short.MaxValue, + Int.MaxValue, + Long.MaxValue, + new JBigDecimal("99999999999999.9999"), + dec38Max, + "日本語😀", + Array[Byte](0, 1, 2, 3), + LocalDate.ofEpochDay(-1), + instant(-1L), + localDateTime(-1L)), + Row( + 0.toByte, + 0.toShort, + 0, + 0L, + JBigDecimal.ZERO.setScale(4), + JBigDecimal.ZERO.setScale(10), + "a", + Array[Byte](0), + LocalDate.ofEpochDay(-365), + instant(-86400000000L), + localDateTime(-86400000000L - 1)), + Row( + (-1).toByte, + (-1).toShort, + -1, + -1L, + new JBigDecimal("-0.0001"), + new JBigDecimal("-0.0000000001"), + "iceberg", + Array[Byte](-1, -1, -1, -1, -1), + LocalDate.ofEpochDay(-366), + instant(-3600000000L - 1), + localDateTime(-3600000000L)), + Row(null, null, null, null, null, null, null, null, null, null, null)) + spark.createDataFrame( + spark.sparkContext.parallelize(randomRows ++ boundaryRows, 3), + sourceSchema) + } + + private def capturePlans(action: => Unit): Seq[SparkPlan] = { + val captured = mutable.Buffer.empty[SparkPlan] + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + captured += qe.executedPlan + } + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + () + } + spark.listenerManager.register(listener) + try { + action + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } finally { + spark.listenerManager.unregister(listener) + } + captured.toSeq + } +} From 3951f0239ca7c2f330d229ac121c156c76cc80b8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 15:28:24 -0600 Subject: [PATCH 02/18] test: simplify wrapped remainder in truncate unit test to satisfy clippy --- native/spark-expr/src/iceberg_funcs/truncate.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs index 579d8f94a63..d875c5f6550 100644 --- a/native/spark-expr/src/iceberg_funcs/truncate.rs +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -315,7 +315,8 @@ mod tests { assert_eq!( longs.as_primitive::(), &Int64Array::from(vec![ - i64::MIN.wrapping_sub(((i64::MIN % 1000) + 1000) % 1000), + // i64::MIN % 1000 == -808, so the wrapped remainder is 192. + i64::MIN.wrapping_sub(192), 9_223_372_036_854_775_000 ]) ); From 8c0d8d61523a48bf41cb20d71fb51e3a67fb402c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 15:51:51 -0600 Subject: [PATCH 03/18] refactor: simplify Iceberg system function kernels, serde, and tests Native: hash tinyint/smallint directly instead of casting first, fold the string/binary bucket arms into one generic helper, compute the minimal two's complement length with leading_ones/leading_zeros, fuse the years/months passes into a single kernel, reuse is_valid_decimal_precision for the decimal overflow check, and share the string/binary truncation helpers with a corrected note on which Arrow kernel overflows. Serde: key CometStaticInvoke's single map by (functionName, class name) so the Iceberg handlers join it instead of a second lookup, and use type sets for the temporal predicates. Tests: move capturePlans into CometIcebergTestBase, write the source parquet once per suite, batch the per-type comparisons into one query per column, and declare boundary values column-wise. --- native/spark-expr/src/iceberg_funcs/bucket.rs | 111 ++++---- .../spark-expr/src/iceberg_funcs/temporal.rs | 49 ++-- .../spark-expr/src/iceberg_funcs/truncate.rs | 69 ++--- .../apache/comet/serde/icebergFunctions.scala | 53 ++-- .../org/apache/comet/serde/statics.scala | 35 ++- .../CometIcebergSystemFunctionSuite.scala | 239 ++++++++---------- .../apache/comet/CometIcebergTestBase.scala | 25 ++ .../comet/CometIcebergWriteActionSuite.scala | 32 +-- 8 files changed, 283 insertions(+), 330 deletions(-) diff --git a/native/spark-expr/src/iceberg_funcs/bucket.rs b/native/spark-expr/src/iceberg_funcs/bucket.rs index e6b332d03fc..aec67f84c04 100644 --- a/native/spark-expr/src/iceberg_funcs/bucket.rs +++ b/native/spark-expr/src/iceberg_funcs/bucket.rs @@ -22,10 +22,10 @@ //! the unscaled value for decimals). use super::{apply_unary, positive_int_param, unsupported_type}; -use arrow::array::{Array, ArrayRef, AsArray, Int32Array}; -use arrow::compute::cast; +use arrow::array::{ArrayRef, AsArray, Int32Array}; use arrow::datatypes::{ - DataType, Date32Type, Decimal128Type, Int32Type, Int64Type, TimeUnit, TimestampMicrosecondType, + DataType, Date32Type, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type, TimeUnit, + TimestampMicrosecondType, }; use datafusion::common::{utils::take_function_args, Result}; use datafusion::logical_expr::{ @@ -80,33 +80,39 @@ fn hash_long(value: i64) -> i32 { } /// `BucketUtil.hash(BigDecimal)`: hashes `unscaledValue().toByteArray()`, the shortest big-endian -/// two's complement encoding of the unscaled value (at least one byte). +/// two's complement encoding of the unscaled value. That keeps exactly one sign bit, so the byte +/// count is one more than the bit length left after the run of leading sign bits. #[inline] fn hash_decimal(unscaled: i128) -> i32 { - let bytes = unscaled.to_be_bytes(); - let mut start = 0; - while start < bytes.len() - 1 { - let redundant_sign_byte = match bytes[start] { - 0x00 => bytes[start + 1] & 0x80 == 0, - 0xFF => bytes[start + 1] & 0x80 != 0, - _ => false, - }; - if !redundant_sign_byte { - break; - } - start += 1; - } - murmur3_32(&bytes[start..]) + let sign_bits = if unscaled < 0 { + unscaled.leading_ones() + } else { + unscaled.leading_zeros() + }; + let skip = ((sign_bits - 1) / 8) as usize; + murmur3_32(&unscaled.to_be_bytes()[skip..]) +} + +/// Buckets string or binary values by their raw bytes, keeping nulls. +fn bucket_bytes<'a, B: AsRef<[u8]> + ?Sized + 'a>( + values: impl Iterator>, + bucket: impl Fn(i32) -> i32, +) -> Int32Array { + values + .map(|v| v.map(|b| bucket(murmur3_32(b.as_ref())))) + .collect() } fn bucket_array(fn_name: &str, array: &ArrayRef, num_buckets: i32) -> Result { let bucket = |hash: i32| (hash & i32::MAX) % num_buckets; let result: Int32Array = match array.data_type() { // Iceberg binds tinyint and smallint inputs to `BucketInt`, hashing them as ints. - DataType::Int8 | DataType::Int16 => { - let widened = cast(array, &DataType::Int32)?; - return bucket_array(fn_name, &widened, num_buckets); - } + DataType::Int8 => array + .as_primitive::() + .unary(|v| bucket(hash_long(v as i64))), + DataType::Int16 => array + .as_primitive::() + .unary(|v| bucket(hash_long(v as i64))), DataType::Int32 => array .as_primitive::() .unary(|v| bucket(hash_long(v as i64))), @@ -122,31 +128,11 @@ fn bucket_array(fn_name: &str, array: &ArrayRef, num_buckets: i32) -> Result array .as_primitive::() .unary(|v| bucket(hash_decimal(v))), - DataType::Utf8 => array - .as_string::() - .iter() - .map(|v| v.map(|s| bucket(murmur3_32(s.as_bytes())))) - .collect(), - DataType::LargeUtf8 => array - .as_string::() - .iter() - .map(|v| v.map(|s| bucket(murmur3_32(s.as_bytes())))) - .collect(), - DataType::Binary => array - .as_binary::() - .iter() - .map(|v| v.map(|b| bucket(murmur3_32(b)))) - .collect(), - DataType::LargeBinary => array - .as_binary::() - .iter() - .map(|v| v.map(|b| bucket(murmur3_32(b)))) - .collect(), - DataType::FixedSizeBinary(_) => array - .as_fixed_size_binary() - .iter() - .map(|v| v.map(|b| bucket(murmur3_32(b)))) - .collect(), + DataType::Utf8 => bucket_bytes(array.as_string::().iter(), bucket), + DataType::LargeUtf8 => bucket_bytes(array.as_string::().iter(), bucket), + DataType::Binary => bucket_bytes(array.as_binary::().iter(), bucket), + DataType::LargeBinary => bucket_bytes(array.as_binary::().iter(), bucket), + DataType::FixedSizeBinary(_) => bucket_bytes(array.as_fixed_size_binary().iter(), bucket), other => return Err(unsupported_type(fn_name, other)), }; Ok(Arc::new(result)) @@ -197,10 +183,9 @@ mod tests { use super::super::test_util::invoke; use super::*; use arrow::array::{ - BinaryArray, Date32Array, Decimal128Array, DictionaryArray, Int32Array, Int64Array, - Int8Array, StringArray, TimestampMicrosecondArray, + Array, BinaryArray, Date32Array, Decimal128Array, DictionaryArray, Int16Array, Int32Array, + Int64Array, Int8Array, StringArray, TimestampMicrosecondArray, }; - use arrow::datatypes::Int8Type; use datafusion::common::ScalarValue; /// Hash values from Appendix B of the Iceberg table spec. @@ -252,13 +237,25 @@ mod tests { #[test] fn buckets_every_supported_type_and_keeps_nulls() { - // bucket(100, 34) -> 79 is the example in Iceberg's function description. - let ints = bucket(100, Arc::new(Int32Array::from(vec![Some(34), None]))); - assert_eq!(ints, Int32Array::from(vec![Some(79), None])); - let longs = bucket(100, Arc::new(Int64Array::from(vec![Some(34), None]))); - assert_eq!(longs, Int32Array::from(vec![Some(79), None])); - let small = bucket(100, Arc::new(Int8Array::from(vec![Some(34), None]))); - assert_eq!(small, Int32Array::from(vec![Some(79), None])); + // bucket(100, 34) -> 79 is the example in Iceberg's function description, and every + // integer width hashes the same 8 little-endian bytes. + let expected_34 = Int32Array::from(vec![Some(79), None]); + assert_eq!( + bucket(100, Arc::new(Int8Array::from(vec![Some(34), None]))), + expected_34 + ); + assert_eq!( + bucket(100, Arc::new(Int16Array::from(vec![Some(34), None]))), + expected_34 + ); + assert_eq!( + bucket(100, Arc::new(Int32Array::from(vec![Some(34), None]))), + expected_34 + ); + assert_eq!( + bucket(100, Arc::new(Int64Array::from(vec![Some(34), None]))), + expected_34 + ); let expected = |hash: i32| (hash & i32::MAX) % 16; let dates = bucket(16, Arc::new(Date32Array::from(vec![Some(17_486), None]))); diff --git a/native/spark-expr/src/iceberg_funcs/temporal.rs b/native/spark-expr/src/iceberg_funcs/temporal.rs index fc25937ccd2..55e4deea8c3 100644 --- a/native/spark-expr/src/iceberg_funcs/temporal.rs +++ b/native/spark-expr/src/iceberg_funcs/temporal.rs @@ -27,7 +27,7 @@ //! `date_part`, which would otherwise shift a `TimestampType` column by the session offset. use super::{apply_unary, unsupported_type}; -use arrow::array::{ArrayRef, AsArray, Date32Array, Int32Array}; +use arrow::array::{ArrayRef, AsArray, Int32Array}; use arrow::datatypes::{DataType, Date32Type, Int32Type, TimeUnit, TimestampMicrosecondType}; use chrono::Datelike; use datafusion::common::{utils::take_function_args, DataFusionError, Result}; @@ -42,7 +42,7 @@ const MICROS_PER_DAY: i64 = 86_400_000_000; const UNIX_EPOCH_YEAR: i32 = 1970; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum TemporalUnit { +pub(crate) enum TemporalUnit { Years, Months, Days, @@ -96,13 +96,17 @@ fn civil_date(days: i32) -> Result { }) } -/// Reduces both supported input types to days since the epoch. -fn to_epoch_days(fn_name: &str, array: &ArrayRef) -> Result { +/// Applies a calendar function of the epoch day to a date or timestamp column in one pass. +fn map_epoch_days( + fn_name: &str, + array: &ArrayRef, + f: impl Fn(i32) -> Result, +) -> Result { match array.data_type() { - DataType::Date32 => Ok(array.as_primitive::().clone()), - DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(array + DataType::Date32 => array.as_primitive::().try_unary(f), + DataType::Timestamp(TimeUnit::Microsecond, _) => array .as_primitive::() - .unary(micros_to_days)), + .try_unary(|micros| f(micros_to_days(micros))), other => Err(unsupported_type(fn_name, other)), } } @@ -110,20 +114,23 @@ fn to_epoch_days(fn_name: &str, array: &ArrayRef) -> Result { fn transform_array(unit: TemporalUnit, array: &ArrayRef) -> Result { let fn_name = unit.fn_name(); let result: ArrayRef = match unit { - TemporalUnit::Years => { - Arc::new(to_epoch_days(fn_name, array)?.try_unary::<_, Int32Type, _>(days_to_years)?) - } - TemporalUnit::Months => { - Arc::new(to_epoch_days(fn_name, array)?.try_unary::<_, Int32Type, _>(days_to_months)?) - } - TemporalUnit::Days => Arc::new(to_epoch_days(fn_name, array)?), + TemporalUnit::Years => Arc::new(map_epoch_days(fn_name, array, days_to_years)?), + TemporalUnit::Months => Arc::new(map_epoch_days(fn_name, array, days_to_months)?), + TemporalUnit::Days => match array.data_type() { + DataType::Date32 => Arc::clone(array), + DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new( + array + .as_primitive::() + .unary::<_, Date32Type>(micros_to_days), + ), + other => return Err(unsupported_type(fn_name, other)), + }, TemporalUnit::Hours => match array.data_type() { - DataType::Timestamp(TimeUnit::Microsecond, _) => { - let hours: Int32Array = array + DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new( + array .as_primitive::() - .unary(micros_to_hours); - Arc::new(hours) - } + .unary::<_, Int32Type>(micros_to_hours), + ), other => return Err(unsupported_type(fn_name, other)), }, }; @@ -139,7 +146,7 @@ pub struct SparkIcebergTemporalTransform { } impl SparkIcebergTemporalTransform { - pub fn new(unit: TemporalUnit) -> Self { + pub(crate) fn new(unit: TemporalUnit) -> Self { Self { unit, signature: Signature::variadic_any(Volatility::Immutable), @@ -186,7 +193,7 @@ impl ScalarUDFImpl for SparkIcebergTemporalTransform { mod tests { use super::super::test_util::invoke; use super::*; - use arrow::array::{Array, TimestampMicrosecondArray}; + use arrow::array::{Array, Date32Array, TimestampMicrosecondArray}; fn transform(unit: TemporalUnit, value: ArrayRef) -> ArrayRef { invoke( diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs index d875c5f6550..97cebf67742 100644 --- a/native/spark-expr/src/iceberg_funcs/truncate.rs +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -20,7 +20,8 @@ //! points of a string, and the first `W` bytes of a binary value. use super::{apply_unary, positive_int_param, unpacked_type, unsupported_type}; -use arrow::array::{ArrayRef, AsArray, Decimal128Array}; +use crate::utils::is_valid_decimal_precision; +use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array, OffsetSizeTrait}; use arrow::compute::kernels::substring::{substring, substring_by_char}; use arrow::datatypes::{DataType, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type}; use datafusion::common::{utils::take_function_args, Result}; @@ -51,6 +52,26 @@ fn truncate_i128(v: i128, w: i128) -> i128 { v - ((v % w) + w) % w } +/// `UTF8String.substring(0, width)` counts code points, not bytes. A width that covers the whole +/// values buffer cannot truncate anything, so the input is returned as is instead of being copied. +fn truncate_string(array: &ArrayRef, width: i32) -> Result { + let strings = array.as_string::(); + if width as usize >= strings.value_data().len() { + return Ok(Arc::clone(array)); + } + Ok(Arc::new(substring_by_char(strings, 0, Some(width as u64))?)) +} + +/// `BinaryUtil.truncateBinaryUnsafe` keeps the first `width` bytes. The whole-buffer shortcut +/// matters here beyond avoiding a copy: Arrow's byte `substring` adds the length to each value's +/// offset without checking for overflow, which panics for a width near `i32::MAX`. +fn truncate_binary(array: &ArrayRef, width: i32) -> Result { + if width as usize >= array.as_binary::().value_data().len() { + return Ok(Arc::clone(array)); + } + Ok(substring(array.as_ref(), 0, Some(width as u64))?) +} + fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result { let result: ArrayRef = match array.data_type() { DataType::Int8 => Arc::new( @@ -78,49 +99,17 @@ fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result().unary_opt(|v| { let truncated = truncate_i128(v, width as i128); - (truncated.abs() <= max_unscaled).then_some(truncated) + is_valid_decimal_precision(truncated, *precision).then_some(truncated) }); Arc::new(truncated.with_precision_and_scale(*precision, *scale)?) } - // `UTF8String.substring(0, width)` counts code points, not bytes. A width that covers - // the whole values buffer cannot truncate anything, so the input is returned as is; that - // also keeps a large width away from Arrow's substring kernels, which add it to the byte - // offsets and overflow on `i32::MAX`. - DataType::Utf8 => { - let strings = array.as_string::(); - if width as usize >= strings.value_data().len() { - Arc::clone(array) - } else { - Arc::new(substring_by_char(strings, 0, Some(width as u64))?) - } - } - DataType::LargeUtf8 => { - let strings = array.as_string::(); - if width as usize >= strings.value_data().len() { - Arc::clone(array) - } else { - Arc::new(substring_by_char(strings, 0, Some(width as u64))?) - } - } - // `BinaryUtil.truncateBinaryUnsafe` keeps the first `width` bytes. - DataType::Binary => { - if width as usize >= array.as_binary::().value_data().len() { - Arc::clone(array) - } else { - substring(array.as_ref(), 0, Some(width as u64))? - } - } - DataType::LargeBinary => { - if width as usize >= array.as_binary::().value_data().len() { - Arc::clone(array) - } else { - substring(array.as_ref(), 0, Some(width as u64))? - } - } + DataType::Utf8 => truncate_string::(array, width)?, + DataType::LargeUtf8 => truncate_string::(array, width)?, + DataType::Binary => truncate_binary::(array, width)?, + DataType::LargeBinary => truncate_binary::(array, width)?, other => return Err(unsupported_type(fn_name, other)), }; Ok(result) @@ -172,9 +161,7 @@ impl ScalarUDFImpl for SparkIcebergTruncate { mod tests { use super::super::test_util::invoke; use super::*; - use arrow::array::{ - Array, BinaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray, - }; + use arrow::array::{BinaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray}; use datafusion::common::ScalarValue; fn truncate(width: i32, value: ArrayRef) -> ArrayRef { diff --git a/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala index 6ff587c134f..4ea5cf5a5b8 100644 --- a/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala +++ b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala @@ -37,24 +37,28 @@ import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctio * projections that users write against hidden partitioning, so routing them through * [[CometStaticInvoke]] covers shuffle, sort, filter, and projection at once. * - * Iceberg is not on Comet's compile classpath, so the handlers are keyed on class name rather - * than `Class[_]`. The list of classes is Iceberg's; `IcebergVersionFunction` is a zero-argument - * constant and is deliberately left out. + * The list of classes is Iceberg's; `IcebergVersionFunction` is a zero-argument constant and is + * deliberately left out. */ object CometIcebergSystemFunctions { private val FunctionsPackage = "org.apache.iceberg.spark.functions." + /** Every Iceberg system function exposes its static magic method under this name. */ + private val MagicMethod = "invoke" + private def implementations( outer: String, handler: CometExpressionSerde[StaticInvoke], - inner: String*): Seq[(String, CometExpressionSerde[StaticInvoke])] = - inner.map(name => s"$FunctionsPackage$outer$$$name" -> handler) + inner: String*): Seq[((String, String), CometExpressionSerde[StaticInvoke])] = + inner.map(name => (MagicMethod, s"$FunctionsPackage$outer$$$name") -> handler) /** - * Handlers keyed by the fully qualified name of the Iceberg class that `StaticInvoke` calls. + * Handlers keyed by `(functionName, class name)` of the Iceberg implementation class that + * `StaticInvoke` calls, the shape [[CometStaticInvoke]] dispatches on. Iceberg is not on + * Comet's compile classpath, which is why the key carries the class name rather than the class. */ - val staticInvokeHandlers: Map[String, CometExpressionSerde[StaticInvoke]] = ( + val staticInvokeHandlers: Map[(String, String), CometExpressionSerde[StaticInvoke]] = ( implementations( "BucketFunction", CometIcebergBucket, @@ -162,10 +166,9 @@ object CometIcebergBucket "iceberg_bucket", "numBuckets", { - case ByteType | ShortType | IntegerType | LongType | DateType | - TimestampType | TimestampNTZType | StringType | BinaryType => + case ByteType | ShortType | IntegerType | LongType | DateType | TimestampType | + TimestampNTZType | StringType | BinaryType | _: DecimalType => true - case _: DecimalType => true case _ => false }) @@ -175,8 +178,9 @@ object CometIcebergTruncate "iceberg_truncate", "width", { - case ByteType | ShortType | IntegerType | LongType | StringType | BinaryType => true - case _: DecimalType => true + case ByteType | ShortType | IntegerType | LongType | StringType | + BinaryType | _: DecimalType => + true case _ => false }) @@ -206,25 +210,20 @@ abstract class CometIcebergTemporalTransform( } } -private object IcebergTemporalTypes { - val dateOrTimestamp: DataType => Boolean = { - case DateType | TimestampType | TimestampNTZType => true - case _ => false - } - val timestampOnly: DataType => Boolean = { - case TimestampType | TimestampNTZType => true - case _ => false - } -} - object CometIcebergYears - extends CometIcebergTemporalTransform("iceberg_years", IcebergTemporalTypes.dateOrTimestamp) + extends CometIcebergTemporalTransform( + "iceberg_years", + Set(DateType, TimestampType, TimestampNTZType)) object CometIcebergMonths - extends CometIcebergTemporalTransform("iceberg_months", IcebergTemporalTypes.dateOrTimestamp) + extends CometIcebergTemporalTransform( + "iceberg_months", + Set(DateType, TimestampType, TimestampNTZType)) object CometIcebergDays - extends CometIcebergTemporalTransform("iceberg_days", IcebergTemporalTypes.dateOrTimestamp) + extends CometIcebergTemporalTransform( + "iceberg_days", + Set(DateType, TimestampType, TimestampNTZType)) object CometIcebergHours - extends CometIcebergTemporalTransform("iceberg_hours", IcebergTemporalTypes.timestampOnly) + extends CometIcebergTemporalTransform("iceberg_hours", Set(TimestampType, TimestampNTZType)) diff --git a/spark/src/main/scala/org/apache/comet/serde/statics.scala b/spark/src/main/scala/org/apache/comet/serde/statics.scala index 6219b335bd8..94beefc0ba5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/statics.scala +++ b/spark/src/main/scala/org/apache/comet/serde/statics.scala @@ -32,33 +32,32 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] { // With Spark 3.4, CharVarcharCodegenUtils.readSidePadding gets called to pad spaces for // char types. // See https://github.com/apache/spark/pull/38151 - private val staticInvokeExpressions - : Map[(String, Class[_]), CometExpressionSerde[StaticInvoke]] = - Map( - ("readSidePadding", classOf[CharVarcharCodegenUtils]) -> CometScalarFunction( + /** + * Handlers keyed by `(functionName, staticObject class name)`. Class names rather than classes + * so that Iceberg's system functions, whose classes are not on Comet's compile classpath, can + * share the map; see [[CometIcebergSystemFunctions]]. + */ + private val staticInvokeExpressions: Map[(String, String), CometExpressionSerde[StaticInvoke]] = + Map[(String, String), CometExpressionSerde[StaticInvoke]]( + ("readSidePadding", classOf[CharVarcharCodegenUtils].getName) -> CometScalarFunction( "read_side_padding"), - ("isLuhnNumber", classOf[ExpressionImplUtils]) -> CometScalarFunction("luhn_check"), - ("encode", UrlCodec.getClass) -> CometUrlEncodeStaticInvoke, - ("decode", UrlCodec.getClass) -> CometUrlDecodeStaticInvoke, - ("aesEncrypt", classOf[ExpressionImplUtils]) -> CometStaticInvokeCodegenDispatch, - ("aesDecrypt", classOf[ExpressionImplUtils]) -> CometStaticInvokeCodegenDispatch, + ("isLuhnNumber", classOf[ExpressionImplUtils].getName) -> CometScalarFunction("luhn_check"), + ("encode", UrlCodec.getClass.getName) -> CometUrlEncodeStaticInvoke, + ("decode", UrlCodec.getClass.getName) -> CometUrlDecodeStaticInvoke, + ("aesEncrypt", classOf[ExpressionImplUtils].getName) -> CometStaticInvokeCodegenDispatch, + ("aesDecrypt", classOf[ExpressionImplUtils].getName) -> CometStaticInvokeCodegenDispatch, // Spark 4.0 lowers `decode(bin, charset)` to `StaticInvoke(StringDecode.decode, ...)` // carrying the `legacyCharsets` / `legacyErrorAction` flags. Routing through the codegen // dispatcher runs Spark's own decoder so both flags are honored. See #4465. - ("decode", classOf[StringDecode]) -> CometStaticInvokeCodegenDispatch, + ("decode", classOf[StringDecode].getName) -> CometStaticInvokeCodegenDispatch, // Spark 3.5+ makes `Base64` RuntimeReplaceable, lowering `base64(bin)` to // `StaticInvoke(Base64.encode, Seq(child, chunkBase64), ...)`. On Spark 3.4 the `Base64` // node survives and is handled directly (see CometBase64). - ("encode", classOf[Base64]) -> CometBase64StaticInvoke) + ("encode", classOf[Base64].getName) -> CometBase64StaticInvoke) ++ + CometIcebergSystemFunctions.staticInvokeHandlers - /** - * Iceberg's system functions (`bucket`, `truncate`, `years`, ...) are matched on class name - * because Iceberg is not on Comet's compile classpath; see [[CometIcebergSystemFunctions]]. - */ private def handlerFor(expr: StaticInvoke): Option[CometExpressionSerde[StaticInvoke]] = - staticInvokeExpressions - .get((expr.functionName, expr.staticObject)) - .orElse(CometIcebergSystemFunctions.staticInvokeHandlers.get(expr.staticObject.getName)) + staticInvokeExpressions.get((expr.functionName, expr.staticObject.getName)) override def getSupportLevel(expr: StaticInvoke): SupportLevel = handlerFor(expr).map(_.getSupportLevel(expr)).getOrElse(Compatible()) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index 6521d9f7b3f..e04fa24ab06 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -19,29 +19,28 @@ package org.apache.comet +import java.io.File import java.math.{BigDecimal => JBigDecimal, BigInteger} +import java.nio.file.Files import java.time.{Instant, LocalDate, LocalDateTime, ZoneOffset} -import scala.collection.mutable import scala.util.Random import org.scalactic.source.Position import org.scalatest.Tag -import org.apache.spark.{CometListenerBusUtils, SparkConf} +import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, Literal} import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSortExec} import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} -import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.spark.sql.util.QueryExecutionListener -import org.apache.comet.serde.{CometIcebergBucket, CometIcebergTruncate, CometStaticInvoke, Compatible, Unsupported} +import org.apache.comet.serde.{CometExpressionSerde, CometIcebergBucket, CometIcebergTruncate, CometStaticInvoke, Compatible, SupportLevel, Unsupported} /** * Native support for Iceberg's system functions (`bucket`, `truncate`, `years`, `months`, `days`, @@ -79,20 +78,37 @@ class CometIcebergSystemFunctionSuite Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin", "dt", "ts", "ts_ntz") private val truncateColumns = Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin") + // The source data is written once per suite; every test reads the same parquet directory. + private var sourceDir: File = _ + private def sourcePath: String = new File(sourceDir, "data").getAbsolutePath + + override def beforeAll(): Unit = { + super.beforeAll() + sourceDir = Files.createTempDirectory("comet-iceberg-system-functions").toFile + sourceData().write.parquet(sourcePath) + } + + override def afterAll(): Unit = { + try deleteRecursively(sourceDir) + finally super.afterAll() + } + test("bucket matches Iceberg for every supported type") { withSourceTable { - for (column <- bucketColumns; numBuckets <- Seq(1, 7, 16, Int.MaxValue)) { - checkSparkAnswerAndOperator( - s"SELECT $column, $catalog.system.bucket($numBuckets, $column) FROM $source") + bucketColumns.foreach { column => + val buckets = + Seq(1, 7, 16, Int.MaxValue).map(n => s"$catalog.system.bucket($n, $column)") + checkSparkAnswerAndOperator(s"SELECT $column, ${buckets.mkString(", ")} FROM $source") } } } test("truncate matches Iceberg for every supported type") { withSourceTable { - for (column <- truncateColumns; width <- Seq(1, 3, 10, 1000, Int.MaxValue)) { - checkSparkAnswerAndOperator( - s"SELECT $column, $catalog.system.truncate($width, $column) FROM $source") + truncateColumns.foreach { column => + val truncated = + Seq(1, 3, 10, 1000, Int.MaxValue).map(w => s"$catalog.system.truncate($w, $column)") + checkSparkAnswerAndOperator(s"SELECT $column, ${truncated.mkString(", ")} FROM $source") } } } @@ -103,13 +119,12 @@ class CometIcebergSystemFunctionSuite // leak into the native result either. for (timezone <- Seq("UTC", "America/Los_Angeles", "Asia/Kathmandu")) { withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { - for (column <- Seq("dt", "ts", "ts_ntz"); function <- Seq("years", "months", "days")) { - checkSparkAnswerAndOperator( - s"SELECT $column, $catalog.system.$function($column) FROM $source") - } - for (column <- Seq("ts", "ts_ntz")) { + Seq("dt", "ts", "ts_ntz").foreach { column => + val functions = Seq("years", "months", "days") ++ (if (column == "dt") Nil + else Seq("hours")) + val transformed = functions.map(f => s"$catalog.system.$f($column)") checkSparkAnswerAndOperator( - s"SELECT $column, $catalog.system.hours($column) FROM $source") + s"SELECT $column, ${transformed.mkString(", ")} FROM $source") } } } @@ -134,11 +149,10 @@ class CometIcebergSystemFunctionSuite s"SELECT i32, str FROM $source " + s"ORDER BY $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), i32, str") checkSparkAnswerAndOperator(df) - assert( - collect(stripAQEPlan(df.queryExecution.executedPlan)) { case s: CometSortExec => - s - }.nonEmpty, - "expected a native sort") + val sorts = collect(stripAQEPlan(df.queryExecution.executedPlan)) { case s: CometSortExec => + s + } + assert(sorts.nonEmpty, "expected a native sort") } } @@ -163,7 +177,7 @@ class CometIcebergSystemFunctionSuite USING iceberg PARTITIONED BY (bucket(4, i32), truncate(2, str), days(ts), months(dt))""") try { - val plans = capturePlans { + val plans = capturePlans(spark) { sql(s"INSERT INTO $table SELECT i32, str, ts, dt FROM $source") } val writePlans = plans.filter(plan => @@ -222,53 +236,31 @@ class CometIcebergSystemFunctionSuite test("support levels follow Iceberg's bind rules") { val value = AttributeReference("v", IntegerType)() - def invoke(cls: Class[_], args: Seq[org.apache.spark.sql.catalyst.expressions.Expression]) = - StaticInvoke(cls, IntegerType, "invoke", args, propagateNull = false) + val float = AttributeReference("f", FloatType)() + val date = AttributeReference("d", DateType)() + def level( + serde: CometExpressionSerde[StaticInvoke], + cls: Class[_], + args: Expression*): SupportLevel = + serde.getSupportLevel(StaticInvoke(cls, IntegerType, "invoke", args, propagateNull = false)) val bucketInt = Class.forName("org.apache.iceberg.spark.functions.BucketFunction$BucketInt") - assert( - CometIcebergBucket.getSupportLevel( - invoke(bucketInt, Seq(Literal(4), value))) == Compatible()) - assert( - CometIcebergBucket - .getSupportLevel(invoke(bucketInt, Seq(Literal(4.toShort), value))) == Compatible()) - assert( - CometIcebergBucket - .getSupportLevel(invoke(bucketInt, Seq(Literal(0), value))) - .isInstanceOf[Unsupported]) - assert( - CometIcebergBucket - .getSupportLevel(invoke(bucketInt, Seq(Literal(-4), value))) - .isInstanceOf[Unsupported]) - assert( - CometIcebergBucket - .getSupportLevel(invoke(bucketInt, Seq(value, value))) - .isInstanceOf[Unsupported]) - assert( - CometIcebergBucket - .getSupportLevel(invoke(bucketInt, Seq(Literal(4), AttributeReference("f", FloatType)()))) - .isInstanceOf[Unsupported]) + assert(level(CometIcebergBucket, bucketInt, Literal(4), value) == Compatible()) + assert(level(CometIcebergBucket, bucketInt, Literal(4.toShort), value) == Compatible()) + assert(level(CometIcebergBucket, bucketInt, Literal(0), value).isInstanceOf[Unsupported]) + assert(level(CometIcebergBucket, bucketInt, Literal(-4), value).isInstanceOf[Unsupported]) + assert(level(CometIcebergBucket, bucketInt, value, value).isInstanceOf[Unsupported]) + assert(level(CometIcebergBucket, bucketInt, Literal(4), float).isInstanceOf[Unsupported]) val truncateInt = Class.forName("org.apache.iceberg.spark.functions.TruncateFunction$TruncateInt") - assert( - CometIcebergTruncate - .getSupportLevel(invoke(truncateInt, Seq(Literal(10), value))) == Compatible()) - assert( - CometIcebergTruncate - .getSupportLevel( - invoke(truncateInt, Seq(Literal(10), AttributeReference("d", DateType)()))) - .isInstanceOf[Unsupported]) + assert(level(CometIcebergTruncate, truncateInt, Literal(10), value) == Compatible()) + assert(level(CometIcebergTruncate, truncateInt, Literal(10), date).isInstanceOf[Unsupported]) - // The dispatch in CometStaticInvoke goes by class name, so the same expressions resolve to - // the Iceberg handlers there too. - assert( - CometStaticInvoke.getSupportLevel( - invoke(bucketInt, Seq(Literal(4), value))) == Compatible()) - assert( - CometStaticInvoke - .getSupportLevel(invoke(bucketInt, Seq(Literal(0), value))) - .isInstanceOf[Unsupported]) + // CometStaticInvoke dispatches on (functionName, class name), so the same expressions reach + // the Iceberg handlers from there too. + assert(level(CometStaticInvoke, bucketInt, Literal(4), value) == Compatible()) + assert(level(CometStaticInvoke, bucketInt, Literal(0), value).isInstanceOf[Unsupported]) } test("fallback reason for an unlisted static invoke names the declaring class") { @@ -293,10 +285,7 @@ class CometIcebergSystemFunctionSuite s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog", s"spark.sql.catalog.$catalog.type" -> "hadoop", s"spark.sql.catalog.$catalog.warehouse" -> warehouseDir.getAbsolutePath) { - withTempPath { dir => - sourceData().write.parquet(dir.getAbsolutePath) - withParquetTable(dir.getAbsolutePath, source)(f) - } + withParquetTable(sourcePath, source)(f) } } @@ -314,6 +303,12 @@ class CometIcebergSystemFunctionSuite StructField("ts", TimestampType), StructField("ts_ntz", TimestampNTZType))) + private def instant(micros: Long): Instant = + Instant.ofEpochSecond(Math.floorDiv(micros, 1000000L), Math.floorMod(micros, 1000000L) * 1000) + + private def localDateTime(micros: Long): LocalDateTime = + LocalDateTime.ofInstant(instant(micros), ZoneOffset.UTC) + /** * Seeded random rows with nulls in every column, followed by the boundary values of each type * (numeric extremes, the epoch and the microsecond before it, empty and multi-byte strings). @@ -335,15 +330,6 @@ class CometIcebergSystemFunctionSuite new JBigDecimal(if (random.nextBoolean()) unscaled else unscaled.negate(), 10) } def randomMicros(): Long = random.nextLong() % 4000000000000000L - def instant(micros: Long): Instant = - Instant.ofEpochSecond( - Math.floorDiv(micros, 1000000L), - Math.floorMod(micros, 1000000L) * 1000) - def localDateTime(micros: Long): LocalDateTime = - LocalDateTime.ofEpochSecond( - Math.floorDiv(micros, 1000000L), - (Math.floorMod(micros, 1000000L) * 1000).toInt, - ZoneOffset.UTC) val randomRows = (0 until 400).map { _ => Row( @@ -359,78 +345,51 @@ class CometIcebergSystemFunctionSuite maybeNull(instant(randomMicros())), maybeNull(localDateTime(randomMicros()))) } + + // One list of boundary values per column, in schema order; transposed into rows below so + // each type's cases sit together (and a list of the wrong length fails loudly). val dec38Max = new JBigDecimal(BigInteger.TEN.pow(38).subtract(BigInteger.ONE), 10) - val boundaryRows = Seq( - Row( - Byte.MinValue, - Short.MinValue, - Int.MinValue, - Long.MinValue, + val boundaryColumns: Seq[Seq[Any]] = Seq( + Seq(Byte.MinValue, Byte.MaxValue, 0.toByte, (-1).toByte, null), + Seq(Short.MinValue, Short.MaxValue, 0.toShort, (-1).toShort, null), + Seq(Int.MinValue, Int.MaxValue, 0, -1, null), + Seq(Long.MinValue, Long.MaxValue, 0L, -1L, null), + Seq( new JBigDecimal("-99999999999999.9999"), - dec38Max.negate(), - "", - Array.empty[Byte], - LocalDate.ofEpochDay(0), - Instant.EPOCH, - LocalDateTime.of(1970, 1, 1, 0, 0)), - Row( - Byte.MaxValue, - Short.MaxValue, - Int.MaxValue, - Long.MaxValue, new JBigDecimal("99999999999999.9999"), - dec38Max, - "日本語😀", - Array[Byte](0, 1, 2, 3), - LocalDate.ofEpochDay(-1), - instant(-1L), - localDateTime(-1L)), - Row( - 0.toByte, - 0.toShort, - 0, - 0L, JBigDecimal.ZERO.setScale(4), + new JBigDecimal("-0.0001"), + null), + Seq( + dec38Max.negate(), + dec38Max, JBigDecimal.ZERO.setScale(10), - "a", + new JBigDecimal("-0.0000000001"), + null), + Seq("", "日本語😀", "a", "iceberg", null), + Seq( + Array.empty[Byte], + Array[Byte](0, 1, 2, 3), Array[Byte](0), + Array.fill[Byte](5)(-1), + null), + Seq( + LocalDate.ofEpochDay(0), + LocalDate.ofEpochDay(-1), LocalDate.ofEpochDay(-365), - instant(-86400000000L), - localDateTime(-86400000000L - 1)), - Row( - (-1).toByte, - (-1).toShort, - -1, - -1L, - new JBigDecimal("-0.0001"), - new JBigDecimal("-0.0000000001"), - "iceberg", - Array[Byte](-1, -1, -1, -1, -1), LocalDate.ofEpochDay(-366), - instant(-3600000000L - 1), - localDateTime(-3600000000L)), - Row(null, null, null, null, null, null, null, null, null, null, null)) + null), + Seq(Instant.EPOCH, instant(-1L), instant(-86400000000L), instant(-3600000000L - 1), null), + Seq( + LocalDateTime.of(1970, 1, 1, 0, 0), + localDateTime(-1L), + localDateTime(-86400000000L - 1), + localDateTime(-3600000000L), + null)) + val boundaryRows = boundaryColumns.transpose.map(Row.fromSeq) + spark.createDataFrame( spark.sparkContext.parallelize(randomRows ++ boundaryRows, 3), sourceSchema) } - - private def capturePlans(action: => Unit): Seq[SparkPlan] = { - val captured = mutable.Buffer.empty[SparkPlan] - val listener = new QueryExecutionListener { - override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { - captured += qe.executedPlan - } - override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = - () - } - spark.listenerManager.register(listener) - try { - action - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - } finally { - spark.listenerManager.unregister(listener) - } - captured.toSeq - } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala index 65d6fac97e9..f6e21f39010 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala @@ -22,8 +22,13 @@ package org.apache.comet import java.io.File import java.nio.file.Files +import scala.collection.mutable + +import org.apache.spark.CometListenerBusUtils import org.apache.spark.sql.SparkSession import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog} +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus import org.apache.comet.iceberg.IcebergReflection @@ -131,4 +136,24 @@ trait CometIcebergTestBase { if (file.isDirectory) file.listFiles().foreach(deleteRecursively) file.delete() } + + /** The executed plan of every query that completes successfully while `action` runs. */ + protected def capturePlans(spark: SparkSession)(action: => Unit): Seq[SparkPlan] = { + val captured = mutable.Buffer.empty[SparkPlan] + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + captured += qe.executedPlan + } + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + () + } + spark.listenerManager.register(listener) + try { + action + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } finally { + spark.listenerManager.unregister(listener) + } + captured.toSeq + } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 514123a2520..83039e54147 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -27,15 +27,14 @@ import scala.concurrent.{Await, Future} import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt -import org.apache.spark.{CometListenerBusUtils, SparkConf} +import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.Row import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog -import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} -import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} @@ -520,7 +519,7 @@ class CometIcebergWriteActionSuite CometConf.COMET_EXEC_ENABLED.key -> "true") { spark.sql("CREATE TABLE testcat.tbl (id INT, region STRING, amount DOUBLE)") try { - val plans = capturePlans { + val plans = capturePlans(spark) { spark.sql("INSERT INTO testcat.tbl VALUES (1, 'us-east', 10.5)") } val (commits, writes) = collectIcebergWriteOps(plans) @@ -1517,14 +1516,14 @@ class CometIcebergWriteActionSuite } } - val ctasPlans = capturePlans { + val ctasPlans = capturePlans(spark) { spark.sql(s"CREATE TABLE $catalog.$ns.ctas_tgt USING iceberg AS SELECT * FROM ctas_src") } assertSplitUsage(ctasPlans, "CTAS") assert(countSnapshots("ctas_tgt") == 1L, "CTAS must land exactly one snapshot") assertRows("ctas_tgt", expectedIds = Seq(1, 2, 3, 4, 5)) - val rtasPlans = capturePlans { + val rtasPlans = capturePlans(spark) { (1 to 2) .map(i => (i, s"r$i", i.toDouble)) .toDF("id", "region", "amount") @@ -1578,28 +1577,9 @@ class CometIcebergWriteActionSuite .append() } - private def capturePlans(action: => Unit): Seq[SparkPlan] = { - val captured = mutable.Buffer.empty[SparkPlan] - val listener = new QueryExecutionListener { - override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { - captured += qe.executedPlan - } - override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = - () - } - spark.listenerManager.register(listener) - try { - action - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - } finally { - spark.listenerManager.unregister(listener) - } - captured.toSeq - } - private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) - val plans = capturePlans(action) + val plans = capturePlans(spark)(action) WriteSnapshot(countSnapshots(tableName) - before, plans) } From 49868b63106ca88ff7e64197c321d314415a771f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 16:03:27 -0600 Subject: [PATCH 04/18] fix: use as_chunks in the Iceberg murmur3 kernel to satisfy clippy 1.98 --- native/spark-expr/src/iceberg_funcs/bucket.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/native/spark-expr/src/iceberg_funcs/bucket.rs b/native/spark-expr/src/iceberg_funcs/bucket.rs index aec67f84c04..cd5b9407900 100644 --- a/native/spark-expr/src/iceberg_funcs/bucket.rs +++ b/native/spark-expr/src/iceberg_funcs/bucket.rs @@ -50,13 +50,11 @@ pub(crate) fn murmur3_32(data: &[u8]) -> i32 { } let mut h1: u32 = 0; - let mut chunks = data.chunks_exact(4); - for chunk in &mut chunks { - let k1 = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); - h1 ^= mix_k1(k1); + let (chunks, tail) = data.as_chunks::<4>(); + for chunk in chunks { + h1 ^= mix_k1(u32::from_le_bytes(*chunk)); h1 = h1.rotate_left(13).wrapping_mul(5).wrapping_add(0xe654_6b64); } - let tail = chunks.remainder(); if !tail.is_empty() { let mut k1: u32 = 0; for (i, byte) in tail.iter().enumerate() { From f8db9b5fdfd5333522620b04c0132572cdb343ec Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 17:16:48 -0600 Subject: [PATCH 05/18] test: keep multi-byte strings and Long.MinValue out of the partitioned write test The CI containers run the JVM with an ASCII platform charset, so raw multi-byte partition directory names written by the pinned iceberg-rust cannot be reopened by iceberg-java for the metrics rebuild; partition on the long column instead. iceberg-rust's truncate transform also overflows on Long.MinValue in debug builds, so that boundary row is not written. --- .../CometIcebergSystemFunctionSuite.scala | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index e04fa24ab06..58f00b66407 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -171,14 +171,21 @@ class CometIcebergSystemFunctionSuite withSourceTable { val table = s"$catalog.db.hidden_partitioning" // No `write.distribution-mode`: Iceberg picks hash distribution for a partitioned table, - // which plans a shuffle and a local sort on the partition transforms. + // which plans a shuffle and a local sort on the partition transforms. The string column is + // deliberately not a partition source: its multi-byte values would land in partition + // directory names, which iceberg-java reads back through the JVM's platform charset. sql(s""" - CREATE TABLE $table (i32 INT, str STRING, ts TIMESTAMP, dt DATE) + CREATE TABLE $table (i32 INT, i64 BIGINT, ts TIMESTAMP, dt DATE) USING iceberg - PARTITIONED BY (bucket(4, i32), truncate(2, str), days(ts), months(dt))""") + PARTITIONED BY (bucket(4, i32), truncate(1000, i64), days(ts), months(dt))""") + // iceberg-rust's own truncate transform, which the writer uses for partition values, does + // `v - ((v % w) + w) % w` without wrapping and overflows on Long.MinValue in debug builds + // (Java wraps), so the boundary row stays out of the written set. + val rows = + s"SELECT i32, i64, ts, dt FROM $source WHERE i64 IS NULL OR i64 <> ${Long.MinValue}" try { val plans = capturePlans(spark) { - sql(s"INSERT INTO $table SELECT i32, str, ts, dt FROM $source") + sql(s"INSERT INTO $table $rows") } val writePlans = plans.filter(plan => collectWithSubqueries(plan) { case w: CometIcebergWriteExec => w }.nonEmpty) @@ -196,16 +203,14 @@ class CometIcebergSystemFunctionSuite s"the distribution shuffle stayed on Spark:\n$plan") } - checkAnswer( - sql(s"SELECT i32, str, ts, dt FROM $table"), - sql(s"SELECT i32, str, ts, dt FROM $source").collect()) + checkAnswer(sql(s"SELECT i32, i64, ts, dt FROM $table"), sql(rows).collect()) // Iceberg's own view of the partitions must match what the JVM transforms compute over // the written rows: one partition per distinct transform tuple. withSQLConf(CometConf.COMET_ENABLED.key -> "false") { val expected = sql(s""" SELECT COUNT(*) FROM ( - SELECT DISTINCT $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), + SELECT DISTINCT $catalog.system.bucket(4, i32), $catalog.system.truncate(1000, i64), $catalog.system.days(ts), $catalog.system.months(dt) FROM $table)""").collect().head.getLong(0) val actual = sql(s"SELECT COUNT(*) FROM $table.partitions").collect().head.getLong(0) From 7de2708a92847e6860824d80dbaa27d3a8b3bcf1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 09:37:31 -0600 Subject: [PATCH 06/18] test: align the system-function corpus write with Spark 4 datetime defaults CometIcebergSystemFunctionSuite aborted in beforeAll on Spark 3.4 and 3.5. Three Spark 3.x defaults differ from Spark 4's and each blocks writing the corpus: - datetimeJava8ApiEnabled is off, so the row encoder rejects the java.time values sourceData supplies (ClassCastException: LocalDate -> java.sql.Date) - datetimeRebaseModeInWrite is EXCEPTION, which rejects the deliberately pre-epoch timestamps (the corpus reaches back to 1843) - outputTimestampType is INT96, which has its own ancient-timestamp check Set all three to Spark 4's values around the write only. Every test reads the data back from parquet, so no result comparison is affected. --- .../CometIcebergSystemFunctionSuite.scala | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index 58f00b66407..c584e07b7bf 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -85,7 +85,27 @@ class CometIcebergSystemFunctionSuite override def beforeAll(): Unit = { super.beforeAll() sourceDir = Files.createTempDirectory("comet-iceberg-system-functions").toFile - sourceData().write.parquet(sourcePath) + // Three Spark 3.x defaults differ from Spark 4's in ways that block writing this corpus. All + // are set to Spark 4's value so one corpus works on every profile, and none affects how a + // result is compared, since every test reads the data back from parquet. + // + // - `datetimeJava8ApiEnabled`: `sourceData` supplies java.time values for the date and + // timestamp columns. Spark 4 resolves those external types; on Spark 3.x the row encoder + // expects java.sql.Date / java.sql.Timestamp instead. The encoder is built here on the + // driver, so setting the flag around the write is enough. + // - `datetimeRebaseModeInWrite`: the timestamp corpus reaches back to 1843, and the corpus + // is deliberately pre-epoch in places, since the temporal transforms go negative before + // 1970. Spark 3.x throws on writing a timestamp before 1900; Spark 4 defaults to + // CORRECTED, which writes the value as-is. + // - `outputTimestampType`: Spark 3.x defaults to INT96, which has its own separate ancient + // timestamp check. Spark 4 defaults to TIMESTAMP_MICROS, which is also what Iceberg + // itself writes. + withSQLConf( + SQLConf.DATETIME_JAVA8API_ENABLED.key -> "true", + SQLConf.PARQUET_REBASE_MODE_IN_WRITE.key -> "CORRECTED", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "TIMESTAMP_MICROS") { + sourceData().write.parquet(sourcePath) + } } override def afterAll(): Unit = { From 71e03beb626404e01d21941a06973b7a7ee700a5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:34:19 -0600 Subject: [PATCH 07/18] fix: cover the whole date domain in the Iceberg temporal kernels `years` and `months` split the calendar with `chrono::NaiveDate`, whose range stops at about year 262143. A Spark `DateType` is an `i32` epoch day (up to year 5881580) and Iceberg's `DateTimeUtil` goes through `LocalDate`, which covers all of it, so `years(DATE)` on a far-future date raised an execution error where the JVM returns a value. Replace the `chrono` round trip with Howard Hinnant's `civil_from_days` in `i64` arithmetic, which is exact over the whole `i32` epoch-day domain and drops the fallible path from the kernels. The pinned expectations for the extremes come from running `DateTimeUtil.convertDays` / `convertMicros` on a JVM, and a new test checks the integer split against `chrono` everywhere `chrono` can represent the date. --- .../spark-expr/src/iceberg_funcs/temporal.rs | 135 +++++++++++++----- 1 file changed, 102 insertions(+), 33 deletions(-) diff --git a/native/spark-expr/src/iceberg_funcs/temporal.rs b/native/spark-expr/src/iceberg_funcs/temporal.rs index 55e4deea8c3..f49c5fc122e 100644 --- a/native/spark-expr/src/iceberg_funcs/temporal.rs +++ b/native/spark-expr/src/iceberg_funcs/temporal.rs @@ -23,14 +23,21 @@ //! `hours` are plain floor division of the epoch value. `days` returns a date (Iceberg's //! `DaysFunction.resultType()` is `DateType`), the other three return an int. //! -//! The kernels work on the raw epoch values rather than going through Arrow's timezone-aware -//! `date_part`, which would otherwise shift a `TimestampType` column by the session offset. +//! The kernels read the raw epoch values instead of Arrow's timezone-aware `date_part`. That is a +//! defensive choice rather than a fix for an observed bug: Comet tags every `TimestampType` array +//! `UTC` and the Iceberg writer casts each batch to a schema that tags `Timestamptz` as `+00:00`, +//! so `date_part` would agree today. It only keeps agreeing for as long as that tagging holds, +//! whereas the epoch arithmetic below is correct for any tag. +//! +//! The calendar split is integer arithmetic rather than a `chrono::NaiveDate`, which covers only +//! about ±262k years. A Spark `DateType` is an `i32` epoch day (up to year 5881580) and Java's +//! `LocalDate`, which Iceberg uses, covers all of it, so going through `chrono` would turn values +//! the JVM handles into execution errors. use super::{apply_unary, unsupported_type}; use arrow::array::{ArrayRef, AsArray, Int32Array}; use arrow::datatypes::{DataType, Date32Type, Int32Type, TimeUnit, TimestampMicrosecondType}; -use chrono::Datelike; -use datafusion::common::{utils::take_function_args, DataFusionError, Result}; +use datafusion::common::{utils::take_function_args, Result}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -67,46 +74,60 @@ impl TemporalUnit { } } -/// `DateTimeUtil.microsToDays`: floor division, so `-1` micros is day `-1`. +/// `DateTimeUtil.microsToDays`: floor division, so `-1` micros is day `-1`. The quotient of the +/// widest `i64` micros is about 1.07e8, so the narrowing is always exact here. #[inline] fn micros_to_days(micros: i64) -> i32 { div_floor(micros, MICROS_PER_DAY) as i32 } -/// `DateTimeUtil.microsToHours`. +/// `DateTimeUtil.microsToHours`. Java narrows the hour count with a plain `(int)` cast, which +/// wraps beyond about 7.7e18 micros; `as i32` truncates the same way. #[inline] fn micros_to_hours(micros: i64) -> i32 { div_floor(micros, MICROS_PER_HOUR) as i32 } /// `DateTimeUtil.daysToYears`: whole calendar years between the epoch and the day, floored. -fn days_to_years(days: i32) -> Result { - Ok(civil_date(days)?.year() - UNIX_EPOCH_YEAR) +fn days_to_years(days: i32) -> i32 { + civil_from_days(days).0 - UNIX_EPOCH_YEAR } /// `DateTimeUtil.daysToMonths`: whole calendar months between the epoch and the day, floored. -fn days_to_months(days: i32) -> Result { - let date = civil_date(days)?; - Ok((date.year() - UNIX_EPOCH_YEAR) * 12 + date.month0() as i32) +fn days_to_months(days: i32) -> i32 { + let (year, month0) = civil_from_days(days); + (year - UNIX_EPOCH_YEAR) * 12 + month0 } -fn civil_date(days: i32) -> Result { - Date32Type::to_naive_date_opt(days).ok_or_else(|| { - DataFusionError::Execution(format!("day {days} is out of the supported date range")) - }) +/// Splits an epoch day into its proleptic Gregorian `(year, month0)`, following Howard Hinnant's +/// `civil_from_days`. The intermediates are `i64` so that every `i32` epoch day is in range, which +/// is what `LocalDate` gives Iceberg; the widest results, at `i32::MIN` and `i32::MAX` days, are +/// years -5877641 and 5881580, so both the year and the month count still fit in an `i32`. +fn civil_from_days(days: i32) -> (i32, i32) { + // Shift the epoch to 0000-03-01 so that the leap day falls at the end of the year. + let shifted = days as i64 + 719_468; + let era = shifted.div_euclid(146_097); + let day_of_era = shifted.rem_euclid(146_097); // [0, 146096] + let year_of_era = + (day_of_era - day_of_era / 1460 + day_of_era / 36_524 - day_of_era / 146_096) / 365; // [0, 399] + let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); // [0, 365] + let shifted_month = (5 * day_of_year + 2) / 153; // [0, 11], March is 0 + let month = if shifted_month < 10 { + shifted_month + 3 + } else { + shifted_month - 9 + }; + let year = era * 400 + year_of_era + i64::from(month <= 2); + (year as i32, (month - 1) as i32) } /// Applies a calendar function of the epoch day to a date or timestamp column in one pass. -fn map_epoch_days( - fn_name: &str, - array: &ArrayRef, - f: impl Fn(i32) -> Result, -) -> Result { +fn map_epoch_days(fn_name: &str, array: &ArrayRef, f: impl Fn(i32) -> i32) -> Result { match array.data_type() { - DataType::Date32 => array.as_primitive::().try_unary(f), - DataType::Timestamp(TimeUnit::Microsecond, _) => array + DataType::Date32 => Ok(array.as_primitive::().unary(f)), + DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(array .as_primitive::() - .try_unary(|micros| f(micros_to_days(micros))), + .unary(|micros| f(micros_to_days(micros)))), other => Err(unsupported_type(fn_name, other)), } } @@ -203,17 +224,25 @@ mod tests { .unwrap() } - // Boundaries around the epoch, as (epoch days, years, months). + // Boundaries around the epoch, as (epoch days, years, months). The values past the epoch + // block are outside `chrono::NaiveDate`'s range but well inside `LocalDate`'s; they come from + // running Iceberg's `DateTimeUtil.convertDays` on a JDK 17 JVM. const DAY_CASES: &[(i32, i32, i32)] = &[ - (17_486, 47, 574), // 2017-11-16, the Iceberg spec example - (0, 0, 0), // 1970-01-01 - (-1, -1, -1), // 1969-12-31 - (-365, -1, -12), // 1969-01-01 - (-366, -2, -13), // 1968-12-31 - (365, 1, 12), // 1971-01-01 - (364, 0, 11), // 1970-12-31 - (31, 0, 1), // 1970-02-01 - (30, 0, 0), // 1970-01-31 + (17_486, 47, 574), // 2017-11-16, the Iceberg spec example + (0, 0, 0), // 1970-01-01 + (-1, -1, -1), // 1969-12-31 + (-365, -1, -12), // 1969-01-01 + (-366, -2, -13), // 1968-12-31 + (365, 1, 12), // 1971-01-01 + (364, 0, 11), // 1970-12-31 + (31, 0, 1), // 1970-02-01 + (30, 0, 0), // 1970-01-31 + (100_000_000, 273_790, 3_285_488), // +275760-09-13 + (-100_000_000, -273_791, -3_285_489), // -271821-04-20 + (1_000_000_000, 2_737_907, 32_854_884), // +2739877-01-03 + (-1_000_000_000, -2_737_908, -32_854_885), // -2735938-12-29 + (i32::MAX, 5_879_610, 70_555_326), // +5881580-07-11 + (i32::MIN, -5_879_611, -70_555_327), // -5877641-06-23 ]; #[test] @@ -266,6 +295,25 @@ mod tests { (-MICROS_PER_DAY - 1, -1, -1, -2, -25), // 1969-12-30T23:59:59.999999 (365 * MICROS_PER_DAY, 1, 12, 365, 8_760), // 1971-01-01T00:00:00 (365 * MICROS_PER_DAY - 1, 0, 11, 364, 8_759), + // The extremes of Spark's timestamp domain, from Iceberg's `DateTimeUtil` on a JVM. + // The hour count is the one conversion that does not fit an `i32` there, and Java's + // `(int)` narrowing wraps it exactly as `as i32` does. + (i64::MAX, 292_277, 3_507_324, 106_751_991, -1_732_919_508), + (i64::MIN, -292_278, -3_507_325, -106_751_992, 1_732_919_507), + ( + 8_000_000_000_000_000_000, + 253_509, + 3_042_118, + 92_592_592, + -2_072_745_074, + ), + ( + -8_000_000_000_000_000_000, + -253_510, + -3_042_119, + -92_592_593, + 2_072_745_073, + ), ]; let micros: Vec> = cases.iter().map(|c| Some(c.0)).chain([None]).collect(); // A non-UTC timezone tag must not change the result. @@ -318,6 +366,27 @@ mod tests { } } + /// The pinned cases above check the endpoints of the domain; this checks the calendar split + /// itself everywhere `chrono` can still represent the date, so the hand-rolled arithmetic + /// cannot drift in between. + #[test] + fn civil_from_days_agrees_with_chrono() { + use chrono::Datelike; + // Every day of one full 400-year Gregorian cycle either side of the epoch, then a stride + // over the rest of the `i32` domain (`to_naive_date_opt` returns `None` past chrono's + // range, which is exactly the region the pinned cases cover). + let dense = -146_097..=146_097; + for days in dense.chain((i32::MIN..=i32::MAX).step_by(999_983)) { + if let Some(date) = Date32Type::to_naive_date_opt(days) { + assert_eq!( + civil_from_days(days), + (date.year(), date.month0() as i32), + "day {days} ({date})" + ); + } + } + } + #[test] fn rejects_unsupported_types() { for unit in [ From 1622503f3f4afbf183905bc9131233e5994eb4be Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:34:32 -0600 Subject: [PATCH 08/18] perf: transform Iceberg dictionary inputs once per distinct value `apply_unary` unpacked a dictionary with `cast` before handing it to the kernel, so a dictionary-encoded string column -- what a Parquet scan hands back for a low cardinality partition column -- paid for a full copy of the values buffer and then hashed or truncated every row instead of every distinct value. Run the kernel over the dictionary's values and expand the result through the keys instead. Over 8192 rows with eight distinct strings (`cargo bench --bench iceberg_transforms`): iceberg_bucket/string_dict 72.9 us -> 3.0 us iceberg_truncate/string_dict 140.9 us -> 35.4 us The decimal comment in `truncate` also now spells out which paths agree with the JVM result and which do not. --- native/spark-expr/src/iceberg_funcs/bucket.rs | 2 +- native/spark-expr/src/iceberg_funcs/mod.rs | 40 +++++++++-------- .../spark-expr/src/iceberg_funcs/truncate.rs | 43 +++++++++++++++++-- 3 files changed, 62 insertions(+), 23 deletions(-) diff --git a/native/spark-expr/src/iceberg_funcs/bucket.rs b/native/spark-expr/src/iceberg_funcs/bucket.rs index cd5b9407900..4e267fdd706 100644 --- a/native/spark-expr/src/iceberg_funcs/bucket.rs +++ b/native/spark-expr/src/iceberg_funcs/bucket.rs @@ -309,7 +309,7 @@ mod tests { } #[test] - fn dictionary_input_is_unpacked() { + fn dictionary_input_is_hashed_once_per_value() { let dict: DictionaryArray = vec![Some("iceberg"), None, Some("iceberg")] .into_iter() .collect(); diff --git a/native/spark-expr/src/iceberg_funcs/mod.rs b/native/spark-expr/src/iceberg_funcs/mod.rs index 241c8d1cbdd..aec13f949b3 100644 --- a/native/spark-expr/src/iceberg_funcs/mod.rs +++ b/native/spark-expr/src/iceberg_funcs/mod.rs @@ -34,21 +34,11 @@ pub use bucket::SparkIcebergBucket; pub use temporal::SparkIcebergTemporalTransform; pub use truncate::SparkIcebergTruncate; -use arrow::array::{Array, ArrayRef}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, AsArray}; +use arrow::compute::take; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result, ScalarValue}; use datafusion::logical_expr::ColumnarValue; -use std::sync::Arc; - -/// Unpacks a dictionary-encoded array to its value type so that the kernels only ever see plain -/// arrays. Any other array is returned unchanged. -fn unpack_dictionary(array: ArrayRef) -> Result { - match array.data_type() { - DataType::Dictionary(_, value_type) => Ok(cast(&array, value_type)?), - _ => Ok(array), - } -} /// The type a kernel sees for an input of type `data_type`, after dictionary unpacking. fn unpacked_type(data_type: &DataType) -> DataType { @@ -64,13 +54,9 @@ fn apply_unary( kernel: impl Fn(&ArrayRef) -> Result, ) -> Result { match value { - ColumnarValue::Array(array) => { - let array = unpack_dictionary(Arc::clone(array))?; - Ok(ColumnarValue::Array(kernel(&array)?)) - } + ColumnarValue::Array(array) => Ok(ColumnarValue::Array(apply_to_array(array, kernel)?)), ColumnarValue::Scalar(scalar) => { - let array = unpack_dictionary(scalar.to_array()?)?; - let result = kernel(&array)?; + let result = apply_to_array(&scalar.to_array()?, kernel)?; Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( &result, 0, )?)) @@ -78,6 +64,24 @@ fn apply_unary( } } +/// Runs `kernel` over `array`, transforming a dictionary one distinct value at a time and then +/// expanding the result through the keys. A Parquet scan hands string partition columns over +/// dictionary-encoded, and unpacking first would hash or truncate every row rather than every +/// distinct value (plus copy the values buffer to do it). +fn apply_to_array( + array: &ArrayRef, + kernel: impl Fn(&ArrayRef) -> Result, +) -> Result { + match array.data_type() { + DataType::Dictionary(_, _) => { + let dictionary = array.as_any_dictionary(); + let values = kernel(dictionary.values())?; + Ok(take(values.as_ref(), dictionary.keys(), None)?) + } + _ => kernel(array), + } +} + /// Reads the `numBuckets` / `width` parameter. The Comet serde only converts these functions when /// the parameter is a positive integer literal, so anything else here is a wiring bug. fn positive_int_param(fn_name: &str, param: &str, value: &ColumnarValue) -> Result { diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs index 97cebf67742..3c5ab0bc729 100644 --- a/native/spark-expr/src/iceberg_funcs/truncate.rs +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -96,9 +96,15 @@ fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result { // Truncating a negative value grows its magnitude by up to `width - 1` units of the - // last digit, so the result can need one more digit than the column allows. Spark's - // `UnsafeRowWriter` writes such a `Decimal` as null (`changePrecision` fails), so - // match that rather than emit a value the column's precision cannot hold. + // last digit, so the result can need one more digit than the column allows. Iceberg's + // `TruncateDecimal.invoke` hands that oversized `Decimal` back to Spark unchanged and + // Spark nulls it only when a row is materialized (`UnsafeRowWriter` calls + // `changePrecision`, which fails). Nulling it here is the same answer for every path + // that writes the value into a row -- a projection, a sort key, a shuffle key, an + // Iceberg partition value -- and it is the only answer available to a kernel that has + // to return a `Decimal128(precision, scale)` array. The two differ where the result + // feeds another expression without being materialized, e.g. `truncate(w, v) IS NULL`; + // see the Iceberg user guide. let truncated: Decimal128Array = array.as_primitive::().unary_opt(|v| { let truncated = truncate_i128(v, width as i128); @@ -161,7 +167,10 @@ impl ScalarUDFImpl for SparkIcebergTruncate { mod tests { use super::super::test_util::invoke; use super::*; - use arrow::array::{BinaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray}; + use arrow::array::{ + BinaryArray, DictionaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray, + }; + use arrow::datatypes::Int8Type; use datafusion::common::ScalarValue; fn truncate(width: i32, value: ArrayRef) -> ArrayRef { @@ -346,6 +355,32 @@ mod tests { ); } + /// A dictionary is truncated once per distinct value and expanded through the keys, so the + /// result is a plain array of the value type, as [`SparkIcebergTruncate::return_type`] says. + #[test] + fn dictionary_input_is_truncated_once_per_value() { + let dict: DictionaryArray = + vec![Some("iceberg"), None, Some("ic"), Some("iceberg")] + .into_iter() + .collect(); + let result = truncate(3, Arc::new(dict)); + assert_eq!(result.data_type(), &DataType::Utf8); + assert_eq!( + result.as_string::(), + &StringArray::from(vec![Some("ice"), None, Some("ic"), Some("ice")]) + ); + // The whole-buffer shortcut has to survive the round trip through the keys too. + let unchanged = truncate(i32::MAX, Arc::new(dict_of(&[Some("ab"), None, Some("ab")]))); + assert_eq!( + unchanged.as_string::(), + &StringArray::from(vec![Some("ab"), None, Some("ab")]) + ); + } + + fn dict_of(values: &[Option<&str>]) -> DictionaryArray { + values.iter().copied().collect() + } + #[test] fn return_type_follows_value_type() { let udf = SparkIcebergTruncate::new(); From 8d898c29e8de588c9168a10a4118802433f4f267 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:34:49 -0600 Subject: [PATCH 09/18] test: pin the Iceberg kernels to iceberg-rust's partition transforms A partitioned write sorts on Comet's kernels and then groups by the partition values iceberg-rust computes, and the clustered writer fails at runtime when the two disagree. Assert they agree over the boundary inputs of every shared type, so an iceberg-rust bump that changes a transform breaks here first. The excluded cases are the interesting ones: iceberg-rust's `truncate` does not wrap like Java's, and its `years` / `months` go through Arrow's `date_part`, which honours the array's timezone tag. The last test pins that tag dependency, since it is the reason `years` and `months` keep local kernels while `bucket`, `days`, and `hours` could in principle be delegated. --- .../src/execution/operators/iceberg_write.rs | 358 ++++++++++++++++++ 1 file changed, 358 insertions(+) diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index 57f39d1b395..e23c856b961 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -1172,3 +1172,361 @@ mod tests { } } } + +/// Pins Comet's Iceberg system-function kernels to iceberg-rust's partition transforms. +/// +/// A partitioned write runs both: the sort in front of [`IcebergWriteExec`] is keyed on the +/// `datafusion-comet-spark-expr` kernels (Iceberg plans the sort as `bucket(...)`, `days(...)`, +/// ... system-function calls), while [`ClusteredWriter`] groups the sorted rows by the partition +/// values that [`PartitionValueCalculator`] computes with iceberg-rust's transforms. The writer +/// requires the two to agree: when they do not it fails at runtime with "The input is not sorted! +/// Cannot write to partition that was previously closed". These tests make an iceberg-rust bump +/// that changes a transform break here first. +#[cfg(test)] +mod iceberg_rust_transform_parity { + use arrow::array::{ + ArrayRef, BinaryArray, Date32Array, Decimal128Array, Int32Array, Int64Array, StringArray, + TimestampMicrosecondArray, + }; + use arrow::datatypes::{DataType, Field}; + use datafusion::common::ScalarValue; + use datafusion::config::ConfigOptions; + use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + use datafusion_comet_spark_expr::{ + SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, + }; + use iceberg::spec::Transform; + use iceberg::transform::create_transform_function; + use std::sync::Arc; + + const MICROS_PER_DAY: i64 = 86_400_000_000; + + /// Runs a Comet kernel over `value`, prepending `parameter` for the two-argument transforms. + fn comet(udf: &dyn ScalarUDFImpl, parameter: Option, value: &ArrayRef) -> ArrayRef { + let mut args: Vec = parameter + .map(|p| ColumnarValue::Scalar(ScalarValue::Int32(Some(p)))) + .into_iter() + .collect(); + args.push(ColumnarValue::Array(Arc::clone(value))); + let arg_fields: Vec<_> = args + .iter() + .enumerate() + .map(|(i, a)| Arc::new(Field::new(format!("arg{i}"), a.data_type(), true))) + .collect(); + let arg_types: Vec = arg_fields.iter().map(|f| f.data_type().clone()).collect(); + let return_type = udf.return_type(&arg_types).unwrap(); + udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: value.len(), + return_field: Arc::new(Field::new(udf.name(), return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + .to_array(value.len()) + .unwrap() + } + + fn iceberg_rust(transform: Transform, value: &ArrayRef) -> ArrayRef { + create_transform_function(&transform) + .unwrap() + .transform(Arc::clone(value)) + .unwrap() + } + + fn assert_agree(label: &str, transform: Transform, udf: &dyn ScalarUDFImpl, value: &ArrayRef) { + let parameter = match transform { + Transform::Bucket(n) => Some(n as i32), + Transform::Truncate(w) => Some(w as i32), + _ => None, + }; + assert_eq!( + comet(udf, parameter, value).as_ref(), + iceberg_rust(transform, value).as_ref(), + "{label} disagrees with iceberg-rust's {transform}" + ); + } + + fn timestamps(micros: Vec>) -> Vec<(&'static str, ArrayRef)> { + // The two tags Comet can produce: `TimestampNTZType` is untagged and `TimestampType` is + // always tagged UTC. + vec![ + ( + "timestamp_ntz", + Arc::new(TimestampMicrosecondArray::from(micros.clone())) as ArrayRef, + ), + ( + "timestamp_utc", + Arc::new(TimestampMicrosecondArray::from(micros).with_timezone("UTC")) as ArrayRef, + ), + ] + } + + /// Every type both sides accept. `Int8` and `Int16` are missing on purpose: Iceberg binds + /// tinyint and smallint to `BucketInt`, iceberg-rust has no arm for them, and Comet's kernel + /// widens them to the same 8 little-endian bytes that the `Int32` case pins here. + #[test] + fn bucket_agrees_with_iceberg_rust() { + let mut inputs: Vec<(&str, ArrayRef)> = vec![ + ( + "int", + Arc::new(Int32Array::from(vec![ + Some(i32::MIN), + Some(-1), + Some(0), + Some(34), + Some(i32::MAX), + None, + ])), + ), + ( + "long", + Arc::new(Int64Array::from(vec![ + Some(i64::MIN), + Some(-1), + Some(0), + Some(34), + Some(i64::MAX), + None, + ])), + ), + ( + "date", + Arc::new(Date32Array::from(vec![ + Some(i32::MIN), + Some(-1), + Some(0), + Some(17_486), + Some(i32::MAX), + None, + ])), + ), + ( + "decimal", + Arc::new( + Decimal128Array::from(vec![ + Some(-(10i128.pow(38) - 1)), + Some(-129), + Some(0), + Some(1420), + Some(10i128.pow(38) - 1), + None, + ]) + .with_precision_and_scale(38, 10) + .unwrap(), + ), + ), + ( + "string", + Arc::new(StringArray::from(vec![ + Some(""), + Some("a"), + Some("iceberg"), + Some("日本語😀"), + None, + ])), + ), + ( + "binary", + Arc::new(BinaryArray::from(vec![ + Some([].as_slice()), + Some([0u8, 1, 2, 3].as_slice()), + Some([0xffu8; 9].as_slice()), + None, + ])), + ), + ]; + inputs.extend(timestamps(vec![ + Some(i64::MIN), + Some(-1), + Some(0), + Some(1_510_871_468_000_000), + Some(i64::MAX), + None, + ])); + + let udf = SparkIcebergBucket::new(); + for num_buckets in [1u32, 7, 16, i32::MAX as u32] { + for (label, input) in &inputs { + assert_agree( + &format!("bucket({num_buckets}, {label})"), + Transform::Bucket(num_buckets), + &udf, + input, + ); + } + } + } + + /// `i32::MIN`, `i64::MIN`, and widths above 2^30 are left out: Java's `TruncateUtil` wraps + /// there and iceberg-rust does not (`truncate_i32` uses `rem_euclid`, `truncate_i64` and the + /// decimal kernel subtract without wrapping and overflow in a debug build). That gap is an + /// iceberg-rust bug affecting the writer's own partition values, independent of these + /// kernels; the wrapping cases are pinned against the JVM in the kernel's own unit tests. + #[test] + fn truncate_agrees_with_iceberg_rust() { + let inputs: Vec<(&str, ArrayRef)> = vec![ + ( + "int", + Arc::new(Int32Array::from(vec![ + Some(i32::MIN + 1_000_000), + Some(-1), + Some(0), + Some(1), + Some(i32::MAX - 1_000_000), + None, + ])), + ), + ( + "long", + Arc::new(Int64Array::from(vec![ + Some(i64::MIN + 1_000_000), + Some(-1), + Some(0), + Some(1), + Some(i64::MAX - 1_000_000), + None, + ])), + ), + ( + "decimal", + Arc::new( + Decimal128Array::from(vec![Some(-1065), Some(0), Some(1065), None]) + .with_precision_and_scale(18, 2) + .unwrap(), + ), + ), + ( + "string", + Arc::new(StringArray::from(vec![ + Some(""), + Some("ic"), + Some("iceberg"), + Some("日本語テキスト"), + Some("a😀b😀c"), + None, + ])), + ), + ( + "binary", + Arc::new(BinaryArray::from(vec![ + Some([].as_slice()), + Some([1u8].as_slice()), + Some([1u8, 2, 3, 4, 5].as_slice()), + None, + ])), + ), + ]; + + let udf = SparkIcebergTruncate::new(); + for width in [1u32, 3, 10, 1000, 1 << 30] { + for (label, input) in &inputs { + assert_agree( + &format!("truncate({width}, {label})"), + Transform::Truncate(width), + &udf, + input, + ); + } + } + } + + /// `days` and `hours` are plain floor division on both sides, so the whole domain agrees. + #[test] + fn days_and_hours_agree_with_iceberg_rust() { + let micros = vec![ + Some(0), + Some(-1), + Some(-MICROS_PER_DAY), + Some(-MICROS_PER_DAY - 1), + Some(1_510_871_468_000_000), + Some(365 * MICROS_PER_DAY - 1), + None, + ]; + let days_udf = SparkIcebergTemporalTransform::days(); + let hours_udf = SparkIcebergTemporalTransform::hours(); + for (label, input) in timestamps(micros) { + assert_agree(&format!("days({label})"), Transform::Day, &days_udf, &input); + assert_agree( + &format!("hours({label})"), + Transform::Hour, + &hours_udf, + &input, + ); + } + let dates: ArrayRef = Arc::new(Date32Array::from(vec![ + Some(i32::MIN), + Some(-366), + Some(0), + Some(17_486), + Some(i32::MAX), + None, + ])); + assert_agree("days(date)", Transform::Day, &days_udf, &dates); + } + + /// `years` and `months` agree over the dates iceberg-rust can represent -- it splits the + /// calendar with `chrono`, so anything past year 262143 errors there while Comet and the JVM + /// keep going (see the kernel's own unit tests for those). + #[test] + fn years_and_months_agree_with_iceberg_rust_within_its_range() { + let years_udf = SparkIcebergTemporalTransform::years(); + let months_udf = SparkIcebergTemporalTransform::months(); + let dates: ArrayRef = Arc::new(Date32Array::from(vec![ + Some(-100_000), + Some(-366), + Some(-365), + Some(-1), + Some(0), + Some(30), + Some(17_486), + Some(100_000), + None, + ])); + assert_agree("years(date)", Transform::Year, &years_udf, &dates); + assert_agree("months(date)", Transform::Month, &months_udf, &dates); + for (label, input) in timestamps(vec![ + Some(-100_000 * MICROS_PER_DAY), + Some(-1), + Some(0), + Some(1_510_871_468_000_000), + None, + ]) { + assert_agree( + &format!("years({label})"), + Transform::Year, + &years_udf, + &input, + ); + assert_agree( + &format!("months({label})"), + Transform::Month, + &months_udf, + &input, + ); + } + } + + /// Why `years` and `months` are not delegated to iceberg-rust even though `bucket`, `days`, + /// and `hours` could be: its kernels go through Arrow's `date_part`, which honours the + /// array's timezone tag, while Iceberg's Java `DateTimeUtil` is always UTC. Comet only ever + /// produces `UTC` and untagged timestamps today, so the parity above holds; this pins the + /// reason the local kernel exists. If this ever fails, iceberg-rust dropped the tag + /// dependency and delegating becomes safe. + #[test] + fn iceberg_rust_years_follow_the_timezone_tag() { + // 1969-12-31T23:59:59.999999Z, which is 1970-01-01T05:44:59.999999 in Kathmandu. + let tagged: ArrayRef = + Arc::new(TimestampMicrosecondArray::from(vec![-1i64]).with_timezone("Asia/Kathmandu")); + let comet_years = comet(&SparkIcebergTemporalTransform::years(), None, &tagged); + let iceberg_years = iceberg_rust(Transform::Year, &tagged); + assert_eq!( + comet_years.as_ref(), + &Int32Array::from(vec![-1]) as &dyn arrow::array::Array + ); + assert_eq!( + iceberg_years.as_ref(), + &Int32Array::from(vec![0]) as &dyn arrow::array::Array + ); + } +} From 19ca6c3bb138a170009d19a88e4049e4f4c50888 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:35:30 -0600 Subject: [PATCH 10/18] docs: describe the decimal truncate difference in the Iceberg guide Iceberg's `TruncateDecimal.invoke` returns a `Decimal` that can exceed the column's precision and Spark nulls it only when the row is materialized, while Comet nulls it in the kernel. Say so, and say which paths that leaves in agreement. --- docs/source/user-guide/latest/iceberg.md | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 3b57a7efd88..3fac24d0760 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -210,11 +210,20 @@ The native kernels reproduce Iceberg's Java semantics exactly rather than approx and timestamps; UTF-8 for strings; raw bytes for binary; the minimal big-endian two's complement of the unscaled value for decimals) with 32-bit Murmur3 and masks the sign bit before taking the modulus. -- `truncate` uses Java's wrapping integer arithmetic, keeps the decimal's precision and scale - (a negative decimal whose truncated value no longer fits the precision becomes null, as it - does in Spark), and counts code points (not bytes) for strings. +- `truncate` uses Java's wrapping integer arithmetic, keeps the decimal's precision and scale, and + counts code points (not bytes) for strings. - `years`, `months`, `days`, and `hours` are evaluated in UTC regardless of the session timezone - and go negative before the epoch; `days` returns a date, the other three an int. + and go negative before the epoch; `days` returns a date, the other three an int. They cover the + whole `DATE` and `TIMESTAMP` domain, as Iceberg's `DateTimeUtil` does. + +One difference is worth calling out. Truncating a negative decimal grows its magnitude, so the +result can need one more digit than the column's precision allows: `truncate(10, v)` on a +`decimal(18,4)` value of `-99999999999999.9999` is `-100000000000000.0000`, which has 19 digits. +Iceberg's Java `TruncateDecimal` hands that oversized value back unchanged, and Spark turns it into +null only when the row is materialized. Comet nulls it in the kernel instead. The two therefore +agree wherever the value is written into a row -- a projection, a sort key, a shuffle key, a +partition value -- and differ only where the truncated decimal feeds another expression without +being materialized, as in `WHERE truncate(10, v) IS NULL`. This matters most for writes. A partitioned table with the default `write.distribution-mode` (`hash`) is planned with a shuffle and a local sort keyed on the partition transforms, and with From fa0158d30484044897e32d6beb0915f825dfc1f8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:35:42 -0600 Subject: [PATCH 11/18] bench: measure the Iceberg system functions natively and against Iceberg's JVM `native/spark-expr/benches/iceberg_transforms.rs` covers every transform over every supported type, with and without nulls, plus the dictionary-encoded string shape a Parquet scan produces. `CometIcebergSystemFunctionBenchmark` runs the same queries with Comet on and off. The Comet-off case is Iceberg's own JVM implementation, since Spark binds each function as a `StaticInvoke` of the class under `org.apache.iceberg.spark.functions`. The data stays in Parquet rather than an Iceberg table so both cases scan identically and only the transform differs. --- native/spark-expr/Cargo.toml | 6 +- .../spark-expr/benches/iceberg_transforms.rs | 218 ++++++++++++++++++ .../CometIcebergSystemFunctionBenchmark.scala | 109 +++++++++ 3 files changed, 332 insertions(+), 1 deletion(-) create mode 100644 native/spark-expr/benches/iceberg_transforms.rs create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 5041d7fe32d..519336c4fec 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -301,4 +301,8 @@ harness = false [[bench]] name = "to_utc_timestamp" -harness = false \ No newline at end of file +harness = false + +[[bench]] +name = "iceberg_transforms" +harness = false diff --git a/native/spark-expr/benches/iceberg_transforms.rs b/native/spark-expr/benches/iceberg_transforms.rs new file mode 100644 index 00000000000..4572ae03920 --- /dev/null +++ b/native/spark-expr/benches/iceberg_transforms.rs @@ -0,0 +1,218 @@ +// 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. + +//! Iceberg's system functions (`bucket`, `truncate`, `years`, `months`, `days`, `hours`) over one +//! input array per supported type, with and without nulls, plus the dictionary-encoded string +//! shape a Parquet scan produces for a partition column. + +use arrow::array::{ + ArrayRef, BinaryArray, Date32Array, Decimal128Array, DictionaryArray, Int32Array, Int64Array, + StringArray, TimestampMicrosecondArray, +}; +use arrow::datatypes::{DataType, Field, Int32Type}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion::common::ScalarValue; +use datafusion::config::ConfigOptions; +use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_comet_spark_expr::{ + SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, +}; +use std::hint::black_box; +use std::sync::Arc; + +const ROWS: usize = 8_192; +const MICROS_PER_DAY: i64 = 86_400_000_000; +/// Every eighth row is null, matching the corpus the correctness suite writes. +const NULL_STRIDE: usize = 8; + +fn maybe_null(i: usize, nulls: bool, value: T) -> Option { + if nulls && i.is_multiple_of(NULL_STRIDE) { + None + } else { + Some(value) + } +} + +/// A deterministic 64-bit value per row; the transforms are data dependent (bucket hashes it, +/// truncate divides by it), so a constant column would not be representative. +fn spread(i: usize) -> i64 { + (i as i64) + .wrapping_mul(6_364_136_223_846_793_005) + .rotate_left(17) +} + +/// The words a low-cardinality string partition column holds; row `i` picks `i % len`. +const WORDS: [&str; 8] = [ + "alpha", + "bravo", + "charlie", + "delta", + "echo", + "foxtrot", + "日本語テキスト", + "a😀b😀c", +]; + +fn inputs(nulls: bool) -> Vec<(&'static str, ArrayRef)> { + let strings: StringArray = (0..ROWS) + .map(|i| maybe_null(i, nulls, WORDS[i % WORDS.len()])) + .collect(); + let dictionary: DictionaryArray = (0..ROWS) + .map(|i| maybe_null(i, nulls, WORDS[i % WORDS.len()])) + .collect(); + let binaries: BinaryArray = (0..ROWS) + .map(|i| maybe_null(i, nulls, spread(i).to_be_bytes())) + .collect(); + vec![ + ( + "int", + Arc::new( + (0..ROWS) + .map(|i| maybe_null(i, nulls, spread(i) as i32)) + .collect::(), + ), + ), + ( + "long", + Arc::new( + (0..ROWS) + .map(|i| maybe_null(i, nulls, spread(i))) + .collect::(), + ), + ), + ( + "decimal38", + Arc::new( + (0..ROWS) + .map(|i| maybe_null(i, nulls, (spread(i) as i128) * 1_000_000_000)) + .collect::() + .with_precision_and_scale(38, 10) + .unwrap(), + ), + ), + ("string", Arc::new(strings)), + ("string_dict", Arc::new(dictionary)), + ("binary", Arc::new(binaries)), + ( + "date", + Arc::new( + (0..ROWS) + .map(|i| maybe_null(i, nulls, (spread(i) % 40_000) as i32)) + .collect::(), + ), + ), + ( + "timestamp", + Arc::new( + (0..ROWS) + .map(|i| maybe_null(i, nulls, spread(i) % (40_000 * MICROS_PER_DAY))) + .collect::() + .with_timezone("UTC"), + ), + ), + ] +} + +fn invoke(udf: &dyn ScalarUDFImpl, args: &[ColumnarValue]) -> ArrayRef { + let arg_fields: Vec<_> = args + .iter() + .enumerate() + .map(|(i, a)| Arc::new(Field::new(format!("arg{i}"), a.data_type(), true))) + .collect(); + let arg_types: Vec = arg_fields.iter().map(|f| f.data_type().clone()).collect(); + let return_type = udf.return_type(&arg_types).unwrap(); + udf.invoke_with_args(ScalarFunctionArgs { + args: args.to_vec(), + arg_fields, + number_rows: ROWS, + return_field: Arc::new(Field::new(udf.name(), return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + .to_array(ROWS) + .unwrap() +} + +/// `true` when `udf` accepts `input`; the transforms are typed the same way Iceberg's `bind` is, +/// so the type matrix is sparse. +fn supported(udf: &dyn ScalarUDFImpl, args: &[ColumnarValue]) -> bool { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let arg_fields: Vec<_> = args + .iter() + .enumerate() + .map(|(i, a)| Arc::new(Field::new(format!("arg{i}"), a.data_type(), true))) + .collect(); + let arg_types: Vec = arg_fields.iter().map(|f| f.data_type().clone()).collect(); + let Ok(return_type) = udf.return_type(&arg_types) else { + return false; + }; + udf.invoke_with_args(ScalarFunctionArgs { + args: args.to_vec(), + arg_fields, + number_rows: ROWS, + return_field: Arc::new(Field::new(udf.name(), return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .is_ok() + })) + .unwrap_or(false) +} + +fn criterion_benchmark(c: &mut Criterion) { + let bucket = SparkIcebergBucket::new(); + let truncate = SparkIcebergTruncate::new(); + let years = SparkIcebergTemporalTransform::years(); + let months = SparkIcebergTemporalTransform::months(); + let days = SparkIcebergTemporalTransform::days(); + let hours = SparkIcebergTemporalTransform::hours(); + + // (name, udf, parameter) -- `bucket` and `truncate` take a literal first argument. + let transforms: Vec<(&str, &dyn ScalarUDFImpl, Option)> = vec![ + ("iceberg_bucket", &bucket, Some(16)), + ("iceberg_truncate", &truncate, Some(4)), + ("iceberg_years", &years, None), + ("iceberg_months", &months, None), + ("iceberg_days", &days, None), + ("iceberg_hours", &hours, None), + ]; + + for (name, udf, parameter) in transforms { + let mut group = c.benchmark_group(name); + group.throughput(Throughput::Elements(ROWS as u64)); + for (nulls, null_tag) in [(false, "no_nulls"), (true, "sparse_nulls")] { + for (type_tag, array) in inputs(nulls) { + let args: Vec = parameter + .map(|p| ColumnarValue::Scalar(ScalarValue::Int32(Some(p)))) + .into_iter() + .chain([ColumnarValue::Array(array)]) + .collect(); + if !supported(udf, &args) { + continue; + } + group.bench_with_input( + BenchmarkId::from_parameter(format!("{type_tag}/{null_tag}")), + &args, + |b, args| b.iter(|| black_box(invoke(udf, black_box(args)))), + ); + } + } + group.finish(); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala new file mode 100644 index 00000000000..43651a3c7d9 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala @@ -0,0 +1,109 @@ +/* + * 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.comet.iceberg.IcebergReflection + +/** + * Benchmark of Iceberg's system functions (`bucket`, `truncate`, `years`, `months`, `days`, + * `hours`) with Comet on and off. The Spark case is Iceberg's own JVM implementation: Spark binds + * each function as a `StaticInvoke` of the matching class under + * `org.apache.iceberg.spark.functions` and whole-stage codegen calls it once per row. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometIcebergSystemFunctionBenchmark + * }}} + * Results will be written to + * "spark/benchmarks/CometIcebergSystemFunctionBenchmark-**results.txt". + */ +object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { + + private val catalog = "benchmark_cat" + + /** + * One case per (transform, input type). `c_str_dict` holds eight distinct values so Parquet + * dictionary-encodes it, which is the shape a string partition column normally arrives in; + * `c_str` is distinct per row. + */ + private def cases: Seq[(String, String)] = { + val bucket = Seq("c_int", "c_long", "c_dec", "c_str_dict", "c_str", "c_bin", "c_date", "c_ts") + .map(column => s"bucket($column)" -> s"select $catalog.system.bucket(16, $column)") + val truncate = Seq("c_int", "c_long", "c_dec", "c_str_dict", "c_str", "c_bin") + .map(column => s"truncate($column)" -> s"select $catalog.system.truncate(4, $column)") + val temporal = Seq("years", "months", "days").flatMap { fn => + Seq("c_date", "c_ts").map(column => + s"$fn($column)" -> s"select $catalog.system.$fn($column)") + } :+ ("hours(c_ts)" -> s"select $catalog.system.hours(c_ts)") + (bucket ++ truncate ++ temporal).map { case (name, select) => + name -> s"$select from parquetV1Table" + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + if (!icebergOnClasspath) { + // scalastyle:off println + println("Iceberg is not on the classpath; skipping. Build with an Iceberg-enabled profile.") + // scalastyle:on println + return + } + // The Iceberg system functions are resolved through a v2 catalog, so one has to be + // registered. No Iceberg table is read: the data stays in Parquet so both cases scan + // identically and the only difference is who evaluates the transform. + withTempPath { warehouse => + spark.conf.set(s"spark.sql.catalog.$catalog", "org.apache.iceberg.spark.SparkCatalog") + spark.conf.set(s"spark.sql.catalog.$catalog.type", "hadoop") + spark.conf.set(s"spark.sql.catalog.$catalog.warehouse", warehouse.getAbsolutePath) + + runBenchmarkWithTable("Iceberg system functions", 1024 * 1024) { v => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s""" + SELECT CAST(value AS INT) AS c_int, + value AS c_long, + CAST(value AS DECIMAL(38,10)) AS c_dec, + CAST(PMOD(value, 8) AS STRING) AS c_str_dict, + REPEAT(CAST(value AS STRING), 3) AS c_str, + CAST(CAST(value AS STRING) AS BINARY) AS c_bin, + DATE_ADD(DATE '1970-01-01', CAST(PMOD(value, 40000) AS INT)) AS c_date, + TIMESTAMP_SECONDS(PMOD(value, 4000000000)) AS c_ts + FROM $tbl""")) + + cases.foreach { case (name, query) => + runBenchmark(name) { + runExpressionBenchmark(name, v, query) + } + } + } + } + } + } + } + + private def icebergOnClasspath: Boolean = + try { + IcebergReflection.loadClass("org.apache.iceberg.spark.functions.BucketFunction") + true + } catch { + case _: ClassNotFoundException => false + } +} From 3dc1a2d5862f5e2cca598888178460eb822600a1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:56:25 -0600 Subject: [PATCH 12/18] docs: correct which paths agree on the truncated decimal `HashPartitioning` hashes the `StaticInvoke` result directly rather than through an `UnsafeRowWriter`, so the shuffle hash is on the differing side, not the agreeing one. The Iceberg partition value does not come from this kernel at all. --- docs/source/user-guide/latest/iceberg.md | 7 ++++--- native/spark-expr/src/iceberg_funcs/truncate.rs | 11 +++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 3fac24d0760..139853516b5 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -221,9 +221,10 @@ result can need one more digit than the column's precision allows: `truncate(10, `decimal(18,4)` value of `-99999999999999.9999` is `-100000000000000.0000`, which has 19 digits. Iceberg's Java `TruncateDecimal` hands that oversized value back unchanged, and Spark turns it into null only when the row is materialized. Comet nulls it in the kernel instead. The two therefore -agree wherever the value is written into a row -- a projection, a sort key, a shuffle key, a -partition value -- and differ only where the truncated decimal feeds another expression without -being materialized, as in `WHERE truncate(10, v) IS NULL`. +agree wherever Spark writes the value into a row -- the output of a projection, and the key a sort +builds -- and differ where the truncated decimal feeds another expression directly, as in +`WHERE truncate(10, v) IS NULL` or the hash behind `DISTRIBUTE BY truncate(10, v)`. Note that the +Iceberg partition value itself is not affected: the writer computes it from the untruncated column. This matters most for writes. A partitioned table with the default `write.distribution-mode` (`hash`) is planned with a shuffle and a local sort keyed on the partition transforms, and with diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs index 3c5ab0bc729..5ec6c6a644c 100644 --- a/native/spark-expr/src/iceberg_funcs/truncate.rs +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -99,12 +99,11 @@ fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result().unary_opt(|v| { let truncated = truncate_i128(v, width as i128); From f4e8e454e925c95b7b87ec64c3d053b07ed8731c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 10:59:52 -0600 Subject: [PATCH 13/18] test: do not assert an unverified cause for the partition-source exclusion The multi-byte string column is out of the write test's partition spec because the test failed on the Linux CI runners with it in, not because the platform-charset mechanism the comment claimed has been demonstrated -- it does not reproduce locally, including with sun.jnu.encoding forced to US-ASCII. --- .../org/apache/comet/CometIcebergSystemFunctionSuite.scala | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index c584e07b7bf..120a7b32a3d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -192,8 +192,11 @@ class CometIcebergSystemFunctionSuite val table = s"$catalog.db.hidden_partitioning" // No `write.distribution-mode`: Iceberg picks hash distribution for a partitioned table, // which plans a shuffle and a local sort on the partition transforms. The string column is - // deliberately not a partition source: its multi-byte values would land in partition - // directory names, which iceberg-java reads back through the JVM's platform charset. + // deliberately not a partition source: with it in the spec this test failed on the Linux + // CI runners with a missing data file, which does not reproduce locally, so the multi-byte + // values it would put in partition directory names stay out of the spec until that is + // understood. `bucket` and `truncate` over strings are covered by the comparison, filter, + // sort, and shuffle tests above. sql(s""" CREATE TABLE $table (i32 INT, i64 BIGINT, ts TIMESTAMP, dt DATE) USING iceberg From ada706be9cfb1ea56ff0f677d6bd8884b4a56e08 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 12:14:38 -0600 Subject: [PATCH 14/18] test: link the upstream iceberg-rust issues from the excluded cases The truncate exclusions in the parity tests and the Long.MinValue filter in the write test now point at apache/iceberg-rust#3141, and the timezone-tag tests at apache/iceberg-rust#3142, instead of describing the workaround inline. --- native/core/src/execution/operators/iceberg_write.rs | 11 ++++++----- .../comet/CometIcebergSystemFunctionSuite.scala | 6 +++--- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index e23c856b961..1d3fb668217 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -1360,9 +1360,10 @@ mod iceberg_rust_transform_parity { /// `i32::MIN`, `i64::MIN`, and widths above 2^30 are left out: Java's `TruncateUtil` wraps /// there and iceberg-rust does not (`truncate_i32` uses `rem_euclid`, `truncate_i64` and the - /// decimal kernel subtract without wrapping and overflow in a debug build). That gap is an + /// decimal kernel subtract without wrapping and overflow in a debug build). That is an /// iceberg-rust bug affecting the writer's own partition values, independent of these - /// kernels; the wrapping cases are pinned against the JVM in the kernel's own unit tests. + /// kernels -- apache/iceberg-rust#3141. The wrapping cases are pinned against the JVM in the + /// kernel's own unit tests; add them here once that issue is fixed. #[test] fn truncate_agrees_with_iceberg_rust() { let inputs: Vec<(&str, ArrayRef)> = vec![ @@ -1467,7 +1468,7 @@ mod iceberg_rust_transform_parity { /// `years` and `months` agree over the dates iceberg-rust can represent -- it splits the /// calendar with `chrono`, so anything past year 262143 errors there while Comet and the JVM - /// keep going (see the kernel's own unit tests for those). + /// keep going (apache/iceberg-rust#3142; see the kernel's own unit tests for those). #[test] fn years_and_months_agree_with_iceberg_rust_within_its_range() { let years_udf = SparkIcebergTemporalTransform::years(); @@ -1511,8 +1512,8 @@ mod iceberg_rust_transform_parity { /// and `hours` could be: its kernels go through Arrow's `date_part`, which honours the /// array's timezone tag, while Iceberg's Java `DateTimeUtil` is always UTC. Comet only ever /// produces `UTC` and untagged timestamps today, so the parity above holds; this pins the - /// reason the local kernel exists. If this ever fails, iceberg-rust dropped the tag - /// dependency and delegating becomes safe. + /// reason the local kernel exists. Reported as apache/iceberg-rust#3142; if this ever fails, + /// iceberg-rust dropped the tag dependency and delegating becomes safe. #[test] fn iceberg_rust_years_follow_the_timezone_tag() { // 1969-12-31T23:59:59.999999Z, which is 1970-01-01T05:44:59.999999 in Kathmandu. diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index 120a7b32a3d..2cd2657666f 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -201,9 +201,9 @@ class CometIcebergSystemFunctionSuite CREATE TABLE $table (i32 INT, i64 BIGINT, ts TIMESTAMP, dt DATE) USING iceberg PARTITIONED BY (bucket(4, i32), truncate(1000, i64), days(ts), months(dt))""") - // iceberg-rust's own truncate transform, which the writer uses for partition values, does - // `v - ((v % w) + w) % w` without wrapping and overflows on Long.MinValue in debug builds - // (Java wraps), so the boundary row stays out of the written set. + // iceberg-rust's own truncate transform, which the writer uses for partition values, + // overflows on Long.MinValue in a debug build where Java wraps + // (apache/iceberg-rust#3141), so the boundary row stays out of the written set. val rows = s"SELECT i32, i64, ts, dt FROM $source WHERE i64 IS NULL OR i64 <> ${Long.MinValue}" try { From 0a57b6aed4182e36ef12908ff479a0dcaf0740e9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 14:45:44 -0600 Subject: [PATCH 15/18] fix: decline Iceberg's decimal truncate instead of nulling early Iceberg's `TruncateDecimal.invoke` returns a `Decimal` that can exceed the column's declared precision, and Spark turns it into null only when the row is materialized. An Arrow `Decimal128(precision, scale)` array has no encoding for that intermediate, so the kernel nulled it during evaluation, which changes what an enclosing predicate or hash sees. Report decimal inputs as `Unsupported` so they stay with Spark. `CometStaticInvoke.getUnsupportedReasons()` now aggregates the per-function handlers' notes, since `GenerateDocs` only asks the serde registered for the expression class and the note would otherwise reach the Iceberg guide but not the compatibility page. Also drops the claim that Comet and Spark agree on the sort key. `SortExec` orders rows with `RowOrdering.create`, which evaluates the sort expression per comparison rather than through an `UnsafeRowWriter`, so Spark's sort sees the oversized decimal. That was wrong in `iceberg.md` and in the kernel comment. --- docs/source/user-guide/latest/iceberg.md | 21 +++++---- .../spark-expr/src/iceberg_funcs/truncate.rs | 15 ++++--- .../apache/comet/serde/icebergFunctions.scala | 44 +++++++++++++++--- .../org/apache/comet/serde/statics.scala | 8 ++++ .../CometIcebergSystemFunctionSuite.scala | 45 ++++++++++++++++++- 5 files changed, 109 insertions(+), 24 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 139853516b5..a693a450f99 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -210,21 +210,20 @@ The native kernels reproduce Iceberg's Java semantics exactly rather than approx and timestamps; UTF-8 for strings; raw bytes for binary; the minimal big-endian two's complement of the unscaled value for decimals) with 32-bit Murmur3 and masks the sign bit before taking the modulus. -- `truncate` uses Java's wrapping integer arithmetic, keeps the decimal's precision and scale, and - counts code points (not bytes) for strings. +- `truncate` uses Java's wrapping integer arithmetic and counts code points (not bytes) for + strings. Decimal inputs are the one case that stays with Spark, see below. - `years`, `months`, `days`, and `hours` are evaluated in UTC regardless of the session timezone and go negative before the epoch; `days` returns a date, the other three an int. They cover the whole `DATE` and `TIMESTAMP` domain, as Iceberg's `DateTimeUtil` does. -One difference is worth calling out. Truncating a negative decimal grows its magnitude, so the -result can need one more digit than the column's precision allows: `truncate(10, v)` on a -`decimal(18,4)` value of `-99999999999999.9999` is `-100000000000000.0000`, which has 19 digits. -Iceberg's Java `TruncateDecimal` hands that oversized value back unchanged, and Spark turns it into -null only when the row is materialized. Comet nulls it in the kernel instead. The two therefore -agree wherever Spark writes the value into a row -- the output of a projection, and the key a sort -builds -- and differ where the truncated decimal feeds another expression directly, as in -`WHERE truncate(10, v) IS NULL` or the hash behind `DISTRIBUTE BY truncate(10, v)`. Note that the -Iceberg partition value itself is not affected: the writer computes it from the untruncated column. +`truncate` on a `decimal` column falls back to Spark. Truncating a negative decimal grows its +magnitude, so the result can need one more digit than the column's precision allows: +`truncate(10, v)` on a `decimal(18,4)` value of `-99999999999999.9999` is +`-100000000000000.0000`, which has 19 digits. Iceberg's `TruncateDecimal` hands that oversized +value back unchanged and Spark turns it into null only when the row is materialized. An Arrow +`Decimal128(precision, scale)` array has no encoding for that intermediate, so a native kernel +would have to null it during evaluation, which changes what an enclosing predicate or hash sees. +Every other `truncate` input type, and `bucket` on decimals, runs natively. This matters most for writes. A partitioned table with the default `write.distribution-mode` (`hash`) is planned with a shuffle and a local sort keyed on the partition transforms, and with diff --git a/native/spark-expr/src/iceberg_funcs/truncate.rs b/native/spark-expr/src/iceberg_funcs/truncate.rs index 5ec6c6a644c..70b638f52bc 100644 --- a/native/spark-expr/src/iceberg_funcs/truncate.rs +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -95,15 +95,16 @@ fn truncate_array(fn_name: &str, array: &ArrayRef, width: i32) -> Result(|v| truncate_i64(v, width as i64)), ), DataType::Decimal128(precision, scale) => { + // Not reachable from a Spark plan: `CometIcebergTruncate` reports decimal inputs as + // `Unsupported` so they stay with Spark. The reason is this arm's only honest answer. // Truncating a negative value grows its magnitude by up to `width - 1` units of the // last digit, so the result can need one more digit than the column allows. Iceberg's - // `TruncateDecimal.invoke` hands that oversized `Decimal` back to Spark unchanged and - // Spark nulls it only when a row is materialized (`UnsafeRowWriter` calls - // `changePrecision`, which fails). Nulling it here is the same answer everywhere Spark - // materializes the value -- the output of a projection, and the key `SortExec` builds - // -- and it is the only answer available to a kernel that has to return a - // `Decimal128(precision, scale)` array. The two differ where the result feeds another - // expression directly, e.g. `truncate(w, v) IS NULL`; see the Iceberg user guide. + // `TruncateDecimal.invoke` hands that oversized `Decimal` back unchanged and Spark + // nulls it only when a row is materialized, but a `Decimal128(precision, scale)` array + // has nowhere to put it, and nulling during evaluation changes what an enclosing + // predicate or hash sees. The arm is kept because it is what iceberg-rust's `Truncate` + // is compared against in the writer's parity tests, and because the choice of gate is + // serde policy that may be revisited. let truncated: Decimal128Array = array.as_primitive::().unary_opt(|v| { let truncated = truncate_i128(v, width as i128); diff --git a/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala index 4ea5cf5a5b8..b9231f46d62 100644 --- a/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala +++ b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala @@ -172,17 +172,51 @@ object CometIcebergBucket case _ => false }) -/** `truncate(width, value)` over the types `TruncateFunction.bind` accepts. */ +/** + * `truncate(width, value)` over the types `TruncateFunction.bind` accepts, minus decimals. + * + * Decimals are declined for a semantic reason rather than a missing kernel; see + * [[CometIcebergTruncate.DecimalNote]]. + */ object CometIcebergTruncate extends CometIcebergParameterizedTransform( "iceberg_truncate", "width", { - case ByteType | ShortType | IntegerType | LongType | StringType | - BinaryType | _: DecimalType => - true + case ByteType | ShortType | IntegerType | LongType | StringType | BinaryType => true case _ => false - }) + }) { + + /** + * Why decimal `truncate` stays with Spark. Truncating a negative decimal grows its magnitude, + * so the result can need one more digit than the column's precision allows. Iceberg's + * `TruncateDecimal.invoke` hands that oversized `Decimal` back unchanged and Spark turns it + * into null only when the row is materialized, whereas an Arrow `Decimal128(precision, scale)` + * array has no encoding for it -- a native kernel would have to null it during evaluation, + * changing what an enclosing predicate or hash sees. + */ + val DecimalNote: String = + "Iceberg's TruncateDecimal returns a Decimal that can exceed the column's declared " + + "precision, and Spark only turns that into null when the row is materialized. An Arrow " + + "Decimal128(precision, scale) array cannot carry that intermediate, so a native kernel " + + "would null it during evaluation and change what an enclosing predicate or hash sees." + + // Ordered after the parameter check so that `truncate(0, decimal_col)` still reports the width + // problem, which is the one that changes whether Iceberg's own ArithmeticException is raised. + override def getSupportLevel(expr: StaticInvoke): SupportLevel = expr.arguments match { + case Seq(parameter, value) + if value.dataType.isInstanceOf[DecimalType] && + CometIcebergSystemFunctions.positiveIntLiteral(parameter).isDefined => + Unsupported(Some(DecimalNote)) + case _ => super.getSupportLevel(expr) + } + + override def getUnsupportedReasons(): Seq[String] = Seq( + "Iceberg's `truncate(width, value)` system function on a `decimal` column. " + DecimalNote + + " Truncating `-99999999999999.9999` in a `decimal(18,4)` column by a width of 10 is one " + + "such value: the result has 19 digits. The other `truncate` input types, and `bucket` on " + + "decimals, are unaffected.") +} /** Shared shape of the single-argument `years`, `months`, `days`, and `hours` transforms. */ abstract class CometIcebergTemporalTransform( diff --git a/spark/src/main/scala/org/apache/comet/serde/statics.scala b/spark/src/main/scala/org/apache/comet/serde/statics.scala index 94beefc0ba5..8944b90835d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/statics.scala +++ b/spark/src/main/scala/org/apache/comet/serde/statics.scala @@ -62,6 +62,14 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] { override def getSupportLevel(expr: StaticInvoke): SupportLevel = handlerFor(expr).map(_.getSupportLevel(expr)).getOrElse(Compatible()) + /** + * `GenerateDocs` only asks the serde registered for the expression class, which is this object, + * so the per-function handlers' notes have to be collected here or they never reach the + * compatibility guide. + */ + override def getUnsupportedReasons(): Seq[String] = + staticInvokeExpressions.values.toSeq.distinct.flatMap(_.getUnsupportedReasons()).distinct + override def convert( expr: StaticInvoke, inputs: Seq[Attribute], diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index 2cd2657666f..6b968e55f14 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -76,7 +76,8 @@ class CometIcebergSystemFunctionSuite private val source = "system_function_source" private val bucketColumns = Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin", "dt", "ts", "ts_ntz") - private val truncateColumns = Seq("i8", "i16", "i32", "i64", "dec18", "dec38", "str", "bin") + private val truncateColumns = Seq("i8", "i16", "i32", "i64", "str", "bin") + private val decimalColumns = Seq("dec18", "dec38") // The source data is written once per suite; every test reads the same parquet directory. private var sourceDir: File = _ @@ -133,6 +134,20 @@ class CometIcebergSystemFunctionSuite } } + test("truncate on a decimal falls back to Spark") { + withSourceTable { + // Iceberg's `TruncateDecimal` can return a value wider than the column's precision, which + // Spark nulls only on materialization and a `Decimal128(p, s)` array cannot represent at + // all; the expression stays with Spark rather than null early. `bucket` on the same columns + // is unaffected and is covered by the bucket test above. + decimalColumns.foreach { column => + checkSparkAnswerAndFallbackReason( + s"SELECT $catalog.system.truncate(10, $column) FROM $source", + "Decimal128(precision, scale) array cannot carry that intermediate") + } + } + } + test("years, months, days, and hours match Iceberg regardless of session timezone") { withSourceTable { // Iceberg evaluates the temporal transforms in UTC; a shifted session timezone must not @@ -285,10 +300,38 @@ class CometIcebergSystemFunctionSuite assert(level(CometIcebergTruncate, truncateInt, Literal(10), value) == Compatible()) assert(level(CometIcebergTruncate, truncateInt, Literal(10), date).isInstanceOf[Unsupported]) + // A decimal value reports its own reason, but only once the width is valid: a zero width has + // to keep reporting the width, since that is what decides whether Iceberg's own + // ArithmeticException is raised. + val decimal = AttributeReference("dec", DecimalType(18, 4))() + val decimalLevel = level(CometIcebergTruncate, truncateInt, Literal(10), decimal) + assert( + decimalLevel == Unsupported(Some(CometIcebergTruncate.DecimalNote)), + s"unexpected support level: $decimalLevel") + val zeroWidth = level(CometIcebergTruncate, truncateInt, Literal(0), decimal) + assert( + zeroWidth.asInstanceOf[Unsupported].notes.exists(_.contains("width must be a positive")), + s"unexpected support level: $zeroWidth") + // `bucket` on a decimal stays native; only `truncate` has the precision problem. + val bucketDecimal = + Class.forName("org.apache.iceberg.spark.functions.BucketFunction$BucketDecimal") + assert(level(CometIcebergBucket, bucketDecimal, Literal(4), decimal) == Compatible()) + + // The reason also has to reach the generated compatibility guide, which only asks the serde + // registered for `StaticInvoke`. + assert( + CometStaticInvoke + .getUnsupportedReasons() + .exists(_.contains(CometIcebergTruncate.DecimalNote)), + "the decimal truncate note is missing from CometStaticInvoke.getUnsupportedReasons") + // CometStaticInvoke dispatches on (functionName, class name), so the same expressions reach // the Iceberg handlers from there too. assert(level(CometStaticInvoke, bucketInt, Literal(4), value) == Compatible()) assert(level(CometStaticInvoke, bucketInt, Literal(0), value).isInstanceOf[Unsupported]) + assert( + level(CometStaticInvoke, truncateInt, Literal(10), decimal) == + Unsupported(Some(CometIcebergTruncate.DecimalNote))) } test("fallback reason for an unlisted static invoke names the declaring class") { From f466a65a96104e00446e17a8bdba43f77b0a8994 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 14:45:44 -0600 Subject: [PATCH 16/18] bench: verify outputs and cover nulls in the Iceberg system function benchmark Every case now runs over a no-null column and a column with one null in eight, and each case's output is compared row by row between the two engines over the same corpus that is then timed, before it is timed. `truncate` on a decimal is dropped since it no longer has a native path. `excludedRulesWith` becomes `protected` so the verification can use the same constant-folding exclusion the timed cases use. --- .../sql/benchmark/CometBenchmarkBase.scala | 2 +- .../CometIcebergSystemFunctionBenchmark.scala | 116 ++++++++++++++---- 2 files changed, 94 insertions(+), 24 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala index 1254da34918..487e4797d12 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala @@ -205,7 +205,7 @@ trait CometBenchmarkBase * Returns the value of `spark.sql.optimizer.excludedRules` with `rule` appended, so that * benchmark-specific exclusions do not clobber exclusions already configured by the caller. */ - private def excludedRulesWith(rule: String): String = + protected def excludedRulesWith(rule: String): String = (Utils.stringToSeq(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, "")) :+ rule).distinct .mkString(",") diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala index 43651a3c7d9..5e3e0a86999 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala @@ -19,6 +19,11 @@ package org.apache.spark.sql.benchmark +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.optimizer.ConstantFolding +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf import org.apache.comet.iceberg.IcebergReflection /** @@ -27,6 +32,10 @@ import org.apache.comet.iceberg.IcebergReflection * each function as a `StaticInvoke` of the matching class under * `org.apache.iceberg.spark.functions` and whole-stage codegen calls it once per row. * + * Every case is run over a no-null column and a column with one null in eight, and every case's + * output is compared between the two engines over the same corpus that is then timed, so a timing + * cannot come from an engine that computed something else. + * * To run this benchmark: * {{{ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometIcebergSystemFunctionBenchmark @@ -38,22 +47,76 @@ object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { private val catalog = "benchmark_cat" + /** One null in eight, matching the null rate of the correctness suite's corpus. */ + private val NullStride = 8 + /** - * One case per (transform, input type). `c_str_dict` holds eight distinct values so Parquet + * Column types each transform accepts. `str_dict` holds eight distinct values so Parquet * dictionary-encodes it, which is the shape a string partition column normally arrives in; - * `c_str` is distinct per row. + * `str` is distinct per row. `truncate` on a decimal is absent because it falls back to Spark + * (see the Iceberg user guide), so there is no native path to measure. */ + private val bucketTypes = Seq("int", "long", "dec", "str_dict", "str", "bin", "date", "ts") + private val truncateTypes = Seq("int", "long", "str_dict", "str", "bin") + + /** (case name, query) for every transform, input type, and null variant. */ private def cases: Seq[(String, String)] = { - val bucket = Seq("c_int", "c_long", "c_dec", "c_str_dict", "c_str", "c_bin", "c_date", "c_ts") - .map(column => s"bucket($column)" -> s"select $catalog.system.bucket(16, $column)") - val truncate = Seq("c_int", "c_long", "c_dec", "c_str_dict", "c_str", "c_bin") - .map(column => s"truncate($column)" -> s"select $catalog.system.truncate(4, $column)") + def variants(types: Seq[String])(select: String => String): Seq[(String, String)] = + for { + t <- types + (suffix, tag) <- Seq("" -> "", "_n" -> ", nulls") + } yield { + val column = s"c_$t$suffix" + s"$t$tag" -> s"select ${select(column)} from parquetV1Table" + } + + val bucket = variants(bucketTypes)(c => s"$catalog.system.bucket(16, $c)") + .map { case (name, query) => s"bucket($name)" -> query } + val truncate = variants(truncateTypes)(c => s"$catalog.system.truncate(4, $c)") + .map { case (name, query) => s"truncate($name)" -> query } val temporal = Seq("years", "months", "days").flatMap { fn => - Seq("c_date", "c_ts").map(column => - s"$fn($column)" -> s"select $catalog.system.$fn($column)") - } :+ ("hours(c_ts)" -> s"select $catalog.system.hours(c_ts)") - (bucket ++ truncate ++ temporal).map { case (name, select) => - name -> s"$select from parquetV1Table" + variants(Seq("date", "ts"))(c => s"$catalog.system.$fn($c)").map { case (name, query) => + s"$fn($name)" -> query + } + } + val hours = variants(Seq("ts"))(c => s"$catalog.system.hours($c)").map { case (name, query) => + s"hours($name)" -> query + } + bucket ++ truncate ++ temporal ++ hours + } + + /** + * Fails if the two engines disagree on `query`. Rows are compared positionally: both cases read + * the same Parquet files with the same partitioning and neither plan shuffles, so the scan + * order is the same. The confs match the ones the benchmark times. + */ + private def verifyOutputsMatch(name: String, query: String): Unit = { + def collect(cometEnabled: Boolean): Array[Row] = + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName), + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString) { + spark.sql(query).collect() + } + + // `Row.equals` compares binary columns by reference, so normalize before comparing. + def comparable(row: Row): String = + row.toSeq + .map { + case bytes: Array[Byte] => bytes.mkString("[", ",", "]") + case other => String.valueOf(other) + } + .mkString("|") + + val sparkRows = collect(false).map(comparable) + val cometRows = collect(true).map(comparable) + assert( + sparkRows.length == cometRows.length, + s"$name: Spark produced ${sparkRows.length} rows, Comet ${cometRows.length}") + val mismatch = sparkRows.indices.find(i => sparkRows(i) != cometRows(i)) + mismatch.foreach { i => + throw new AssertionError( + s"$name: row $i differs -- Spark ${sparkRows(i)}, Comet ${cometRows(i)}") } } @@ -75,20 +138,10 @@ object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { runBenchmarkWithTable("Iceberg system functions", 1024 * 1024) { v => withTempPath { dir => withTempTable("parquetV1Table") { - prepareTable( - dir, - spark.sql(s""" - SELECT CAST(value AS INT) AS c_int, - value AS c_long, - CAST(value AS DECIMAL(38,10)) AS c_dec, - CAST(PMOD(value, 8) AS STRING) AS c_str_dict, - REPEAT(CAST(value AS STRING), 3) AS c_str, - CAST(CAST(value AS STRING) AS BINARY) AS c_bin, - DATE_ADD(DATE '1970-01-01', CAST(PMOD(value, 40000) AS INT)) AS c_date, - TIMESTAMP_SECONDS(PMOD(value, 4000000000)) AS c_ts - FROM $tbl""")) + prepareTable(dir, spark.sql(corpusQuery)) cases.foreach { case (name, query) => + verifyOutputsMatch(name, query) runBenchmark(name) { runExpressionBenchmark(name, v, query) } @@ -99,6 +152,23 @@ object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { } } + /** Every column, followed by a `_n` twin carrying one null in [[NullStride]]. */ + private def corpusQuery: String = { + val columns = Seq( + "c_int" -> "CAST(value AS INT)", + "c_long" -> "value", + "c_dec" -> "CAST(value AS DECIMAL(38,10))", + "c_str_dict" -> "CAST(PMOD(value, 8) AS STRING)", + "c_str" -> "REPEAT(CAST(value AS STRING), 3)", + "c_bin" -> "CAST(CAST(value AS STRING) AS BINARY)", + "c_date" -> "DATE_ADD(DATE '1970-01-01', CAST(PMOD(value, 40000) AS INT))", + "c_ts" -> "TIMESTAMP_SECONDS(PMOD(value, 4000000000))") + val projections = columns.flatMap { case (name, expr) => + Seq(s"$expr AS $name", s"IF(PMOD(value, $NullStride) = 0, NULL, $expr) AS ${name}_n") + } + s"SELECT ${projections.mkString(", ")} FROM $tbl" + } + private def icebergOnClasspath: Boolean = try { IcebergReflection.loadClass("org.apache.iceberg.spark.functions.BucketFunction") From f2f0c704b4e6175b9c17d52634c2b121b4b74a86 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 18:26:26 -0600 Subject: [PATCH 17/18] fix: keep the Iceberg benchmark's row collection compatible with Spark 3.x Spark 3.4 and 3.5 declare `SQLHelper.withSQLConf` as `(pairs: (String, String)*)(f: => Unit): Unit`; only Spark 4 has the generic result-returning form. Returning the wrapper from a method typed `Array[Row]` therefore failed to compile on both 3.x profiles and on the two Celeborn jobs, before any test could run. Collect into a local inside the configuration scope and return it after the block. Also drop the claim that the two arms scan identically. Enabling Comet replaces the Parquet scan and the projection with native operators as well as the transform, so each ratio is a query-level scan-plus-projection result, not a measurement of the transform on its own; the criterion benchmark in native/spark-expr/benches is what isolates the kernel. --- .../CometIcebergSystemFunctionBenchmark.scala | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala index 5e3e0a86999..7369a532304 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala @@ -36,6 +36,11 @@ import org.apache.comet.iceberg.IcebergReflection * output is compared between the two engines over the same corpus that is then timed, so a timing * cannot come from an engine that computed something else. * + * The two cases differ in more than the transform: enabling Comet also replaces the Parquet scan + * and the projection with native operators. Each ratio here is therefore a query-level + * scan-plus-projection result and does not isolate the cost of the transform. + * `native/spark-expr/benches/iceberg_transforms.rs` is the kernel-level measurement. + * * To run this benchmark: * {{{ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometIcebergSystemFunctionBenchmark @@ -91,13 +96,19 @@ object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { * order is the same. The confs match the ones the benchmark times. */ private def verifyOutputsMatch(name: String, query: String): Unit = { - def collect(cometEnabled: Boolean): Array[Row] = + // The rows are assigned to a local rather than returned from the `withSQLConf` block, because + // Spark 3.4 and 3.5 declare `SQLHelper.withSQLConf` as returning `Unit`; only Spark 4 has the + // generic result-returning form. + def collect(cometEnabled: Boolean): Array[Row] = { + var rows: Array[Row] = Array.empty withSQLConf( SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName), CometConf.COMET_ENABLED.key -> cometEnabled.toString, CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString) { - spark.sql(query).collect() + rows = spark.sql(query).collect() } + rows + } // `Row.equals` compares binary columns by reference, so normalize before comparing. def comparable(row: Row): String = @@ -128,8 +139,8 @@ object CometIcebergSystemFunctionBenchmark extends CometBenchmarkBase { return } // The Iceberg system functions are resolved through a v2 catalog, so one has to be - // registered. No Iceberg table is read: the data stays in Parquet so both cases scan - // identically and the only difference is who evaluates the transform. + // registered. No Iceberg table is read: the data stays in Parquet, which keeps the Iceberg + // reader out of the Spark case and lets both cases read the same files. withTempPath { warehouse => spark.conf.set(s"spark.sql.catalog.$catalog", "org.apache.iceberg.spark.SparkCatalog") spark.conf.set(s"spark.sql.catalog.$catalog.type", "hadoop") From 8c6a0645ae4ef699d93422c89a9935b3c8d80b24 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 18:26:37 -0600 Subject: [PATCH 18/18] test: partition the Iceberg write test on the string column again The multi-byte string column came out of the partitioned write test's spec because the test failed on the Linux CI runners with a missing data file, which never reproduced locally. The cause was iceberg-rust writing partition directory names raw where iceberg-java percent-encodes them, so a name containing multi-byte characters could not be reopened under the runners' ASCII platform charset. apache/iceberg-rust#2875 fixed that upstream and #5651 bumped the pin, which is now merged into this branch, so put the column back. The test then also covers the escaping. Dropping `truncate(1000, i64)` from the spec removes the reason for the Long.MinValue row filter as well, so the whole corpus is written again. --- .../CometIcebergSystemFunctionSuite.scala | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala index 6b968e55f14..d551a7cd0c3 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -207,20 +207,15 @@ class CometIcebergSystemFunctionSuite val table = s"$catalog.db.hidden_partitioning" // No `write.distribution-mode`: Iceberg picks hash distribution for a partitioned table, // which plans a shuffle and a local sort on the partition transforms. The string column is - // deliberately not a partition source: with it in the spec this test failed on the Linux - // CI runners with a missing data file, which does not reproduce locally, so the multi-byte - // values it would put in partition directory names stay out of the spec until that is - // understood. `bucket` and `truncate` over strings are covered by the comparison, filter, - // sort, and shuffle tests above. + // a partition source on purpose: its values include multi-byte characters and a surrogate + // pair, which end up in partition directory names. Those names were written raw until + // apache/iceberg-rust#2875, which #5651 picked up, so this also covers the URL escaping + // iceberg-java's `PartitionSpec.partitionToPath` applies. sql(s""" - CREATE TABLE $table (i32 INT, i64 BIGINT, ts TIMESTAMP, dt DATE) + CREATE TABLE $table (i32 INT, str STRING, ts TIMESTAMP, dt DATE) USING iceberg - PARTITIONED BY (bucket(4, i32), truncate(1000, i64), days(ts), months(dt))""") - // iceberg-rust's own truncate transform, which the writer uses for partition values, - // overflows on Long.MinValue in a debug build where Java wraps - // (apache/iceberg-rust#3141), so the boundary row stays out of the written set. - val rows = - s"SELECT i32, i64, ts, dt FROM $source WHERE i64 IS NULL OR i64 <> ${Long.MinValue}" + PARTITIONED BY (bucket(4, i32), truncate(2, str), days(ts), months(dt))""") + val rows = s"SELECT i32, str, ts, dt FROM $source" try { val plans = capturePlans(spark) { sql(s"INSERT INTO $table $rows") @@ -241,14 +236,14 @@ class CometIcebergSystemFunctionSuite s"the distribution shuffle stayed on Spark:\n$plan") } - checkAnswer(sql(s"SELECT i32, i64, ts, dt FROM $table"), sql(rows).collect()) + checkAnswer(sql(s"SELECT i32, str, ts, dt FROM $table"), sql(rows).collect()) // Iceberg's own view of the partitions must match what the JVM transforms compute over // the written rows: one partition per distinct transform tuple. withSQLConf(CometConf.COMET_ENABLED.key -> "false") { val expected = sql(s""" SELECT COUNT(*) FROM ( - SELECT DISTINCT $catalog.system.bucket(4, i32), $catalog.system.truncate(1000, i64), + SELECT DISTINCT $catalog.system.bucket(4, i32), $catalog.system.truncate(2, str), $catalog.system.days(ts), $catalog.system.months(dt) FROM $table)""").collect().head.getLong(0) val actual = sql(s"SELECT COUNT(*) FROM $table.partitions").collect().head.getLong(0)