Allow multiple root plans in plan modifier (#193)

koralium · Dec 14, 2023 · cb0b5a6 · cb0b5a6
1 parent d0b21f6
commit cb0b5a6
Showing 1 changed file with 26 additions and 22 deletions.
diff --git a/src/FlowtideDotNet.Substrait/PlanModifier.cs b/src/FlowtideDotNet.Substrait/PlanModifier.cs
@@ -18,12 +18,13 @@ namespace FlowtideDotNet.Substrait
 {
     public class PlanModifier
     {
-        private Plan? _rootPlan;
+        private List<Plan> _rootPlans;
         private Dictionary<string, Plan> _subplans;
         private List<string> _writeToTables;
 
         public PlanModifier()
         {
+            _rootPlans = new List<Plan>();
             _subplans = new Dictionary<string, Plan>();
             _writeToTables = new List<string>();
         }
@@ -44,14 +45,11 @@ public PlanModifier AddPlanAsView(string viewName, Plan plan)
 
         public PlanModifier AddRootPlan(Plan plan)
         {
-            if (_rootPlan != null)
-            {
-                throw new InvalidOperationException("A root plan has already been added.");
-            }
-            _rootPlan = plan;
+            _rootPlans.Add(plan);
             return this;
         }
 
+        [Obsolete("Use inserts with SQL instead")]
         public PlanModifier WriteToTable(string tableName) 
         {
             _writeToTables.Add(tableName);
@@ -60,7 +58,7 @@ public PlanModifier WriteToTable(string tableName)
 
         public Plan Modify()
         {
-            if (_rootPlan == null)
+            if (_rootPlans.Count == 0)
             {
                 throw new InvalidOperationException("No root plan has been added.");
             }
@@ -124,27 +122,33 @@ public Plan Modify()
                 }
             }
 
+
             var rootRelationId = -1;
             RootRelation? oldRootRel = null;
-            for (int i = 0; i < _rootPlan.Relations.Count; i++)
+            foreach (var rootPlan in _rootPlans)
             {
-                var relation = _rootPlan.Relations[i];
-                if (relation is RootRelation rootRelation)
+
+                for (int i = 0; i < rootPlan.Relations.Count; i++)
                 {
-                    oldRootRel = rootRelation;
-                    var modified = modifierVisitor.Visit(rootRelation.Input, default);
-                    //var modified = rootRelation.Input;
-                    rootRelationId = newPlan.Relations.Count;
-                    newPlan.Relations.Add(modified);
-                }
-                else
-                {
-                    var modified = modifierVisitor.Visit(relation, default);
-                    newPlan.Relations.Add(modified);
+                    var relation = rootPlan.Relations[i];
+                    if (relation is RootRelation rootRelation)
+                    {
+                        oldRootRel = rootRelation;
+                        var modified = modifierVisitor.Visit(rootRelation.Input, default);
+                        //var modified = rootRelation.Input;
+                        rootRelationId = newPlan.Relations.Count;
+                        newPlan.Relations.Add(modified);
+                    }
+                    else
+                    {
+                        var modified = modifierVisitor.Visit(relation, default);
+                        newPlan.Relations.Add(modified);
+                    }
                 }
             }
 
-            foreach(var write in _writeToTables)
+
+            foreach (var write in _writeToTables)
             {
                 newPlan.Relations.Add(new RootRelation()
                 {
@@ -180,7 +184,7 @@ public Plan Modify()
                     }
                 });
             }
-            
+
             return newPlan;
         }
     }