You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2024/04/18 11:41:25 UTC
(arrow-datafusion) branch main updated: Avoid cloning in `log::simplify` and `power::simplify` (#10086)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 0f88a2d8eb Avoid cloning in `log::simplify` and `power::simplify` (#10086)
0f88a2d8eb is described below
commit 0f88a2d8eb72edce1b6c728574190ef620ca80a1
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Thu Apr 18 07:41:20 2024 -0400
Avoid cloning in `log::simplify` and `power::simplify` (#10086)
* Avoid cloning in power::simplify
* Avoid cloning in log::simplify
* Apply suggestions from code review
Co-authored-by: comphead <co...@users.noreply.github.com>
---------
Co-authored-by: comphead <co...@users.noreply.github.com>
---
datafusion/functions/src/math/log.rs | 88 ++++++++++++++++++++++------------
datafusion/functions/src/math/power.rs | 59 +++++++++++++----------
2 files changed, 92 insertions(+), 55 deletions(-)
diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs
index 2131b6aa67..0a29e1ecfc 100644
--- a/datafusion/functions/src/math/log.rs
+++ b/datafusion/functions/src/math/log.rs
@@ -18,10 +18,15 @@
//! Math function: `log()`.
use arrow::datatypes::DataType;
-use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_common::{
+ exec_err, internal_err, plan_datafusion_err, plan_err, DataFusionError, Result,
+ ScalarValue,
+};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
-use datafusion_expr::{ColumnarValue, Expr, FuncMonotonicity, ScalarFunctionDefinition};
+use datafusion_expr::{
+ lit, ColumnarValue, Expr, FuncMonotonicity, ScalarFunctionDefinition,
+};
use arrow::array::{ArrayRef, Float32Array, Float64Array};
use datafusion_expr::TypeSignature::*;
@@ -146,44 +151,54 @@ impl ScalarUDFImpl for LogFunc {
/// 3. Log(a, a) ===> 1
fn simplify(
&self,
- args: Vec<Expr>,
+ mut args: Vec<Expr>,
info: &dyn SimplifyInfo,
) -> Result<ExprSimplifyResult> {
- let mut number = &args[0];
- let mut base =
- &Expr::Literal(ScalarValue::new_ten(&info.get_data_type(number)?)?);
- if args.len() == 2 {
- base = &args[0];
- number = &args[1];
+ // Args are either
+ // log(number)
+ // log(base, number)
+ let num_args = args.len();
+ if num_args > 2 {
+ return plan_err!("Expected log to have 1 or 2 arguments, got {num_args}");
}
+ let number = args.pop().ok_or_else(|| {
+ plan_datafusion_err!("Expected log to have 1 or 2 arguments, got 0")
+ })?;
+ let number_datatype = info.get_data_type(&number)?;
+ // default to base 10
+ let base = if let Some(base) = args.pop() {
+ base
+ } else {
+ lit(ScalarValue::new_ten(&number_datatype)?)
+ };
match number {
- Expr::Literal(value)
- if value == &ScalarValue::new_one(&info.get_data_type(number)?)? =>
- {
- Ok(ExprSimplifyResult::Simplified(Expr::Literal(
- ScalarValue::new_zero(&info.get_data_type(base)?)?,
- )))
+ Expr::Literal(value) if value == ScalarValue::new_one(&number_datatype)? => {
+ Ok(ExprSimplifyResult::Simplified(lit(ScalarValue::new_zero(
+ &info.get_data_type(&base)?,
+ )?)))
}
- Expr::ScalarFunction(ScalarFunction {
- func_def: ScalarFunctionDefinition::UDF(fun),
- args,
- }) if base == &args[0]
- && fun
- .as_ref()
- .inner()
- .as_any()
- .downcast_ref::<PowerFunc>()
- .is_some() =>
+ Expr::ScalarFunction(ScalarFunction { func_def, mut args })
+ if is_pow(&func_def) && args.len() == 2 && base == args[0] =>
{
- Ok(ExprSimplifyResult::Simplified(args[1].clone()))
+ let b = args.pop().unwrap(); // length checked above
+ Ok(ExprSimplifyResult::Simplified(b))
}
- _ => {
+ number => {
if number == base {
- Ok(ExprSimplifyResult::Simplified(Expr::Literal(
- ScalarValue::new_one(&info.get_data_type(number)?)?,
- )))
+ Ok(ExprSimplifyResult::Simplified(lit(ScalarValue::new_one(
+ &number_datatype,
+ )?)))
} else {
+ let args = match num_args {
+ 1 => vec![number],
+ 2 => vec![number, base],
+ _ => {
+ return internal_err!(
+ "Unexpected number of arguments in log::simplify"
+ )
+ }
+ };
Ok(ExprSimplifyResult::Original(args))
}
}
@@ -191,6 +206,19 @@ impl ScalarUDFImpl for LogFunc {
}
}
+/// Returns true if the function is `PowerFunc`
+fn is_pow(func_def: &ScalarFunctionDefinition) -> bool {
+ if let ScalarFunctionDefinition::UDF(fun) = func_def {
+ fun.as_ref()
+ .inner()
+ .as_any()
+ .downcast_ref::<PowerFunc>()
+ .is_some()
+ } else {
+ false
+ }
+}
+
#[cfg(test)]
mod tests {
use datafusion_common::cast::{as_float32_array, as_float64_array};
diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs
index 8e3b2cf024..29caa63a94 100644
--- a/datafusion/functions/src/math/power.rs
+++ b/datafusion/functions/src/math/power.rs
@@ -18,7 +18,9 @@
//! Math function: `power()`.
use arrow::datatypes::DataType;
-use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_common::{
+ exec_err, plan_datafusion_err, DataFusionError, Result, ScalarValue,
+};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionDefinition};
@@ -118,43 +120,50 @@ impl ScalarUDFImpl for PowerFunc {
/// 3. Power(a, Log(a, b)) ===> b
fn simplify(
&self,
- args: Vec<Expr>,
+ mut args: Vec<Expr>,
info: &dyn SimplifyInfo,
) -> Result<ExprSimplifyResult> {
- let base = &args[0];
- let exponent = &args[1];
-
+ let exponent = args.pop().ok_or_else(|| {
+ plan_datafusion_err!("Expected power to have 2 arguments, got 0")
+ })?;
+ let base = args.pop().ok_or_else(|| {
+ plan_datafusion_err!("Expected power to have 2 arguments, got 1")
+ })?;
+
+ let exponent_type = info.get_data_type(&exponent)?;
match exponent {
- Expr::Literal(value)
- if value == &ScalarValue::new_zero(&info.get_data_type(exponent)?)? =>
- {
+ Expr::Literal(value) if value == ScalarValue::new_zero(&exponent_type)? => {
Ok(ExprSimplifyResult::Simplified(Expr::Literal(
- ScalarValue::new_one(&info.get_data_type(base)?)?,
+ ScalarValue::new_one(&info.get_data_type(&base)?)?,
)))
}
- Expr::Literal(value)
- if value == &ScalarValue::new_one(&info.get_data_type(exponent)?)? =>
- {
- Ok(ExprSimplifyResult::Simplified(base.clone()))
+ Expr::Literal(value) if value == ScalarValue::new_one(&exponent_type)? => {
+ Ok(ExprSimplifyResult::Simplified(base))
}
- Expr::ScalarFunction(ScalarFunction {
- func_def: ScalarFunctionDefinition::UDF(fun),
- args,
- }) if base == &args[0]
- && fun
- .as_ref()
- .inner()
- .as_any()
- .downcast_ref::<LogFunc>()
- .is_some() =>
+ Expr::ScalarFunction(ScalarFunction { func_def, mut args })
+ if is_log(&func_def) && args.len() == 2 && base == args[0] =>
{
- Ok(ExprSimplifyResult::Simplified(args[1].clone()))
+ let b = args.pop().unwrap(); // length checked above
+ Ok(ExprSimplifyResult::Simplified(b))
}
- _ => Ok(ExprSimplifyResult::Original(args)),
+ _ => Ok(ExprSimplifyResult::Original(vec![base, exponent])),
}
}
}
+/// Return true if this function call is a call to `Log`
+fn is_log(func_def: &ScalarFunctionDefinition) -> bool {
+ if let ScalarFunctionDefinition::UDF(fun) = func_def {
+ fun.as_ref()
+ .inner()
+ .as_any()
+ .downcast_ref::<LogFunc>()
+ .is_some()
+ } else {
+ false
+ }
+}
+
#[cfg(test)]
mod tests {
use datafusion_common::cast::{as_float64_array, as_int64_array};