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 8d428c386c feat: BuiltinScalarFunction::Cbrt (#5839)
8d428c386c is described below
commit 8d428c386cf113a4f0c2ac4f9b997ac549e49485
Author: Igor Izvekov <[email protected]>
AuthorDate: Wed Apr 5 15:50:28 2023 +0300
feat: BuiltinScalarFunction::Cbrt (#5839)
* feat: BuiltinScalarFunction::Cbrt
* fix: "ScalarFunction::Cbrt" covered in "datafusion.proto" and "pbjson.rs"
* feat: add documentation for "cbrt" function
* fix: use prettier
* fix: use regen.sh
* fix: remove useless docstring
---
datafusion/core/tests/sql/expr.rs | 13 +++++++++++++
datafusion/expr/src/built_in_function.rs | 4 ++++
datafusion/expr/src/expr_fn.rs | 2 ++
datafusion/expr/src/function.rs | 1 +
datafusion/physical-expr/src/functions.rs | 1 +
datafusion/physical-expr/src/math_expressions.rs | 1 +
datafusion/proto/proto/datafusion.proto | 1 +
datafusion/proto/src/generated/pbjson.rs | 3 +++
datafusion/proto/src/generated/prost.rs | 3 +++
datafusion/proto/src/logical_plan/from_proto.rs | 4 +++-
datafusion/proto/src/logical_plan/to_proto.rs | 1 +
docs/source/user-guide/expressions.md | 1 +
docs/source/user-guide/sql/scalar_functions.md | 14 ++++++++++++++
13 files changed, 48 insertions(+), 1 deletion(-)
diff --git a/datafusion/core/tests/sql/expr.rs
b/datafusion/core/tests/sql/expr.rs
index aceb09c299..d4930d631d 100644
--- a/datafusion/core/tests/sql/expr.rs
+++ b/datafusion/core/tests/sql/expr.rs
@@ -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");
@@ -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();
diff --git a/datafusion/expr/src/built_in_function.rs
b/datafusion/expr/src/built_in_function.rs
index 8ea96f185c..8280309087 100644
--- a/datafusion/expr/src/built_in_function.rs
+++ b/datafusion/expr/src/built_in_function.rs
@@ -66,6 +66,8 @@ pub enum BuiltinScalarFunction {
Sin,
/// sqrt
Sqrt,
+ /// cbrt
+ Cbrt,
/// tan
Tan,
/// trunc
@@ -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,
@@ -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,
diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
index abf0e1be98..9912cecde1 100644
--- a/datafusion/expr/src/expr_fn.rs
+++ b/datafusion/expr/src/expr_fn.rs
@@ -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");
@@ -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);
diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
index 03932ccc4d..81740f10ff 100644
--- a/datafusion/expr/src/function.rs
+++ b/datafusion/expr/src/function.rs
@@ -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),
diff --git a/datafusion/physical-expr/src/functions.rs
b/datafusion/physical-expr/src/functions.rs
index 035e72136f..f04b192bf0 100644
--- a/datafusion/physical-expr/src/functions.rs
+++ b/datafusion/physical-expr/src/functions.rs
@@ -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 => {
diff --git a/datafusion/physical-expr/src/math_expressions.rs
b/datafusion/physical-expr/src/math_expressions.rs
index ea89c2cd4b..2527e858e2 100644
--- a/datafusion/physical-expr/src/math_expressions.rs
+++ b/datafusion/physical-expr/src/math_expressions.rs
@@ -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);
diff --git a/datafusion/proto/proto/datafusion.proto
b/datafusion/proto/proto/datafusion.proto
index 76ec5b0017..4abee6d900 100644
--- a/datafusion/proto/proto/datafusion.proto
+++ b/datafusion/proto/proto/datafusion.proto
@@ -523,6 +523,7 @@ enum ScalarFunction {
CurrentDate = 70;
CurrentTime = 71;
Uuid = 72;
+ Cbrt = 73;
}
message ScalarFunctionNode {
diff --git a/datafusion/proto/src/generated/pbjson.rs
b/datafusion/proto/src/generated/pbjson.rs
index 406d9ee27a..97375681c3 100644
--- a/datafusion/proto/src/generated/pbjson.rs
+++ b/datafusion/proto/src/generated/pbjson.rs
@@ -17083,6 +17083,7 @@ impl serde::Serialize for ScalarFunction {
Self::CurrentDate => "CurrentDate",
Self::CurrentTime => "CurrentTime",
Self::Uuid => "Uuid",
+ Self::Cbrt => "Cbrt",
};
serializer.serialize_str(variant)
}
@@ -17167,6 +17168,7 @@ impl<'de> serde::Deserialize<'de> for ScalarFunction {
"CurrentDate",
"CurrentTime",
"Uuid",
+ "Cbrt",
];
struct GeneratedVisitor;
@@ -17282,6 +17284,7 @@ impl<'de> serde::Deserialize<'de> for ScalarFunction {
"CurrentDate" => Ok(ScalarFunction::CurrentDate),
"CurrentTime" => Ok(ScalarFunction::CurrentTime),
"Uuid" => Ok(ScalarFunction::Uuid),
+ "Cbrt" => Ok(ScalarFunction::Cbrt),
_ => Err(serde::de::Error::unknown_variant(value, FIELDS)),
}
}
diff --git a/datafusion/proto/src/generated/prost.rs
b/datafusion/proto/src/generated/prost.rs
index e5b4534f60..f4f48dfe41 100644
--- a/datafusion/proto/src/generated/prost.rs
+++ b/datafusion/proto/src/generated/prost.rs
@@ -2126,6 +2126,7 @@ pub enum ScalarFunction {
CurrentDate = 70,
CurrentTime = 71,
Uuid = 72,
+ Cbrt = 73,
}
impl ScalarFunction {
/// String value of the enum field names used in the ProtoBuf definition.
@@ -2207,6 +2208,7 @@ impl ScalarFunction {
ScalarFunction::CurrentDate => "CurrentDate",
ScalarFunction::CurrentTime => "CurrentTime",
ScalarFunction::Uuid => "Uuid",
+ ScalarFunction::Cbrt => "Cbrt",
}
}
/// Creates an enum from field names used in the ProtoBuf definition.
@@ -2285,6 +2287,7 @@ impl ScalarFunction {
"CurrentDate" => Some(Self::CurrentDate),
"CurrentTime" => Some(Self::CurrentTime),
"Uuid" => Some(Self::Uuid),
+ "Cbrt" => Some(Self::Cbrt),
_ => None,
}
}
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs
b/datafusion/proto/src/logical_plan/from_proto.rs
index 755b942e43..5d198d47a6 100644
--- a/datafusion/proto/src/logical_plan/from_proto.rs
+++ b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -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},
@@ -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,
@@ -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)?)),
diff --git a/datafusion/proto/src/logical_plan/to_proto.rs
b/datafusion/proto/src/logical_plan/to_proto.rs
index e8570cf3c7..c53fdf8aab 100644
--- a/datafusion/proto/src/logical_plan/to_proto.rs
+++ b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -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,
diff --git a/docs/source/user-guide/expressions.md
b/docs/source/user-guide/expressions.md
index 7717366d25..85918a4518 100644
--- a/docs/source/user-guide/expressions.md
+++ b/docs/source/user-guide/expressions.md
@@ -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 |
diff --git a/docs/source/user-guide/sql/scalar_functions.md
b/docs/source/user-guide/sql/scalar_functions.md
index 89fb567146..ea109664c5 100644
--- a/docs/source/user-guide/sql/scalar_functions.md
+++ b/docs/source/user-guide/sql/scalar_functions.md
@@ -39,6 +39,7 @@
- [signum](#signum)
- [sin](#sin)
- [sqrt](#sqrt)
+- [cbrt](#cbrt)
- [tan](#tan)
- [trunc](#trunc)
@@ -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.
