Skip to content

Commit

Permalink
refactor: move acos() to function crate (#9297)
Browse files Browse the repository at this point in the history
  • Loading branch information
SteveLauC committed Feb 28, 2024
1 parent 544b3d9 commit 935ebca
Show file tree
Hide file tree
Showing 14 changed files with 240 additions and 145 deletions.
193 changes: 92 additions & 101 deletions datafusion-cli/Cargo.lock

Large diffs are not rendered by default.

10 changes: 2 additions & 8 deletions datafusion/expr/src/built_in_function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,6 @@ use strum_macros::EnumIter;
#[derive(Debug, Clone, PartialEq, Eq, Hash, EnumIter, Copy)]
pub enum BuiltinScalarFunction {
// math functions
/// acos
Acos,
/// asin
Asin,
/// atan
Expand Down Expand Up @@ -362,7 +360,6 @@ impl BuiltinScalarFunction {
pub fn volatility(&self) -> Volatility {
match self {
// Immutable scalar builtins
BuiltinScalarFunction::Acos => Volatility::Immutable,
BuiltinScalarFunction::Asin => Volatility::Immutable,
BuiltinScalarFunction::Atan => Volatility::Immutable,
BuiltinScalarFunction::Atan2 => Volatility::Immutable,
Expand Down Expand Up @@ -873,8 +870,7 @@ impl BuiltinScalarFunction {
utf8_to_int_type(&input_expr_types[0], "levenshtein")
}

BuiltinScalarFunction::Acos
| BuiltinScalarFunction::Asin
BuiltinScalarFunction::Asin
| BuiltinScalarFunction::Atan
| BuiltinScalarFunction::Acosh
| BuiltinScalarFunction::Asinh
Expand Down Expand Up @@ -1346,8 +1342,7 @@ impl BuiltinScalarFunction {
vec![Exact(vec![Utf8, Utf8]), Exact(vec![LargeUtf8, LargeUtf8])],
self.volatility(),
),
BuiltinScalarFunction::Acos
| BuiltinScalarFunction::Asin
BuiltinScalarFunction::Asin
| BuiltinScalarFunction::Atan
| BuiltinScalarFunction::Acosh
| BuiltinScalarFunction::Asinh
Expand Down Expand Up @@ -1438,7 +1433,6 @@ impl BuiltinScalarFunction {
/// Returns all names that can be used to call this function
pub fn aliases(&self) -> &'static [&'static str] {
match self {
BuiltinScalarFunction::Acos => &["acos"],
BuiltinScalarFunction::Acosh => &["acosh"],
BuiltinScalarFunction::Asin => &["asin"],
BuiltinScalarFunction::Asinh => &["asinh"],
Expand Down
2 changes: 0 additions & 2 deletions datafusion/expr/src/expr_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -531,7 +531,6 @@ scalar_expr!(Sinh, sinh, num, "hyperbolic sine");
scalar_expr!(Cosh, cosh, num, "hyperbolic cosine");
scalar_expr!(Tanh, tanh, num, "hyperbolic tangent");
scalar_expr!(Asin, asin, num, "inverse sine");
scalar_expr!(Acos, acos, num, "inverse cosine");
scalar_expr!(Atan, atan, num, "inverse tangent");
scalar_expr!(Asinh, asinh, num, "inverse hyperbolic sine");
scalar_expr!(Acosh, acosh, num, "inverse hyperbolic cosine");
Expand Down Expand Up @@ -1339,7 +1338,6 @@ mod test {
test_unary_scalar_expr!(Cosh, cosh);
test_unary_scalar_expr!(Tanh, tanh);
test_unary_scalar_expr!(Asin, asin);
test_unary_scalar_expr!(Acos, acos);
test_unary_scalar_expr!(Atan, atan);
test_unary_scalar_expr!(Asinh, asinh);
test_unary_scalar_expr!(Acosh, acosh);
Expand Down
110 changes: 110 additions & 0 deletions datafusion/functions/src/math/acos.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

//! Math function: `acos()`.

use arrow::array::{ArrayRef, Float32Array, Float64Array};
use arrow::datatypes::DataType;
use datafusion_common::{exec_err, plan_datafusion_err, DataFusionError, Result};
use datafusion_expr::ColumnarValue;
use datafusion_expr::{
utils::generate_signature_error_msg, ScalarUDFImpl, Signature, Volatility,
};
use std::any::Any;
use std::sync::Arc;

#[derive(Debug)]
pub struct AcosFunc {
signature: Signature,
}

impl AcosFunc {
pub fn new() -> Self {
use DataType::*;
Self {
signature: Signature::uniform(
1,
vec![Float64, Float32],
Volatility::Immutable,
),
}
}
}

impl ScalarUDFImpl for AcosFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"acos"
}

fn signature(&self) -> &Signature {
&self.signature
}

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
if arg_types.len() != 1 {
return Err(plan_datafusion_err!(
"{}",
generate_signature_error_msg(
self.name(),
self.signature().clone(),
arg_types,
)
));
}

let arg_type = &arg_types[0];

match arg_type {
DataType::Float64 => Ok(DataType::Float64),
DataType::Float32 => Ok(DataType::Float32),

// For other types (possible values null/int), use Float 64
_ => Ok(DataType::Float64),
}
}

fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
let args = ColumnarValue::values_to_arrays(args)?;

let arr: ArrayRef = match args[0].data_type() {
DataType::Float64 => Arc::new(make_function_scalar_inputs_return_type!(
&args[0],
self.name(),
Float64Array,
Float64Array,
{ f64::acos }
)),
DataType::Float32 => Arc::new(make_function_scalar_inputs_return_type!(
&args[0],
self.name(),
Float32Array,
Float32Array,
{ f32::acos }
)),
other => {
return exec_err!(
"Unsupported data type {other:?} for function {}",
self.name()
)
}
};
Ok(ColumnarValue::Array(arr))
}
}
19 changes: 15 additions & 4 deletions datafusion/functions/src/math/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,26 @@

//! "math" DataFusion functions

mod nans;
mod abs;
mod acos;
mod nans;

// create UDFs
make_udf_function!(nans::IsNanFunc, ISNAN, isnan);
make_udf_function!(abs::AbsFunc, ABS, abs);
make_udf_function!(acos::AcosFunc, ACOS, acos);

// Export the functions out of this package, both as expr_fn as well as a list of functions
export_functions!(
(isnan, num, "returns true if a given number is +NaN or -NaN otherwise returns false"),
(abs, num, "returns the absolute value of a given number")
);
(
isnan,
num,
"returns true if a given number is +NaN or -NaN otherwise returns false"
),
(abs, num, "returns the absolute value of a given number"),
(
acos,
num,
"returns the arc cosine or inverse cosine of a number"
)
);
2 changes: 1 addition & 1 deletion datafusion/functions/src/math/nans.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.

//! Encoding expressions
//! Math function: `isnan()`.

use arrow::datatypes::DataType;
use datafusion_common::{exec_err, DataFusionError, Result};
Expand Down
4 changes: 2 additions & 2 deletions datafusion/optimizer/src/analyzer/type_coercion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -889,14 +889,14 @@ mod test {
// test that automatic argument type coercion for scalar functions work
let empty = empty();
let lit_expr = lit(10i64);
let fun: BuiltinScalarFunction = BuiltinScalarFunction::Acos;
let fun: BuiltinScalarFunction = BuiltinScalarFunction::Floor;
let scalar_function_expr =
Expr::ScalarFunction(ScalarFunction::new(fun, vec![lit_expr]));
let plan = LogicalPlan::Projection(Projection::try_new(
vec![scalar_function_expr],
empty,
)?);
let expected = "Projection: acos(CAST(Int64(10) AS Float64))\n EmptyRelation";
let expected = "Projection: floor(CAST(Int64(10) AS Float64))\n EmptyRelation";
assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), &plan, expected)
}

Expand Down
1 change: 0 additions & 1 deletion datafusion/physical-expr/src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,6 @@ pub fn create_physical_fun(
) -> Result<ScalarFunctionImplementation> {
Ok(match fun {
// math functions
BuiltinScalarFunction::Acos => Arc::new(math_expressions::acos),
BuiltinScalarFunction::Asin => Arc::new(math_expressions::asin),
BuiltinScalarFunction::Atan => Arc::new(math_expressions::atan),
BuiltinScalarFunction::Acosh => Arc::new(math_expressions::acosh),
Expand Down
2 changes: 1 addition & 1 deletion datafusion/proto/proto/datafusion.proto
Original file line number Diff line number Diff line change
Expand Up @@ -548,7 +548,7 @@ enum ScalarFunction {
// 0 was Abs before
// The first enum value must be zero for open enums
unknown = 0;
Acos = 1;
// 1 was Acos
Asin = 2;
Atan = 3;
Ascii = 4;
Expand Down
3 changes: 0 additions & 3 deletions datafusion/proto/src/generated/pbjson.rs

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

4 changes: 1 addition & 3 deletions datafusion/proto/src/generated/prost.rs

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

20 changes: 9 additions & 11 deletions datafusion/proto/src/logical_plan/from_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,15 @@ use datafusion_common::{
use datafusion_expr::expr::Unnest;
use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by};
use datafusion_expr::{
acos, acosh, array, array_append, array_concat, array_dims, array_distinct,
array_element, array_empty, array_except, array_has, array_has_all, array_has_any,
array_intersect, array_length, array_ndims, array_pop_back, array_pop_front,
array_position, array_positions, array_prepend, array_remove, array_remove_all,
array_remove_n, array_repeat, array_replace, array_replace_all, array_replace_n,
array_resize, array_slice, array_sort, array_union, arrow_typeof, ascii, asin, asinh,
atan, atan2, atanh, bit_length, btrim, cardinality, cbrt, ceil, character_length,
chr, coalesce, concat_expr, concat_ws_expr, cos, cosh, cot, current_date,
current_time, date_bin, date_part, date_trunc, degrees, digest, ends_with, exp,
acosh, array, array_append, array_concat, array_dims, array_distinct, array_element,
array_empty, array_except, array_has, array_has_all, array_has_any, array_intersect,
array_length, array_ndims, array_pop_back, array_pop_front, array_position,
array_positions, array_prepend, array_remove, array_remove_all, array_remove_n,
array_repeat, array_replace, array_replace_all, array_replace_n, array_resize,
array_slice, array_sort, array_union, arrow_typeof, ascii, asin, asinh, atan, atan2,
atanh, bit_length, btrim, cardinality, cbrt, ceil, character_length, chr, coalesce,
concat_expr, concat_ws_expr, cos, cosh, cot, current_date, current_time, date_bin,
date_part, date_trunc, degrees, digest, ends_with, exp,
expr::{self, InList, Sort, WindowFunction},
factorial, find_in_set, flatten, floor, from_unixtime, gcd, gen_range, initcap,
instr, iszero, lcm, left, levenshtein, ln, log, log10, log2,
Expand Down Expand Up @@ -450,7 +450,6 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
ScalarFunction::Tan => Self::Tan,
ScalarFunction::Cot => Self::Cot,
ScalarFunction::Asin => Self::Asin,
ScalarFunction::Acos => Self::Acos,
ScalarFunction::Atan => Self::Atan,
ScalarFunction::Sinh => Self::Sinh,
ScalarFunction::Cosh => Self::Cosh,
Expand Down Expand Up @@ -1362,7 +1361,6 @@ pub fn parse_expr(
match scalar_function {
ScalarFunction::Unknown => Err(proto_error("Unknown scalar function")),
ScalarFunction::Asin => Ok(asin(parse_expr(&args[0], registry)?)),
ScalarFunction::Acos => Ok(acos(parse_expr(&args[0], registry)?)),
ScalarFunction::Asinh => Ok(asinh(parse_expr(&args[0], registry)?)),
ScalarFunction::Acosh => Ok(acosh(parse_expr(&args[0], registry)?)),
ScalarFunction::Array => Ok(array(
Expand Down
1 change: 0 additions & 1 deletion datafusion/proto/src/logical_plan/to_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1432,7 +1432,6 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
BuiltinScalarFunction::Cosh => Self::Cosh,
BuiltinScalarFunction::Tanh => Self::Tanh,
BuiltinScalarFunction::Asin => Self::Asin,
BuiltinScalarFunction::Acos => Self::Acos,
BuiltinScalarFunction::Atan => Self::Atan,
BuiltinScalarFunction::Asinh => Self::Asinh,
BuiltinScalarFunction::Acosh => Self::Acosh,
Expand Down
14 changes: 7 additions & 7 deletions datafusion/proto/tests/cases/roundtrip_physical_plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,6 @@
// under the License.

use arrow::csv::WriterBuilder;
use std::ops::Deref;
use std::sync::Arc;
use std::vec;

use datafusion::arrow::array::ArrayRef;
use datafusion::arrow::compute::kernels::sort::SortOptions;
use datafusion::arrow::datatypes::{DataType, Field, Fields, IntervalUnit, Schema};
Expand Down Expand Up @@ -52,6 +48,7 @@ use datafusion::physical_plan::expressions::{
StringAgg, Sum,
};
use datafusion::physical_plan::filter::FilterExec;
use datafusion::physical_plan::functions;
use datafusion::physical_plan::insert::FileSinkExec;
use datafusion::physical_plan::joins::{
HashJoinExec, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode,
Expand All @@ -66,7 +63,7 @@ use datafusion::physical_plan::windows::{
BuiltInWindowExpr, PlainAggregateWindowExpr, WindowAggExec,
};
use datafusion::physical_plan::{
functions, udaf, AggregateExpr, ExecutionPlan, Partitioning, PhysicalExpr, Statistics,
udaf, AggregateExpr, ExecutionPlan, Partitioning, PhysicalExpr, Statistics,
};
use datafusion::prelude::SessionContext;
use datafusion::scalar::ScalarValue;
Expand All @@ -82,6 +79,9 @@ use datafusion_expr::{
};
use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
use datafusion_proto::protobuf;
use std::ops::Deref;
use std::sync::Arc;
use std::vec;

/// Perform a serde roundtrip and assert that the string representation of the before and after plans
/// are identical. Note that this often isn't sufficient to guarantee that no information is
Expand Down Expand Up @@ -600,10 +600,10 @@ fn roundtrip_builtin_scalar_function() -> Result<()> {
let execution_props = ExecutionProps::new();

let fun_expr =
functions::create_physical_fun(&BuiltinScalarFunction::Acos, &execution_props)?;
functions::create_physical_fun(&BuiltinScalarFunction::Sin, &execution_props)?;

let expr = ScalarFunctionExpr::new(
"acos",
"sin",
fun_expr,
vec![col("a", &schema)?],
DataType::Float64,
Expand Down

0 comments on commit 935ebca

Please sign in to comment.