perf: optimize log in datafusion-functions

datafusion/functions/Cargo.toml

@@ -103,6 +103,11 @@ harness = false
 name = "ascii"
 required-features = ["string_expressions"]
 
+[[bench]]
+harness = false
+name = "log"
+required-features = ["math_expressions"]
+
 [[bench]]
 harness = false
 name = "concat"

datafusion/functions/benches/log.rs

@@ -0,0 +1,132 @@
+// 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.
+
+extern crate criterion;
+
+use arrow::array::{ArrayRef, Float32Array, Float64Array};
+use arrow::datatypes::{DataType, Field};
+use arrow::util::bench_util::create_primitive_array;
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_functions::math::log::LogFunc;
+use std::hint::black_box;
+use std::sync::Arc;
+
+fn run(c: &mut Criterion, name: &str, args: Vec<ColumnarValue>, return_type: DataType) {
+    let log = LogFunc::new();
+    let size = args
+        .iter()
+        .filter_map(|arg| match arg {
+            ColumnarValue::Array(array) => Some(array.len()),
+            ColumnarValue::Scalar(_) => None,
+        })
+        .max()
+        .unwrap();
+    let arg_fields: Vec<_> = args
+        .iter()
+        .enumerate()
+        .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
+        .collect();
+    let return_field = Arc::new(Field::new("f", return_type, true));
+    let config_options = Arc::new(ConfigOptions::default());
+
+    c.bench_function(name, |b| {
+        b.iter(|| {
+            black_box(
+                log.invoke_with_args(ScalarFunctionArgs {
+                    args: args.clone(),
+                    arg_fields: arg_fields.clone(),
+                    number_rows: size,
+                    return_field: Arc::clone(&return_field),
+                    config_options: Arc::clone(&config_options),
+                })
+                .unwrap(),
+            )
+        })
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let size = 8192;
+    let f64_array = Arc::new(create_primitive_array::<arrow::datatypes::Float64Type>(
+        size, 0.1,
+    )) as ArrayRef;
+    let f32_array = Arc::new(create_primitive_array::<arrow::datatypes::Float32Type>(
+        size, 0.1,
+    )) as ArrayRef;
+
+    // log(x), which defaults to base 10
+    run(
+        c,
+        "log_f64_default_base",
+        vec![ColumnarValue::Array(Arc::clone(&f64_array))],
+        DataType::Float64,
+    );
+
+    // log(base, x) with a scalar base
+    run(
+        c,
+        "log_f64_scalar_base",
+        vec![
+            ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+            ColumnarValue::Array(Arc::clone(&f64_array)),
+        ],
+        DataType::Float64,
+    );
+
+    run(
+        c,
+        "log_f32_scalar_base",
+        vec![
+            ColumnarValue::Scalar(ScalarValue::Float32(Some(2.0))),
+            ColumnarValue::Array(Arc::clone(&f32_array)),
+        ],
+        DataType::Float32,
+    );
+
+    // log(base, x) with a base column, which cannot hoist the base logarithm
+    let f64_bases = Arc::new(Float64Array::from_iter_values(
+        (0..size).map(|i| 2.0 + (i % 8) as f64),
+    )) as ArrayRef;
+    run(
+        c,
+        "log_f64_array_base",
+        vec![
+            ColumnarValue::Array(f64_bases),
+            ColumnarValue::Array(Arc::clone(&f64_array)),
+        ],
+        DataType::Float64,
+    );
+
+    let f32_bases = Arc::new(Float32Array::from_iter_values(
+        (0..size).map(|i| 2.0 + (i % 8) as f32),
+    )) as ArrayRef;
+    run(
+        c,
+        "log_f32_array_base",
+        vec![
+            ColumnarValue::Array(f32_bases),
+            ColumnarValue::Array(Arc::clone(&f32_array)),
+        ],
+        DataType::Float32,
+    );
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/functions/src/math/log.rs

@@ -22,8 +22,8 @@ use super::power::PowerFunc;
 use crate::utils::calculate_binary_math;
 use arrow::array::{Array, ArrayRef};
 use arrow::datatypes::{
-    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Float16Type,
-    Float32Type, Float64Type,
+    ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, Decimal128Type,
+    Decimal256Type, Float16Type, Float32Type, Float64Type,
 };
 use arrow::error::ArrowError;
 use arrow_buffer::i256;
@@ -185,6 +185,46 @@ fn log_decimal256(value: i256, scale: i8, base: f64) -> Result<f64, ArrowError>
     }
 }
 
+/// If `base` is a non-null scalar, returns its natural logarithm as the native
+/// type of `T`. The cast mirrors the one [`calculate_binary_math`] applies to
+/// the base, so the value seen here is the one the kernel would have used.
+fn scalar_base_ln<T: ArrowPrimitiveType>(base: &ColumnarValue) -> Option<T::Native>
+where
+    T::Native: Float + TryFrom<ScalarValue>,
+{
+    let ColumnarValue::Scalar(scalar) = base else {
+        return None;
+    };
+    if scalar.is_null() {
+        return None;
+    }
+    let scalar = scalar.cast_to(&T::DATA_TYPE).ok()?;
+    T::Native::try_from(scalar).ok().map(T::Native::ln)
+}
+
+/// `log(base, value)` over a float array.
+///
+/// `x.log(base)` is `x.ln() / base.ln()`, so a scalar base contributes the same
+/// logarithm to every row. Hoisting it out of the loop halves the transcendental
+/// calls per row without changing the formula or its operands.
+fn log_float<T>(value: &dyn Array, base: &ColumnarValue) -> Result<ArrayRef>
+where
+    T: ArrowPrimitiveType,
+    T::Native: Float + TryFrom<ScalarValue>,
+{
+    let base_ln = scalar_base_ln::<T>(base);
+    Ok(calculate_binary_math::<T, T, T, _>(
+        value,
+        base,
+        |value, base| {
+            Ok(match base_ln {
+                Some(base_ln) => value.ln() / base_ln,
+                None => value.log(base),
+            })
+        },
+    )?)
+}
+
 impl ScalarUDFImpl for LogFunc {
     fn name(&self) -> &str {
         "log"
@@ -254,20 +294,8 @@ impl ScalarUDFImpl for LogFunc {
                     |value, base| Ok(value.log(base)),
                 )?
             }
-            DataType::Float32 => {
-                calculate_binary_math::<Float32Type, Float32Type, Float32Type, _>(
-                    &value,
-                    &base,
-                    |value, base| Ok(value.log(base)),
-                )?
-            }
-            DataType::Float64 => {
-                calculate_binary_math::<Float64Type, Float64Type, Float64Type, _>(
-                    &value,
-                    &base,
-                    |value, base| Ok(value.log(base)),
-                )?
-            }
+            DataType::Float32 => log_float::<Float32Type>(&value, &base)?,
+            DataType::Float64 => log_float::<Float64Type>(&value, &base)?,
             DataType::Decimal32(_, scale) => {
                 calculate_binary_math::<Decimal32Type, Float64Type, Float64Type, _>(
                     &value,