Skip to content

Implement monotonicity for ScalarUDF #8799

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

Merged
merged 2 commits into from
Jan 11, 2024
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
8 changes: 7 additions & 1 deletion datafusion-examples/examples/advanced_udf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ use arrow::datatypes::Float64Type;
use datafusion::error::Result;
use datafusion::prelude::*;
use datafusion_common::{internal_err, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature};
use datafusion_expr::{
ColumnarValue, FuncMonotonicity, ScalarUDF, ScalarUDFImpl, Signature,
};
use std::sync::Arc;

/// This example shows how to use the full ScalarUDFImpl API to implement a user
Expand Down Expand Up @@ -184,6 +186,10 @@ impl ScalarUDFImpl for PowUdf {
fn aliases(&self) -> &[String] {
&self.aliases
}

fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
Ok(Some(vec![Some(true)]))
}
}

/// In this example we register `PowUdf` as a user defined function
Expand Down
15 changes: 14 additions & 1 deletion datafusion/expr/src/udf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@
//! [`ScalarUDF`]: Scalar User Defined Functions

use crate::{
ColumnarValue, Expr, ReturnTypeFunction, ScalarFunctionImplementation, Signature,
ColumnarValue, Expr, FuncMonotonicity, ReturnTypeFunction,
ScalarFunctionImplementation, Signature,
};
use arrow::datatypes::DataType;
use datafusion_common::Result;
Expand Down Expand Up @@ -164,6 +165,13 @@ impl ScalarUDF {
let captured = self.inner.clone();
Arc::new(move |args| captured.invoke(args))
}

/// This function specifies monotonicity behaviors for User defined scalar functions.
///
/// See [`ScalarUDFImpl::monotonicity`] for more details.
pub fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
self.inner.monotonicity()
}
}

impl<F> From<F> for ScalarUDF
Expand Down Expand Up @@ -271,6 +279,11 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
fn aliases(&self) -> &[String] {
&[]
}

/// This function specifies monotonicity behaviors for User defined scalar functions.
fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
Ok(None)
}
}

/// ScalarUDF that adds an alias to the underlying function. It is better to
Expand Down
75 changes: 74 additions & 1 deletion datafusion/physical-expr/src/udf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,79 @@ pub fn create_physical_expr(
fun.fun(),
input_phy_exprs.to_vec(),
fun.return_type(&input_exprs_types)?,
None,
fun.monotonicity()?,
)))
}

#[cfg(test)]
mod tests {
use arrow::datatypes::Schema;
use arrow_schema::DataType;
use datafusion_common::Result;
use datafusion_expr::{
ColumnarValue, FuncMonotonicity, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};

use crate::ScalarFunctionExpr;

use super::create_physical_expr;

#[test]
fn test_functions() -> Result<()> {
#[derive(Debug, Clone)]
struct TestScalarUDF {
signature: Signature,
}

impl TestScalarUDF {
fn new() -> Self {
let signature =
Signature::exact(vec![DataType::Float64], Volatility::Immutable);

Self { signature }
}
}

impl ScalarUDFImpl for TestScalarUDF {
fn as_any(&self) -> &dyn std::any::Any {
self
}

fn name(&self) -> &str {
"my_fn"
}

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

fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Float64)
}

fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
unimplemented!("my_fn is not implemented")
}

fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
Ok(Some(vec![Some(true)]))
}
}

// create and register the udf
let udf = ScalarUDF::from(TestScalarUDF::new());

let p_expr = create_physical_expr(&udf, &[], &Schema::empty())?;

assert_eq!(
p_expr
.as_any()
.downcast_ref::<ScalarFunctionExpr>()
.unwrap()
.monotonicity(),
&Some(vec![Some(true)])
);

Ok(())
}
}