Skip to content

Commit ecdfd02

Browse files
authored
Implement monotonicity for ScalarUDF (#8799)
* Implement monotonicity for ScalarUDF * add unit test case
1 parent 1f13607 commit ecdfd02

File tree

3 files changed

+95
-3
lines changed

3 files changed

+95
-3
lines changed

datafusion-examples/examples/advanced_udf.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@ use arrow::datatypes::Float64Type;
3131
use datafusion::error::Result;
3232
use datafusion::prelude::*;
3333
use datafusion_common::{internal_err, ScalarValue};
34-
use datafusion_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature};
34+
use datafusion_expr::{
35+
ColumnarValue, FuncMonotonicity, ScalarUDF, ScalarUDFImpl, Signature,
36+
};
3537
use std::sync::Arc;
3638

3739
/// This example shows how to use the full ScalarUDFImpl API to implement a user
@@ -184,6 +186,10 @@ impl ScalarUDFImpl for PowUdf {
184186
fn aliases(&self) -> &[String] {
185187
&self.aliases
186188
}
189+
190+
fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
191+
Ok(Some(vec![Some(true)]))
192+
}
187193
}
188194

189195
/// In this example we register `PowUdf` as a user defined function

datafusion/expr/src/udf.rs

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
//! [`ScalarUDF`]: Scalar User Defined Functions
1919
2020
use crate::{
21-
ColumnarValue, Expr, ReturnTypeFunction, ScalarFunctionImplementation, Signature,
21+
ColumnarValue, Expr, FuncMonotonicity, ReturnTypeFunction,
22+
ScalarFunctionImplementation, Signature,
2223
};
2324
use arrow::datatypes::DataType;
2425
use datafusion_common::Result;
@@ -164,6 +165,13 @@ impl ScalarUDF {
164165
let captured = self.inner.clone();
165166
Arc::new(move |args| captured.invoke(args))
166167
}
168+
169+
/// This function specifies monotonicity behaviors for User defined scalar functions.
170+
///
171+
/// See [`ScalarUDFImpl::monotonicity`] for more details.
172+
pub fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
173+
self.inner.monotonicity()
174+
}
167175
}
168176

169177
impl<F> From<F> for ScalarUDF
@@ -271,6 +279,11 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
271279
fn aliases(&self) -> &[String] {
272280
&[]
273281
}
282+
283+
/// This function specifies monotonicity behaviors for User defined scalar functions.
284+
fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
285+
Ok(None)
286+
}
274287
}
275288

276289
/// ScalarUDF that adds an alias to the underlying function. It is better to

datafusion/physical-expr/src/udf.rs

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,79 @@ pub fn create_physical_expr(
3939
fun.fun(),
4040
input_phy_exprs.to_vec(),
4141
fun.return_type(&input_exprs_types)?,
42-
None,
42+
fun.monotonicity()?,
4343
)))
4444
}
45+
46+
#[cfg(test)]
47+
mod tests {
48+
use arrow::datatypes::Schema;
49+
use arrow_schema::DataType;
50+
use datafusion_common::Result;
51+
use datafusion_expr::{
52+
ColumnarValue, FuncMonotonicity, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
53+
};
54+
55+
use crate::ScalarFunctionExpr;
56+
57+
use super::create_physical_expr;
58+
59+
#[test]
60+
fn test_functions() -> Result<()> {
61+
#[derive(Debug, Clone)]
62+
struct TestScalarUDF {
63+
signature: Signature,
64+
}
65+
66+
impl TestScalarUDF {
67+
fn new() -> Self {
68+
let signature =
69+
Signature::exact(vec![DataType::Float64], Volatility::Immutable);
70+
71+
Self { signature }
72+
}
73+
}
74+
75+
impl ScalarUDFImpl for TestScalarUDF {
76+
fn as_any(&self) -> &dyn std::any::Any {
77+
self
78+
}
79+
80+
fn name(&self) -> &str {
81+
"my_fn"
82+
}
83+
84+
fn signature(&self) -> &Signature {
85+
&self.signature
86+
}
87+
88+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
89+
Ok(DataType::Float64)
90+
}
91+
92+
fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
93+
unimplemented!("my_fn is not implemented")
94+
}
95+
96+
fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
97+
Ok(Some(vec![Some(true)]))
98+
}
99+
}
100+
101+
// create and register the udf
102+
let udf = ScalarUDF::from(TestScalarUDF::new());
103+
104+
let p_expr = create_physical_expr(&udf, &[], &Schema::empty())?;
105+
106+
assert_eq!(
107+
p_expr
108+
.as_any()
109+
.downcast_ref::<ScalarFunctionExpr>()
110+
.unwrap()
111+
.monotonicity(),
112+
&Some(vec![Some(true)])
113+
);
114+
115+
Ok(())
116+
}
117+
}

0 commit comments

Comments
 (0)