diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 07c1086624..8eac031540 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -84,15 +84,16 @@ incompatibility details. ## Configs for Inspecting Plans and Fallback -Comet provides four configs for understanding what is happening in a plan. +Comet provides five configs for understanding what is happening in a plan. They serve different purposes and produce output in different places. -| Config | Output destination | What you see | -| ---------------------------------------- | ---------------------------------- | --------------------------------------------------------------------------------------------- | -| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | -| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | -| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | -| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | +| Config | Output destination | What you see | +| ---------------------------------------- | ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | +| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | +| `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A single `[COMET-INFO: JVM codegen dispatcher: , , ...]` segment naming each expression on the operator that was routed through the JVM codegen dispatcher. | +| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | +| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | ### `spark.comet.explainFallback.enabled` @@ -108,6 +109,48 @@ when you want to see all reasons, including ones that not include the surrounding plan, so it is best for accumulating diagnostics across many queries. +### `spark.comet.explainCodegen.enabled` + +Disabled by default. When enabled, Comet annotates the surrounding Comet +operator (`CometProject`, `CometFilter`, etc.) with a single combined +`[COMET-INFO: JVM codegen dispatcher: , , ...]` segment listing +every expression on that operator that was routed through the JVM codegen +dispatcher (Spark's own `doGenCode` running inside a Comet kernel), gated by +`spark.comet.exec.scalaUDF.codegen.enabled=true`. The projection stays +Comet-native — this is informational only, useful for distinguishing "runs +natively in DataFusion" from "runs Spark's generated code inside a Comet +kernel". The annotation only appears for expressions Comet actually routed +through the dispatcher on this plan — either because no native DataFusion +implementation exists (`CometCodegenDispatch` serdes such as `hypot`, +`levenshtein`, `pmod`), or because a native path exists but declined this +specific input (`CodegenDispatchFallback` mixins). It is not a general "this +expression ran here" marker. + +Example: + +```scala +spark.conf.set("spark.comet.exec.scalaUDF.codegen.enabled", "true") +spark.conf.set("spark.comet.explainCodegen.enabled", "true") + +val df = spark.sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") +println(new org.apache.comet.ExtendedExplainInfo() + .generateExtendedInfo(df.queryExecution.executedPlan)) +``` + +Output: + +``` +CometNativeColumnarToRow ++- CometProject [COMET-INFO: JVM codegen dispatcher: hypot, levenshtein] + +- CometNativeScan parquet spark_catalog.default.t + +Comet accelerated 2 out of 2 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. +``` + +Note that the operator is still `CometProject` (Comet-accelerated); only the +per-expression evaluation for `hypot` and `levenshtein` takes the JVM codegen +path. + ### `spark.comet.explain.format` This config is read by `org.apache.comet.ExtendedExplainInfo`, which Spark diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..2d5b0c0f1b 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -139,3 +139,11 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "floor" +harness = false + +[[bench]] +name = "ceil" +harness = false diff --git a/native/spark-expr/benches/ceil.rs b/native/spark-expr/benches/ceil.rs new file mode 100644 index 0000000000..0436860f11 --- /dev/null +++ b/native/spark-expr/benches/ceil.rs @@ -0,0 +1,65 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Decimal128Array; +use arrow::datatypes::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_ceil; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: i128 = 8192; + +fn decimal_args(precision: u8, scale: i8, spread: i128) -> Vec { + let values: Vec = (0..NUM_ROWS) + .map(|i| (i - NUM_ROWS / 2) * spread + i % 7) + .collect(); + let array = Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .unwrap(); + vec![ColumnarValue::Array(Arc::new(array))] +} + +fn criterion_benchmark(c: &mut Criterion) { + // decimal(18, 4): unscaled values fit comfortably in 64 bits + let narrow = decimal_args(18, 4, 1_000_003); + // decimal(38, 6): unscaled values exceed the 64-bit range + let wide = decimal_args(38, 6, 1_000_000_000_000_000_003); + + let mut group = c.benchmark_group("ceil"); + group.bench_function("ceil_decimal_18_4", |b| { + b.iter(|| { + black_box(spark_ceil( + black_box(&narrow), + black_box(&DataType::Decimal128(18, 0)), + )) + }) + }); + group.bench_function("ceil_decimal_38_6", |b| { + b.iter(|| { + black_box(spark_ceil( + black_box(&wide), + black_box(&DataType::Decimal128(38, 0)), + )) + }) + }); + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/floor.rs b/native/spark-expr/benches/floor.rs new file mode 100644 index 0000000000..5005c40f00 --- /dev/null +++ b/native/spark-expr/benches/floor.rs @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, Decimal128Array, Float64Array}; +use arrow::datatypes::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_floor; +use std::hint::black_box; +use std::sync::Arc; + +const ROWS: usize = 8192; + +/// Unscaled decimal values that fit in 64 bits, which is the common case, with every 10th row +/// null. +fn decimal_array(precision: u8, scale: i8) -> ArrayRef { + let values: Vec> = (0..ROWS) + .map(|i| { + if i % 10 == 0 { + None + } else { + let v = (i as i128) * 987_654_321 % 100_000_000_000_000_000; + Some(if i % 2 == 0 { v } else { -v }) + } + }) + .collect(); + Arc::new( + Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) +} + +/// Unscaled decimal values too wide for 64 bits, exercising the 128-bit fallback. +fn wide_decimal_array(precision: u8, scale: i8) -> ArrayRef { + let values: Vec> = (0..ROWS) + .map(|i| { + if i % 10 == 0 { + None + } else { + let v = (i as i128 + 1) * 10_000_000_000_000_000_000_000_000_000; + Some(if i % 2 == 0 { v } else { -v }) + } + }) + .collect(); + Arc::new( + Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) +} + +fn float_array() -> ArrayRef { + Arc::new(Float64Array::from( + (0..ROWS) + .map(|i| i as f64 * 1.5 - 4096.0) + .collect::>(), + )) +} + +fn criterion_benchmark(c: &mut Criterion) { + let dec_2 = decimal_array(38, 2); + c.bench_function("spark_floor: decimal128(38,2)", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&dec_2))]; + b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(37, 0)).unwrap())) + }); + + let dec_18 = decimal_array(38, 18); + c.bench_function("spark_floor: decimal128(38,18)", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&dec_18))]; + b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(21, 0)).unwrap())) + }); + + let wide = wide_decimal_array(38, 6); + c.bench_function("spark_floor: decimal128(38,6) wide values", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&wide))]; + b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(33, 0)).unwrap())) + }); + + let floats = float_array(); + c.bench_function("spark_floor: float64", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&floats))]; + b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Float64).unwrap())) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index bba34d0368..ce2fe515b6 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -255,53 +255,36 @@ macro_rules! cast_int_to_int_macro { $eval_mode:expr, $from_arrow_primitive_type: ty, $to_arrow_primitive_type: ty, - $from_data_type: expr, $to_native_type: ty, $spark_from_data_type_name: expr, - $spark_to_data_type_name: expr + $spark_to_data_type_name: expr, + $spark_int_literal_suffix: expr ) => {{ let cast_array = $array .as_any() .downcast_ref::>() .unwrap(); - let spark_int_literal_suffix = match $from_data_type { - &DataType::Int64 => "L", - &DataType::Int16 => "S", - &DataType::Int8 => "T", - _ => "", - }; - let output_array = match $eval_mode { - EvalMode::Legacy => cast_array - .iter() - .map(|value| match value { - Some(value) => { - Ok::, SparkError>(Some(value as $to_native_type)) - } - _ => Ok(None), - }) - .collect::, _>>(), - _ => cast_array - .iter() - .map(|value| match value { - Some(value) => { - let res = <$to_native_type>::try_from(value); - if res.is_err() { - Err(cast_overflow( - &(value.to_string() + spark_int_literal_suffix), - $spark_from_data_type_name, - $spark_to_data_type_name, - )) - } else { - Ok::, SparkError>(Some(res.unwrap())) - } - } - _ => Ok(None), + let output_array: PrimitiveArray<$to_arrow_primitive_type> = match $eval_mode { + // Legacy narrowing keeps the low-order bits of the source value, which is what a + // wrapping `as` conversion does. Being infallible, it maps the values buffer in a + // single pass and carries the null buffer over untouched. + EvalMode::Legacy => { + cast_array.unary::<_, $to_arrow_primitive_type>(|value| value as $to_native_type) + } + // `try_unary` only applies the conversion to non-null slots, so an out-of-range + // value sitting under a null cannot raise a spurious overflow error. + _ => cast_array.try_unary::<_, $to_arrow_primitive_type, SparkError>(|value| { + <$to_native_type>::try_from(value).map_err(|_| { + cast_overflow( + &(value.to_string() + $spark_int_literal_suffix), + $spark_from_data_type_name, + $spark_to_data_type_name, + ) }) - .collect::, _>>(), - }?; - let result: SparkResult = Ok(Arc::new(output_array) as ArrayRef); - result + })?, + }; + Ok(Arc::new(output_array) as ArrayRef) }}; } @@ -780,22 +763,22 @@ pub(crate) fn spark_cast_int_to_int( ) -> SparkResult { match (from_type, to_type) { (DataType::Int64, DataType::Int32) => cast_int_to_int_macro!( - array, eval_mode, Int64Type, Int32Type, from_type, i32, "BIGINT", "INT" + array, eval_mode, Int64Type, Int32Type, i32, "BIGINT", "INT", "L" ), (DataType::Int64, DataType::Int16) => cast_int_to_int_macro!( - array, eval_mode, Int64Type, Int16Type, from_type, i16, "BIGINT", "SMALLINT" + array, eval_mode, Int64Type, Int16Type, i16, "BIGINT", "SMALLINT", "L" ), (DataType::Int64, DataType::Int8) => cast_int_to_int_macro!( - array, eval_mode, Int64Type, Int8Type, from_type, i8, "BIGINT", "TINYINT" + array, eval_mode, Int64Type, Int8Type, i8, "BIGINT", "TINYINT", "L" ), (DataType::Int32, DataType::Int16) => cast_int_to_int_macro!( - array, eval_mode, Int32Type, Int16Type, from_type, i16, "INT", "SMALLINT" - ), - (DataType::Int32, DataType::Int8) => cast_int_to_int_macro!( - array, eval_mode, Int32Type, Int8Type, from_type, i8, "INT", "TINYINT" + array, eval_mode, Int32Type, Int16Type, i16, "INT", "SMALLINT", "" ), + (DataType::Int32, DataType::Int8) => { + cast_int_to_int_macro!(array, eval_mode, Int32Type, Int8Type, i8, "INT", "TINYINT", "") + } (DataType::Int16, DataType::Int8) => cast_int_to_int_macro!( - array, eval_mode, Int16Type, Int8Type, from_type, i8, "SMALLINT", "TINYINT" + array, eval_mode, Int16Type, Int8Type, i8, "SMALLINT", "TINYINT", "S" ), _ => unreachable!( "{}", diff --git a/native/spark-expr/src/math_funcs/ceil.rs b/native/spark-expr/src/math_funcs/ceil.rs index 324a118fcb..292ef117e9 100644 --- a/native/spark-expr/src/math_funcs/ceil.rs +++ b/native/spark-expr/src/math_funcs/ceil.rs @@ -16,7 +16,9 @@ // under the License. use crate::downcast_compute_op; -use crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar}; +use crate::math_funcs::utils::{ + dispatch_pow10, get_precision_scale, make_decimal_array, make_decimal_scalar, +}; use arrow::array::{Array, ArrowNativeTypeOp}; use arrow::array::{Float32Array, Float64Array, Int64Array}; use arrow::datatypes::DataType; @@ -45,10 +47,13 @@ pub fn spark_ceil( let result = array.as_any().downcast_ref::().unwrap(); Ok(ColumnarValue::Array(Arc::new(result.clone()))) } - DataType::Decimal128(_, scale) if *scale > 0 => { - let f = decimal_ceil_f(scale); + DataType::Decimal128(_, input_scale) if *input_scale > 0 => { let (precision, scale) = get_precision_scale(data_type); - make_decimal_array(array, precision, scale, &f) + dispatch_pow10!( + *input_scale, + EXP => make_decimal_array(array, precision, scale, decimal_ceil_pow10::), + make_decimal_array(array, precision, scale, decimal_ceil_f(*input_scale)) + ) } other => Err(DataFusionError::Internal(format!( "Unsupported data type {other:?} for function ceil", @@ -62,10 +67,10 @@ pub fn spark_ceil( a.map(|x| x.ceil() as i64), ))), ScalarValue::Int64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))), - ScalarValue::Decimal128(a, _, scale) if *scale > 0 => { - let f = decimal_ceil_f(scale); + ScalarValue::Decimal128(a, _, input_scale) if *input_scale > 0 => { + let f = decimal_ceil_f(*input_scale); let (precision, scale) = get_precision_scale(data_type); - make_decimal_scalar(a, precision, scale, &f) + make_decimal_scalar(a, precision, scale, f) } _ => Err(DataFusionError::Internal(format!( "Unsupported data type {:?} for function ceil", @@ -76,11 +81,32 @@ pub fn spark_ceil( } #[inline] -fn decimal_ceil_f(scale: &i8) -> impl Fn(i128) -> i128 { - let div = 10_i128.pow_wrapping(*scale as u32); +fn decimal_ceil_f(scale: i8) -> impl Fn(i128) -> i128 { + let div = 10_i128.pow_wrapping(scale as u32); move |x: i128| div_ceil(x, div) } +/// Ceiling-divides an unscaled decimal by `10^EXP`. +/// +/// `EXP` is a compile-time constant so that the divisor is folded in and the division lowered to a +/// multiply-and-shift. A 128-bit division is always a libcall, even by a constant, so values that +/// fit in 64 bits take a 64-bit path; unscaled decimals rarely exceed that range. +#[inline] +fn decimal_ceil_pow10(x: i128) -> i128 { + match i64::try_from(x) { + Ok(x) => div_ceil(x, const { 10_i64.pow(EXP) }) as i128, + Err(_) => decimal_ceil_wide(x, const { 10_i128.pow(EXP) }), + } +} + +/// Kept out of line so that the libcall and its stack frame stay out of the loop body of every +/// [`decimal_ceil_pow10`] instantiation. +#[cold] +#[inline(never)] +fn decimal_ceil_wide(x: i128, div: i128) -> i128 { + div_ceil(x, div) +} + #[cfg(test)] mod test { use crate::spark_ceil; @@ -187,6 +213,36 @@ mod test { Ok(()) } + #[test] + fn test_ceil_decimal128_wide_array() -> Result<()> { + // Unscaled values that exceed i64::MAX (~9.22e18) exercise the wide fallback in + // decimal_ceil_pow10. Values chosen at scale 6 targeting Decimal128(38, 0): + // 20_000_000_000_000_000_000 / 10^6 = 20_000_000_000_000 (exact multiple) + // 20_000_000_000_000_000_001 / 10^6 -> ceil 20_000_000_000_001 (positive remainder) + // -20_000_000_000_000_000_001 / 10^6 -> ceil -20_000_000_000_000 (negative remainder) + let array = Decimal128Array::from(vec![ + Some(20_000_000_000_000_000_000_i128), + Some(20_000_000_000_000_000_001_i128), + Some(-20_000_000_000_000_000_001_i128), + None, + ]) + .with_precision_and_scale(38, 6)?; + let args = vec![ColumnarValue::Array(Arc::new(array))]; + let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Decimal128(38, 0))? else { + unreachable!() + }; + let expected = Decimal128Array::from(vec![ + Some(20_000_000_000_000_i128), + Some(20_000_000_000_001_i128), + Some(-20_000_000_000_000_i128), + None, + ]) + .with_precision_and_scale(38, 0)?; + let actual = result.as_any().downcast_ref::().unwrap(); + assert_eq!(actual, &expected); + Ok(()) + } + #[test] fn test_ceil_f32_scalar() -> Result<()> { let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.2345)))]; diff --git a/native/spark-expr/src/math_funcs/floor.rs b/native/spark-expr/src/math_funcs/floor.rs index c7c1c9e904..973bbab708 100644 --- a/native/spark-expr/src/math_funcs/floor.rs +++ b/native/spark-expr/src/math_funcs/floor.rs @@ -16,7 +16,9 @@ // under the License. use crate::downcast_compute_op; -use crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar}; +use crate::math_funcs::utils::{ + dispatch_pow10, get_precision_scale, make_decimal_array, make_decimal_scalar, +}; use arrow::array::{Array, ArrowNativeTypeOp}; use arrow::array::{Float32Array, Float64Array, Int64Array}; use arrow::datatypes::DataType; @@ -45,10 +47,13 @@ pub fn spark_floor( let result = array.as_any().downcast_ref::().unwrap(); Ok(ColumnarValue::Array(Arc::new(result.clone()))) } - DataType::Decimal128(_, scale) if *scale > 0 => { - let f = decimal_floor_f(scale); + DataType::Decimal128(_, input_scale) if *input_scale > 0 => { let (precision, scale) = get_precision_scale(data_type); - make_decimal_array(array, precision, scale, &f) + dispatch_pow10!( + *input_scale, + EXP => make_decimal_array(array, precision, scale, decimal_floor_pow10::), + make_decimal_array(array, precision, scale, decimal_floor_f(*input_scale)) + ) } other => Err(DataFusionError::Internal(format!( "Unsupported data type {other:?} for function floor", @@ -62,10 +67,10 @@ pub fn spark_floor( a.map(|x| x.floor() as i64), ))), ScalarValue::Int64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))), - ScalarValue::Decimal128(a, _, scale) if *scale > 0 => { - let f = decimal_floor_f(scale); + ScalarValue::Decimal128(a, _, input_scale) if *input_scale > 0 => { + let f = decimal_floor_f(*input_scale); let (precision, scale) = get_precision_scale(data_type); - make_decimal_scalar(a, precision, scale, &f) + make_decimal_scalar(a, precision, scale, f) } _ => Err(DataFusionError::Internal(format!( "Unsupported data type {:?} for function floor", @@ -76,11 +81,32 @@ pub fn spark_floor( } #[inline] -fn decimal_floor_f(scale: &i8) -> impl Fn(i128) -> i128 { - let div = 10_i128.pow_wrapping(*scale as u32); +fn decimal_floor_f(scale: i8) -> impl Fn(i128) -> i128 { + let div = 10_i128.pow_wrapping(scale as u32); move |x: i128| div_floor(x, div) } +/// Floor-divides an unscaled decimal by `10^EXP`. +/// +/// `EXP` is a compile-time constant so that the divisor is folded in and the division lowered to a +/// multiply-and-shift. A 128-bit division is always a libcall, even by a constant, so values that +/// fit in 64 bits take a 64-bit path; unscaled decimals rarely exceed that range. +#[inline] +fn decimal_floor_pow10(x: i128) -> i128 { + match i64::try_from(x) { + Ok(x) => div_floor(x, const { 10_i64.pow(EXP) }) as i128, + Err(_) => decimal_floor_wide(x, const { 10_i128.pow(EXP) }), + } +} + +/// Kept out of line so that the libcall and its stack frame stay out of the loop body of every +/// [`decimal_floor_pow10`] instantiation. +#[cold] +#[inline(never)] +fn decimal_floor_wide(x: i128, div: i128) -> i128 { + div_floor(x, div) +} + #[cfg(test)] mod test { use crate::spark_floor; diff --git a/native/spark-expr/src/math_funcs/utils.rs b/native/spark-expr/src/math_funcs/utils.rs index 6109c8d161..23e7b336ef 100644 --- a/native/spark-expr/src/math_funcs/utils.rs +++ b/native/spark-expr/src/math_funcs/utils.rs @@ -41,23 +41,113 @@ macro_rules! downcast_compute_op { }}; } +/// Evaluates `$body` with `$exp` bound as a compile-time constant equal to `$scale`, for every +/// decimal scale whose divisor `10^scale` fits in an `i64`. Larger scales evaluate `$fallback`, +/// which must derive the divisor at runtime. +/// +/// Specializing on the scale lets the divisor be folded into the generated code, turning the +/// division by `10^scale` into a multiply-and-shift. +macro_rules! dispatch_pow10 { + ($scale:expr, $exp:ident => $body:expr, $fallback:expr) => { + match $scale { + 1 => { + const $exp: u32 = 1; + $body + } + 2 => { + const $exp: u32 = 2; + $body + } + 3 => { + const $exp: u32 = 3; + $body + } + 4 => { + const $exp: u32 = 4; + $body + } + 5 => { + const $exp: u32 = 5; + $body + } + 6 => { + const $exp: u32 = 6; + $body + } + 7 => { + const $exp: u32 = 7; + $body + } + 8 => { + const $exp: u32 = 8; + $body + } + 9 => { + const $exp: u32 = 9; + $body + } + 10 => { + const $exp: u32 = 10; + $body + } + 11 => { + const $exp: u32 = 11; + $body + } + 12 => { + const $exp: u32 = 12; + $body + } + 13 => { + const $exp: u32 = 13; + $body + } + 14 => { + const $exp: u32 = 14; + $body + } + 15 => { + const $exp: u32 = 15; + $body + } + 16 => { + const $exp: u32 = 16; + $body + } + 17 => { + const $exp: u32 = 17; + $body + } + 18 => { + const $exp: u32 = 18; + $body + } + _ => $fallback, + } + }; +} + +pub(crate) use dispatch_pow10; + #[inline] pub(crate) fn make_decimal_scalar( a: &Option, precision: u8, scale: i8, - f: &dyn Fn(i128) -> i128, + f: impl Fn(i128) -> i128, ) -> Result { let result = ScalarValue::Decimal128(a.map(f), precision, scale); Ok(ColumnarValue::Scalar(result)) } +/// Generic over the operation so that it is monomorphized into the element loop rather than +/// called through a vtable for every value. #[inline] pub(crate) fn make_decimal_array( array: &ArrayRef, precision: u8, scale: i8, - f: &dyn Fn(i128) -> i128, + f: impl Fn(i128) -> i128, ) -> Result { let array = array.as_primitive::(); let result: Decimal128Array = arrow::compute::kernels::arity::unary(array, f); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 1351a391c0..119493ecfb 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -686,6 +686,15 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXPLAIN_CODEGEN_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.explainCodegen.enabled") + .category(CATEGORY_EXEC_EXPLAIN) + .doc("When enabled, Comet annotates the surrounding Comet operator with a `[COMET-INFO: " + + "JVM codegen dispatcher: ]` segment listing every expression it routed through " + + "the JVM codegen dispatcher. Disabled by default.") + .booleanConf + .createWithDefault(false) + val COMET_ONHEAP_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.onHeap.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index f117372b3c..202c52b158 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -383,4 +383,19 @@ object CometSparkSessionExtensions extends Logging { node } + /** + * Record that `node` (typically an `Expression`) is routing through the JVM codegen dispatcher. + * `CometExecRule.rollUpInfoMessages` collects the names across an operator's expression trees + * and emits one combined `[COMET-INFO: ...]` segment. + */ + def withCodegenDispatchExpr[T <: TreeNode[_]](node: T, name: String): T = { + if (name != null && name.nonEmpty) { + val existing = node + .getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS) + .getOrElse(Set.empty[String]) + node.setTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS, existing + name) + } + node + } + } diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 5318c3f87c..0a7a8f4e30 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -224,6 +224,11 @@ object CometExplainInfo { val FALLBACK_REASONS = new TreeNodeTag[Set[String]]("CometFallbackReasons") val EXTENSION_INFO = new TreeNodeTag[Set[String]]("CometExtensionInfo") + // Expression names the serde routed through the JVM codegen dispatcher. Rolled up per + // operator by `CometExecRule.rollUpInfoMessages` into one combined `[COMET-INFO: ...]`. + val CODEGEN_DISPATCH_EXPRS = + new TreeNodeTag[Set[String]]("CometCodegenDispatchExprs") + def getActualPlan(node: TreeNode[_]): TreeNode[_] = { node match { case p: AdaptiveSparkPlanExec => getActualPlan(p.executedPlan) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index dc8268950e..70982bdb26 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -729,14 +729,25 @@ case class CometExecRule(session: SparkSession) * Lift informational (non-fallback) messages tagged on an operator and its expressions onto the * converted Comet plan node so they appear in verbose extended explain output. Expression-level * hints would otherwise be invisible because explain only traverses plan nodes, not - * expressions. + * expressions. `CODEGEN_DISPATCH_EXPRS` names across the tree are aggregated into one combined + * info line. */ private def rollUpInfoMessages(op: SparkPlan, exec: SparkPlan): Unit = { - val fromOp = op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]) - val fromExprs = op.expressions - .flatMap(_.collect { case e: Expression => e }) - .flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String])) - (fromOp ++ fromExprs).foreach(msg => withInfo(exec, msg)) + val allExprs = op.expressions.flatMap(_.collect { case e: Expression => e }) + + val infos = + op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]) ++ + allExprs.flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO)).flatten + infos.foreach(msg => withInfo(exec, msg)) + + val routedNames = allExprs + .flatMap(_.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS)) + .flatten + .distinct + .sorted + if (routedNames.nonEmpty) { + withInfo(exec, s"JVM codegen dispatcher: ${routedNames.mkString(", ")}") + } } private def isOperatorEnabled( diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 3df275f057..e04f680b3b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -19,12 +19,14 @@ package org.apache.comet.serde +import java.util.Locale + import org.apache.spark.SparkEnv import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, withFallbackReason} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} @@ -55,6 +57,19 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = emitJvmCodegenDispatch(expr, inputs, binding) + /** + * Canonical name for annotation. `BinaryMathExpression` (Hypot, Pow, ...) overrides + * `prettyName` to raw uppercase; `ScalaUDF.prettyName` collapses to `"scalaudf"` for every user + * UDF. Prefer `ScalaUDF.udfName` when set, then lowercase to normalize. + */ + private def displayName(expr: Expression): String = { + val raw = expr match { + case s: ScalaUDF => s.udfName.getOrElse(s.prettyName) + case other => other.prettyName + } + raw.toLowerCase(Locale.ROOT) + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -70,11 +85,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val exprName = displayName(expr) if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { withFallbackReason( expr, - s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has no native " + - "path so the plan falls back to Spark") + s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " + + "no native path so the plan falls back to Spark") return None } @@ -97,7 +113,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // at execute. CometBatchKernelCodegen.canHandle(boundExpr) match { case Some(reason) => - withFallbackReason(expr, reason) + withFallbackReason(expr, s"$exprName: $reason") return None case None => } @@ -110,12 +126,26 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val buffer = serializer.serialize(boundExpr) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) - val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding) - .getOrElse(return None) + val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + + "expression payload") + return None + } - val dataArgs = - attrs.map(a => exprToProtoInternal(a, inputs, binding).getOrElse(return None)) - val returnTypeProto = serializeDataType(expr.dataType).getOrElse(return None) + val dataArgs = attrs.map { a => + exprToProtoInternal(a, inputs, binding).getOrElse { + withFallbackReason(expr, s"$exprName: codegen dispatch: could not serialize data arg $a") + return None + } + } + val returnTypeProto = serializeDataType(expr.dataType).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") + return None + } val udfBuilder = ExprOuterClass.JvmScalarUdf .newBuilder() @@ -125,6 +155,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { udfBuilder .setReturnType(returnTypeProto) .setReturnNullable(expr.nullable) + // Opt-in dispatch annotation for extended explain. Rolled up per operator by + // `CometExecRule.rollUpInfoMessages` into a single `[COMET-INFO: JVM codegen dispatcher: + // ...]` line. Informational only - does not trigger fallback. + if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { + withCodegenDispatchExpr(expr, exprName) + } Some( ExprOuterClass.Expr .newBuilder() diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 76ffbc8c99..aa610d0751 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -131,6 +131,72 @@ class CometCodegenSuite } } + test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") { + // With the opt-in flag on, `hypot` and `levenshtein` (both `CometCodegenDispatch`) roll up + // into one `[COMET-INFO: JVM codegen dispatcher: hypot, levenshtein]` line on the + // `CometProject`. With the flag off (default), no such line appears. + withTable("t") { + sql("CREATE TABLE t (a DOUBLE, b DOUBLE, s1 STRING, s2 STRING) USING parquet") + sql("INSERT INTO t VALUES (3.0, 4.0, 'kitten', 'sitting')") + + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + explain.contains("[COMET-INFO:"), + s"expected a [COMET-INFO: segment, got:\n$explain") + // Names appear alphabetically via `.distinct.sorted` in rollUpInfoMessages. + assert( + explain.contains("JVM codegen dispatcher: hypot, levenshtein"), + s"expected combined codegen-dispatch info, got:\n$explain") + } + + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + !explain.contains("JVM codegen dispatcher"), + s"expected NO codegen-dispatch info with the flag off, got:\n$explain") + } + } + } + + test("codegen dispatch fallback reasons name the expression") { + // Flag-off short-circuit tags the expression `: ` so distinct expressions + // don't collapse in the `Set[String]` roll-up. + withTable("t") { + sql("CREATE TABLE t (a DOUBLE, b DOUBLE) USING parquet") + sql("INSERT INTO t VALUES (3.0, 4.0)") + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b) FROM t") + checkSparkAnswer(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + explain.contains("hypot:") && + explain.contains(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false"), + s"expected 'hypot:' prefix and disabled-flag reason, got:\n$explain") + } + } + } + test("dispatcher caches the compiled kernel across batches of one query") { // Within a single query, the dispatcher compiles a kernel for the (expression, schema) pair // once and reuses it across every subsequent batch of the same shape. Force multiple batches