Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid cloning in log::simplify and power::simplify #10086

Merged
merged 5 commits into from
Apr 18, 2024
Merged

Avoid cloning in log::simplify and power::simplify #10086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d05dc54
Avoid cloning in power::simplify
alamb Apr 8, 2024
c017057
Avoid cloning in log::simplify
alamb Apr 8, 2024
4f9a3f1
Merge remote-tracking branch 'apache/main' into alamb/less_cloning
alamb Apr 17, 2024
2ef63bc
Apply suggestions from code review
alamb Apr 17, 2024
9ca3408
Merge branch 'alamb/less_cloning' of github.com:alamb/arrow-datafusio…
alamb Apr 17, 2024
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
88 changes: 58 additions & 30 deletions datafusion/functions/src/math/log.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down Expand Up @@ -146,51 +151,74 @@ 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}");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be an exec error to match similar patterns elsewhere?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lets fix it in follow up. tbh we should unify error messaging for all builtin functions

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the idea here is that simplify gets called as part of planning (not execution) so planning makes sense to me. Though I am not sure what practical difference it makes in this case 🤔

}
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))
}
}
}
}
}

/// 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};
Expand Down
59 changes: 34 additions & 25 deletions datafusion/functions/src/math/power.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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};
Expand Down