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.

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

Fix coalesce, struct and named_strct expr_fn function to take multiple arguments #10321

Merged
merged 3 commits into from
May 7, 2024

Conversation

alamb
Copy link
Contributor

@alamb alamb commented Apr 30, 2024

Which issue does this PR close?

Closes #10320

Rationale for this change

There was a regression in the coalesce function's signature due to lack of testing

What changes are included in this PR?

  1. Fix signature
  2. Add test coverage

Are these changes tested?

Yes, by existing CI and new tests

Are there any user-facing changes?

@github-actions github-actions bot added the core Core DataFusion crate label Apr 30, 2024
@@ -161,6 +161,181 @@ async fn test_fn_btrim_with_chars() -> Result<()> {
Ok(())
}

#[tokio::test]
Copy link
Contributor Author

Choose a reason for hiding this comment

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

While adding a test for coalesce I realized the other functions in core weren't covered so I added new tests for them too (except for get_field which needs a map / struct and I was too lazy to add one to the test fixture

(get_field, arg_1 arg_2, "Returns the value of the field with the given name from the struct"),
(coalesce, args, "Returns `coalesce(args...)`, which evaluates to the value of the first expr which is not NULL")
);
pub mod expr_fn {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

It is not possible to create an expr_fn that takes a Vec<Expr> using the export_functions macro, so follow the model @Omega359 used in the math module

pub mod expr_fn {
use datafusion_expr::Expr;
#[doc = "returns the absolute value of a given number"]
pub fn abs(num: Expr) -> Expr {
super::abs().call(vec![num])
}
#[doc = "returns the arc cosine or inverse cosine of a number"]
pub fn acos(num: Expr) -> Expr {
super::acos().call(vec![num])
}
#[doc = "returns inverse hyperbolic cosine"]
pub fn acosh(num: Expr) -> Expr {
super::acosh().call(vec![num])
}
#[doc = "returns the arc sine or inverse sine of a number"]
pub fn asin(num: Expr) -> Expr {
super::asin().call(vec![num])
}
#[doc = "returns inverse hyperbolic sine"]
pub fn asinh(num: Expr) -> Expr {
super::asinh().call(vec![num])
}
#[doc = "returns inverse tangent"]
pub fn atan(num: Expr) -> Expr {
super::atan().call(vec![num])
}
#[doc = "returns inverse tangent of a division given in the argument"]
pub fn atan2(y: Expr, x: Expr) -> Expr {
super::atan2().call(vec![y, x])
}
#[doc = "returns inverse hyperbolic tangent"]
pub fn atanh(num: Expr) -> Expr {
super::atanh().call(vec![num])
}
#[doc = "cube root of a number"]
pub fn cbrt(num: Expr) -> Expr {
super::cbrt().call(vec![num])
}
#[doc = "nearest integer greater than or equal to argument"]
pub fn ceil(num: Expr) -> Expr {
super::ceil().call(vec![num])
}
#[doc = "cosine"]
pub fn cos(num: Expr) -> Expr {
super::cos().call(vec![num])
}
#[doc = "hyperbolic cosine"]
pub fn cosh(num: Expr) -> Expr {
super::cosh().call(vec![num])
}
#[doc = "cotangent of a number"]
pub fn cot(num: Expr) -> Expr {
super::cot().call(vec![num])
}
#[doc = "converts radians to degrees"]
pub fn degrees(num: Expr) -> Expr {
super::degrees().call(vec![num])
}
#[doc = "exponential"]
pub fn exp(num: Expr) -> Expr {
super::exp().call(vec![num])
}
#[doc = "factorial"]
pub fn factorial(num: Expr) -> Expr {
super::factorial().call(vec![num])
}
#[doc = "nearest integer less than or equal to argument"]
pub fn floor(num: Expr) -> Expr {
super::floor().call(vec![num])
}
#[doc = "greatest common divisor"]
pub fn gcd(x: Expr, y: Expr) -> Expr {
super::gcd().call(vec![x, y])
}
#[doc = "returns true if a given number is +NaN or -NaN otherwise returns false"]
pub fn isnan(num: Expr) -> Expr {
super::isnan().call(vec![num])
}
#[doc = "returns true if a given number is +0.0 or -0.0 otherwise returns false"]
pub fn iszero(num: Expr) -> Expr {
super::iszero().call(vec![num])
}
#[doc = "least common multiple"]
pub fn lcm(x: Expr, y: Expr) -> Expr {
super::lcm().call(vec![x, y])
}
#[doc = "natural logarithm (base e) of a number"]
pub fn ln(num: Expr) -> Expr {
super::ln().call(vec![num])
}
#[doc = "logarithm of a number for a particular `base`"]
pub fn log(base: Expr, num: Expr) -> Expr {
super::log().call(vec![base, num])
}
#[doc = "base 2 logarithm of a number"]
pub fn log2(num: Expr) -> Expr {
super::log2().call(vec![num])
}
#[doc = "base 10 logarithm of a number"]
pub fn log10(num: Expr) -> Expr {
super::log10().call(vec![num])
}
#[doc = "returns x if x is not NaN otherwise returns y"]
pub fn nanvl(x: Expr, y: Expr) -> Expr {
super::nanvl().call(vec![x, y])
}
#[doc = "Returns an approximate value of π"]
pub fn pi() -> Expr {
super::pi().call(vec![])
}
#[doc = "`base` raised to the power of `exponent`"]
pub fn power(base: Expr, exponent: Expr) -> Expr {
super::power().call(vec![base, exponent])
}
#[doc = "converts degrees to radians"]
pub fn radians(num: Expr) -> Expr {
super::radians().call(vec![num])
}
#[doc = "Returns a random value in the range 0.0 <= x < 1.0"]
pub fn random() -> Expr {
super::random().call(vec![])
}
#[doc = "round to nearest integer"]
pub fn round(args: Vec<Expr>) -> Expr {
super::round().call(args)
}
#[doc = "sign of the argument (-1, 0, +1)"]
pub fn signum(num: Expr) -> Expr {
super::signum().call(vec![num])
}
#[doc = "sine"]
pub fn sin(num: Expr) -> Expr {
super::sin().call(vec![num])
}
#[doc = "hyperbolic sine"]
pub fn sinh(num: Expr) -> Expr {
super::sinh().call(vec![num])
}
#[doc = "square root of a number"]
pub fn sqrt(num: Expr) -> Expr {
super::sqrt().call(vec![num])
}
#[doc = "returns the tangent of a number"]
pub fn tan(num: Expr) -> Expr {
super::tan().call(vec![num])
}
#[doc = "returns the hyperbolic tangent of a number"]
pub fn tanh(num: Expr) -> Expr {
super::tanh().call(vec![num])
}
#[doc = "truncate toward zero, with optional precision"]
pub fn trunc(args: Vec<Expr>) -> Expr {
super::trunc().call(args)
}
}

Copy link
Contributor

Choose a reason for hiding this comment

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

I think it is possible to take Vec like what functions-array macro does

($UDF:ty, $EXPR_FN:ident, $DOC:expr , $SCALAR_UDF_FN:ident) => {
        paste::paste! {
            // "fluent expr_fn" style function
            #[doc = $DOC]
            pub fn $EXPR_FN(arg: Vec<Expr>) -> Expr {
                Expr::ScalarFunction(ScalarFunction::new_udf(
                    $SCALAR_UDF_FN(),
                    arg,
                ))
            }

Copy link
Contributor Author

Choose a reason for hiding this comment

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

What I could't figure out how to do was make the macro take both syntaxes

It would have to look something like

export_functions!(
    // create a function with arg_1 and arg_2 Expr arguments
    (arrow_cast, arg_1 arg_2, "returns arg_1 cast to the `arrow_type` given the second argument. This can be used to cast to a specific `arrow_type`."),
   /// create a function with a single  Vec<Expr> arg
    (coalesce, args, "Returns `coalesce(args...)`, which evaluates to the value of the first expr which is not NULL")
...
}

the macro definition is

macro_rules! export_functions {
($(($FUNC:ident, $($arg:ident)*, $DOC:expr)),*) => {
pub mod expr_fn {
$(
#[doc = $DOC]
/// Return $name(arg)
pub fn $FUNC($($arg: datafusion_expr::Expr),*) -> datafusion_expr::Expr {
super::$FUNC().call(vec![$($arg),*],)
}
)*
}
/// Return a list of all functions in this package
pub fn functions() -> Vec<std::sync::Arc<datafusion_expr::ScalarUDF>> {
vec![
$(
$FUNC(),
)*
]
}
};
}

there may be some way to do this in rust, but I could't figure it out

Copy link
Contributor

Choose a reason for hiding this comment

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

I spent a whole day trying to make it work as well and got lost in the depths of macro programming. I'm not sure it's worth the effort of trying to get it to work tbh - the alternative is not exactly hard work.

Copy link
Contributor

Choose a reason for hiding this comment

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

macro_rules! make_udf_function {
($UDF:ty, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr , $SCALAR_UDF_FN:ident) => {
paste::paste! {
// "fluent expr_fn" style function
#[doc = $DOC]
pub fn $EXPR_FN($($arg: Expr),*) -> Expr {
Expr::ScalarFunction(ScalarFunction::new_udf(
$SCALAR_UDF_FN(),
vec![$($arg),*],
))
}
/// Singleton instance of [`$UDF`], ensures the UDF is only created once
/// named STATIC_$(UDF). For example `STATIC_ArrayToString`
#[allow(non_upper_case_globals)]
static [< STATIC_ $UDF >]: std::sync::OnceLock<std::sync::Arc<datafusion_expr::ScalarUDF>> =
std::sync::OnceLock::new();
/// ScalarFunction that returns a [`ScalarUDF`] for [`$UDF`]
///
/// [`ScalarUDF`]: datafusion_expr::ScalarUDF
pub fn $SCALAR_UDF_FN() -> std::sync::Arc<datafusion_expr::ScalarUDF> {
[< STATIC_ $UDF >]
.get_or_init(|| {
std::sync::Arc::new(datafusion_expr::ScalarUDF::new_from_impl(
<$UDF>::new(),
))
})
.clone()
}
}
};
($UDF:ty, $EXPR_FN:ident, $DOC:expr , $SCALAR_UDF_FN:ident) => {
paste::paste! {
// "fluent expr_fn" style function
#[doc = $DOC]
pub fn $EXPR_FN(arg: Vec<Expr>) -> Expr {
Expr::ScalarFunction(ScalarFunction::new_udf(
$SCALAR_UDF_FN(),
arg,
))
}
/// Singleton instance of [`$UDF`], ensures the UDF is only created once
/// named STATIC_$(UDF). For example `STATIC_ArrayToString`
#[allow(non_upper_case_globals)]
static [< STATIC_ $UDF >]: std::sync::OnceLock<std::sync::Arc<datafusion_expr::ScalarUDF>> =
std::sync::OnceLock::new();
/// ScalarFunction that returns a [`ScalarUDF`] for [`$UDF`]
///
/// [`ScalarUDF`]: datafusion_expr::ScalarUDF
pub fn $SCALAR_UDF_FN() -> std::sync::Arc<datafusion_expr::ScalarUDF> {
[< STATIC_ $UDF >]
.get_or_init(|| {
std::sync::Arc::new(datafusion_expr::ScalarUDF::new_from_impl(
<$UDF>::new(),
))
})
.clone()
}
}
};
}

In array macro, we support two syntaxes, which is surprisingly easy

Copy link
Contributor

Choose a reason for hiding this comment

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

I was thinking of alamb#19.

However, it does not reuse the macro export_functions that accepts multiple functions with one macro call.

export_functions_single accepts a single function with each macro call.

Since this requires introducing another macro, it is also not an ideal solution.

I'm fine to move on with the current PR.

Copy link
Contributor

@jayzhan211 jayzhan211 May 3, 2024

Choose a reason for hiding this comment

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

https://users.rust-lang.org/t/macro-repetition-with-multiple-rules/110816/2?u=jayzhan

I got a probably better solution (I had not tried it) from rust forum. To anyone that is interested in

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thank you @jayzhan211 -- I will give it a try

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 couldn't make this work (though I didn't try very hard) so I filed #10397 as a follow on task for this

Copy link
Contributor

Choose a reason for hiding this comment

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

I think we can merge this first and release 38.

}

/// Returns `coalesce(args...)`, which evaluates to the value of the first expr which is not NULL
pub fn coalesce(args: Vec<Expr>) -> Expr {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is the actual fix, though I suspect that struct and named_struct are also affected

@alamb alamb changed the title Fix coalesce expr_fn function to take multiple arguments Fix coalesce, struct and named_strct expr_fn function to take multiple arguments Apr 30, 2024
Copy link
Contributor

@appletreeisyellow appletreeisyellow left a comment

Choose a reason for hiding this comment

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

🧹 Thank you for fixing them! With the added tests, we have more confidence with future changes to these functions

@alamb alamb mentioned this pull request May 3, 2024
@alamb alamb merged commit 40a2055 into apache:main May 7, 2024
23 checks passed
@alamb
Copy link
Contributor Author

alamb commented May 7, 2024

Thanks @jayzhan211

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
core Core DataFusion crate
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Regression with coalesce expr_fn function: can't take multiple arguments
4 participants