Support Substrait functions and_not, xor, and between in consumer built-in expression builder (#16984)

* Added support for functions with consumer's build in expression builder

* Format fix

* Support for logb

* Use built in functions for binary expressions

* Added a comment

* Quick fix

* Simplified producer logic

* quick fix

Author: lorenarosati

datafusion/substrait/src/extensions.rs

@@ -45,6 +45,7 @@ impl Extensions {
         // Rename those to match the Substrait extensions for interoperability
         let function_name = match function_name.as_str() {
             "substr" => "substring".to_string(),
+            "log" => "logb".to_string(),
             "isnan" => "is_nan".to_string(),
             _ => function_name,
         };

datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs

@@ -21,10 +21,9 @@ use datafusion::common::{
     not_impl_err, plan_err, substrait_err, DFSchema, DataFusionError, ScalarValue,
 };
 use datafusion::execution::FunctionRegistry;
-use datafusion::logical_expr::{expr, BinaryExpr, Expr, Like, Operator};
+use datafusion::logical_expr::{expr, Between, BinaryExpr, Expr, Like, Operator};
 use std::vec::Drain;
 use substrait::proto::expression::ScalarFunction;
-use substrait::proto::function_argument::ArgType;
 
 pub async fn from_scalar_function(
     consumer: &impl SubstraitConsumer,
@@ -70,7 +69,7 @@ pub async fn from_scalar_function(
         // In those cases we build a balanced tree of BinaryExprs
         arg_list_to_binary_op_tree(op, args)
     } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) {
-        builder.build(consumer, f, input_schema).await
+        builder.build(consumer, f, args).await
     } else {
         not_impl_err!("Unsupported function name: {fn_name:?}")
     }
@@ -180,7 +179,8 @@ impl BuiltinExprBuilder {
         match name {
             "not" | "like" | "ilike" | "is_null" | "is_not_null" | "is_true"
             | "is_false" | "is_not_true" | "is_not_false" | "is_unknown"
-            | "is_not_unknown" | "negative" | "negate" => Some(Self {
+            | "is_not_unknown" | "negative" | "negate" | "and_not" | "xor"
+            | "between" | "logb" => Some(Self {
                 expr_name: name.to_string(),
             }),
             _ => None,
@@ -191,37 +191,30 @@ impl BuiltinExprBuilder {
         self,
         consumer: &impl SubstraitConsumer,
         f: &ScalarFunction,
-        input_schema: &DFSchema,
+        args: Vec<Expr>,
     ) -> Result<Expr> {
         match self.expr_name.as_str() {
-            "like" => Self::build_like_expr(consumer, false, f, input_schema).await,
-            "ilike" => Self::build_like_expr(consumer, true, f, input_schema).await,
+            "like" => Self::build_like_expr(false, f, args).await,
+            "ilike" => Self::build_like_expr(true, f, args).await,
             "not" | "negative" | "negate" | "is_null" | "is_not_null" | "is_true"
             | "is_false" | "is_not_true" | "is_not_false" | "is_unknown"
-            | "is_not_unknown" => {
-                Self::build_unary_expr(consumer, &self.expr_name, f, input_schema).await
+            | "is_not_unknown" => Self::build_unary_expr(&self.expr_name, args).await,
+            "and_not" | "xor" => Self::build_binary_expr(&self.expr_name, args).await,
+            "between" => Self::build_between_expr(&self.expr_name, args).await,
+            "logb" => {
+                Self::build_custom_handling_expr(consumer, &self.expr_name, args).await
             }
             _ => {
                 not_impl_err!("Unsupported builtin expression: {}", self.expr_name)
             }
         }
     }
 
-    async fn build_unary_expr(
-        consumer: &impl SubstraitConsumer,
-        fn_name: &str,
-        f: &ScalarFunction,
-        input_schema: &DFSchema,
-    ) -> Result<Expr> {
-        if f.arguments.len() != 1 {
-            return substrait_err!("Expect one argument for {fn_name} expr");
-        }
-        let Some(ArgType::Value(expr_substrait)) = &f.arguments[0].arg_type else {
-            return substrait_err!("Invalid arguments type for {fn_name} expr");
+    async fn build_unary_expr(fn_name: &str, args: Vec<Expr>) -> Result<Expr> {
+        let [arg] = match args.try_into() {
+            Ok(args_arr) => args_arr,
+            Err(_) => return substrait_err!("Expected one argument for {fn_name} expr"),
         };
-        let arg = consumer
-            .consume_expression(expr_substrait, input_schema)
-            .await?;
         let arg = Box::new(arg);
 
         let expr = match fn_name {
@@ -242,40 +235,29 @@ impl BuiltinExprBuilder {
     }
 
     async fn build_like_expr(
-        consumer: &impl SubstraitConsumer,
         case_insensitive: bool,
         f: &ScalarFunction,
-        input_schema: &DFSchema,
+        args: Vec<Expr>,
     ) -> Result<Expr> {
         let fn_name = if case_insensitive { "ILIKE" } else { "LIKE" };
-        if f.arguments.len() != 2 && f.arguments.len() != 3 {
+        if args.len() != 2 && args.len() != 3 {
             return substrait_err!("Expect two or three arguments for `{fn_name}` expr");
         }
 
-        let Some(ArgType::Value(expr_substrait)) = &f.arguments[0].arg_type else {
-            return substrait_err!("Invalid arguments type for `{fn_name}` expr");
+        let mut args_iter = args.into_iter();
+        let Some(expr) = args_iter.next() else {
+            return substrait_err!("Missing first argument for {fn_name} expression");
         };
-        let expr = consumer
-            .consume_expression(expr_substrait, input_schema)
-            .await?;
-        let Some(ArgType::Value(pattern_substrait)) = &f.arguments[1].arg_type else {
-            return substrait_err!("Invalid arguments type for `{fn_name}` expr");
+        let Some(pattern) = args_iter.next() else {
+            return substrait_err!("Missing second argument for {fn_name} expression");
         };
-        let pattern = consumer
-            .consume_expression(pattern_substrait, input_schema)
-            .await?;
 
         // Default case: escape character is Literal(Utf8(None))
         let escape_char = if f.arguments.len() == 3 {
-            let Some(ArgType::Value(escape_char_substrait)) = &f.arguments[2].arg_type
-            else {
-                return substrait_err!("Invalid arguments type for `{fn_name}` expr");
+            let Some(escape_char_expr) = args_iter.next() else {
+                return substrait_err!("Missing third argument for {fn_name} expression");
             };
 
-            let escape_char_expr = consumer
-                .consume_expression(escape_char_substrait, input_schema)
-                .await?;
-
             match escape_char_expr {
                 Expr::Literal(ScalarValue::Utf8(escape_char_string), _) => {
                     // Convert Option<String> to Option<char>
@@ -299,6 +281,80 @@ impl BuiltinExprBuilder {
             case_insensitive,
         }))
     }
+
+    async fn build_binary_expr(fn_name: &str, args: Vec<Expr>) -> Result<Expr> {
+        let [a, b] = match args.try_into() {
+            Ok(args_arr) => args_arr,
+            Err(_) => {
+                return substrait_err!("Expected two arguments for `{fn_name}` expr")
+            }
+        };
+        match fn_name {
+            "and_not" => Ok(Self::build_and_not_expr(a, b)),
+            "xor" => Ok(Self::build_xor_expr(a, b)),
+            _ => not_impl_err!("Unsupported builtin expression: {}", fn_name),
+        }
+    }
+
+    fn build_and_not_expr(a: Expr, b: Expr) -> Expr {
+        a.and(Expr::Not(Box::new(b)))
+    }
+
+    fn build_xor_expr(a: Expr, b: Expr) -> Expr {
+        let or_expr = a.clone().or(b.clone());
+        let and_expr = a.and(b);
+        Self::build_and_not_expr(or_expr, and_expr)
+    }
+
+    async fn build_between_expr(fn_name: &str, args: Vec<Expr>) -> Result<Expr> {
+        let [expression, low, high] = match args.try_into() {
+            Ok(args_arr) => args_arr,
+            Err(_) => {
+                return substrait_err!("Expected three arguments for `{fn_name}` expr")
+            }
+        };
+
+        Ok(Expr::Between(Between {
+            expr: Box::new(expression),
+            negated: false,
+            low: Box::new(low),
+            high: Box::new(high),
+        }))
+    }
+
+    //This handles any functions that require custom handling
+    async fn build_custom_handling_expr(
+        consumer: &impl SubstraitConsumer,
+        fn_name: &str,
+        args: Vec<Expr>,
+    ) -> Result<Expr> {
+        match fn_name {
+            "logb" => Self::build_logb_expr(consumer, args).await,
+            _ => not_impl_err!("Unsupported custom handled expression: {}", fn_name),
+        }
+    }
+
+    async fn build_logb_expr(
+        consumer: &impl SubstraitConsumer,
+        args: Vec<Expr>,
+    ) -> Result<Expr> {
+        if args.len() != 2 {
+            return substrait_err!("Expect two arguments for logb function");
+        }
+
+        let mut args = args;
+        args.swap(0, 1);
+
+        //The equivalent of logb in DataFusion is the log function (which has its arguments in reverse order)
+        if let Ok(func) = consumer.get_function_registry().udf("log") {
+            Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
+                func.to_owned(),
+                args,
+            )))
+        } else {
+            not_impl_err!("Unsupported function name: logb")
+        }
+    }
 }
 
 #[cfg(test)]

datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs

@@ -34,6 +34,8 @@ pub fn from_scalar_function(
         });
     }
 
+    let arguments = custom_argument_handler(fun.name(), arguments);
+
     let function_anchor = producer.register_function(fun.name().to_string());
     #[allow(deprecated)]
     Ok(Expression {
@@ -47,6 +49,25 @@ pub fn from_scalar_function(
     })
 }
 
+// Handle functions that require custom handling for their arguments (e.g. log)
+pub fn custom_argument_handler(
+    name: &str,
+    args: Vec<FunctionArgument>,
+) -> Vec<FunctionArgument> {
+    match name {
+        "log" => {
+            if args.len() == 2 {
+                let mut args = args;
+                args.swap(0, 1);
+                args
+            } else {
+                args
+            }
+        }
+        _ => args,
+    }
+}
+
 pub fn from_unary_expr(
     producer: &mut impl SubstraitProducer,
     expr: &Expr,

datafusion/substrait/tests/cases/builtin_expr_semantics_tests.rs

@@ -0,0 +1,124 @@
+// 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.
+
+//! There are some Substrait functions that are semantically equivalent to nested built-in expressions, such as xor:bool_bool and and_not:bool_bool
+//! This module tests that the semantics of these functions are correct roundtripped
+
+#[cfg(test)]
+mod tests {
+    use crate::utils::test::add_plan_schemas_to_ctx;
+    use datafusion::arrow::util::pretty;
+    use datafusion::common::Result;
+    use datafusion::prelude::DataFrame;
+    use datafusion::prelude::SessionContext;
+    use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+    use datafusion_substrait::logical_plan::producer::to_substrait_plan;
+    use std::fs::File;
+    use std::io::BufReader;
+    use substrait::proto::Plan;
+
+    // Helper function to test scalar function semantics and roundtrip conversion
+    async fn test_scalar_fn_semantics(
+        file_path: &str,
+        expected_results: Vec<&str>,
+    ) -> Result<()> {
+        let path = format!("tests/testdata/test_plans/{file_path}");
+        let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
+            File::open(path).expect("file not found"),
+        ))
+        .expect("failed to parse json");
+
+        let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto)?;
+        let plan = from_substrait_plan(&ctx.state(), &proto).await?;
+
+        // Test correct semantics of function
+        let df = DataFrame::new(ctx.state().clone(), plan.clone());
+        let results = df.collect().await?;
+        let pretty_results = pretty::pretty_format_batches(&results)?.to_string();
+        assert_eq!(
+            pretty_results.trim().lines().collect::<Vec<_>>(),
+            expected_results
+        );
+
+        // Test roundtrip semantics
+        let proto = to_substrait_plan(&plan, &ctx.state())?;
+        let plan2 = from_substrait_plan(&ctx.state(), &proto).await?;
+        let df2 = DataFrame::new(ctx.state().clone(), plan2.clone());
+        let results2 = df2.collect().await?;
+        let pretty_results2 = pretty::pretty_format_batches(&results2)?.to_string();
+        assert_eq!(
+            pretty_results2.trim().lines().collect::<Vec<_>>(),
+            expected_results
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_xor_semantics() -> Result<()> {
+        let expected = vec![
+            "+-------+-------+--------+",
+            "| a     | b     | result |",
+            "+-------+-------+--------+",
+            "| true  | true  | false  |",
+            "| true  | false | true   |",
+            "| false | true  | true   |",
+            "| false | false | false  |",
+            "+-------+-------+--------+",
+        ];
+
+        test_scalar_fn_semantics(
+            "scalar_fn_to_built_in_binary_expr_xor.substrait.json",
+            expected,
+        )
+        .await
+    }
+
+    #[tokio::test]
+    async fn test_and_not_semantics() -> Result<()> {
+        let expected = vec![
+            "+-------+-------+--------+",
+            "| a     | b     | result |",
+            "+-------+-------+--------+",
+            "| true  | true  | false  |",
+            "| true  | false | true   |",
+            "| false | true  | false  |",
+            "| false | false | false  |",
+            "+-------+-------+--------+",
+        ];
+
+        test_scalar_fn_semantics(
+            "scalar_fn_to_built_in_binary_expr_and_not.substrait.json",
+            expected,
+        )
+        .await
+    }
+
+    #[tokio::test]
+    async fn test_logb_semantics() -> Result<()> {
+        let expected = vec![
+            "+-------+------+--------+",
+            "| x     | base | result |",
+            "+-------+------+--------+",
+            "| 1.0   | 10.0 | 0.0    |",
+            "| 100.0 | 10.0 | 2.0    |",
+            "+-------+------+--------+",
+        ];
+
+        test_scalar_fn_semantics("scalar_fn_logb_expr.substrait.json", expected).await
+    }
+}