You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/17 17:37:31 UTC
(datafusion-comet) 01/01: refactor: replace remaining hand-rolled loops with Arrow kernels (#5367)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5367-74725611130396969b77e7f44ebe64ef20d4c452
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 3b942e593ebb7715db32ea3d09f6e6bac314cdf7
Author: ChenChen Lai <72...@users.noreply.github.com>
AuthorDate: Thu Sep 17 17:37:08 2026 +0000
refactor: replace remaining hand-rolled loops with Arrow kernels (#5367)
* refactor: replace remaining hand-rolled loops with Arrow kernels
* add all-null fast path and regression test
* fix spark_pow dense-null
* perf: dispatch binary spark_pow on the combined null mask
Base the dense-null dispatch for array/array pow on the union of both
input null masks and reuse that union for the result. With independent
null patterns each operand can be under the threshold while most output
rows are null, so the per-operand max picked the full evaluation path.
Add a pow(a + 2.5D, b + 2.5D) benchmark with independent null masks.
* perf: skip null slots in spark_pow at any null density
Replace the 75% null-density threshold with skipping null slots whenever
the output has nulls. A spark_powf call costs far more than walking the
null bitmap, so the null-skipping loop matched or beat evaluating every
slot from 5% nulls upward, both for zero-filled null slots and for slots
carrying a payload from an upstream Arrow op. The threshold left inputs
between roughly 50% and 75% nulls slower than the row-wise loop on main.
Also short-circuit a null Float64 scalar before datafusion's apply(),
which materializes each scalar as a one-element array.
---------
Co-authored-by: Chao Sun <su...@apache.org>
---
native/spark-expr/Cargo.toml | 8 +
.../spark-expr/benches/cast_decimal_to_boolean.rs | 82 ++++
native/spark-expr/benches/spark_pow.rs | 276 +++++++++++++
native/spark-expr/src/array_funcs/array_insert.rs | 75 +++-
native/spark-expr/src/conversion_funcs/numeric.rs | 131 ++++++-
native/spark-expr/src/math_funcs/pow.rs | 431 +++++++++++++++++----
.../org/apache/comet/CometExpressionSuite.scala | 34 ++
7 files changed, 941 insertions(+), 96 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 6ce2835c60..380656f107 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -403,6 +403,14 @@ harness = false
name = "iceberg_transforms"
harness = false
+[[bench]]
+name = "spark_pow"
+harness = false
+
+[[bench]]
+name = "cast_decimal_to_boolean"
+harness = false
+
[[bench]]
name = "dayofweek_weekday"
harness = false
diff --git a/native/spark-expr/benches/cast_decimal_to_boolean.rs b/native/spark-expr/benches/cast_decimal_to_boolean.rs
new file mode 100644
index 0000000000..0d1bfe255c
--- /dev/null
+++ b/native/spark-expr/benches/cast_decimal_to_boolean.rs
@@ -0,0 +1,82 @@
+// 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, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions};
+use std::hint::black_box;
+use std::sync::Arc;
+
+const PRECISION: u8 = 20;
+const SCALE: i8 = 2;
+
+/// Build a Decimal128(20, 2) column of `rows` rows. Every `null_every`-th row is null
+/// (`null_every == 0` means no nulls). Values alternate between 0 and non-zero so the
+/// boolean result is a realistic mix.
+fn create_batch(rows: usize, null_every: usize) -> RecordBatch {
+ let arr: Decimal128Array = (0..rows)
+ .map(|i| {
+ if null_every != 0 && i % null_every == 0 {
+ None
+ } else if i % 3 == 0 {
+ Some(0i128)
+ } else {
+ Some(((i % 100) as i128) * 100)
+ }
+ })
+ .collect::<Decimal128Array>()
+ .with_precision_and_scale(PRECISION, SCALE)
+ .unwrap();
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ DataType::Decimal128(PRECISION, SCALE),
+ true,
+ )]));
+ RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let rows = 8192;
+ let expr = Arc::new(Column::new("a", 0));
+ let cast_to_bool = Cast::new(
+ expr,
+ DataType::Boolean,
+ SparkCastOptions::new(EvalMode::Legacy, "UTC", false),
+ None,
+ None,
+ );
+
+ let no_nulls = create_batch(rows, 0);
+ let sparse_nulls = create_batch(rows, 10);
+ let dense_nulls = create_batch(rows, 2);
+
+ let mut bench = |name: &str, batch: &RecordBatch| {
+ c.bench_function(name, |b| {
+ b.iter(|| black_box(cast_to_bool.evaluate(black_box(batch)).unwrap()))
+ });
+ };
+ bench("cast_decimal_to_boolean: no nulls", &no_nulls);
+ bench("cast_decimal_to_boolean: sparse nulls", &sparse_nulls);
+ bench("cast_decimal_to_boolean: dense nulls", &dense_nulls);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/benches/spark_pow.rs b/native/spark-expr/benches/spark_pow.rs
new file mode 100644
index 0000000000..6a63bb306c
--- /dev/null
+++ b/native/spark-expr/benches/spark_pow.rs
@@ -0,0 +1,276 @@
+// 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, Datum, Float64Array, Scalar};
+use arrow::buffer::NullBuffer;
+use arrow::compute::kernels::numeric::add;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::spark_pow;
+use std::hint::black_box;
+use std::sync::Arc;
+
+/// Build a Float64 column of `rows` rows, with every `null_every`-th row null
+/// (`null_every == 0` means no nulls). Values stay in [0.5, 5.0] so `powf` is finite.
+fn create_f64_array(rows: usize, null_every: usize) -> ArrayRef {
+ let arr: Float64Array = (0..rows)
+ .map(|i| {
+ if null_every != 0 && i % null_every == 0 {
+ None
+ } else {
+ Some(0.5 + ((i % 10) as f64) * 0.5)
+ }
+ })
+ .collect();
+ Arc::new(arr)
+}
+
+/// Build a Float64 column of `rows` rows with approximately `null_pct`% nulls striped
+/// across the batch (`i % 100 < null_pct` marks a null). Payload in null slots is
+/// the default 0. Meant as input for a downstream Arrow op (like `add`) whose
+/// output is what feeds `spark_pow`.
+fn create_f64_array_with_null_pct(rows: usize, null_pct: usize) -> ArrayRef {
+ let arr: Float64Array = (0..rows)
+ .map(|i| {
+ if null_pct != 0 && i % 100 < null_pct {
+ None
+ } else {
+ Some(0.5 + ((i % 10) as f64) * 0.5)
+ }
+ })
+ .collect();
+ Arc::new(arr)
+}
+
+/// Build a Float64 column whose null slots still carry a real (non-zero) payload,
+/// simulating the output of `a + 2.5D` where the addition preserves null bits but
+/// writes a real value into every slot. `null_pct` is the target fraction of null rows
+/// in `0..=100`. Nulls are striped across the batch modulo 100
+/// (`i % 100 < null_pct` marks a null).
+fn create_f64_array_with_payload_in_nulls(rows: usize, null_pct: usize) -> ArrayRef {
+ let values: Vec<f64> = (0..rows).map(|i| 2.5 + ((i % 10) as f64) * 0.1).collect();
+ let nulls = if null_pct == 0 {
+ None
+ } else {
+ Some(NullBuffer::from(
+ (0..rows)
+ .map(|i| i % 100 >= null_pct)
+ .collect::<Vec<bool>>(),
+ ))
+ };
+ Arc::new(Float64Array::new(values.into(), nulls))
+}
+
+/// Build a Float64 column with approximately `null_pct`% nulls placed by a seeded hash, so
+/// two columns built with different seeds have independent null masks. Payload in null
+/// slots is the default 0, matching a nullable column read before any arithmetic.
+fn create_f64_array_with_hashed_nulls(rows: usize, null_pct: u64, seed: u64) -> ArrayRef {
+ let arr: Float64Array = (0..rows as u64)
+ .map(|i| {
+ // splitmix64 finalizer: cheap, deterministic, and well mixed across bits.
+ let mut h = i.wrapping_add(seed).wrapping_mul(0x9E37_79B9_7F4A_7C15);
+ h = (h ^ (h >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+ h = (h ^ (h >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+ h ^= h >> 31;
+ if h % 100 < null_pct {
+ None
+ } else {
+ Some(0.5 + ((i % 10) as f64) * 0.5)
+ }
+ })
+ .collect();
+ Arc::new(arr)
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let rows = 8192;
+ let no_nulls_a = create_f64_array(rows, 0);
+ let no_nulls_b = create_f64_array(rows, 0);
+ let sparse_a = create_f64_array(rows, 10);
+ let sparse_b = create_f64_array(rows, 10);
+ let dense_a = create_f64_array(rows, 2);
+ let dense_b = create_f64_array(rows, 2);
+
+ // Array/array: exercises `binary` over spark_powf.
+ let mut bench_arr_arr = |name: &str, a: &ArrayRef, b: &ArrayRef| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(a)),
+ ColumnarValue::Array(Arc::clone(b)),
+ ];
+ c.bench_function(name, move |bencher| {
+ bencher.iter(|| black_box(spark_pow(black_box(&args)).unwrap()))
+ });
+ };
+ bench_arr_arr("spark_pow: array/array no nulls", &no_nulls_a, &no_nulls_b);
+ bench_arr_arr("spark_pow: array/array sparse nulls", &sparse_a, &sparse_b);
+ bench_arr_arr("spark_pow: array/array dense nulls", &dense_a, &dense_b);
+
+ // Scalar/array: exercises `unary` with the base captured.
+ let mut bench_scalar_arr = |name: &str, exp: &ArrayRef| {
+ let args = vec![
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.5))),
+ ColumnarValue::Array(Arc::clone(exp)),
+ ];
+ c.bench_function(name, move |bencher| {
+ bencher.iter(|| black_box(spark_pow(black_box(&args)).unwrap()))
+ });
+ };
+ bench_scalar_arr("spark_pow: scalar/array no nulls", &no_nulls_b);
+ bench_scalar_arr("spark_pow: scalar/array sparse nulls", &sparse_b);
+ bench_scalar_arr("spark_pow: scalar/array dense nulls", &dense_b);
+
+ // Array/scalar: exercises `unary` with the exponent captured.
+ let mut bench_arr_scalar = |name: &str, base: &ArrayRef| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(base)),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0))),
+ ];
+ c.bench_function(name, move |bencher| {
+ bencher.iter(|| black_box(spark_pow(black_box(&args)).unwrap()))
+ });
+ };
+ bench_arr_scalar("spark_pow: array/scalar no nulls", &no_nulls_a);
+ bench_arr_scalar("spark_pow: array/scalar sparse nulls", &sparse_a);
+ bench_arr_scalar("spark_pow: array/scalar dense nulls", &dense_a);
+
+ // Null-scalar short-circuit: whole output is null, no work per row.
+ let null_scalar_args = vec![
+ ColumnarValue::Scalar(ScalarValue::Float64(None)),
+ ColumnarValue::Array(Arc::clone(&no_nulls_b)),
+ ];
+ c.bench_function("spark_pow: null scalar short-circuit", |b| {
+ b.iter(|| black_box(spark_pow(black_box(&null_scalar_args)).unwrap()))
+ });
+
+ // Composed-nullable: models `pow(a + 2.5D, b)` where null slots carry a real
+ // payload (Arrow arithmetic preserves the null bit but overwrites the value).
+ // Sweeps null density from sparse to nearly all-null.
+ for null_pct in [10usize, 30, 50, 70, 90, 99] {
+ let a = create_f64_array_with_payload_in_nulls(rows, null_pct);
+ let b = create_f64_array_with_payload_in_nulls(rows, null_pct);
+
+ let arr_arr_args = vec![
+ ColumnarValue::Array(Arc::clone(&a)),
+ ColumnarValue::Array(Arc::clone(&b)),
+ ];
+ c.bench_function(
+ &format!("spark_pow: array/array composed nulls {null_pct}%"),
+ move |bencher| bencher.iter(|| black_box(spark_pow(black_box(&arr_arr_args)).unwrap())),
+ );
+
+ let a = create_f64_array_with_payload_in_nulls(rows, null_pct);
+ let scalar_arr_args = vec![
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.5))),
+ ColumnarValue::Array(Arc::clone(&a)),
+ ];
+ c.bench_function(
+ &format!("spark_pow: scalar/array composed nulls {null_pct}%"),
+ move |bencher| {
+ bencher.iter(|| black_box(spark_pow(black_box(&scalar_arr_args)).unwrap()))
+ },
+ );
+
+ let a = create_f64_array_with_payload_in_nulls(rows, null_pct);
+ let arr_scalar_args = vec![
+ ColumnarValue::Array(Arc::clone(&a)),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0))),
+ ];
+ c.bench_function(
+ &format!("spark_pow: array/scalar composed nulls {null_pct}%"),
+ move |bencher| {
+ bencher.iter(|| black_box(spark_pow(black_box(&arr_scalar_args)).unwrap()))
+ },
+ );
+ }
+
+ // End-to-end pipeline. The `add` step preserves null bits but overwrites the
+ // underlying payload, which is the exact shape the reviewer flagged for a null-skipping
+ // kernel. Timing includes both the Arrow `add` and the `spark_pow` call so it reflects
+ // the real query cost, not a pre-materialised intermediate. Two shapes:
+ // 1. `pow(a + 2.5D, 3)` — array/scalar dispatch (`pow_array_scalar`)
+ // 2. `pow(a + 2.5D, b)` — array/array dispatch (`pow_binary`), with
+ // nullable `a` and a non-null array of finite fractional exponents.
+ let exp_arg = ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0)));
+ let two_point_five: Arc<dyn Datum> = Arc::new(Scalar::new(Float64Array::from(vec![2.5])));
+ let fractional_exponents: ArrayRef = Arc::new(Float64Array::from(
+ (0..rows)
+ .map(|i| 1.25 + (i % 10) as f64 * 0.1)
+ .collect::<Vec<_>>(),
+ ));
+ for null_pct in [10usize, 30, 50, 70, 90, 99] {
+ let base: ArrayRef = create_f64_array_with_null_pct(rows, null_pct);
+ let scalar = Arc::clone(&two_point_five);
+ let exp = exp_arg.clone();
+ c.bench_function(
+ &format!("spark_pow: pipeline pow(a + 2.5D, 3) nulls {null_pct}%"),
+ move |bencher| {
+ bencher.iter(|| {
+ let composed =
+ add(black_box(&base.as_ref()), black_box(scalar.as_ref())).unwrap();
+ let args = [ColumnarValue::Array(composed), exp.clone()];
+ black_box(spark_pow(black_box(&args)).unwrap())
+ })
+ },
+ );
+
+ let base: ArrayRef = create_f64_array_with_null_pct(rows, null_pct);
+ let scalar = Arc::clone(&two_point_five);
+ let exp = Arc::clone(&fractional_exponents);
+ c.bench_function(
+ &format!("spark_pow: pipeline pow(a + 2.5D, b) nulls {null_pct}%"),
+ move |bencher| {
+ bencher.iter(|| {
+ let composed_base =
+ add(black_box(&base.as_ref()), black_box(scalar.as_ref())).unwrap();
+ let args = [
+ ColumnarValue::Array(composed_base),
+ ColumnarValue::Array(Arc::clone(&exp)),
+ ];
+ black_box(spark_pow(black_box(&args)).unwrap())
+ })
+ },
+ );
+ }
+ // Independent null masks on both operands: `pow(a + 2.5D, b + 2.5D)`. The output null
+ // density is roughly `1 - (1 - p)^2` (about 91% at 70% per operand), so dispatch must
+ // use the combined mask rather than either operand alone. Both additions are timed.
+ for null_pct in [10u64, 30, 50, 70] {
+ let a: ArrayRef = create_f64_array_with_hashed_nulls(rows, null_pct, 1);
+ let b: ArrayRef = create_f64_array_with_hashed_nulls(rows, null_pct, 2);
+ let scalar = Arc::clone(&two_point_five);
+ c.bench_function(
+ &format!("spark_pow: pipeline pow(a + 2.5D, b + 2.5D) independent nulls {null_pct}%"),
+ move |bencher| {
+ bencher.iter(|| {
+ let composed_base =
+ add(black_box(&a.as_ref()), black_box(scalar.as_ref())).unwrap();
+ let composed_exp =
+ add(black_box(&b.as_ref()), black_box(scalar.as_ref())).unwrap();
+ let args = [
+ ColumnarValue::Array(composed_base),
+ ColumnarValue::Array(composed_exp),
+ ];
+ black_box(spark_pow(black_box(&args)).unwrap())
+ })
+ },
+ );
+ }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/array_funcs/array_insert.rs b/native/spark-expr/src/array_funcs/array_insert.rs
index 16954f1243..6d996052a0 100644
--- a/native/spark-expr/src/array_funcs/array_insert.rs
+++ b/native/spark-expr/src/array_funcs/array_insert.rs
@@ -15,9 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{
- make_array, Array, ArrayRef, BooleanArray, GenericListArray, Int32Array, OffsetSizeTrait,
-};
+use arrow::array::{make_array, Array, ArrayRef, GenericListArray, Int32Array, OffsetSizeTrait};
+use arrow::compute::{and, is_not_null};
use arrow::datatypes::{DataType, Schema};
use arrow::{
array::{as_primitive_array, Capacities, MutableArrayData},
@@ -111,7 +110,11 @@ impl PhysicalExpr for ArrayInsert {
// 2. pos only when src is non-null
// 3. item only when src and pos are non-null
- // Check that src array is actually an array and get it's value type
+ // Check that src array is actually an array and get it's value type.
+ // `into_array(batch.num_rows())` broadcasts scalar results to the batch length, so
+ // `src_value.len() == batch.num_rows()`. `is_not_null(&src_value)` below inherits
+ // that length; without the broadcast the mask would be short and the downstream
+ // `and` / `evaluate_selection` would fail on length mismatch.
let src_value = self
.src_array_expr
.evaluate(batch)?
@@ -123,11 +126,7 @@ impl PhysicalExpr for ArrayInsert {
_ => unreachable!(),
};
- let evaluate_pos = BooleanArray::from(
- (0..batch.num_rows())
- .map(|row| src_value.is_valid(row))
- .collect::<Vec<_>>(),
- );
+ let evaluate_pos = is_not_null(&src_value)?;
let pos_value = self
.pos_expr
@@ -143,11 +142,7 @@ impl PhysicalExpr for ArrayInsert {
)));
}
- let evaluate_item = BooleanArray::from(
- (0..batch.num_rows())
- .map(|row| src_value.is_valid(row) && pos_value.is_valid(row))
- .collect::<Vec<_>>(),
- );
+ let evaluate_item = and(&evaluate_pos, &is_not_null(&pos_value)?)?;
// Check that inserted value has the same type as an array
let item_value = self
@@ -505,4 +500,56 @@ mod test {
assert_eq!(&result.to_data(), &expected.to_data());
Ok(())
}
+
+ // Pins the Spark evaluation-order contract: `pos` is evaluated only on rows where
+ // `src` is non-null, and `item` only on rows where both `src` and `pos` are non-null.
+ // Each of `src`, `pos`, and `item` is null in a different row, and the fourth row
+ // has all three non-null, so any regression that evaluates a column on a row where
+ // one of its guards is null produces a different output.
+ #[test]
+ fn test_array_insert_evaluate_cross_null_patterns() -> Result<()> {
+ use arrow::datatypes::{Field, Int32Type, Schema};
+ use datafusion::physical_expr::expressions::col;
+
+ let src = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(1), Some(2), Some(3)]), // src ok, pos ok, item ok
+ Some(vec![Some(4), Some(5)]), // src ok, pos NULL
+ None, // src NULL, pos ok
+ Some(vec![Some(6), Some(7)]), // src ok, pos ok, item NULL
+ ]);
+ let positions = Int32Array::from(vec![Some(2), None, Some(1), Some(1)]);
+ let items = Int32Array::from(vec![Some(99), Some(99), Some(99), None]);
+
+ let list_field = match src.data_type() {
+ DataType::List(f) => Arc::clone(f),
+ _ => unreachable!(),
+ };
+ let schema = Schema::new(vec![
+ Field::new("src", DataType::List(list_field), true),
+ Field::new("pos", DataType::Int32, true),
+ Field::new("item", DataType::Int32, true),
+ ]);
+ let schema_ref = Arc::new(schema);
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema_ref),
+ vec![Arc::new(src), Arc::new(positions), Arc::new(items)],
+ )?;
+
+ let expr = ArrayInsert::new(
+ col("src", &schema_ref)?,
+ col("pos", &schema_ref)?,
+ col("item", &schema_ref)?,
+ false,
+ );
+ let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?;
+
+ let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(1), Some(99), Some(2), Some(3)]),
+ None,
+ None,
+ Some(vec![None, Some(6), Some(7)]),
+ ]);
+ assert_eq!(&result.to_data(), &expected.to_data());
+ Ok(())
+ }
}
diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs
index 868d059552..81972425e8 100644
--- a/native/spark-expr/src/conversion_funcs/numeric.rs
+++ b/native/spark-expr/src/conversion_funcs/numeric.rs
@@ -20,15 +20,17 @@ use crate::conversion_funcs::utils::cast_overflow;
use crate::conversion_funcs::utils::MICROS_PER_SECOND;
use crate::{EvalMode, SparkError, SparkResult};
use arrow::array::{
- Array, ArrayRef, AsArray, BooleanBuilder, Decimal128Array, Float32Array, Float64Array,
+ Array, ArrayRef, AsArray, BooleanArray, Decimal128Array, Float32Array, Float64Array,
GenericStringBuilder, Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait,
- PrimitiveArray, StringBuilder, TimestampMicrosecondBuilder,
+ PrimitiveArray, Scalar, StringBuilder, TimestampMicrosecondBuilder,
};
+use arrow::buffer::BooleanBuffer;
+use arrow::compute::kernels::cmp::neq;
use arrow::datatypes::{
i256, is_validate_decimal_precision, ArrowPrimitiveType, DataType, Decimal128Type, Float32Type,
Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
};
-use num::{cast::AsPrimitive, Float, ToPrimitive, Zero};
+use num::{cast::AsPrimitive, Float, ToPrimitive};
use std::fmt::{self, Write};
use std::sync::Arc;
@@ -83,8 +85,8 @@ pub(crate) fn is_df_cast_from_decimal_spark_compatible(to_type: &DataType) -> bo
| DataType::Utf8
)
// Note: Boolean is intentionally absent. Decimal-to-boolean uses a dedicated
- // spark_cast_decimal_to_boolean function (in cast.rs) that checks the raw i128
- // value, bypassing the DataFusion cast kernel entirely.
+ // spark_cast_decimal_to_boolean function that compares against a zero decimal of
+ // the same precision/scale, bypassing the DataFusion cast kernel entirely.
}
macro_rules! cast_float_to_timestamp_impl {
@@ -952,15 +954,33 @@ pub(crate) fn spark_cast_int_to_int(
pub(crate) fn spark_cast_decimal_to_boolean(array: &dyn Array) -> SparkResult<ArrayRef> {
let decimal_array = array.as_primitive::<Decimal128Type>();
- let mut result = BooleanBuilder::with_capacity(decimal_array.len());
- for i in 0..decimal_array.len() {
- if decimal_array.is_null(i) {
- result.append_null()
- } else {
- result.append_value(!decimal_array.value(i).is_zero());
- }
+ // All-null fast path: skips the zero-scalar construction for a batch whose output
+ // is trivially all-null. Also covers all-null `Decimal128(0, 0)`, which the default
+ // path cannot handle (see the precision-zero fast path below).
+ if decimal_array.null_count() == decimal_array.len() {
+ return Ok(Arc::new(BooleanArray::new_null(decimal_array.len())));
+ }
+ // Precision-zero fast path. Arrow rejects `precision == 0` in `with_precision_and_scale`,
+ // so the default `neq`-against-zero path cannot round-trip a `Decimal128(0, 0)` batch
+ // even when it has valid slots. Spark reaches this shape via JVM-side writers that do
+ // not validate precision (e.g. a UDF returning `BigInteger.ZERO` written through
+ // `DecimalVector.setSafe(long)`). The type contract says only 0 is representable, but
+ // we still cast the raw i128 payload (`v != 0`) rather than hard-coding `false`, so an
+ // out-of-contract non-zero value in a valid slot still round-trips correctly.
+ if decimal_array.precision() == 0 {
+ let values: BooleanBuffer = decimal_array.values().iter().map(|&v| v != 0).collect();
+ return Ok(Arc::new(BooleanArray::new(
+ values,
+ decimal_array.nulls().cloned(),
+ )));
}
- Ok(Arc::new(result.finish()))
+ // Arrow has no Decimal-to-Boolean cast. `neq` against a zero of the same
+ // precision/scale is exactly `!value.is_zero()`, including null handling.
+ let zero = Scalar::new(
+ Decimal128Array::from(vec![0i128])
+ .with_precision_and_scale(decimal_array.precision(), decimal_array.scale())?,
+ );
+ Ok(Arc::new(neq(decimal_array, &zero)?))
}
/// Powers of ten that are exactly representable as `f64` (`10^n = 2^n * 5^n` and `5^22 < 2^53`);
@@ -1821,6 +1841,91 @@ mod tests {
assert!(bool_array.value(1)); // 100 -> true
assert!(bool_array.value(2)); // -100 -> true
assert!(bool_array.is_null(3)); // null -> null
+
+ // A different precision/scale must still compare against a matching zero scalar.
+ let array: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![Some(0), Some(1)])
+ .with_precision_and_scale(38, 0)
+ .unwrap(),
+ );
+ let result = spark_cast_decimal_to_boolean(&array).unwrap();
+ let bool_array = result.as_boolean();
+ assert!(!bool_array.value(0));
+ assert!(bool_array.value(1));
+
+ // All-null Decimal128(0, 0) is reachable via Spark's RDD row-to-Arrow path;
+ // `with_precision_and_scale(0, 0)` rejects precision 0, so the all-null fast path
+ // must yield an all-null boolean array without constructing the zero scalar.
+ // SAFETY: builds a well-formed Decimal128 ArrayData:
+ // - len = 2 slots
+ // - values buffer = 32 bytes = 2 * 16 (Decimal128 slot width), zero-initialized
+ // - null buffer = 1 byte, all bits clear, covering >= len bits as required
+ // Precision 0 skips ArrowError validation but is otherwise a legal DataType tag;
+ // no non-null slot is ever read, so precision-range invariants are vacuous.
+ let array_data = unsafe {
+ arrow::array::ArrayData::builder(DataType::Decimal128(0, 0))
+ .len(2)
+ .null_bit_buffer(Some(arrow::buffer::Buffer::from(&[0u8])))
+ .add_buffer(arrow::buffer::Buffer::from(&[0u8; 32]))
+ .build_unchecked()
+ };
+ let array: ArrayRef = Arc::new(Decimal128Array::from(array_data));
+ assert_eq!(array.data_type(), &DataType::Decimal128(0, 0));
+ let result = spark_cast_decimal_to_boolean(&array).unwrap();
+ let bool_array = result.as_boolean();
+ assert_eq!(bool_array.len(), 2);
+ assert!(bool_array.is_null(0));
+ assert!(bool_array.is_null(1));
+
+ // Empty Decimal128(0, 0) input: `null_count() == len()` (0 == 0) must still take
+ // the fast path. Precision 0 makes this discriminating — without the fast path,
+ // building the zero scalar would fail regardless of the empty length.
+ // SAFETY: len = 0, so both the (empty) values buffer and the absent null buffer
+ // trivially cover every slot; the precision-range invariant is vacuous.
+ let array_data = unsafe {
+ arrow::array::ArrayData::builder(DataType::Decimal128(0, 0))
+ .len(0)
+ .add_buffer(arrow::buffer::Buffer::from(&[] as &[u8]))
+ .build_unchecked()
+ };
+ let array: ArrayRef = Arc::new(Decimal128Array::from(array_data));
+ // The load-bearing check is that `spark_cast_decimal_to_boolean` returns Ok — without
+ // the fast path, precision 0 would fail during zero-scalar construction. Length is
+ // asserted for completeness; null_count is trivially 0 on any empty array.
+ let result = spark_cast_decimal_to_boolean(&array).unwrap();
+ assert_eq!(result.as_boolean().len(), 0);
+
+ // Mixed valid + null Decimal128(0, 0). Reachable via JVM-side writers that do not
+ // validate precision (e.g. a Java UDF returning `BigInteger.ZERO` declared as
+ // `DecimalType(0, 0)`, written through `DecimalVector.setSafe(long)`). The all-null
+ // fast path does not cover this: `null_count() < len()` yet the batch still cannot
+ // round-trip through a zero scalar. We build a 3-row batch where slot 0 is a valid
+ // zero, slot 1 is null, and slot 2 is an out-of-contract non-zero i128; the path
+ // must read the raw value (not hard-code `false`) so slot 2 comes back as `true`.
+ // Expected output: `[false, null, true]`.
+ // SAFETY: len = 3; values buffer = 48 bytes = 3 * 16, with i128(0), i128(0),
+ // i128(7) laid out little-endian; null bit buffer 0b0000_0101 = slots 0 and 2
+ // valid, slot 1 null.
+ let mut vals = [0u8; 48];
+ vals[32..48].copy_from_slice(&7i128.to_le_bytes()); // slot 2 = 7
+ let array_data = unsafe {
+ arrow::array::ArrayData::builder(DataType::Decimal128(0, 0))
+ .len(3)
+ .null_bit_buffer(Some(arrow::buffer::Buffer::from(&[0b0000_0101u8])))
+ .add_buffer(arrow::buffer::Buffer::from(&vals))
+ .build_unchecked()
+ };
+ let array: ArrayRef = Arc::new(Decimal128Array::from(array_data));
+ assert_eq!(array.data_type(), &DataType::Decimal128(0, 0));
+ assert_eq!(array.null_count(), 1);
+ let result = spark_cast_decimal_to_boolean(&array).unwrap();
+ let bool_array = result.as_boolean();
+ assert_eq!(bool_array.len(), 3);
+ assert!(!bool_array.is_null(0));
+ assert!(!bool_array.value(0)); // valid zero -> false
+ assert!(bool_array.is_null(1)); // null -> null
+ assert!(!bool_array.is_null(2));
+ assert!(bool_array.value(2)); // valid non-zero i128 -> true (raw-value cast)
}
#[test]
diff --git a/native/spark-expr/src/math_funcs/pow.rs b/native/spark-expr/src/math_funcs/pow.rs
index 61168f7d18..a7f277c6fa 100644
--- a/native/spark-expr/src/math_funcs/pow.rs
+++ b/native/spark-expr/src/math_funcs/pow.rs
@@ -15,8 +15,12 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::Float64Array;
-use datafusion::common::{DataFusionError, ScalarValue};
+use arrow::array::{Array, ArrayRef, Datum, Float64Array};
+use arrow::buffer::{NullBuffer, ScalarBuffer};
+use arrow::compute::kernels::arity::unary;
+use arrow::error::ArrowError;
+use datafusion::common::{utils::take_function_args, DataFusionError, ScalarValue};
+use datafusion::physical_expr_common::datum::apply;
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
@@ -42,86 +46,127 @@ fn spark_powf(base: f64, exp: f64) -> f64 {
/// Unlike DataFusion's `power`, `pow(0, -1)` returns `Infinity` rather than erroring. Only null
/// inputs produce null; otherwise every result is the `spark_powf` value.
pub fn spark_pow(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
- if args.len() != 2 {
- return Err(DataFusionError::Internal(format!(
- "spark_pow requires 2 arguments, got {}",
- args.len()
- )));
+ let [base, exp] = take_function_args("spark_pow", args)?;
+ // A null scalar on either side makes the whole result null. Handle it before `apply`,
+ // which would otherwise materialize each scalar as a one-element array first.
+ if is_null_f64_scalar(base) || is_null_f64_scalar(exp) {
+ return match (base, exp) {
+ (ColumnarValue::Array(array), _) | (_, ColumnarValue::Array(array)) => {
+ let len = as_f64_array(array.as_ref())?.len();
+ Ok(ColumnarValue::Array(Arc::new(Float64Array::new_null(len))))
+ }
+ _ => Ok(ColumnarValue::Scalar(ScalarValue::Float64(None))),
+ };
}
+ apply(base, exp, spark_pow_kernel)
+}
+
+fn is_null_f64_scalar(value: &ColumnarValue) -> bool {
+ matches!(value, ColumnarValue::Scalar(ScalarValue::Float64(None)))
+}
+
+fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> {
+ array
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .ok_or_else(|| {
+ ArrowError::ComputeError(format!(
+ "spark_pow expected Float64, got {:?}",
+ array.data_type()
+ ))
+ })
+}
- fn as_f64_array(
- value: &Arc<dyn arrow::array::Array>,
- ) -> Result<&Float64Array, DataFusionError> {
- value
- .as_any()
- .downcast_ref::<Float64Array>()
- .ok_or_else(|| {
- DataFusionError::Internal(format!(
- "spark_pow expected Float64, got {:?}",
- value.data_type()
- ))
- })
+/// Whenever the output has nulls, `spark_powf` runs only on the valid indices. Masked
+/// slots can carry a real payload from an upstream Arrow op (e.g. `a + 2.5D` preserves
+/// the null bit but overwrites the value), and a `spark_powf` call costs far more than
+/// walking the null bitmap, so skipping pays off at any null density. Inputs without
+/// nulls use a plain loop over the raw values. Scalars are never broadcast.
+fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<ArrayRef, ArrowError> {
+ let (left, left_is_scalar) = lhs.get();
+ let (right, right_is_scalar) = rhs.get();
+ let left = as_f64_array(left)?;
+ let right = as_f64_array(right)?;
+
+ let result = match (left_is_scalar, right_is_scalar) {
+ // Null scalars are handled in `spark_pow` before `apply`.
+ (true, false) => pow_scalar_array(left.value(0), right),
+ (false, true) => pow_array_scalar(left, right.value(0)),
+ _ => pow_binary(left, right)?,
+ };
+ Ok(Arc::new(result))
+}
+
+/// Returns the null buffer when it masks at least one slot, so callers can take the
+/// null-skipping loop only when there is something to skip.
+#[inline]
+fn nulls_to_skip(nulls: Option<&NullBuffer>) -> Option<&NullBuffer> {
+ nulls.filter(|n| n.null_count() > 0)
+}
+
+fn pow_scalar_array(base: f64, exp: &Float64Array) -> Float64Array {
+ let Some(nulls) = nulls_to_skip(exp.nulls()) else {
+ return unary(exp, |e| spark_powf(base, e));
+ };
+ let exp_values = exp.values();
+ let mut out = vec![0.0f64; exp.len()];
+ for i in nulls.valid_indices() {
+ out[i] = spark_powf(base, exp_values[i]);
}
+ Float64Array::new(out.into(), Some(nulls.clone()))
+}
- fn as_f64_scalar(scalar: &ScalarValue) -> Result<Option<f64>, DataFusionError> {
- match scalar {
- ScalarValue::Float64(v) => Ok(*v),
- _ => Err(DataFusionError::Internal(format!(
- "spark_pow expected Float64 scalar, got {scalar:?}",
- ))),
- }
+fn pow_array_scalar(base: &Float64Array, exp: f64) -> Float64Array {
+ let Some(nulls) = nulls_to_skip(base.nulls()) else {
+ return unary(base, |b| spark_powf(b, exp));
+ };
+ let base_values = base.values();
+ let mut out = vec![0.0f64; base.len()];
+ for i in nulls.valid_indices() {
+ out[i] = spark_powf(base_values[i], exp);
}
+ Float64Array::new(out.into(), Some(nulls.clone()))
+}
- match (&args[0], &args[1]) {
- (ColumnarValue::Array(base_arr), ColumnarValue::Array(exp_arr)) => {
- let bases = as_f64_array(base_arr)?;
- let exps = as_f64_array(exp_arr)?;
- let result: Float64Array = bases
- .iter()
- .zip(exps.iter())
- .map(|(b, e)| match (b, e) {
- (Some(base), Some(exp)) => Some(spark_powf(base, exp)),
- _ => None,
- })
- .collect();
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- (ColumnarValue::Scalar(base_scalar), ColumnarValue::Array(exp_arr)) => {
- let exps = as_f64_array(exp_arr)?;
- let result: Float64Array = match as_f64_scalar(base_scalar)? {
- Some(base) => exps
- .iter()
- .map(|e| e.map(|exp| spark_powf(base, exp)))
- .collect(),
- None => Float64Array::new_null(exp_arr.len()),
- };
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- (ColumnarValue::Array(base_arr), ColumnarValue::Scalar(exp_scalar)) => {
- let bases = as_f64_array(base_arr)?;
- let result: Float64Array = match as_f64_scalar(exp_scalar)? {
- Some(exp) => bases
- .iter()
- .map(|b| b.map(|base| spark_powf(base, exp)))
- .collect(),
- None => Float64Array::new_null(base_arr.len()),
- };
- Ok(ColumnarValue::Array(Arc::new(result)))
- }
- (ColumnarValue::Scalar(base_scalar), ColumnarValue::Scalar(exp_scalar)) => {
- let result = match (as_f64_scalar(base_scalar)?, as_f64_scalar(exp_scalar)?) {
- (Some(base), Some(exp)) => ScalarValue::Float64(Some(spark_powf(base, exp))),
- _ => ScalarValue::Float64(None),
- };
- Ok(ColumnarValue::Scalar(result))
- }
+/// Array/array power. The output null mask is the union of both input masks; it is
+/// computed once and reused both to pick the null-skipping loop and as the result's null
+/// buffer. Deciding on either input's mask alone would miss rows that are null only on
+/// the other side.
+fn pow_binary(base: &Float64Array, exp: &Float64Array) -> Result<Float64Array, ArrowError> {
+ // Match arrow's `binary` contract: reject mismatched lengths with an error. This must
+ // precede `NullBuffer::union`, which panics on mismatched lengths.
+ if base.len() != exp.len() {
+ return Err(ArrowError::ComputeError(format!(
+ "spark_pow: arrays have different lengths: {} vs {}",
+ base.len(),
+ exp.len()
+ )));
}
+ let nulls = NullBuffer::union(base.nulls(), exp.nulls());
+ let base_values = base.values();
+ let exp_values = exp.values();
+ let values: ScalarBuffer<f64> = match nulls_to_skip(nulls.as_ref()) {
+ Some(n) => {
+ let mut out = vec![0.0f64; base.len()];
+ for i in n.valid_indices() {
+ out[i] = spark_powf(base_values[i], exp_values[i]);
+ }
+ out.into()
+ }
+ None => base_values
+ .iter()
+ .zip(exp_values.iter())
+ .map(|(&b, &e)| spark_powf(b, e))
+ .collect(),
+ };
+ Ok(Float64Array::new(values, nulls))
}
#[cfg(test)]
mod test {
use super::*;
use arrow::array::Array;
+ use datafusion::common::ScalarValue;
#[test]
fn test_spark_pow_basic() {
@@ -311,4 +356,252 @@ mod test {
panic!("expected array result");
}
}
+
+ /// The null-aware binary path must reject length mismatches the same way `binary`
+ /// does — returning an `ArrowError`, not panicking on out-of-bounds indexing.
+ #[test]
+ fn test_spark_pow_null_aware_binary_length_mismatch() {
+ use arrow::buffer::NullBuffer;
+ // Nulls on both sides so the null-skipping path is taken.
+ let rows_a = 100;
+ let rows_b = 90;
+ let make = |rows: usize| -> Float64Array {
+ let vals: Vec<f64> = vec![2.5; rows];
+ let nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::<Vec<bool>>());
+ Float64Array::new(vals.into(), Some(nulls))
+ };
+ let a = make(rows_a);
+ let b = make(rows_b);
+ let err = spark_pow(&[
+ ColumnarValue::Array(Arc::new(a)),
+ ColumnarValue::Array(Arc::new(b)),
+ ])
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("different lengths"),
+ "expected length-mismatch error, got: {err}"
+ );
+ }
+
+ /// A null buffer that masks nothing must take the plain loop and still produce the
+ /// same values as an input without a null buffer, on every dispatch shape.
+ #[test]
+ fn test_spark_pow_all_valid_null_buffer() {
+ use arrow::buffer::NullBuffer;
+ let rows = 16;
+ let with_buffer = || {
+ Float64Array::new(
+ (0..rows)
+ .map(|i| 1.0 + i as f64 * 0.25)
+ .collect::<Vec<_>>()
+ .into(),
+ Some(NullBuffer::new_valid(rows)),
+ )
+ };
+ let expected: Vec<f64> = (0..rows)
+ .map(|i| spark_powf(1.0 + i as f64 * 0.25, 1.0 + i as f64 * 0.25))
+ .collect();
+ let cases = [
+ (
+ ColumnarValue::Array(Arc::new(with_buffer())),
+ ColumnarValue::Array(Arc::new(with_buffer())),
+ expected.clone(),
+ ),
+ (
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+ ColumnarValue::Array(Arc::new(with_buffer())),
+ (0..rows)
+ .map(|i| spark_powf(2.0, 1.0 + i as f64 * 0.25))
+ .collect(),
+ ),
+ (
+ ColumnarValue::Array(Arc::new(with_buffer())),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+ (0..rows)
+ .map(|i| spark_powf(1.0 + i as f64 * 0.25, 2.0))
+ .collect(),
+ ),
+ ];
+ for (lhs, rhs, want) in cases {
+ let ColumnarValue::Array(arr) = spark_pow(&[lhs, rhs]).unwrap() else {
+ panic!("expected array result");
+ };
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ assert_eq!(arr.null_count(), 0);
+ assert_eq!(arr.values().as_ref(), want.as_slice());
+ }
+ }
+
+ /// Independent null masks: each operand is 70% null and the union is 91% null. Null
+ /// slots carry a real payload, and the result must still be null exactly where either
+ /// input is null and `spark_powf` everywhere else.
+ #[test]
+ fn test_spark_pow_binary_independent_null_masks() {
+ use arrow::buffer::NullBuffer;
+ let rows = 100;
+ // Valid when `i % 10 < 3` (base) and when `(i / 10) % 10 < 3` (exp): 30% valid each,
+ // independent, 9% valid in the union.
+ let base_valid: Vec<bool> = (0..rows).map(|i| i % 10 < 3).collect();
+ let exp_valid: Vec<bool> = (0..rows).map(|i| (i / 10) % 10 < 3).collect();
+ let base = Float64Array::new(
+ vec![2.5; rows].into(),
+ Some(NullBuffer::from(base_valid.clone())),
+ );
+ let exp = Float64Array::new(
+ vec![1.5; rows].into(),
+ Some(NullBuffer::from(exp_valid.clone())),
+ );
+
+ let result = spark_pow(&[
+ ColumnarValue::Array(Arc::new(base)),
+ ColumnarValue::Array(Arc::new(exp)),
+ ])
+ .unwrap();
+ let ColumnarValue::Array(arr) = result else {
+ panic!("expected array result");
+ };
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ assert_eq!(arr.len(), rows);
+ assert_eq!(arr.null_count(), 91);
+ for i in 0..rows {
+ if base_valid[i] && exp_valid[i] {
+ assert!(!arr.is_null(i), "row {i} should be valid");
+ assert_eq!(arr.value(i), 2.5f64.powf(1.5));
+ } else {
+ assert!(arr.is_null(i), "row {i} should be null");
+ }
+ }
+ }
+
+ /// Simulates the output of an upstream Arrow op (e.g. `a + 2.5D`) that preserves the
+ /// null bit but writes a real payload into the underlying value. The null-aware path
+ /// must ignore that payload and still return null for masked slots.
+ #[test]
+ fn test_spark_pow_payload_in_null_slots() {
+ use arrow::buffer::NullBuffer;
+ let rows = 100;
+ // Every slot has value 2.5; only every 10th slot is valid (90% null).
+ let base_values: Vec<f64> = vec![2.5; rows];
+ let base_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::<Vec<bool>>());
+ let base = Float64Array::new(base_values.into(), Some(base_nulls));
+
+ let result = spark_pow(&[
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0))),
+ ColumnarValue::Array(Arc::new(base)),
+ ])
+ .unwrap();
+ let ColumnarValue::Array(arr) = result else {
+ panic!("expected array result");
+ };
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ assert_eq!(arr.len(), rows);
+ for i in 0..rows {
+ if i % 10 == 0 {
+ assert!(!arr.is_null(i), "row {i} should be valid");
+ assert!((arr.value(i) - 3.0f64.powf(2.5)).abs() < 1e-10);
+ } else {
+ assert!(arr.is_null(i), "row {i} should be null");
+ }
+ }
+
+ // Same setup, but with a scalar exponent so the array-scalar null-aware path runs.
+ let base_values: Vec<f64> = vec![2.5; rows];
+ let base_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::<Vec<bool>>());
+ let base = Float64Array::new(base_values.into(), Some(base_nulls));
+ let result = spark_pow(&[
+ ColumnarValue::Array(Arc::new(base)),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0))),
+ ])
+ .unwrap();
+ let ColumnarValue::Array(arr) = result else {
+ panic!("expected array result");
+ };
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ for i in 0..rows {
+ if i % 10 == 0 {
+ assert!(!arr.is_null(i));
+ assert!((arr.value(i) - 2.5f64.powf(3.0)).abs() < 1e-10);
+ } else {
+ assert!(arr.is_null(i));
+ }
+ }
+
+ // And the binary path: both sides carry payload in their null slots.
+ let a_values: Vec<f64> = vec![2.5; rows];
+ let a_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::<Vec<bool>>());
+ let a = Float64Array::new(a_values.into(), Some(a_nulls));
+ let b_values: Vec<f64> = vec![3.0; rows];
+ let b_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::<Vec<bool>>());
+ let b = Float64Array::new(b_values.into(), Some(b_nulls));
+ let result = spark_pow(&[
+ ColumnarValue::Array(Arc::new(a)),
+ ColumnarValue::Array(Arc::new(b)),
+ ])
+ .unwrap();
+ let ColumnarValue::Array(arr) = result else {
+ panic!("expected array result");
+ };
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ for i in 0..rows {
+ if i % 10 == 0 {
+ assert!(!arr.is_null(i));
+ assert!((arr.value(i) - 2.5f64.powf(3.0)).abs() < 1e-10);
+ } else {
+ assert!(arr.is_null(i));
+ }
+ }
+ }
+
+ #[test]
+ fn test_spark_pow_null_scalar() {
+ let exps = Float64Array::from(vec![Some(3.0), Some(2.0)]);
+ let result = spark_pow(&[
+ ColumnarValue::Scalar(ScalarValue::Float64(None)),
+ ColumnarValue::Array(Arc::new(exps)),
+ ])
+ .unwrap();
+ if let ColumnarValue::Array(arr) = result {
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ assert!(arr.is_null(0));
+ assert!(arr.is_null(1));
+ } else {
+ panic!("expected array result");
+ }
+
+ let bases = Float64Array::from(vec![Some(2.0), Some(3.0)]);
+ let result = spark_pow(&[
+ ColumnarValue::Array(Arc::new(bases)),
+ ColumnarValue::Scalar(ScalarValue::Float64(None)),
+ ])
+ .unwrap();
+ if let ColumnarValue::Array(arr) = result {
+ let arr = arr.as_any().downcast_ref::<Float64Array>().unwrap();
+ assert!(arr.is_null(0));
+ assert!(arr.is_null(1));
+ } else {
+ panic!("expected array result");
+ }
+
+ let scalar_result = spark_pow(&[
+ ColumnarValue::Scalar(ScalarValue::Float64(None)),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+ ])
+ .unwrap();
+ assert!(matches!(
+ scalar_result,
+ ColumnarValue::Scalar(ScalarValue::Float64(None))
+ ));
+
+ // The null-scalar short-circuit must still reject a non-Float64 array.
+ let ints = arrow::array::Int32Array::from(vec![Some(1), Some(2)]);
+ let err = spark_pow(&[
+ ColumnarValue::Scalar(ScalarValue::Float64(None)),
+ ColumnarValue::Array(Arc::new(ints)),
+ ])
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("expected Float64"),
+ "unexpected error: {err}"
+ );
+ }
}
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index 7ec40e4dc9..edb610f9d9 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -3514,6 +3514,40 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ test("regression: cast Decimal(0, 0) to Boolean via Java UDF") {
+ // Reachable via a Java UDF that declares return type `DecimalType(0, 0)` and returns
+ // either `BigInteger.ZERO` or `null`. Spark accepts the schema; Arrow rejects
+ // `precision == 0` inside `Decimal128Array::with_precision_and_scale`, so the native
+ // cast kernel used to error. Enabling the ScalaUDF codegen dispatcher routes the UDF
+ // output straight into the native cast, so this test exercises the actual native path
+ // rather than falling back to Spark. The mixed valid + null batch covers the
+ // precision-zero fast path that reads the raw i128 payload; the all-null batch covers
+ // the earlier all-null shortcut.
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") {
+ spark.udf.register(
+ "zero_decimal",
+ new org.apache.spark.sql.api.java.UDF1[java.lang.Long, java.math.BigInteger] {
+ override def call(id: java.lang.Long): java.math.BigInteger =
+ if (id == 0L) java.math.BigInteger.ZERO else null
+ },
+ DecimalType(0, 0))
+ // id = 0 -> BigInteger.ZERO -> Decimal(0,0) 0 -> false
+ // id = 1 -> null -> null
+ val mixed = spark.range(0, 2).selectExpr("CAST(zero_decimal(id) AS BOOLEAN) AS b")
+ checkSparkAnswerAndOperator(mixed)
+ checkAnswer(mixed, Seq(Row(false), Row(null)))
+
+ // All-null batch: exercises the earlier `null_count() == len()` fast path that
+ // bypasses the zero-scalar construction entirely.
+ val allNull = spark.range(1, 3).selectExpr("CAST(zero_decimal(id) AS BOOLEAN) AS b")
+ checkSparkAnswerAndOperator(allNull)
+ checkAnswer(allNull, Seq(Row(null), Row(null)))
+ }
+ }
+
test("NativeOptIn message and Compatible field") {
import org.apache.comet.serde.{Compatible, NativeOptIn}
val key = "spark.comet.expression.RLike.allowIncompatible"
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org