diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0d62838cd27..bedaafae878 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 f784dc0e510..5d99f5e5bf8 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 e9227dea8bf..49f35009202 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..a693a450f99 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/docs/source/user-guide/latest/iceberg.md @@ -196,6 +196,41 @@ 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 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. + +`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 +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/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index 57f39d1b395..1d3fb668217 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -1172,3 +1172,362 @@ 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 is an + /// iceberg-rust bug affecting the writer's own partition values, independent of these + /// 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![ + ( + "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 (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(); + 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. 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. + 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 + ); + } +} 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/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..4e267fdd706 --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/bucket.rs @@ -0,0 +1,371 @@ +// 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::{ArrayRef, AsArray, Int32Array}; +use arrow::datatypes::{ + DataType, Date32Type, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type, 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 (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); + } + 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. 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 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 => 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))), + 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 => 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)) +} + +/// `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::{ + Array, BinaryArray, Date32Array, Decimal128Array, DictionaryArray, Int16Array, Int32Array, + Int64Array, Int8Array, StringArray, TimestampMicrosecondArray, + }; + 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, 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]))); + 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_hashed_once_per_value() { + 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..aec13f949b3 --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/mod.rs @@ -0,0 +1,139 @@ +// 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, AsArray}; +use arrow::compute::take; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, Result, ScalarValue}; +use datafusion::logical_expr::ColumnarValue; + +/// 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) => Ok(ColumnarValue::Array(apply_to_array(array, kernel)?)), + ColumnarValue::Scalar(scalar) => { + let result = apply_to_array(&scalar.to_array()?, kernel)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } +} + +/// 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 { + 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..f49c5fc122e --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/temporal.rs @@ -0,0 +1,408 @@ +// 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 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 datafusion::common::{utils::take_function_args, 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(crate) 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`. 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`. 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) -> 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) -> i32 { + let (year, month0) = civil_from_days(days); + (year - UNIX_EPOCH_YEAR) * 12 + month0 +} + +/// 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) -> i32) -> Result { + match array.data_type() { + DataType::Date32 => Ok(array.as_primitive::().unary(f)), + DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(array + .as_primitive::() + .unary(|micros| f(micros_to_days(micros)))), + 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(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, _) => Arc::new( + array + .as_primitive::() + .unary::<_, Int32Type>(micros_to_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(crate) 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, Date32Array, 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). 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 + (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] + 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), + // 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. + 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) + ); + } + } + + /// 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 [ + 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..70b638f52bc --- /dev/null +++ b/native/spark-expr/src/iceberg_funcs/truncate.rs @@ -0,0 +1,427 @@ +// 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 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}; +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 +} + +/// `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( + 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) => { + // 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 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); + is_valid_decimal_precision(truncated, *precision).then_some(truncated) + }); + Arc::new(truncated.with_precision_and_scale(*precision, *scale)?) + } + 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) +} + +/// `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::{ + BinaryArray, DictionaryArray, Int16Array, Int32Array, Int64Array, Int8Array, StringArray, + }; + use arrow::datatypes::Int8Type; + 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 % 1000 == -808, so the wrapped remainder is 192. + i64::MIN.wrapping_sub(192), + 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]) + ); + } + + /// 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(); + 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..b9231f46d62 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/icebergFunctions.scala @@ -0,0 +1,263 @@ +/* + * 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. + * + * 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, String), CometExpressionSerde[StaticInvoke])] = + inner.map(name => (MagicMethod, s"$FunctionsPackage$outer$$$name") -> handler) + + /** + * 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, 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 | _: DecimalType => + true + case _ => false + }) + +/** + * `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 => 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( + 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) + } +} + +object CometIcebergYears + extends CometIcebergTemporalTransform( + "iceberg_years", + Set(DateType, TimestampType, TimestampNTZType)) + +object CometIcebergMonths + extends CometIcebergTemporalTransform( + "iceberg_months", + Set(DateType, TimestampType, TimestampNTZType)) + +object CometIcebergDays + extends CometIcebergTemporalTransform( + "iceberg_days", + Set(DateType, TimestampType, TimestampNTZType)) + +object CometIcebergHours + 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 6501a1d46cb..8944b90835d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/statics.scala +++ b/spark/src/main/scala/org/apache/comet/serde/statics.scala @@ -32,36 +32,57 @@ 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 + + private def handlerFor(expr: StaticInvoke): Option[CometExpressionSerde[StaticInvoke]] = + staticInvokeExpressions.get((expr.functionName, expr.staticObject.getName)) + + 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], 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..d551a7cd0c3 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometIcebergSystemFunctionSuite.scala @@ -0,0 +1,461 @@ +/* + * 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.io.File +import java.math.{BigDecimal => JBigDecimal, BigInteger} +import java.nio.file.Files +import java.time.{Instant, LocalDate, LocalDateTime, ZoneOffset} + +import scala.util.Random + +import org.scalactic.source.Position +import org.scalatest.Tag + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +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.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.comet.serde.{CometExpressionSerde, CometIcebergBucket, CometIcebergTruncate, CometStaticInvoke, Compatible, SupportLevel, 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", "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 = _ + private def sourcePath: String = new File(sourceDir, "data").getAbsolutePath + + override def beforeAll(): Unit = { + super.beforeAll() + sourceDir = Files.createTempDirectory("comet-iceberg-system-functions").toFile + // 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 = { + try deleteRecursively(sourceDir) + finally super.afterAll() + } + + test("bucket matches Iceberg for every supported type") { + withSourceTable { + 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 { + 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") + } + } + } + + 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 + // leak into the native result either. + for (timezone <- Seq("UTC", "America/Los_Angeles", "Asia/Kathmandu")) { + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { + 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, ${transformed.mkString(", ")} 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) + val sorts = collect(stripAQEPlan(df.queryExecution.executedPlan)) { case s: CometSortExec => + s + } + assert(sorts.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. The string column is + // 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, str STRING, ts TIMESTAMP, dt DATE) + USING iceberg + 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") + } + 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(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), + $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)() + 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(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(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") { + 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) { + withParquetTable(sourcePath, 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))) + + 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). + */ + 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 + + 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()))) + } + + // 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 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"), + new JBigDecimal("99999999999999.9999"), + JBigDecimal.ZERO.setScale(4), + new JBigDecimal("-0.0001"), + null), + Seq( + dec38Max.negate(), + dec38Max, + JBigDecimal.ZERO.setScale(10), + 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), + LocalDate.ofEpochDay(-366), + 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) + } +} 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 b964bab56cd..66ea7216786 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) @@ -1570,14 +1569,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") @@ -1631,28 +1630,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) } 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 new file mode 100644 index 00000000000..7369a532304 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergSystemFunctionBenchmark.scala @@ -0,0 +1,190 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.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 + +/** + * 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. + * + * 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. + * + * 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 + * }}} + * Results will be written to + * "spark/benchmarks/CometIcebergSystemFunctionBenchmark-**results.txt". + */ +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 + + /** + * 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; + * `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)] = { + 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 => + 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 = { + // 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) { + rows = spark.sql(query).collect() + } + rows + } + + // `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)}") + } + } + + 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, 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") + 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(corpusQuery)) + + cases.foreach { case (name, query) => + verifyOutputsMatch(name, query) + runBenchmark(name) { + runExpressionBenchmark(name, v, query) + } + } + } + } + } + } + } + + /** 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") + true + } catch { + case _: ClassNotFoundException => false + } +}