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

feat: BuiltinScalarFunction::Cbrt #5839

Merged
merged 6 commits into from
Apr 5, 2023
Merged
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
13 changes: 13 additions & 0 deletions datafusion/core/tests/sql/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -676,6 +676,7 @@ async fn test_boolean_expressions() -> Result<()> {
#[tokio::test]
async fn test_mathematical_expressions_with_null() -> Result<()> {
test_expression!("sqrt(NULL)", "NULL");
test_expression!("cbrt(NULL)", "NULL");
test_expression!("sin(NULL)", "NULL");
test_expression!("cos(NULL)", "NULL");
test_expression!("tan(NULL)", "NULL");
Expand Down Expand Up @@ -1629,6 +1630,18 @@ async fn csv_query_sqrt_sqrt() -> Result<()> {
Ok(())
}

#[tokio::test]
async fn csv_query_cbrt_cbrt() -> Result<()> {
let ctx = create_ctx();
register_aggregate_csv(&ctx).await?;
let sql = "SELECT cbrt(cbrt(c12)) FROM aggregate_test_100 LIMIT 1";
let actual = execute(&ctx, sql).await;
// cbrt(cbrt(c12=0.9294097332465232)) = 0.9918990366780552
let expected = vec![vec!["0.9918990366780552"]];
assert_float_eq(&expected, &actual);
Ok(())
}

#[tokio::test]
async fn nested_subquery() -> Result<()> {
let ctx = SessionContext::new();
Expand Down
4 changes: 4 additions & 0 deletions datafusion/expr/src/built_in_function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ pub enum BuiltinScalarFunction {
Sin,
/// sqrt
Sqrt,
/// cbrt
Cbrt,
/// tan
Tan,
/// trunc
Expand Down Expand Up @@ -212,6 +214,7 @@ impl BuiltinScalarFunction {
BuiltinScalarFunction::Signum => Volatility::Immutable,
BuiltinScalarFunction::Sin => Volatility::Immutable,
BuiltinScalarFunction::Sqrt => Volatility::Immutable,
BuiltinScalarFunction::Cbrt => Volatility::Immutable,
BuiltinScalarFunction::Tan => Volatility::Immutable,
BuiltinScalarFunction::Trunc => Volatility::Immutable,
BuiltinScalarFunction::MakeArray => Volatility::Immutable,
Expand Down Expand Up @@ -304,6 +307,7 @@ impl FromStr for BuiltinScalarFunction {
"signum" => BuiltinScalarFunction::Signum,
"sin" => BuiltinScalarFunction::Sin,
"sqrt" => BuiltinScalarFunction::Sqrt,
"cbrt" => BuiltinScalarFunction::Cbrt,
"tan" => BuiltinScalarFunction::Tan,
"trunc" => BuiltinScalarFunction::Trunc,

Expand Down
2 changes: 2 additions & 0 deletions datafusion/expr/src/expr_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -446,6 +446,7 @@ macro_rules! nary_scalar_expr {

// math functions
scalar_expr!(Sqrt, sqrt, num, "square root of a number");
scalar_expr!(Cbrt, cbrt, num, "cube root of a number");
scalar_expr!(Sin, sin, num, "sine");
scalar_expr!(Cos, cos, num, "cosine");
scalar_expr!(Tan, tan, num, "tangent");
Expand Down Expand Up @@ -758,6 +759,7 @@ mod test {
#[test]
fn scalar_function_definitions() {
test_unary_scalar_expr!(Sqrt, sqrt);
test_unary_scalar_expr!(Cbrt, cbrt);
test_unary_scalar_expr!(Sin, sin);
test_unary_scalar_expr!(Cos, cos);
test_unary_scalar_expr!(Tan, tan);
Expand Down
1 change: 1 addition & 0 deletions datafusion/expr/src/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ pub fn return_type(
| BuiltinScalarFunction::Signum
| BuiltinScalarFunction::Sin
| BuiltinScalarFunction::Sqrt
| BuiltinScalarFunction::Cbrt
| BuiltinScalarFunction::Tan
| BuiltinScalarFunction::Trunc => match input_expr_types[0] {
DataType::Float32 => Ok(DataType::Float32),
Expand Down
1 change: 1 addition & 0 deletions datafusion/physical-expr/src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,7 @@ pub fn create_physical_fun(
BuiltinScalarFunction::Signum => Arc::new(math_expressions::signum),
BuiltinScalarFunction::Sin => Arc::new(math_expressions::sin),
BuiltinScalarFunction::Sqrt => Arc::new(math_expressions::sqrt),
BuiltinScalarFunction::Cbrt => Arc::new(math_expressions::cbrt),
BuiltinScalarFunction::Tan => Arc::new(math_expressions::tan),
BuiltinScalarFunction::Trunc => Arc::new(math_expressions::trunc),
BuiltinScalarFunction::Power => {
Expand Down
1 change: 1 addition & 0 deletions datafusion/physical-expr/src/math_expressions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ macro_rules! make_function_inputs2 {
}

math_unary_function!("sqrt", sqrt);
math_unary_function!("cbrt", cbrt);
math_unary_function!("sin", sin);
math_unary_function!("cos", cos);
math_unary_function!("tan", tan);
Expand Down
1 change: 1 addition & 0 deletions datafusion/proto/proto/datafusion.proto
Original file line number Diff line number Diff line change
Expand Up @@ -523,6 +523,7 @@ enum ScalarFunction {
CurrentDate = 70;
CurrentTime = 71;
Uuid = 72;
Cbrt = 73;
}

message ScalarFunctionNode {
Expand Down
3 changes: 3 additions & 0 deletions datafusion/proto/src/generated/pbjson.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions datafusion/proto/src/generated/prost.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 3 additions & 1 deletion datafusion/proto/src/logical_plan/from_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use datafusion_common::{
ScalarValue,
};
use datafusion_expr::{
abs, acos, array, ascii, asin, atan, atan2, bit_length, btrim, ceil,
abs, acos, array, ascii, asin, atan, atan2, bit_length, btrim, cbrt, ceil,
character_length, chr, coalesce, concat_expr, concat_ws_expr, cos, date_bin,
date_part, date_trunc, digest, exp,
expr::{self, Sort, WindowFunction},
Expand Down Expand Up @@ -400,6 +400,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
use protobuf::ScalarFunction;
match f {
ScalarFunction::Sqrt => Self::Sqrt,
ScalarFunction::Cbrt => Self::Cbrt,
ScalarFunction::Sin => Self::Sin,
ScalarFunction::Cos => Self::Cos,
ScalarFunction::Tan => Self::Tan,
Expand Down Expand Up @@ -1129,6 +1130,7 @@ pub fn parse_expr(
.collect::<Result<Vec<_>, _>>()?,
)),
ScalarFunction::Sqrt => Ok(sqrt(parse_expr(&args[0], registry)?)),
ScalarFunction::Cbrt => Ok(cbrt(parse_expr(&args[0], registry)?)),
ScalarFunction::Sin => Ok(sin(parse_expr(&args[0], registry)?)),
ScalarFunction::Cos => Ok(cos(parse_expr(&args[0], registry)?)),
ScalarFunction::Tan => Ok(tan(parse_expr(&args[0], registry)?)),
Expand Down
1 change: 1 addition & 0 deletions datafusion/proto/src/logical_plan/to_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1251,6 +1251,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
fn try_from(scalar: &BuiltinScalarFunction) -> Result<Self, Self::Error> {
let scalar_function = match scalar {
BuiltinScalarFunction::Sqrt => Self::Sqrt,
BuiltinScalarFunction::Cbrt => Self::Cbrt,
BuiltinScalarFunction::Sin => Self::Sin,
BuiltinScalarFunction::Cos => Self::Cos,
BuiltinScalarFunction::Tan => Self::Tan,
Expand Down
1 change: 1 addition & 0 deletions docs/source/user-guide/expressions.md
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ expressions such as `col("a") + col("b")` to be used.
| signum(x) | sign of the argument (-1, 0, +1) |
| sin(x) | sine |
| sqrt(x) | square root |
| cbrt(x) | cube root |
| tan(x) | tangent |
| trunc(x) | truncate toward zero |

Expand Down
14 changes: 14 additions & 0 deletions docs/source/user-guide/sql/scalar_functions.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
- [signum](#signum)
- [sin](#sin)
- [sqrt](#sqrt)
- [cbrt](#cbrt)
- [tan](#tan)
- [trunc](#trunc)

Expand Down Expand Up @@ -275,6 +276,19 @@ sqrt(numeric_expression)

#### Arguments

- **numeric_expression**: Numeric expression to operate on.
Can be a constant, column, or function, and any combination of arithmetic operators.

### `cbrt`

Returns the cube root of a number.

```
cbrt(numeric_expression)
```

#### Arguments

- **numeric_expression**: Numeric expression to operate on.
Can be a constant, column, or function, and any combination of arithmetic operators.

Expand Down