Fix tests in split_actor_pool_projects

Eventual-Inc · Aug 24, 2024 · 3d77075 · 3d77075
1 parent 16e77db
commit 3d77075
Showing 1 changed file with 3 additions and 10 deletions.
diff --git a/src/daft-plan/src/logical_optimization/rules/split_actor_pool_projects.rs b/src/daft-plan/src/logical_optimization/rules/split_actor_pool_projects.rs
@@ -834,18 +834,13 @@ mod tests {
         let expected = LogicalPlan::ActorPoolProject(ActorPoolProject::try_new(
             expected,
             vec![
-                col("a"),
                 create_stateful_udf(vec![col("a")])
                     .clone()
                     .alias(intermediate_name),
+                col("a"),
             ],
         )?)
         .arced();
-        let expected = LogicalPlan::Project(Project::try_new(
-            expected,
-            vec![col(intermediate_name), col("a")],
-        )?)
-        .arced();
         let expected = LogicalPlan::ActorPoolProject(ActorPoolProject::try_new(
             expected,
             vec![
@@ -933,11 +928,10 @@ mod tests {
         let expected = LogicalPlan::ActorPoolProject(ActorPoolProject::try_new(
             expected,
             vec![
-                col("a"), // TODO: This should be able to be pruned as well, but it seems Projection Pushdown isn't working as intended
-                col("b"),
                 create_stateful_udf(vec![col("a")])
                     .clone()
                     .alias(intermediate_name_0),
+                col("b"),
             ],
         )?)
         .arced();
@@ -1049,11 +1043,10 @@ mod tests {
         let expected = LogicalPlan::ActorPoolProject(ActorPoolProject::try_new(
             expected,
             vec![
-                col("a"), // TODO: This should be pruned by Projection Pushdown, but isn't for some reason
-                col("b"),
                 create_stateful_udf(vec![col("a")])
                     .clone()
                     .alias(intermediate_name_0),
+                col("b"),
             ],
         )?)
         .arced();