alamb commented on code in PR #9304:
URL: https://github.com/apache/arrow-datafusion/pull/9304#discussion_r1510046782
##########
datafusion/optimizer/src/simplify_expressions/mod.rs:
##########
@@ -15,14 +15,15 @@
// specific language governing permissions and limitations
// under the License.
-pub mod context;
pub mod expr_simplifier;
mod guarantees;
mod inlist_simplifier;
mod regex;
pub mod simplify_exprs;
mod utils;
-pub use context::*;
+// backwards compatibility
+pub use datafusion_expr::simplify::{SimplifyContext, SimplifyInfo};
Review Comment:
WIth this change I think the PR is still backwards compatible
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -514,6 +519,97 @@ async fn deregister_udf() -> Result<()> {
Ok(())
}
+#[derive(Debug)]
+struct CastToI64UDF {
+ signature: Signature,
+}
+
+impl CastToI64UDF {
+ fn new() -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for CastToI64UDF {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "cast_to_i64"
+ }
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+ // Wrap with Expr::Cast() to Int64
+ fn simplify(
+ &self,
+ mut args: Vec<Expr>,
+ info: &dyn SimplifyInfo,
+ ) -> Result<ExprSimplifyResult> {
+ // Note that Expr::cast_to requires an ExprSchema but simplify gets a
Review Comment:
I couldn't figure out how to do `cast_to` but I think this way is OK too.
##########
datafusion/expr/src/udf.rs:
##########
@@ -338,6 +350,33 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
Ok(None)
}
+
+ /// Optionally apply per-UDF simplification / rewrite rules.
+ ///
+ /// This can be used to apply function specific simplification rules during
+ /// optimization (e.g. `arrow_cast` --> `Expr::Cast`). The default
+ /// implementation does nothing.
+ ///
+ /// Note that DataFusion handles simplifying arguments and "constant
+ /// folding" (replacing a function call with constant arguments such as
+ /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such
+ /// optimizations manually for specific UDFs.
+ ///
+ /// # Arguments
+ /// * 'args': The arguments of the function
+ /// * 'schema': The schema of the function
+ ///
+ /// # Returns
+ /// [`ExprSimplifyResult`] indicating the result of the simplification NOTE
+ /// if the function cannot be simplified, the arguments *MUST* be returned
+ /// unmodified
+ fn simplify(
Review Comment:
Here is the new API -- it takes owned args
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]