fix: LogFunc simplify swaps arguments

apache · May 2, 2024 · 7ce7f93 · 7ce7f93
1 parent 3b245ff
commit 7ce7f93
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 2 deletions.
diff --git a/datafusion/expr/src/simplify.rs b/datafusion/expr/src/simplify.rs
@@ -109,6 +109,7 @@ impl<'a> SimplifyInfo for SimplifyContext<'a> {
 }
 
 /// Was the expression simplified?
+#[derive(Debug)]
 pub enum ExprSimplifyResult {
     /// The function call was simplified to an entirely new Expr
     Simplified(Expr),

diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs
@@ -192,7 +192,7 @@ impl ScalarUDFImpl for LogFunc {
                 } else {
                     let args = match num_args {
                         1 => vec![number],
-                        2 => vec![number, base],
+                        2 => vec![base, number],
                         _ => {
                             return internal_err!(
                                 "Unexpected number of arguments in log::simplify"
@@ -220,7 +220,13 @@ fn is_pow(func_def: &ScalarFunctionDefinition) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use datafusion_common::cast::{as_float32_array, as_float64_array};
+    use std::collections::HashMap;
+
+    use datafusion_common::{
+        cast::{as_float32_array, as_float64_array},
+        DFSchema,
+    };
+    use datafusion_expr::{execution_props::ExecutionProps, simplify::SimplifyContext};
 
     use super::*;
 
@@ -283,4 +289,52 @@ mod tests {
             }
         }
     }
+    #[test]
+    // Test log() simplification errors
+    fn test_log_simplify_errors() {
+        let props = ExecutionProps::new();
+        let schema =
+            Arc::new(DFSchema::new_with_metadata(vec![], HashMap::new()).unwrap());
+        let context = SimplifyContext::new(&props).with_schema(schema);
+        // Expect 0 args to error
+        let _ = LogFunc::new().simplify(vec![], &context).unwrap_err();
+        // Expect 3 args to error
+        let _ = LogFunc::new()
+            .simplify(vec![lit(1), lit(2), lit(3)], &context)
+            .unwrap_err();
+    }
+
+    #[test]
+    // Test that non-simplifiable log() expressions are unchanged after simplification
+    fn test_log_simplify_original() {
+        let props = ExecutionProps::new();
+        let schema =
+            Arc::new(DFSchema::new_with_metadata(vec![], HashMap::new()).unwrap());
+        let context = SimplifyContext::new(&props).with_schema(schema);
+        // One argument with no simplifications
+        let result = LogFunc::new().simplify(vec![lit(2)], &context).unwrap();
+        match result {
+            ExprSimplifyResult::Simplified(_) => {
+                panic!("Expected ExprSimplifyResult::Original")
+            }
+            ExprSimplifyResult::Original(args) => {
+                assert_eq!(args.len(), 1);
+                assert_eq!(args[0], lit(2));
+            }
+        }
+        // Two arguments with no simplifications
+        let result = LogFunc::new()
+            .simplify(vec![lit(2), lit(3)], &context)
+            .unwrap();
+        match result {
+            ExprSimplifyResult::Simplified(_) => {
+                panic!("Expected ExprSimplifyResult::Original")
+            }
+            ExprSimplifyResult::Original(args) => {
+                assert_eq!(args.len(), 2);
+                assert_eq!(args[0], lit(2));
+                assert_eq!(args[1], lit(3));
+            }
+        }
+    }
 }