diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..5d587c7d4746f 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -98,6 +98,10 @@ env_logger = { workspace = true } rand = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync"] } +[[bench]] +harness = false +name = "greatest" + [[bench]] harness = false name = "ascii" diff --git a/datafusion/functions/benches/greatest.rs b/datafusion/functions/benches/greatest.rs new file mode 100644 index 0000000000000..cee5803a0d84e --- /dev/null +++ b/datafusion/functions/benches/greatest.rs @@ -0,0 +1,124 @@ +// 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, ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::Field; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions::core::greatest; +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const SIZE: usize = 1024; + +fn int64_array(seed: u64, null_density: f32) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(seed); + let values: Int64Array = (0..SIZE) + .map(|_| { + if rng.random::() < null_density { + None + } else { + Some(rng.random_range(0i64..1000)) + } + }) + .collect(); + Arc::new(values) +} + +fn string_array(seed: u64, null_density: f32) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(seed); + let values: StringArray = (0..SIZE) + .map(|_| { + if rng.random::() < null_density { + None + } else { + Some(format!("value_{}", rng.random_range(0..1000))) + } + }) + .collect(); + Arc::new(values) +} + +fn bench_greatest(c: &mut Criterion, func: &ScalarUDF, name: &str, args: Vec) { + let return_type = args[0].data_type().clone(); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type().clone(), true).into() + }) + .collect::>(); + let args = args + .into_iter() + .map(ColumnarValue::Array) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { + b.iter(|| { + black_box( + func.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: SIZE, + return_field: Arc::new(Field::new("f", return_type.clone(), true)), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let greatest_func = greatest(); + + bench_greatest( + c, + &greatest_func, + "greatest_i64_nullable", + vec![int64_array(1, 0.2), int64_array(2, 0.2)], + ); + bench_greatest( + c, + &greatest_func, + "greatest_i64_no_nulls", + vec![int64_array(3, 0.0), int64_array(4, 0.0)], + ); + bench_greatest( + c, + &greatest_func, + "greatest_i64_nullable_3_args", + vec![ + int64_array(5, 0.2), + int64_array(6, 0.2), + int64_array(7, 0.2), + ], + ); + bench_greatest( + c, + &greatest_func, + "greatest_utf8_nullable", + vec![string_array(8, 0.2), string_array(9, 0.2)], + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/core/greatest.rs b/datafusion/functions/src/core/greatest.rs index 64eaefb9b887d..7aad0e488e8ea 100644 --- a/datafusion/functions/src/core/greatest.rs +++ b/datafusion/functions/src/core/greatest.rs @@ -21,6 +21,7 @@ use arrow::buffer::BooleanBuffer; use arrow::compute::SortOptions; use arrow::compute::kernels::cmp; use arrow::datatypes::DataType; +use datafusion_common::types::NativeType; use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err}; use datafusion_doc::Documentation; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; @@ -95,25 +96,18 @@ impl GreatestLeastOperator for GreatestFunc { /// Return boolean array where `arr[i] = lhs[i] >= rhs[i]` for all i, where `arr` is the result array /// Nulls are always considered smaller than any other value fn get_indexes_to_keep(lhs: &dyn Array, rhs: &dyn Array) -> Result { - // Fast path: - // If both arrays are not nested, have the same length and no nulls, we can use the faster vectorized kernel - // - If both arrays are not nested: Nested types, such as lists, are not supported as the null semantics are not well-defined. - // - both array does not have any nulls: cmp::gt_eq will return null if any of the input is null while we want to return false in that case - if !lhs.data_type().is_nested() - && lhs.logical_null_count() == 0 - && rhs.logical_null_count() == 0 - { - return cmp::gt_eq(&lhs, &rhs).map_err(|e| e.into()); - } - - let cmp = make_comparator(lhs, rhs, SORT_OPTIONS)?; - assert_eq_or_internal_err!( lhs.len(), rhs.len(), "All arrays should have the same length for greatest comparison" ); + if let Some(values) = vectorized_indexes_to_keep(lhs, rhs) { + return Ok(BooleanArray::new(values, None)); + } + + let cmp = make_comparator(lhs, rhs, SORT_OPTIONS)?; + let values = BooleanBuffer::collect_bool(lhs.len(), |i| cmp(i, i).is_ge()); // No nulls as we only want to keep the values that are larger, its either true or false @@ -121,6 +115,52 @@ impl GreatestLeastOperator for GreatestFunc { } } +/// Computes `lhs[i] >= rhs[i]` with nulls ordered before every other value, +/// using the vectorized comparison kernel rather than a per-row comparator. +/// +/// Returns `None` when the inputs are not a good fit for the kernel and the +/// caller should fall back to [`make_comparator`]: +/// - nested types, such as lists, for which the kernel's null semantics are not +/// well-defined, and +/// - nullable floating point values, because `cmp::gt_eq` compares them with +/// IEEE semantics (every comparison against `NaN` is false) while the +/// comparator uses the total order in which `NaN` is the largest value. +fn vectorized_indexes_to_keep(lhs: &dyn Array, rhs: &dyn Array) -> Option { + if lhs.data_type().is_nested() { + return None; + } + + let (values, both_valid) = cmp::gt_eq(&lhs, &rhs).ok()?.into_parts(); + + // The kernel's validity is the intersection of both inputs', so an absent or + // fully set buffer means there is nothing to fix up: the values are already + // the answer. + let Some(both_valid) = both_valid.filter(|nulls| nulls.null_count() > 0) else { + return Some(values); + }; + + // `NativeType` looks through the encodings that wrap a value type, so this also + // catches dictionary and run-end encoded floats. + if NativeType::from(lhs.data_type()).is_float() { + return None; + } + + // `cmp::gt_eq` yields null wherever either side is null, but a null must compare + // as smaller than any other value rather than propagate. So keep the left value + // where both sides are valid and it compares greater or equal, and wherever the + // right side is null (which includes both sides being null, where the values are + // equivalent and the left one is kept). + let both_valid_and_ge = &values & both_valid.inner(); + Some(match rhs.logical_nulls() { + // `!rhs_valid` is all zeroes when the right side has no nulls, so the OR + // would be a no-op. + Some(rhs_nulls) if rhs_nulls.null_count() > 0 => { + &both_valid_and_ge | &!rhs_nulls.inner() + } + _ => both_valid_and_ge, + }) +} + impl ScalarUDFImpl for GreatestFunc { fn name(&self) -> &str { "greatest"