Fixes for working with functions in dataframes

datafusion/src/logical_plan/expr.rs

@@ -1564,7 +1564,7 @@ pub fn approx_distinct(expr: Expr) -> Expr {
 /// Create an convenience function representing a unary scalar function
 macro_rules! unary_scalar_expr {
     ($ENUM:ident, $FUNC:ident) => {
-        #[doc = "this scalar function is not documented yet"]
+        #[doc = concat!("Unary scalar function definition for ", stringify!($FUNC) ) ]
         pub fn $FUNC(e: Expr) -> Expr {
             Expr::ScalarFunction {
                 fun: functions::BuiltinScalarFunction::$ENUM,
@@ -1574,17 +1574,28 @@ macro_rules! unary_scalar_expr {
     };
 }
 
-/// Create an convenience function representing a binary scalar function
-macro_rules! binary_scalar_expr {
+macro_rules! scalar_expr {
+    ($ENUM:ident, $FUNC:ident, $($arg:ident),*) => {
+        #[doc = concat!("Scalar function definition for ", stringify!($FUNC) ) ]
+        pub fn $FUNC($($arg: Expr),*) -> Expr {            
+            Expr::ScalarFunction {
+                fun: functions::BuiltinScalarFunction::$ENUM,
+                args: vec![$($arg),*],
+            }
+        }
+    };  
+}
+
+macro_rules! nary_scalar_expr {
     ($ENUM:ident, $FUNC:ident) => {
-        #[doc = "this scalar function is not documented yet"]
-        pub fn $FUNC(arg1: Expr, arg2: Expr) -> Expr {
+        #[doc = concat!("Scalar function definition for ", stringify!($FUNC) ) ]
+        pub fn $FUNC(args: Vec<Expr>) -> Expr {            
             Expr::ScalarFunction {
                 fun: functions::BuiltinScalarFunction::$ENUM,
-                args: vec![arg1, arg2],
+                args: args,
             }
         }
-    };
+    };  
 }
 
 // generate methods for creating the supported unary/binary expressions
@@ -1610,44 +1621,46 @@ unary_scalar_expr!(Log10, log10);
 unary_scalar_expr!(Ln, ln);
 
 // string functions
-unary_scalar_expr!(Ascii, ascii);
-unary_scalar_expr!(BitLength, bit_length);
-unary_scalar_expr!(Btrim, btrim);
-unary_scalar_expr!(CharacterLength, character_length);
-unary_scalar_expr!(CharacterLength, length);
-unary_scalar_expr!(Chr, chr);
-unary_scalar_expr!(InitCap, initcap);
-unary_scalar_expr!(Left, left);
-unary_scalar_expr!(Lower, lower);
-unary_scalar_expr!(Lpad, lpad);
-unary_scalar_expr!(Ltrim, ltrim);
-unary_scalar_expr!(MD5, md5);
-unary_scalar_expr!(OctetLength, octet_length);
-unary_scalar_expr!(RegexpMatch, regexp_match);
-unary_scalar_expr!(RegexpReplace, regexp_replace);
-unary_scalar_expr!(Replace, replace);
-unary_scalar_expr!(Repeat, repeat);
-unary_scalar_expr!(Reverse, reverse);
-unary_scalar_expr!(Right, right);
-unary_scalar_expr!(Rpad, rpad);
-unary_scalar_expr!(Rtrim, rtrim);
-unary_scalar_expr!(SHA224, sha224);
-unary_scalar_expr!(SHA256, sha256);
-unary_scalar_expr!(SHA384, sha384);
-unary_scalar_expr!(SHA512, sha512);
-unary_scalar_expr!(SplitPart, split_part);
-unary_scalar_expr!(StartsWith, starts_with);
-unary_scalar_expr!(Strpos, strpos);
-unary_scalar_expr!(Substr, substr);
-unary_scalar_expr!(ToHex, to_hex);
-unary_scalar_expr!(Translate, translate);
-unary_scalar_expr!(Trim, trim);
-unary_scalar_expr!(Upper, upper);
+scalar_expr!(Ascii, ascii, string);
+scalar_expr!(BitLength, bit_length, string);
+nary_scalar_expr!(Btrim, btrim);
+// scalar_expr!(Btrim, btrim, string);
+// scalar_expr!(Btrim, btrim_chars, string, characters);
+scalar_expr!(CharacterLength, character_length, string);
+scalar_expr!(CharacterLength, length, string);
+scalar_expr!(Chr, chr, string);
+scalar_expr!(Digest, digest, string, algorithm);
+scalar_expr!(InitCap, initcap, string);
+scalar_expr!(Left, left, string, count);
+scalar_expr!(Lower, lower, string);
+nary_scalar_expr!(Lpad, lpad);
+scalar_expr!(Ltrim, ltrim, string);
+scalar_expr!(MD5, md5, string);
+scalar_expr!(OctetLength, octet_length, string);
+nary_scalar_expr!(RegexpMatch, regexp_match);
+nary_scalar_expr!(RegexpReplace, regexp_replace);
+scalar_expr!(Replace, replace, string, from, to);
+scalar_expr!(Repeat, repeat, string, count);
+scalar_expr!(Reverse, reverse, string);
+scalar_expr!(Right, right, string, count);
+nary_scalar_expr!(Rpad, rpad);
+scalar_expr!(Rtrim, rtrim, string);
+scalar_expr!(SHA224, sha224, string);
+scalar_expr!(SHA256, sha256, string);
+scalar_expr!(SHA384, sha384, string);
+scalar_expr!(SHA512, sha512, string);
+scalar_expr!(SplitPart, split_part, expr, delimiter, index);
+scalar_expr!(StartsWith, starts_with, string, characters);
+scalar_expr!(Strpos, strpos, string, substring);
+scalar_expr!(Substr, substr, string, position);
+scalar_expr!(ToHex, to_hex, string);
+scalar_expr!(Translate, translate, string, from, to);
+scalar_expr!(Trim, trim, string);
+scalar_expr!(Upper, upper, string);
 
 // date functions
-binary_scalar_expr!(DatePart, date_part);
-binary_scalar_expr!(DateTrunc, date_trunc);
-binary_scalar_expr!(Digest, digest);
+scalar_expr!(DatePart, date_part, part, date);
+scalar_expr!(DateTrunc, date_trunc, part, date);
 
 /// returns an array of fixed size with each argument on it.
 pub fn array(args: Vec<Expr>) -> Expr {
@@ -2217,6 +2230,44 @@ mod tests {
         }};
     }
 
+    macro_rules! test_scalar_expr {
+        ($ENUM:ident, $FUNC:ident, $($arg:ident),*) => {
+            let expected = vec![$(stringify!($arg)),*];            
+            let result = $FUNC(
+                $(
+                    col(stringify!($arg.to_string()))
+                ),*    
+            );
+            if let Expr::ScalarFunction { fun, args } = result {
+                let name = functions::BuiltinScalarFunction::$ENUM;
+                assert_eq!(name, fun);
+                assert_eq!(expected.len(), args.len());
+            } else {
+                assert!(false, "unexpected");
+            }               
+        };  
+    }    
+
+    macro_rules! test_nary_scalar_expr {
+        ($ENUM:ident, $FUNC:ident, $($arg:ident),*) => {
+            let expected = vec![$(stringify!($arg)),*]; 
+            let result = $FUNC(
+                vec![
+                    $(
+                        col(stringify!($arg.to_string()))
+                    ),*    
+                ]
+            );
+            if let Expr::ScalarFunction { fun, args } = result {
+                let name = functions::BuiltinScalarFunction::$ENUM;
+                assert_eq!(name, fun);
+                assert_eq!(expected.len(), args.len());
+            } else {
+                assert!(false, "unexpected");
+            }               
+        };  
+    }    
+
     #[test]
     fn digest_function_definitions() {
         if let Expr::ScalarFunction { fun, args } = digest(col("tableA.a"), lit("md5")) {
@@ -2248,39 +2299,49 @@ mod tests {
         test_unary_scalar_expr!(Log2, log2);
         test_unary_scalar_expr!(Log10, log10);
         test_unary_scalar_expr!(Ln, ln);
-        test_unary_scalar_expr!(Ascii, ascii);
-        test_unary_scalar_expr!(BitLength, bit_length);
-        test_unary_scalar_expr!(Btrim, btrim);
-        test_unary_scalar_expr!(CharacterLength, character_length);
-        test_unary_scalar_expr!(CharacterLength, length);
-        test_unary_scalar_expr!(Chr, chr);
-        test_unary_scalar_expr!(InitCap, initcap);
-        test_unary_scalar_expr!(Left, left);
-        test_unary_scalar_expr!(Lower, lower);
-        test_unary_scalar_expr!(Lpad, lpad);
-        test_unary_scalar_expr!(Ltrim, ltrim);
-        test_unary_scalar_expr!(MD5, md5);
-        test_unary_scalar_expr!(OctetLength, octet_length);
-        test_unary_scalar_expr!(RegexpMatch, regexp_match);
-        test_unary_scalar_expr!(RegexpReplace, regexp_replace);
-        test_unary_scalar_expr!(Replace, replace);
-        test_unary_scalar_expr!(Repeat, repeat);
-        test_unary_scalar_expr!(Reverse, reverse);
-        test_unary_scalar_expr!(Right, right);
-        test_unary_scalar_expr!(Rpad, rpad);
-        test_unary_scalar_expr!(Rtrim, rtrim);
-        test_unary_scalar_expr!(SHA224, sha224);
-        test_unary_scalar_expr!(SHA256, sha256);
-        test_unary_scalar_expr!(SHA384, sha384);
-        test_unary_scalar_expr!(SHA512, sha512);
-        test_unary_scalar_expr!(SplitPart, split_part);
-        test_unary_scalar_expr!(StartsWith, starts_with);
-        test_unary_scalar_expr!(Strpos, strpos);
-        test_unary_scalar_expr!(Substr, substr);
-        test_unary_scalar_expr!(ToHex, to_hex);
-        test_unary_scalar_expr!(Translate, translate);
-        test_unary_scalar_expr!(Trim, trim);
-        test_unary_scalar_expr!(Upper, upper);
+
+        test_scalar_expr!(Ascii, ascii, input);
+        test_scalar_expr!(BitLength, bit_length, string);
+        test_nary_scalar_expr!(Btrim, btrim, string);
+        test_nary_scalar_expr!(Btrim, btrim, string, characters);
+        test_scalar_expr!(CharacterLength, character_length, string);
+        test_scalar_expr!(CharacterLength, length, string);
+        test_scalar_expr!(Chr, chr, string);
+        test_scalar_expr!(Digest, digest, string, algorithm);
+        test_scalar_expr!(InitCap, initcap, string);
+        test_scalar_expr!(Left, left, string, count);
+        test_scalar_expr!(Lower, lower, string);
+        test_nary_scalar_expr!(Lpad, lpad, string, count);
+        test_nary_scalar_expr!(Lpad, lpad, string, count, characters);
+        test_scalar_expr!(Ltrim, ltrim, string);
+        test_scalar_expr!(MD5, md5, string);
+        test_scalar_expr!(OctetLength, octet_length, string);
+        test_nary_scalar_expr!(RegexpMatch, regexp_match, string, pattern);
+        test_nary_scalar_expr!(RegexpMatch, regexp_match, string, pattern, flags);
+        test_nary_scalar_expr!(RegexpReplace, regexp_replace, string, pattern, replacement);
+        test_nary_scalar_expr!(RegexpReplace, regexp_replace, string, pattern, replacement, flags);
+        test_scalar_expr!(Replace, replace, string, from, to);
+        test_scalar_expr!(Repeat, repeat, string, count);
+        test_scalar_expr!(Reverse, reverse, string);
+        test_scalar_expr!(Right, right, string, count);
+        test_nary_scalar_expr!(Rpad, rpad, string, count);
+        test_nary_scalar_expr!(Rpad, rpad, string, count, characters);
+        test_scalar_expr!(Rtrim, rtrim, string);
+        test_scalar_expr!(SHA224, sha224, string);
+        test_scalar_expr!(SHA256, sha256, string);
+        test_scalar_expr!(SHA384, sha384, string);
+        test_scalar_expr!(SHA512, sha512, string);
+        test_scalar_expr!(SplitPart, split_part, expr, delimiter, index);
+        test_scalar_expr!(StartsWith, starts_with, string, characters);
+        test_scalar_expr!(Strpos, strpos, string, substring);
+        test_scalar_expr!(Substr, substr, string, position);
+        test_scalar_expr!(ToHex, to_hex, string);
+        test_scalar_expr!(Translate, translate, string, from, to);
+        test_scalar_expr!(Trim, trim, string);
+        test_scalar_expr!(Upper, upper, string);
+
+        test_scalar_expr!(DatePart, date_part, part, date);
+        test_scalar_expr!(DateTrunc, date_trunc, part, date);
     }
 
     #[test]

datafusion/src/physical_plan/functions.rs

@@ -2575,6 +2575,19 @@ mod tests {
             Int32,
             Int32Array
         );
+        // #[cfg(feature = "regex_expressions")]
+        // test_function!(
+        //     RegexpMatch,
+        //     &[
+        //         lit(ScalarValue::Utf8(Some("abc".to_string()))),
+        //         lit(ScalarValue::Utf8(Some("a..".to_string()))),
+        //         lit(ScalarValue::Utf8(Some("g".to_string()))),
+        //     ],
+        //     Ok(Some("abc")),
+        //     &str,
+        //     Utf8,
+        //     StringArray
+        // );        
         #[cfg(feature = "regex_expressions")]
         test_function!(
             RegexpReplace,

datafusion/src/prelude.rs

@@ -32,7 +32,7 @@ pub use crate::execution::options::{CsvReadOptions, NdJsonReadOptions};
 pub use crate::logical_plan::{
     array, ascii, avg, bit_length, btrim, character_length, chr, col, concat, concat_ws,
     count, create_udf, date_part, date_trunc, digest, in_list, initcap, left, length,
-    lit, lower, lpad, ltrim, max, md5, min, now, octet_length, random, regexp_replace,
+    lit, lower, lpad, ltrim, max, md5, min, now, octet_length, random, regexp_match, regexp_replace,
     repeat, replace, reverse, right, rpad, rtrim, sha224, sha256, sha384, sha512,
     split_part, starts_with, strpos, substr, sum, to_hex, translate, trim, upper, Column,
     JoinType, Partitioning,