Add missing Substrait to DataFusion function name mappings (#16950)

* Add missing function mappings

* Added roundtrip test

* Quick fix

* Tests for all mappings, refactored logic

* Quick fix to log test

* Removed logb and sign from name mappings changes

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(),
+            "isnan" => "is_nan".to_string(),
             _ => function_name,
         };
 

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

@@ -41,12 +41,21 @@ pub async fn from_scalar_function(
             f.function_reference
         );
     };
+
     let fn_name = substrait_fun_name(fn_signature);
     let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?;
 
+    let udf_func = consumer.get_function_registry().udf(fn_name).or_else(|e| {
+        if let Some(alt_name) = substrait_to_df_name(fn_name) {
+            consumer.get_function_registry().udf(alt_name).or(Err(e))
+        } else {
+            Err(e)
+        }
+    });
+
     // try to first match the requested function into registered udfs, then built-in ops
     // and finally built-in expressions
-    if let Ok(func) = consumer.get_function_registry().udf(fn_name) {
+    if let Ok(func) = udf_func {
         Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
             func.to_owned(),
             args,
@@ -113,6 +122,13 @@ pub fn name_to_op(name: &str) -> Option<Operator> {
     }
 }
 
+pub fn substrait_to_df_name(name: &str) -> Option<&str> {
+    match name {
+        "is_nan" => Some("isnan"),
+        _ => None,
+    }
+}
+
 /// Build a balanced tree of binary operations from a binary operator and a list of arguments.
 ///
 /// For example, `OR` `(a, b, c, d, e)` will be converted to: `OR(OR(a, OR(b, c)), OR(d, e))`.
@@ -369,4 +385,39 @@ mod tests {
         assert_snapshot!(expr.to_string(), @"Int64(1) OR Int64(2) OR Int64(3) OR Int64(4)");
         Ok(())
     }
+
+    //Test that DataFusion can consume scalar functions that have a different name in Substrait
+    #[tokio::test]
+    async fn test_substrait_to_df_name_mapping() -> Result<()> {
+        // Build substrait extensions (we are using only one function)
+        let mut extensions = Extensions::default();
+        //is_nan is one of the functions that has a different name in Substrait (mapping is in substrait_to_df_name())
+        extensions.functions.insert(0, String::from("is_nan:fp32"));
+        // Build substrait consumer
+        let consumer = DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE);
+
+        // Build arguments for the function call
+        let arg = FunctionArgument {
+            arg_type: Some(ArgType::Value(Expression {
+                rex_type: Some(RexType::Literal(Literal {
+                    nullable: false,
+                    type_variation_reference: 0,
+                    literal_type: Some(LiteralType::Fp32(1.0)),
+                })),
+            })),
+        };
+        let arguments = vec![arg];
+        let func = ScalarFunction {
+            function_reference: 0,
+            arguments,
+            ..Default::default()
+        };
+        // Trivial input schema
+        let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]);
+        let df_schema = DFSchema::try_from(schema).unwrap();
+
+        // Consume the expression and ensure we don't get an error
+        let _ = consumer.consume_scalar_function(&func, &df_schema).await?;
+        Ok(())
+    }
 }

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -426,6 +426,41 @@ async fn simple_scalar_function_substr() -> Result<()> {
     roundtrip("SELECT SUBSTR(f, 1, 3) FROM data").await
 }
 
+// Test that DataFusion functions gets correctly mapped to Substrait names (when the names are diferent)
+// Follows the same structure as existing roundtrip tests, but more explicitly tests for name mappings
+async fn test_substrait_to_df_name_mapping(
+    substrait_name: &str,
+    sql: &str,
+) -> Result<()> {
+    let ctx = create_context().await?;
+    let df = ctx.sql(sql).await?;
+    let plan = df.into_optimized_plan()?;
+    let proto = to_substrait_plan(&plan, &ctx.state())?;
+
+    let function_name = match proto.extensions[0].mapping_type.as_ref().unwrap() {
+        MappingType::ExtensionFunction(ext_f) => &ext_f.name,
+        _ => unreachable!("Expected function extension"),
+    };
+
+    assert_eq!(function_name, substrait_name);
+
+    let plan2 = from_substrait_plan(&ctx.state(), &proto).await?;
+    let plan2 = ctx.state().optimize(&plan2)?;
+
+    let plan1str = format!("{plan}");
+    let plan2str = format!("{plan2}");
+    assert_eq!(plan1str, plan2str);
+
+    assert_eq!(plan.schema(), plan2.schema());
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn scalar_function_is_nan_mapping() -> Result<()> {
+    test_substrait_to_df_name_mapping("is_nan", "SELECT ISNAN(a) FROM data").await
+}
+
 #[tokio::test]
 async fn simple_scalar_function_is_null() -> Result<()> {
     roundtrip("SELECT * FROM data WHERE a IS NULL").await