Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions datafusion/core/tests/expr_api/simplification.rs
Original file line number Diff line number Diff line change
Expand Up @@ -648,13 +648,19 @@ fn test_simplify_power() {
let expected = col("c3_non_null");
test_simplify(expr, expected)
}
// Power(c3, Log(c3, c4)) ===> c4
// Power(c3, Log(c3, c4)) ===> cast(c4 AS Int64)
// The simplifier rewrites `power(b, log(b, x))` to `x`, but the
// rewritten expression must keep the same type as the original
// `power` call. `power`'s declared return type follows its base
// argument (c3 = Int64), so the UInt32 c4 has to be cast to Int64
// to preserve the output schema the optimizer already committed to.
{
let expr = power(
col("c3_non_null"),
log(col("c3_non_null"), col("c4_non_null")),
);
let expected = col("c4_non_null");
let expected =
Expr::Cast(Cast::new(Box::new(col("c4_non_null")), DataType::Int64));
test_simplify(expr, expected)
}
// Power(c3, c4) ===> Power(c3, c4)
Expand Down
5 changes: 5 additions & 0 deletions datafusion/functions/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,11 @@ harness = false
name = "atan2"
required-features = ["math_expressions"]

[[bench]]
harness = false
name = "power"
required-features = ["math_expressions"]

[[bench]]
harness = false
name = "substr_index"
Expand Down
140 changes: 140 additions & 0 deletions datafusion/functions/benches/power.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
// 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.

//! Microbenchmark for `power(decimal_array, int_*)`.
//!
//! Covers both array- and scalar-shaped integer exponents on a Decimal
//! base. Both shapes are dispatched to the native per-row decimal kernel;
//! the bench guards against any future change that routes either shape
//! through a Float64 round-trip, which is measurably slower than the
//! decimal kernel for the cases the kernel can handle.

extern crate criterion;

use arrow::array::{Decimal128Array, Int64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use criterion::{Criterion, criterion_group, criterion_main};
use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF};
use datafusion_functions::math::power;
use std::hint::black_box;
use std::sync::Arc;

fn make_decimal_array(size: usize, precision: u8, scale: i8) -> Decimal128Array {
// Use a fixed unscaled value (250) so the bench is independent of `scale`.
// The four-arm dispatch in `power` only cares about the Decimal variant
// and the exponent's shape, not the numeric value.
let arr = Decimal128Array::from(vec![250i128; size]);
arr.with_precision_and_scale(precision, scale).unwrap()
}

fn make_int_array(size: usize, value: i64) -> Int64Array {
Int64Array::from(vec![value; size])
}

fn run_power(
power_fn: &ScalarUDF,
args: &[ColumnarValue],
arg_fields: &[FieldRef],
return_field: &FieldRef,
config_options: &Arc<ConfigOptions>,
num_rows: usize,
) {
black_box(
power_fn
.invoke_with_args(ScalarFunctionArgs {
args: args.to_vec(),
arg_fields: arg_fields.to_vec(),
number_rows: num_rows,
return_field: Arc::clone(return_field),
config_options: Arc::clone(config_options),
})
.unwrap(),
);
}

fn criterion_benchmark(c: &mut Criterion) {
let power_fn = power();
let config_options = Arc::new(ConfigOptions::default());
let precision: u8 = 20;
let scale: i8 = 2;
let decimal_ty = DataType::Decimal128(precision, scale);

// Exponents are bounded by what the native decimal kernel can handle
// without overflowing the i128 intermediate; see
// <https://github.com/apache/datafusion/issues/22480>
let exponents = [2i64, 4, 8];

for size in [1024usize, 8192] {
let base_arr = Arc::new(make_decimal_array(size, precision, scale));
let base_field: FieldRef = Field::new("base", decimal_ty.clone(), true).into();
let exp_field: FieldRef = Field::new("exp", DataType::Int64, true).into();
let return_field: FieldRef = Field::new("r", decimal_ty.clone(), true).into();
let arg_fields = vec![base_field, exp_field];

for &exp in &exponents {
let exp_arr = Arc::new(make_int_array(size, exp));
let array_args = vec![
ColumnarValue::Array(base_arr.clone()),
ColumnarValue::Array(exp_arr),
];
c.bench_function(
&format!(
"power decimal({precision},{scale}) array x int array, exp={exp}, n={size}"
),
|b| {
b.iter(|| {
run_power(
&power_fn,
&array_args,
&arg_fields,
&return_field,
&config_options,
size,
)
})
},
);

let scalar_args = vec![
ColumnarValue::Array(base_arr.clone()),
ColumnarValue::Scalar(ScalarValue::Int64(Some(exp))),
];
c.bench_function(
&format!(
"power decimal({precision},{scale}) array x int scalar, exp={exp}, n={size}"
),
|b| {
b.iter(|| {
run_power(
&power_fn,
&scalar_args,
&arg_fields,
&return_field,
&config_options,
size,
)
})
},
);
}
}
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
Loading
Loading