From b993851bf478290a488da0f70fc5805d277b8ee8 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Wed, 13 May 2026 10:27:09 +0530 Subject: [PATCH 1/2] bench: run array_replace kernels in benchmark --- .../benches/array_expression.rs | 146 +++++++++++++++--- 1 file changed, 128 insertions(+), 18 deletions(-) diff --git a/datafusion/functions-nested/benches/array_expression.rs b/datafusion/functions-nested/benches/array_expression.rs index 71bb939238aa..b1acb4c8e808 100644 --- a/datafusion/functions-nested/benches/array_expression.rs +++ b/datafusion/functions-nested/benches/array_expression.rs @@ -15,33 +15,143 @@ // specific language governing permissions and limitations // under the License. -use criterion::{Criterion, criterion_group, criterion_main}; -use datafusion_expr::lit; -use datafusion_functions_nested::expr_fn::{array_replace_all, make_array}; +use arrow::array::{ArrayRef, Int64Builder, ListArray}; +use arrow::datatypes::{DataType, Field}; +use criterion::{ + BenchmarkGroup, Criterion, criterion_group, criterion_main, measurement::WallTime, +}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions_nested::replace::{ + array_replace_all_udf, array_replace_n_udf, array_replace_udf, +}; +use rand::seq::IndexedRandom; +use rand::{Rng, SeedableRng, rngs::StdRng}; use std::hint::black_box; +use std::sync::Arc; + +// (num_rows, list_size) — tuned so benches finish in a few seconds. At +// NEEDLE_DENSITY = 0.1, list_size = 100/500 reliably exercise the +// match-and-replace path (and the counter cutoff for `array_replace_n`); +// list_size = 10 exercises the null-row and no-match early-return paths. +const SIZES: &[(usize, usize)] = &[(4_000, 10), (10_000, 100), (10_000, 500)]; +const SEED: u64 = 42; +const HAYSTACK_NULL_DENSITY: f64 = 0.1; +const NEEDLE_DENSITY: f64 = 0.1; +const NEEDLE: i64 = 0; +const REPLACEMENT: i64 = -1; fn criterion_benchmark(c: &mut Criterion) { - // Construct large arrays for benchmarking + let list_data_type = + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))); + let return_field: Arc = + Field::new("result", list_data_type.clone(), true).into(); + let config_options = Arc::new(ConfigOptions::default()); - let array_len = 100_000; + let array_field: Arc = Field::new("array", list_data_type, true).into(); + let from_field: Arc = Field::new("from", DataType::Int64, true).into(); + let to_field: Arc = Field::new("to", DataType::Int64, true).into(); + let max_field: Arc = Field::new("max", DataType::Int64, true).into(); - let array = (0..array_len).map(|_| lit(2_i64)).collect::>(); - let list_array = make_array(vec![make_array(array); 3]); - let from_array = make_array(vec![lit(2_i64); 3]); - let to_array = make_array(vec![lit(-2_i64); 3]); + let three_arg_fields = + vec![array_field.clone(), from_field.clone(), to_field.clone()]; + let four_arg_fields = vec![array_field, from_field, to_field, max_field]; - // Benchmark array functions + let mut group = c.benchmark_group("array_replace_int64"); + let udf = array_replace_udf(); + for &(num_rows, list_size) in SIZES { + let fn_args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(create_list_array(num_rows, list_size)), + ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), + ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), + ], + arg_fields: three_arg_fields.clone(), + number_rows: num_rows, + return_field: return_field.clone(), + config_options: config_options.clone(), + }; + bench_case(&mut group, &udf, list_size, num_rows, &fn_args); + } + group.finish(); - c.bench_function("array_replace", |b| { - b.iter(|| { - black_box(array_replace_all( - list_array.clone(), - from_array.clone(), - to_array.clone(), - )) - }) + let mut group = c.benchmark_group("array_replace_n_int64"); + let udf = array_replace_n_udf(); + // n = 2 makes the counter cutoff fire on any row with ≥ 2 needles. + let n = 2_i64; + for &(num_rows, list_size) in SIZES { + let fn_args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(create_list_array(num_rows, list_size)), + ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), + ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), + ColumnarValue::Scalar(ScalarValue::from(n)), + ], + arg_fields: four_arg_fields.clone(), + number_rows: num_rows, + return_field: return_field.clone(), + config_options: config_options.clone(), + }; + bench_case(&mut group, &udf, list_size, num_rows, &fn_args); + } + group.finish(); + + let mut group = c.benchmark_group("array_replace_all_int64"); + let udf = array_replace_all_udf(); + for &(num_rows, list_size) in SIZES { + let fn_args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(create_list_array(num_rows, list_size)), + ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), + ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), + ], + arg_fields: three_arg_fields.clone(), + number_rows: num_rows, + return_field: return_field.clone(), + config_options: config_options.clone(), + }; + bench_case(&mut group, &udf, list_size, num_rows, &fn_args); + } + group.finish(); +} + +fn bench_case( + group: &mut BenchmarkGroup<'_, WallTime>, + udf: &Arc, + list_size: usize, + num_rows: usize, + fn_args: &ScalarFunctionArgs, +) { + let name = format!("list_size: {list_size}, num_rows: {num_rows}"); + group.bench_function(name, |b| { + b.iter(|| black_box(udf.invoke_with_args(fn_args.clone()).unwrap())) }); } +fn create_list_array(num_rows: usize, list_size: usize) -> ArrayRef { + let filler_values = [None, Some(1_i64), Some(2), Some(3), Some(4), Some(5)]; + let mut rng = StdRng::seed_from_u64(SEED); + let values = (0..num_rows) + .map(|_| { + if rng.random_bool(HAYSTACK_NULL_DENSITY) { + None + } else { + let list = (0..list_size) + .map(|_| { + if rng.random_bool(NEEDLE_DENSITY) { + Some(NEEDLE) + } else { + *filler_values.choose(&mut rng).unwrap() + } + }) + .collect::>(); + Some(list) + } + }) + .collect::>(); + Arc::new(ListArray::from_nested_iter::(values)) +} + criterion_group!(benches, criterion_benchmark); criterion_main!(benches); From 00e80e9f57f5eaa74dccac97523cb9b7e0e11cf5 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Thu, 14 May 2026 13:37:16 +0530 Subject: [PATCH 2/2] remove stale array_expression benchmark --- datafusion/functions-nested/Cargo.toml | 4 - .../benches/array_expression.rs | 157 ------------------ 2 files changed, 161 deletions(-) delete mode 100644 datafusion/functions-nested/benches/array_expression.rs diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index 83bed152f69d..ed5a89b8e3e7 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -78,10 +78,6 @@ name = "array_concat" harness = false name = "array_min_max" -[[bench]] -harness = false -name = "array_expression" - [[bench]] harness = false name = "arrays_zip" diff --git a/datafusion/functions-nested/benches/array_expression.rs b/datafusion/functions-nested/benches/array_expression.rs deleted file mode 100644 index b1acb4c8e808..000000000000 --- a/datafusion/functions-nested/benches/array_expression.rs +++ /dev/null @@ -1,157 +0,0 @@ -// 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, Int64Builder, ListArray}; -use arrow::datatypes::{DataType, Field}; -use criterion::{ - BenchmarkGroup, Criterion, criterion_group, criterion_main, measurement::WallTime, -}; -use datafusion_common::ScalarValue; -use datafusion_common::config::ConfigOptions; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; -use datafusion_functions_nested::replace::{ - array_replace_all_udf, array_replace_n_udf, array_replace_udf, -}; -use rand::seq::IndexedRandom; -use rand::{Rng, SeedableRng, rngs::StdRng}; -use std::hint::black_box; -use std::sync::Arc; - -// (num_rows, list_size) — tuned so benches finish in a few seconds. At -// NEEDLE_DENSITY = 0.1, list_size = 100/500 reliably exercise the -// match-and-replace path (and the counter cutoff for `array_replace_n`); -// list_size = 10 exercises the null-row and no-match early-return paths. -const SIZES: &[(usize, usize)] = &[(4_000, 10), (10_000, 100), (10_000, 500)]; -const SEED: u64 = 42; -const HAYSTACK_NULL_DENSITY: f64 = 0.1; -const NEEDLE_DENSITY: f64 = 0.1; -const NEEDLE: i64 = 0; -const REPLACEMENT: i64 = -1; - -fn criterion_benchmark(c: &mut Criterion) { - let list_data_type = - DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))); - let return_field: Arc = - Field::new("result", list_data_type.clone(), true).into(); - let config_options = Arc::new(ConfigOptions::default()); - - let array_field: Arc = Field::new("array", list_data_type, true).into(); - let from_field: Arc = Field::new("from", DataType::Int64, true).into(); - let to_field: Arc = Field::new("to", DataType::Int64, true).into(); - let max_field: Arc = Field::new("max", DataType::Int64, true).into(); - - let three_arg_fields = - vec![array_field.clone(), from_field.clone(), to_field.clone()]; - let four_arg_fields = vec![array_field, from_field, to_field, max_field]; - - let mut group = c.benchmark_group("array_replace_int64"); - let udf = array_replace_udf(); - for &(num_rows, list_size) in SIZES { - let fn_args = ScalarFunctionArgs { - args: vec![ - ColumnarValue::Array(create_list_array(num_rows, list_size)), - ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), - ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), - ], - arg_fields: three_arg_fields.clone(), - number_rows: num_rows, - return_field: return_field.clone(), - config_options: config_options.clone(), - }; - bench_case(&mut group, &udf, list_size, num_rows, &fn_args); - } - group.finish(); - - let mut group = c.benchmark_group("array_replace_n_int64"); - let udf = array_replace_n_udf(); - // n = 2 makes the counter cutoff fire on any row with ≥ 2 needles. - let n = 2_i64; - for &(num_rows, list_size) in SIZES { - let fn_args = ScalarFunctionArgs { - args: vec![ - ColumnarValue::Array(create_list_array(num_rows, list_size)), - ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), - ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), - ColumnarValue::Scalar(ScalarValue::from(n)), - ], - arg_fields: four_arg_fields.clone(), - number_rows: num_rows, - return_field: return_field.clone(), - config_options: config_options.clone(), - }; - bench_case(&mut group, &udf, list_size, num_rows, &fn_args); - } - group.finish(); - - let mut group = c.benchmark_group("array_replace_all_int64"); - let udf = array_replace_all_udf(); - for &(num_rows, list_size) in SIZES { - let fn_args = ScalarFunctionArgs { - args: vec![ - ColumnarValue::Array(create_list_array(num_rows, list_size)), - ColumnarValue::Scalar(ScalarValue::from(NEEDLE)), - ColumnarValue::Scalar(ScalarValue::from(REPLACEMENT)), - ], - arg_fields: three_arg_fields.clone(), - number_rows: num_rows, - return_field: return_field.clone(), - config_options: config_options.clone(), - }; - bench_case(&mut group, &udf, list_size, num_rows, &fn_args); - } - group.finish(); -} - -fn bench_case( - group: &mut BenchmarkGroup<'_, WallTime>, - udf: &Arc, - list_size: usize, - num_rows: usize, - fn_args: &ScalarFunctionArgs, -) { - let name = format!("list_size: {list_size}, num_rows: {num_rows}"); - group.bench_function(name, |b| { - b.iter(|| black_box(udf.invoke_with_args(fn_args.clone()).unwrap())) - }); -} - -fn create_list_array(num_rows: usize, list_size: usize) -> ArrayRef { - let filler_values = [None, Some(1_i64), Some(2), Some(3), Some(4), Some(5)]; - let mut rng = StdRng::seed_from_u64(SEED); - let values = (0..num_rows) - .map(|_| { - if rng.random_bool(HAYSTACK_NULL_DENSITY) { - None - } else { - let list = (0..list_size) - .map(|_| { - if rng.random_bool(NEEDLE_DENSITY) { - Some(NEEDLE) - } else { - *filler_values.choose(&mut rng).unwrap() - } - }) - .collect::>(); - Some(list) - } - }) - .collect::>(); - Arc::new(ListArray::from_nested_iter::(values)) -} - -criterion_group!(benches, criterion_benchmark); -criterion_main!(benches);