From 94b4d38c4e32289fbf8650abca871cb396197aeb Mon Sep 17 00:00:00 2001
From: Kev Wang <kevin@eventualcomputing.com>
Date: Mon, 5 Aug 2024 18:01:23 -0700
Subject: [PATCH] [BUG] Fix Expr::with_new_children for ScalarFunction (#2624)

---
 src/daft-dsl/src/expr.rs | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/daft-dsl/src/expr.rs b/src/daft-dsl/src/expr.rs
index e59d837a23..e96bd001ad 100644
--- a/src/daft-dsl/src/expr.rs
+++ b/src/daft-dsl/src/expr.rs
@@ -618,7 +618,10 @@ impl Expr {
         use Expr::*;
         match self {
             // no children
-            Column(..) | Literal(..) => self.clone(),
+            Column(..) | Literal(..) => {
+                assert!(children.is_empty(), "Should have no children");
+                self.clone()
+            }
             // 1 child
             Not(..) => Not(children.first().expect("Should have 1 child").clone()),
             Alias(.., name) => Alias(
@@ -658,11 +661,30 @@ impl Expr {
             },
             // N-ary
             Agg(agg_expr) => Agg(agg_expr.with_new_children(children)),
-            Function { func, .. } => Function {
-                func: func.clone(),
-                inputs: children,
-            },
-            ScalarFunction(sf) => ScalarFunction(sf.clone()),
+            Function {
+                func,
+                inputs: old_children,
+            } => {
+                assert!(
+                    children.len() == old_children.len(),
+                    "Should have same number of children"
+                );
+                Function {
+                    func: func.clone(),
+                    inputs: children,
+                }
+            }
+            ScalarFunction(sf) => {
+                assert!(
+                    children.len() == sf.inputs.len(),
+                    "Should have same number of children"
+                );
+
+                ScalarFunction(crate::functions::ScalarFunction {
+                    udf: sf.udf.clone(),
+                    inputs: children,
+                })
+            }
         }
     }