perf: optimize spark_ceil (3x faster)

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: <name1>, <name2>, ...]` 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: <name1>, <name2>, ...]` 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

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

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<ColumnarValue> {
+    let values: Vec<i128> = (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);

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<Option<i128>> = (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<Option<i128>> = (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::<Vec<_>>(),
+    ))
+}
+
+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);

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::<PrimitiveArray<$from_arrow_primitive_type>>()
             .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::<Option<$to_native_type>, SparkError>(Some(value as $to_native_type))
-                    }
-                    _ => Ok(None),
-                })
-                .collect::<Result<PrimitiveArray<$to_arrow_primitive_type>, _>>(),
-            _ => 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::<Option<$to_native_type>, 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::<Result<PrimitiveArray<$to_arrow_primitive_type>, _>>(),
-        }?;
-        let result: SparkResult<ArrayRef> = 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<ArrayRef> {
     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!(
             "{}",