refactor: replace remaining hand-rolled loops with Arrow kernels

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

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<Int32>, 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::<Int32Type, _, _>(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::<Vec<_>>(),
+    );
+    let items = Int32Array::from((0..rows).map(|i| Some(i as i32)).collect::<Vec<_>>());
+
+    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);

native/spark-expr/benches/cast_decimal_to_boolean.rs

@@ -0,0 +1,82 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Decimal128Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions};
+use std::hint::black_box;
+use std::sync::Arc;
+
+const PRECISION: u8 = 20;
+const SCALE: i8 = 2;
+
+/// Build a Decimal128(20, 2) column of `rows` rows. Every `null_every`-th row is null
+/// (`null_every == 0` means no nulls). Values alternate between 0 and non-zero so the
+/// boolean result is a realistic mix.
+fn create_batch(rows: usize, null_every: usize) -> RecordBatch {
+    let arr: Decimal128Array = (0..rows)
+        .map(|i| {
+            if null_every != 0 && i % null_every == 0 {
+                None
+            } else if i % 3 == 0 {
+                Some(0i128)
+            } else {
+                Some(((i % 100) as i128) * 100)
+            }
+        })
+        .collect::<Decimal128Array>()
+        .with_precision_and_scale(PRECISION, SCALE)
+        .unwrap();
+
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "a",
+        DataType::Decimal128(PRECISION, SCALE),
+        true,
+    )]));
+    RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let rows = 8192;
+    let expr = Arc::new(Column::new("a", 0));
+    let cast_to_bool = Cast::new(
+        expr,
+        DataType::Boolean,
+        SparkCastOptions::new(EvalMode::Legacy, "UTC", false),
+        None,
+        None,
+    );
+
+    let no_nulls = create_batch(rows, 0);
+    let sparse_nulls = create_batch(rows, 10);
+    let dense_nulls = create_batch(rows, 2);
+
+    let mut bench = |name: &str, batch: &RecordBatch| {
+        c.bench_function(name, |b| {
+            b.iter(|| black_box(cast_to_bool.evaluate(black_box(batch)).unwrap()))
+        });
+    };
+    bench("cast_decimal_to_boolean: no nulls", &no_nulls);
+    bench("cast_decimal_to_boolean: sparse nulls", &sparse_nulls);
+    bench("cast_decimal_to_boolean: dense nulls", &dense_nulls);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

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);