assert first physical operator is DataSourcePhysicalOp & use source_operator.get_datasource()

Author: mdr223

src/palimpzest/query/execution/parallel_execution_strategy.py

@@ -8,11 +8,12 @@
 from palimpzest.core.lib.schemas import SourceRecord
 from palimpzest.query.execution.execution_strategy import ExecutionStrategy
 from palimpzest.query.operators.aggregate import AggregateOp
-from palimpzest.query.operators.datasource import MarshalAndScanDataOp
+from palimpzest.query.operators.datasource import DataSourcePhysicalOp
 from palimpzest.query.operators.limit import LimitScanOp
 from palimpzest.query.operators.physical import PhysicalOperator
 from palimpzest.query.optimizer.plan import PhysicalPlan
 
+g
 
 class PipelinedParallelExecutionStrategy(ExecutionStrategy):
     """
@@ -84,12 +85,9 @@ def execute_plan(self, plan: PhysicalPlan, num_samples: int | float = float("inf
 
         # get handle to DataSource and pre-compute its op_id and size
         source_operator = plan.operators[0]
+        assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
         source_op_id = source_operator.get_op_id()
-        datasource = (
-            source_operator.get_datasource()
-            if isinstance(source_operator, MarshalAndScanDataOp)
-            else self.datadir.get_cached_result(source_operator.dataset_id)
-        )
+        datasource = source_operator.get_datasource()
         datasource_len = len(datasource)
 
         # get limit of final limit operator (if one exists)

src/palimpzest/query/execution/single_threaded_execution_strategy.py

@@ -5,7 +5,7 @@
 from palimpzest.core.lib.schemas import SourceRecord
 from palimpzest.query.execution.execution_strategy import ExecutionStrategy
 from palimpzest.query.operators.aggregate import AggregateOp
-from palimpzest.query.operators.datasource import DataSourcePhysicalOp, MarshalAndScanDataOp
+from palimpzest.query.operators.datasource import DataSourcePhysicalOp
 from palimpzest.query.operators.filter import FilterOp
 from palimpzest.query.operators.limit import LimitScanOp
 from palimpzest.query.optimizer.plan import PhysicalPlan
@@ -48,11 +48,8 @@ def execute_plan(self, plan: PhysicalPlan, num_samples: int | float = float("inf
 
         # get handle to DataSource and pre-compute its size
         source_operator = plan.operators[0]
-        datasource = (
-            source_operator.get_datasource()
-            if isinstance(source_operator, MarshalAndScanDataOp)
-            else self.datadir.get_cached_result(source_operator.dataset_id)
-        )
+        assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
+        datasource = source_operator.get_datasource()
         datasource_len = len(datasource)
 
         # initialize processing queues for each operation
@@ -186,11 +183,8 @@ def execute_plan(self, plan: PhysicalPlan, num_samples: int | float = float("inf
 
         # get handle to DataSource and pre-compute its size
         source_operator = plan.operators[0]
-        datasource = (
-            source_operator.get_datasource()
-            if isinstance(source_operator, MarshalAndScanDataOp)
-            else self.datadir.get_cached_result(source_operator.dataset_id)
-        )
+        assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
+        datasource = source_operator.get_datasource()
         datasource_len = len(datasource)
 
         # initialize processing queues for each operation

src/palimpzest/query/processor/nosentinel_processor.py

@@ -92,11 +92,8 @@ def execute_plan(self, plan: PhysicalPlan, num_samples: int | float = float("inf
 
         # get handle to DataSource and pre-compute its size
         source_operator = plan.operators[0]
-        datasource = (
-            source_operator.get_datasource()
-            if isinstance(source_operator, MarshalAndScanDataOp)
-            else self.datadir.get_cached_result(source_operator.dataset_id)
-        )
+        assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
+        datasource = source_operator.get_datasource()
         datasource_len = len(datasource)
 
         # Calculate total work units - each record needs to go through each operator
@@ -272,11 +269,8 @@ def execute_plan(self, plan: PhysicalPlan, num_samples: int | float = float("inf
 
         # get handle to DataSource and pre-compute its size
         source_operator = plan.operators[0]
-        datasource = (
-            source_operator.get_datasource()
-            if isinstance(source_operator, MarshalAndScanDataOp)
-            else self.datadir.get_cached_result(source_operator.dataset_id)
-        )
+        assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
+        datasource = source_operator.get_datasource()
         datasource_len = len(datasource)
 
         # Calculate total work units - each record needs to go through each operator
@@ -468,11 +462,8 @@ def __init__(self, *args, **kwargs):
 
     #     # get handle to DataSource and pre-compute its size
     #     source_operator = plan.operators[0]
-    #     datasource = (
-    #         source_operator.get_datasource()
-    #         if isinstance(source_operator, MarshalAndScanDataOp)
-    #         else self.datadir.get_cached_result(source_operator.dataset_id)
-    #     )
+    #     assert isinstance(source_operator, DataSourcePhysicalOp), "First operator in physical plan must be a DataSourcePhysicalOp"
+    #     datasource = source_operator.get_datasource()
     #     datasource_len = len(datasource)
 
     #     # Calculate total work units - each record needs to go through each operator