From 353eab5f78eb4f31a1a3250c3865bcf16880f1e6 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Sat, 15 Aug 2026 23:43:55 +0800 Subject: [PATCH 1/5] refactor: replace remaining hand-rolled loops with Arrow kernels --- native/spark-expr/Cargo.toml | 8 + .../benches/cast_decimal_to_boolean.rs | 82 ++++++++++ native/spark-expr/benches/spark_pow.rs | 103 +++++++++++++ .../src/array_funcs/array_insert.rs | 69 +++++++-- .../src/conversion_funcs/numeric.rs | 40 +++-- native/spark-expr/src/math_funcs/pow.rs | 141 +++++++++--------- 6 files changed, 344 insertions(+), 99 deletions(-) create mode 100644 native/spark-expr/benches/cast_decimal_to_boolean.rs create mode 100644 native/spark-expr/benches/spark_pow.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 98e5a999044..77037904931 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -390,3 +390,11 @@ harness = false [[bench]] name = "iceberg_transforms" harness = false + +[[bench]] +name = "spark_pow" +harness = false + +[[bench]] +name = "cast_decimal_to_boolean" +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 00000000000..0d1bfe255c3 --- /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::() + .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 00000000000..073b5bd0b06 --- /dev/null +++ b/native/spark-expr/benches/spark_pow.rs @@ -0,0 +1,103 @@ +// 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, Float64Array}; +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) +} + +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())) + }); +} + +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 16954f12437..4693f2ebf26 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}, @@ -123,11 +122,7 @@ impl PhysicalExpr for ArrayInsert { _ => unreachable!(), }; - let evaluate_pos = BooleanArray::from( - (0..batch.num_rows()) - .map(|row| src_value.is_valid(row)) - .collect::>(), - ); + let evaluate_pos = is_not_null(&src_value)?; let pos_value = self .pos_expr @@ -143,11 +138,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::>(), - ); + 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 +496,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::(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::(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 61db16d973c..c61056c909a 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -20,15 +20,16 @@ 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, - GenericStringBuilder, Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, - PrimitiveArray, StringBuilder, TimestampMicrosecondBuilder, + Array, ArrayRef, AsArray, Decimal128Array, Float32Array, Float64Array, GenericStringBuilder, + Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, PrimitiveArray, Scalar, + StringBuilder, TimestampMicrosecondBuilder, }; +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, ToPrimitive, Zero}; +use num::{cast::AsPrimitive, ToPrimitive}; use std::sync::Arc; /// Check if DataFusion cast from integer types is Spark compatible @@ -82,8 +83,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 { @@ -872,15 +873,13 @@ pub(crate) fn spark_cast_int_to_int( pub(crate) fn spark_cast_decimal_to_boolean(array: &dyn Array) -> SparkResult { let decimal_array = array.as_primitive::(); - 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()); - } - } - 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`); @@ -1741,6 +1740,17 @@ 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)); } #[test] diff --git a/native/spark-expr/src/math_funcs/pow.rs b/native/spark-expr/src/math_funcs/pow.rs index 61168f7d181..2837f74d02f 100644 --- a/native/spark-expr/src/math_funcs/pow.rs +++ b/native/spark-expr/src/math_funcs/pow.rs @@ -15,8 +15,11 @@ // 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::compute::kernels::arity::{binary, unary}; +use arrow::error::ArrowError; +use datafusion::common::{utils::take_function_args, DataFusionError}; +use datafusion::physical_expr_common::datum::apply; use datafusion::physical_plan::ColumnarValue; use std::sync::Arc; @@ -42,86 +45,55 @@ 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 { - 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)?; + apply(base, exp, spark_pow_kernel) +} - fn as_f64_array( - value: &Arc, - ) -> Result<&Float64Array, DataFusionError> { - value - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "spark_pow expected Float64, got {:?}", - value.data_type() - )) - }) - } +fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> { + array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::ComputeError(format!( + "spark_pow expected Float64, got {:?}", + array.data_type() + )) + }) +} - fn as_f64_scalar(scalar: &ScalarValue) -> Result, DataFusionError> { - match scalar { - ScalarValue::Float64(v) => Ok(*v), - _ => Err(DataFusionError::Internal(format!( - "spark_pow expected Float64 scalar, got {scalar:?}", - ))), - } - } +/// Array/array uses [`binary`] over [`spark_powf`]. Scalar/array uses [`unary`] so the +/// scalar is not broadcast. A null scalar short-circuits to an all-null array. +fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { + 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)?; - 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))) + let result = match (left_is_scalar, right_is_scalar) { + (true, false) => { + if left.is_null(0) { + Float64Array::new_null(right.len()) + } else { + unary(right, |exp| spark_powf(left.value(0), exp)) + } } - (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)) + (false, true) => { + if right.is_null(0) { + Float64Array::new_null(left.len()) + } else { + unary(left, |base| spark_powf(base, right.value(0))) + } } - } + _ => binary(left, right, spark_powf)?, + }; + Ok(Arc::new(result)) } #[cfg(test)] mod test { use super::*; use arrow::array::Array; + use datafusion::common::ScalarValue; #[test] fn test_spark_pow_basic() { @@ -311,4 +283,31 @@ mod test { panic!("expected array result"); } } + + #[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::().unwrap(); + assert!(arr.is_null(0)); + assert!(arr.is_null(1)); + } else { + panic!("expected array result"); + } + + let both_null = spark_pow(&[ + ColumnarValue::Scalar(ScalarValue::Float64(None)), + ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))), + ]) + .unwrap(); + assert!(matches!( + both_null, + ColumnarValue::Scalar(ScalarValue::Float64(None)) + )); + } } From 82759dac2f70aaecda5ee742852821f903c5d24d Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Sat, 22 Aug 2026 15:42:17 +0800 Subject: [PATCH 2/5] add all-null fast path and regression test --- .../src/conversion_funcs/numeric.rs | 53 +++++++++++++++++-- 1 file changed, 50 insertions(+), 3 deletions(-) diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index c61056c909a..27160c8a444 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -20,9 +20,9 @@ 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, Decimal128Array, Float32Array, Float64Array, GenericStringBuilder, - Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, PrimitiveArray, Scalar, - StringBuilder, TimestampMicrosecondBuilder, + Array, ArrayRef, AsArray, BooleanArray, Decimal128Array, Float32Array, Float64Array, + GenericStringBuilder, Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, + PrimitiveArray, Scalar, StringBuilder, TimestampMicrosecondBuilder, }; use arrow::compute::kernels::cmp::neq; use arrow::datatypes::{ @@ -873,6 +873,11 @@ pub(crate) fn spark_cast_int_to_int( pub(crate) fn spark_cast_decimal_to_boolean(array: &dyn Array) -> SparkResult { let decimal_array = array.as_primitive::(); + // All-null fast path avoids constructing a zero scalar with the input's precision/scale, + // which would fail validation for Decimal128(0, 0) inputs that Spark accepts as nullable. + if decimal_array.null_count() == decimal_array.len() { + return Ok(Arc::new(BooleanArray::new_null(decimal_array.len()))); + } // 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( @@ -1751,6 +1756,48 @@ mod tests { 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); } #[test] From e9106a6f10a442900efd00c16ace4f7f00dfc916 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Fri, 28 Aug 2026 23:46:43 +0800 Subject: [PATCH 3/5] fix spark_pow dense-null --- native/spark-expr/benches/spark_pow.rs | 131 ++++++++++- .../src/array_funcs/array_insert.rs | 6 +- .../src/conversion_funcs/numeric.rs | 52 ++++- native/spark-expr/src/math_funcs/pow.rs | 220 +++++++++++++++++- .../apache/comet/CometExpressionSuite.scala | 34 +++ 5 files changed, 434 insertions(+), 9 deletions(-) diff --git a/native/spark-expr/benches/spark_pow.rs b/native/spark-expr/benches/spark_pow.rs index 073b5bd0b06..3d9743db394 100644 --- a/native/spark-expr/benches/spark_pow.rs +++ b/native/spark-expr/benches/spark_pow.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, Float64Array}; +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; @@ -38,6 +40,42 @@ fn create_f64_array(rows: usize, null_every: usize) -> ArrayRef { 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 = (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::>(), + )) + }; + Arc::new(Float64Array::new(values.into(), nulls)) +} + fn criterion_benchmark(c: &mut Criterion) { let rows = 8192; let no_nulls_a = create_f64_array(rows, 0); @@ -97,6 +135,97 @@ fn criterion_benchmark(c: &mut Criterion) { 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 to locate the crossover between the raw-buffer kernels + // (unary/binary) and a null-skipping path. + for null_pct in [50usize, 70, 80, 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_null_aware`) + // 2. `pow(a + 2.5D, b)` — array/array dispatch (`pow_binary_null_aware`), 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 = 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::>(), + )); + for null_pct in [50usize, 70, 80, 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()) + }) + }, + ); + } } criterion_group!(benches, criterion_benchmark); diff --git a/native/spark-expr/src/array_funcs/array_insert.rs b/native/spark-expr/src/array_funcs/array_insert.rs index 4693f2ebf26..6d996052a02 100644 --- a/native/spark-expr/src/array_funcs/array_insert.rs +++ b/native/spark-expr/src/array_funcs/array_insert.rs @@ -110,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)? diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index 27160c8a444..a25c6aefcb9 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -24,6 +24,7 @@ use arrow::array::{ GenericStringBuilder, Int16Array, Int32Array, Int64Array, Int8Array, OffsetSizeTrait, 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, @@ -873,11 +874,26 @@ pub(crate) fn spark_cast_int_to_int( pub(crate) fn spark_cast_decimal_to_boolean(array: &dyn Array) -> SparkResult { let decimal_array = array.as_primitive::(); - // All-null fast path avoids constructing a zero scalar with the input's precision/scale, - // which would fail validation for Decimal128(0, 0) inputs that Spark accepts as nullable. + // 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(), + ))); + } // 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( @@ -1798,6 +1814,38 @@ mod tests { // 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 2837f74d02f..b3329218665 100644 --- a/native/spark-expr/src/math_funcs/pow.rs +++ b/native/spark-expr/src/math_funcs/pow.rs @@ -16,6 +16,7 @@ // under the License. use arrow::array::{Array, ArrayRef, Datum, Float64Array}; +use arrow::buffer::NullBuffer; use arrow::compute::kernels::arity::{binary, unary}; use arrow::error::ArrowError; use datafusion::common::{utils::take_function_args, DataFusionError}; @@ -23,6 +24,20 @@ use datafusion::physical_expr_common::datum::apply; use datafusion::physical_plan::ColumnarValue; use std::sync::Arc; +/// When null density exceeds this fraction (numerator / denominator), the raw-buffer +/// kernels (`unary`/`binary`) waste `spark_powf` calls on masked-out slots that carry +/// non-zero payload (e.g. `pow(a + 2.5D, b)` after Arrow addition preserves null bits +/// but overwrites the value). Above the threshold we iterate valid indices instead. +/// +/// See `benches/spark_pow.rs::spark_pow: * composed nulls *` for the crossover. +const NULL_SKIP_THRESHOLD_NUM: usize = 3; +const NULL_SKIP_THRESHOLD_DEN: usize = 4; + +#[inline] +fn is_dense_null(null_count: usize, len: usize) -> bool { + null_count * NULL_SKIP_THRESHOLD_DEN > len * NULL_SKIP_THRESHOLD_NUM +} + /// Spark-compatible scalar power matching Java's `Math.pow`. /// /// Rust's `f64::powf` follows C99 `pow` semantics, which agree with `Math.pow` on almost every @@ -61,8 +76,11 @@ fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> { }) } -/// Array/array uses [`binary`] over [`spark_powf`]. Scalar/array uses [`unary`] so the -/// scalar is not broadcast. A null scalar short-circuits to an all-null array. +/// Array/array uses [`binary`] over [`spark_powf`]. Scalar/array and array/scalar use +/// [`unary`] so the scalar is not broadcast. A null scalar on either side short-circuits +/// to an all-null array. When null density exceeds the threshold (see [`is_dense_null`]) +/// we skip masked slots to avoid running `spark_powf` on carried-over payload from an +/// upstream Arrow op. fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { let (left, left_is_scalar) = lhs.get(); let (right, right_is_scalar) = rhs.get(); @@ -73,6 +91,8 @@ fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { if left.is_null(0) { Float64Array::new_null(right.len()) + } else if is_dense_null(right.null_count(), right.len()) { + pow_scalar_array_null_aware(left.value(0), right) } else { unary(right, |exp| spark_powf(left.value(0), exp)) } @@ -80,15 +100,86 @@ fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { if right.is_null(0) { Float64Array::new_null(left.len()) + } else if is_dense_null(left.null_count(), left.len()) { + pow_array_scalar_null_aware(left, right.value(0)) } else { unary(left, |base| spark_powf(base, right.value(0))) } } - _ => binary(left, right, spark_powf)?, + _ => { + // Effective null count of the output is bounded below by max(left, right). + // Using the max avoids a full NullBuffer::union scan just for the density + // check; the union still happens if we actually take the null-aware path. + let approx_nulls = left.null_count().max(right.null_count()); + if is_dense_null(approx_nulls, left.len()) { + pow_binary_null_aware(left, right)? + } else { + binary(left, right, spark_powf)? + } + } }; Ok(Arc::new(result)) } +/// `unary` runs `spark_powf` over every value in the raw buffer and copies the null +/// buffer through. When most slots are null, only running `spark_powf` on the valid +/// indices beats that. Output payload in masked slots remains initialized to zero. +fn pow_scalar_array_null_aware(base: f64, exp: &Float64Array) -> Float64Array { + let Some(nulls) = 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 pow_array_scalar_null_aware(base: &Float64Array, exp: f64) -> Float64Array { + let Some(nulls) = 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())) +} + +fn pow_binary_null_aware( + base: &Float64Array, + exp: &Float64Array, +) -> Result { + // Match arrow's `binary` contract: reject mismatched lengths with an error rather + // than panicking on out-of-bounds indexing below. + if base.len() != exp.len() { + return Err(ArrowError::ComputeError(format!( + "spark_pow: arrays have different lengths: {} vs {}", + base.len(), + exp.len() + ))); + } + let combined = NullBuffer::union(base.nulls(), exp.nulls()); + let base_values = base.values(); + let exp_values = exp.values(); + let mut out = vec![0.0f64; base.len()]; + match &combined { + Some(nulls) => { + for i in nulls.valid_indices() { + out[i] = spark_powf(base_values[i], exp_values[i]); + } + } + None => { + for i in 0..base.len() { + out[i] = spark_powf(base_values[i], exp_values[i]); + } + } + } + Ok(Float64Array::new(out.into(), combined)) +} + #[cfg(test)] mod test { use super::*; @@ -284,6 +375,111 @@ mod test { } } + /// 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; + // 90% null on both sides so the dense-null dispatch fires. + let rows_a = 100; + let rows_b = 90; + let make = |rows: usize| -> Float64Array { + let vals: Vec = vec![2.5; rows]; + let nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::>()); + 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}" + ); + } + + /// 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 = vec![2.5; rows]; + let base_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::>()); + 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::().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 = vec![2.5; rows]; + let base_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::>()); + 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::().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 = vec![2.5; rows]; + let a_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::>()); + let a = Float64Array::new(a_values.into(), Some(a_nulls)); + let b_values: Vec = vec![3.0; rows]; + let b_nulls = NullBuffer::from((0..rows).map(|i| i % 10 == 0).collect::>()); + 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::().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)]); @@ -300,13 +496,27 @@ mod test { panic!("expected array result"); } - let both_null = spark_pow(&[ + 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::().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!( - both_null, + scalar_result, ColumnarValue::Scalar(ScalarValue::Float64(None)) )); } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 4ab86b4f373..5ce576dd4f4 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -3459,6 +3459,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" From b9197b34eef96c30679097bd2f82fa927c6e6937 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Wed, 16 Sep 2026 22:18:31 +0300 Subject: [PATCH 4/5] 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. --- native/spark-expr/benches/spark_pow.rs | 48 ++++++++++- native/spark-expr/src/math_funcs/pow.rs | 107 ++++++++++++++++-------- 2 files changed, 119 insertions(+), 36 deletions(-) diff --git a/native/spark-expr/benches/spark_pow.rs b/native/spark-expr/benches/spark_pow.rs index 3d9743db394..cfc9961217b 100644 --- a/native/spark-expr/benches/spark_pow.rs +++ b/native/spark-expr/benches/spark_pow.rs @@ -76,6 +76,27 @@ fn create_f64_array_with_payload_in_nulls(rows: usize, null_pct: usize) -> Array 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); @@ -183,7 +204,7 @@ fn criterion_benchmark(c: &mut Criterion) { // 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_null_aware`) - // 2. `pow(a + 2.5D, b)` — array/array dispatch (`pow_binary_null_aware`), with + // 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 = Arc::new(Scalar::new(Float64Array::from(vec![2.5]))); @@ -226,6 +247,31 @@ fn criterion_benchmark(c: &mut Criterion) { }, ); } + // Independent null masks on both operands: `pow(a + 2.5D, b + 2.5D)`. Each operand is + // below the dense-null threshold on its own, but the output null density is roughly + // `1 - (1 - p)^2` (about 91% at 70% per operand), so the dispatch must look at the + // combined mask rather than either operand alone. Both additions are timed. + for null_pct in [50u64, 70, 74, 80] { + 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); diff --git a/native/spark-expr/src/math_funcs/pow.rs b/native/spark-expr/src/math_funcs/pow.rs index b3329218665..1faa3b0a5de 100644 --- a/native/spark-expr/src/math_funcs/pow.rs +++ b/native/spark-expr/src/math_funcs/pow.rs @@ -16,8 +16,8 @@ // under the License. use arrow::array::{Array, ArrayRef, Datum, Float64Array}; -use arrow::buffer::NullBuffer; -use arrow::compute::kernels::arity::{binary, unary}; +use arrow::buffer::{NullBuffer, ScalarBuffer}; +use arrow::compute::kernels::arity::unary; use arrow::error::ArrowError; use datafusion::common::{utils::take_function_args, DataFusionError}; use datafusion::physical_expr_common::datum::apply; @@ -25,9 +25,10 @@ use datafusion::physical_plan::ColumnarValue; use std::sync::Arc; /// When null density exceeds this fraction (numerator / denominator), the raw-buffer -/// kernels (`unary`/`binary`) waste `spark_powf` calls on masked-out slots that carry -/// non-zero payload (e.g. `pow(a + 2.5D, b)` after Arrow addition preserves null bits -/// but overwrites the value). Above the threshold we iterate valid indices instead. +/// loops (`unary`, or the zipped loop in `pow_binary`) waste `spark_powf` calls on +/// masked-out slots that carry non-zero payload (e.g. `pow(a + 2.5D, b)` after Arrow +/// addition preserves null bits but overwrites the value). Above the threshold we iterate +/// valid indices instead. For array/array the density is that of the combined mask. /// /// See `benches/spark_pow.rs::spark_pow: * composed nulls *` for the crossover. const NULL_SKIP_THRESHOLD_NUM: usize = 3; @@ -76,7 +77,7 @@ fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> { }) } -/// Array/array uses [`binary`] over [`spark_powf`]. Scalar/array and array/scalar use +/// Array/array uses [`pow_binary`]. Scalar/array and array/scalar use /// [`unary`] so the scalar is not broadcast. A null scalar on either side short-circuits /// to an all-null array. When null density exceeds the threshold (see [`is_dense_null`]) /// we skip masked slots to avoid running `spark_powf` on carried-over payload from an @@ -106,17 +107,7 @@ fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { - // Effective null count of the output is bounded below by max(left, right). - // Using the max avoids a full NullBuffer::union scan just for the density - // check; the union still happens if we actually take the null-aware path. - let approx_nulls = left.null_count().max(right.null_count()); - if is_dense_null(approx_nulls, left.len()) { - pow_binary_null_aware(left, right)? - } else { - binary(left, right, spark_powf)? - } - } + _ => pow_binary(left, right)?, }; Ok(Arc::new(result)) } @@ -148,12 +139,13 @@ fn pow_array_scalar_null_aware(base: &Float64Array, exp: f64) -> Float64Array { Float64Array::new(out.into(), Some(nulls.clone())) } -fn pow_binary_null_aware( - base: &Float64Array, - exp: &Float64Array, -) -> Result { - // Match arrow's `binary` contract: reject mismatched lengths with an error rather - // than panicking on out-of-bounds indexing below. +/// Array/array power. The output null mask is the union of both input masks, and the +/// dense-null dispatch must be based on that union: with independent null patterns each +/// operand can be under the threshold while most output rows are null. The union is +/// computed once and reused for both the density check and the result. +fn pow_binary(base: &Float64Array, exp: &Float64Array) -> Result { + // 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 {}", @@ -161,23 +153,24 @@ fn pow_binary_null_aware( exp.len() ))); } - let combined = NullBuffer::union(base.nulls(), exp.nulls()); + let nulls = NullBuffer::union(base.nulls(), exp.nulls()); let base_values = base.values(); let exp_values = exp.values(); - let mut out = vec![0.0f64; base.len()]; - match &combined { - Some(nulls) => { - for i in nulls.valid_indices() { + let values: ScalarBuffer = match &nulls { + Some(n) if is_dense_null(n.null_count(), base.len()) => { + 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 => { - for i in 0..base.len() { - out[i] = spark_powf(base_values[i], exp_values[i]); - } - } - } - Ok(Float64Array::new(out.into(), combined)) + _ => base_values + .iter() + .zip(exp_values.iter()) + .map(|(&b, &e)| spark_powf(b, e)) + .collect(), + }; + Ok(Float64Array::new(values, nulls)) } #[cfg(test)] @@ -401,6 +394,50 @@ mod test { ); } + /// Independent null masks: each operand is 70% null (under the dense-null threshold on + /// its own) but the union is 91% null, so the dispatch takes the null-skipping path. + /// 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 = (0..rows).map(|i| i % 10 < 3).collect(); + let exp_valid: Vec = (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())), + ); + assert!(!is_dense_null(base.null_count(), rows)); + assert!(!is_dense_null(exp.null_count(), rows)); + + 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::().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. From 5be638d03d55a21a38f910fe62c98fda65c09770 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Thu, 17 Sep 2026 13:45:25 +0300 Subject: [PATCH 5/5] 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. --- native/spark-expr/benches/spark_pow.rs | 18 ++- native/spark-expr/src/math_funcs/pow.rs | 167 +++++++++++++++--------- 2 files changed, 115 insertions(+), 70 deletions(-) diff --git a/native/spark-expr/benches/spark_pow.rs b/native/spark-expr/benches/spark_pow.rs index cfc9961217b..6a63bb306cb 100644 --- a/native/spark-expr/benches/spark_pow.rs +++ b/native/spark-expr/benches/spark_pow.rs @@ -159,9 +159,8 @@ fn criterion_benchmark(c: &mut Criterion) { // 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 to locate the crossover between the raw-buffer kernels - // (unary/binary) and a null-skipping path. - for null_pct in [50usize, 70, 80, 90, 99] { + // 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); @@ -203,7 +202,7 @@ fn criterion_benchmark(c: &mut Criterion) { // 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_null_aware`) + // 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))); @@ -213,7 +212,7 @@ fn criterion_benchmark(c: &mut Criterion) { .map(|i| 1.25 + (i % 10) as f64 * 0.1) .collect::>(), )); - for null_pct in [50usize, 70, 80, 90, 99] { + 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(); @@ -247,11 +246,10 @@ fn criterion_benchmark(c: &mut Criterion) { }, ); } - // Independent null masks on both operands: `pow(a + 2.5D, b + 2.5D)`. Each operand is - // below the dense-null threshold on its own, but the output null density is roughly - // `1 - (1 - p)^2` (about 91% at 70% per operand), so the dispatch must look at the - // combined mask rather than either operand alone. Both additions are timed. - for null_pct in [50u64, 70, 74, 80] { + // 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); diff --git a/native/spark-expr/src/math_funcs/pow.rs b/native/spark-expr/src/math_funcs/pow.rs index 1faa3b0a5de..a7f277c6fa7 100644 --- a/native/spark-expr/src/math_funcs/pow.rs +++ b/native/spark-expr/src/math_funcs/pow.rs @@ -19,26 +19,11 @@ 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}; +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; -/// When null density exceeds this fraction (numerator / denominator), the raw-buffer -/// loops (`unary`, or the zipped loop in `pow_binary`) waste `spark_powf` calls on -/// masked-out slots that carry non-zero payload (e.g. `pow(a + 2.5D, b)` after Arrow -/// addition preserves null bits but overwrites the value). Above the threshold we iterate -/// valid indices instead. For array/array the density is that of the combined mask. -/// -/// See `benches/spark_pow.rs::spark_pow: * composed nulls *` for the crossover. -const NULL_SKIP_THRESHOLD_NUM: usize = 3; -const NULL_SKIP_THRESHOLD_DEN: usize = 4; - -#[inline] -fn is_dense_null(null_count: usize, len: usize) -> bool { - null_count * NULL_SKIP_THRESHOLD_DEN > len * NULL_SKIP_THRESHOLD_NUM -} - /// Spark-compatible scalar power matching Java's `Math.pow`. /// /// Rust's `f64::powf` follows C99 `pow` semantics, which agree with `Math.pow` on almost every @@ -62,9 +47,24 @@ fn spark_powf(base: f64, exp: f64) -> f64 { /// inputs produce null; otherwise every result is the `spark_powf` value. pub fn spark_pow(args: &[ColumnarValue]) -> Result { 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() @@ -77,11 +77,11 @@ fn as_f64_array(array: &dyn Array) -> Result<&Float64Array, ArrowError> { }) } -/// Array/array uses [`pow_binary`]. Scalar/array and array/scalar use -/// [`unary`] so the scalar is not broadcast. A null scalar on either side short-circuits -/// to an all-null array. When null density exceeds the threshold (see [`is_dense_null`]) -/// we skip masked slots to avoid running `spark_powf` on carried-over payload from an -/// upstream Arrow op. +/// 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 { let (left, left_is_scalar) = lhs.get(); let (right, right_is_scalar) = rhs.get(); @@ -89,34 +89,23 @@ fn spark_pow_kernel(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { - if left.is_null(0) { - Float64Array::new_null(right.len()) - } else if is_dense_null(right.null_count(), right.len()) { - pow_scalar_array_null_aware(left.value(0), right) - } else { - unary(right, |exp| spark_powf(left.value(0), exp)) - } - } - (false, true) => { - if right.is_null(0) { - Float64Array::new_null(left.len()) - } else if is_dense_null(left.null_count(), left.len()) { - pow_array_scalar_null_aware(left, right.value(0)) - } else { - unary(left, |base| spark_powf(base, right.value(0))) - } - } + // 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)) } -/// `unary` runs `spark_powf` over every value in the raw buffer and copies the null -/// buffer through. When most slots are null, only running `spark_powf` on the valid -/// indices beats that. Output payload in masked slots remains initialized to zero. -fn pow_scalar_array_null_aware(base: f64, exp: &Float64Array) -> Float64Array { - let Some(nulls) = exp.nulls() else { +/// 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(); @@ -127,8 +116,8 @@ fn pow_scalar_array_null_aware(base: f64, exp: &Float64Array) -> Float64Array { Float64Array::new(out.into(), Some(nulls.clone())) } -fn pow_array_scalar_null_aware(base: &Float64Array, exp: f64) -> Float64Array { - let Some(nulls) = base.nulls() else { +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(); @@ -139,10 +128,10 @@ fn pow_array_scalar_null_aware(base: &Float64Array, exp: f64) -> Float64Array { Float64Array::new(out.into(), Some(nulls.clone())) } -/// Array/array power. The output null mask is the union of both input masks, and the -/// dense-null dispatch must be based on that union: with independent null patterns each -/// operand can be under the threshold while most output rows are null. The union is -/// computed once and reused for both the density check and the 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 { // Match arrow's `binary` contract: reject mismatched lengths with an error. This must // precede `NullBuffer::union`, which panics on mismatched lengths. @@ -156,15 +145,15 @@ fn pow_binary(base: &Float64Array, exp: &Float64Array) -> Result = match &nulls { - Some(n) if is_dense_null(n.null_count(), base.len()) => { + let values: ScalarBuffer = 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() } - _ => base_values + None => base_values .iter() .zip(exp_values.iter()) .map(|(&b, &e)| spark_powf(b, e)) @@ -373,7 +362,7 @@ mod test { #[test] fn test_spark_pow_null_aware_binary_length_mismatch() { use arrow::buffer::NullBuffer; - // 90% null on both sides so the dense-null dispatch fires. + // Nulls on both sides so the null-skipping path is taken. let rows_a = 100; let rows_b = 90; let make = |rows: usize| -> Float64Array { @@ -394,10 +383,58 @@ mod test { ); } - /// Independent null masks: each operand is 70% null (under the dense-null threshold on - /// its own) but the union is 91% null, so the dispatch takes the null-skipping path. - /// Null slots carry a real payload, and the result must still be null exactly where - /// either input is null and `spark_powf` everywhere else. + /// 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::>() + .into(), + Some(NullBuffer::new_valid(rows)), + ) + }; + let expected: Vec = (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::().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; @@ -414,8 +451,6 @@ mod test { vec![1.5; rows].into(), Some(NullBuffer::from(exp_valid.clone())), ); - assert!(!is_dense_null(base.null_count(), rows)); - assert!(!is_dense_null(exp.null_count(), rows)); let result = spark_pow(&[ ColumnarValue::Array(Arc::new(base)), @@ -556,5 +591,17 @@ mod test { 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}" + ); } }