perf: Split projections with expressions that need granular batching (#4329)

Author: srilman

src/daft-logical-plan/src/builder/mod.rs

@@ -776,6 +776,7 @@ impl LogicalPlanBuilder {
                     |builder| builder.reorder_joins(),
                 )
                 .simplify_expressions()
+                .split_granular_projections()
                 .build();
 
             let optimized_plan = optimizer.optimize(
@@ -830,6 +831,7 @@ impl LogicalPlanBuilder {
                 |builder| builder.reorder_joins(),
             )
             .simplify_expressions()
+            .split_granular_projections()
             .build();
 
         let optimized_plan = optimizer.optimize(

src/daft-logical-plan/src/optimization/optimizer.rs

@@ -11,7 +11,7 @@ use super::{
         MaterializeScans, OptimizerRule, PushDownAntiSemiJoin, PushDownFilter,
         PushDownJoinPredicate, PushDownLimit, PushDownProjection, ReorderJoins,
         SimplifyExpressionsRule, SimplifyNullFilteredJoin, SplitActorPoolProjects,
-        UnnestPredicateSubquery, UnnestScalarSubquery,
+        SplitGranularProjection, UnnestPredicateSubquery, UnnestScalarSubquery,
     },
 };
 use crate::LogicalPlan;
@@ -208,6 +208,14 @@ impl OptimizerBuilder {
         self
     }
 
+    pub fn split_granular_projections(mut self) -> Self {
+        self.rule_batches.push(RuleBatch::new(
+            vec![Box::new(SplitGranularProjection::new())],
+            RuleExecutionStrategy::Once,
+        ));
+        self
+    }
+
     pub fn with_optimizer_config(mut self, config: OptimizerConfig) -> Self {
         self.config = config;
         self