diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6faa9fec4e..fabf7220e2 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -222,4 +222,16 @@ harness = false [[bench]] name = "cast_int_to_decimal" +harness = false + +[[bench]] +name = "spark_pow" +harness = false + +[[bench]] +name = "cast_decimal_to_boolean" +harness = false + +[[bench]] +name = "array_insert" harness = false \ No newline at end of file diff --git a/native/spark-expr/benches/array_insert.rs b/native/spark-expr/benches/array_insert.rs new file mode 100644 index 0000000000..21f5b78345 --- /dev/null +++ b/native/spark-expr/benches/array_insert.rs @@ -0,0 +1,125 @@ +// 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::{Array, Int32Array, ListArray, RecordBatch}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_expr::expressions::col; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::ArrayInsert; +use std::hint::black_box; +use std::sync::Arc; + +/// Build a `RecordBatch` with columns (src List, pos Int32, item Int32) for +/// `rows` rows. Every `src_null_every`-th src row is null, every `pos_null_every`-th +/// pos is null (`_ == 0` disables that pattern). `item` is always non-null. The +/// exercised code path is the one refactored in the PR: `is_not_null` on src and pos, +/// then `and` for the item mask. +fn create_batch(rows: usize, src_null_every: usize, pos_null_every: usize) -> RecordBatch { + let src_iter = (0..rows).map(|i| { + if src_null_every != 0 && i % src_null_every == 0 { + None + } else { + Some(vec![ + Some(i as i32), + Some((i + 1) as i32), + Some((i + 2) as i32), + ]) + } + }); + let src = ListArray::from_iter_primitive::(src_iter); + + let positions = Int32Array::from( + (0..rows) + .map(|i| { + if pos_null_every != 0 && i % pos_null_every == 0 { + None + } else { + Some(((i as i32) % 3) + 1) + } + }) + .collect::>(), + ); + let items = Int32Array::from((0..rows).map(|i| Some(i as i32)).collect::>()); + + let list_field = match src.data_type() { + DataType::List(f) => Arc::clone(f), + _ => unreachable!(), + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::List(Arc::clone(&list_field)), true), + Field::new("pos", DataType::Int32, true), + Field::new("item", DataType::Int32, true), + ])); + RecordBatch::try_new( + schema, + vec![Arc::new(src), Arc::new(positions), Arc::new(items)], + ) + .unwrap() +} + +fn make_expr(batch: &RecordBatch) -> ArrayInsert { + let schema = batch.schema(); + ArrayInsert::new( + col("src", &schema).unwrap(), + col("pos", &schema).unwrap(), + col("item", &schema).unwrap(), + false, + ) +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + + let no_nulls = create_batch(rows, 0, 0); + let sparse_src_nulls = create_batch(rows, 10, 0); + let dense_src_nulls = create_batch(rows, 2, 0); + let mixed_nulls = create_batch(rows, 10, 7); + + let no_nulls_expr = make_expr(&no_nulls); + let sparse_src_expr = make_expr(&sparse_src_nulls); + let dense_src_expr = make_expr(&dense_src_nulls); + let mixed_expr = make_expr(&mixed_nulls); + + c.bench_function("array_insert: no nulls", |b| { + b.iter(|| black_box(no_nulls_expr.evaluate(black_box(&no_nulls)).unwrap())) + }); + c.bench_function("array_insert: sparse src nulls", |b| { + b.iter(|| { + black_box( + sparse_src_expr + .evaluate(black_box(&sparse_src_nulls)) + .unwrap(), + ) + }) + }); + c.bench_function("array_insert: dense src nulls", |b| { + b.iter(|| { + black_box( + dense_src_expr + .evaluate(black_box(&dense_src_nulls)) + .unwrap(), + ) + }) + }); + c.bench_function("array_insert: mixed src+pos nulls", |b| { + b.iter(|| black_box(mixed_expr.evaluate(black_box(&mixed_nulls)).unwrap())) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); 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::() + .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..073b5bd0b0 --- /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 e056c108e0..1b5e43d866 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 029a38e3e7..439b046fa2 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 @@ -80,8 +81,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 { @@ -852,15 +853,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)?)) } pub(crate) fn cast_float64_to_decimal128( @@ -1295,6 +1294,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 61168f7d18..2837f74d02 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)) + )); + } }