pixeltable 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/exec/exec_node.py

@@ -1,13 +1,26 @@
 from __future__ import annotations
-from typing import Iterable, Optional, List
+
 import abc
+from typing import TYPE_CHECKING, Iterable, Iterator, List, Optional
+
+import pixeltable.exprs as exprs
 
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
-import pixeltable.exprs as exprs
+
+if TYPE_CHECKING:
+    from pixeltable import exec
 
 class ExecNode(abc.ABC):
     """Base class of all execution nodes"""
+    output_exprs: Iterable[exprs.Expr]
+    row_builder: exprs.RowBuilder
+    input: Optional[ExecNode]
+    flushed_img_slots: list[int]  # idxs of image slots of our output_exprs dependencies
+    stored_img_cols: list[exprs.ColumnSlotIdx]
+    ctx: Optional[ExecContext]
+    __iter: Optional[Iterator[DataRowBatch]]
+
     def __init__(
             self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr],
             input_exprs: Iterable[exprs.Expr], input: Optional[ExecNode] = None):
@@ -21,8 +34,9 @@ class ExecNode(abc.ABC):
             e.slot_idx for e in output_dependencies
             if e.col_type.is_image_type() and e.slot_idx not in output_slot_idxs
         ]
-        self.stored_img_cols: List[exprs.ColumnSlotIdx] = []
-        self.ctx: Optional[ExecContext] = None  # all nodes of a tree share the same context
+        self.stored_img_cols = []
+        self.ctx = None  # all nodes of a tree share the same context
+        self.__iter = None
 
     def set_ctx(self, ctx: ExecContext) -> None:
         self.ctx = ctx
@@ -35,12 +49,15 @@ class ExecNode(abc.ABC):
         if self.input is not None:
             self.input.set_stored_img_cols(stored_img_cols)
 
-    def __iter__(self):
+    # TODO: make this an abstractmethod when __next__() is removed
+    def __iter__(self) -> Iterator[DataRowBatch]:
         return self
 
-    @abc.abstractmethod
+    # TODO: remove this and switch every subclass over to implementing __iter__
     def __next__(self) -> DataRowBatch:
-        pass
+        if self.__iter is None:
+            self.__iter = iter(self)
+        return next(self.__iter)
 
     def open(self) -> None:
         """Bottom-up initialization of nodes for execution. Must be called before __next__."""
@@ -60,3 +77,15 @@ class ExecNode(abc.ABC):
     def _close(self) -> None:
         pass
 
+    def get_sql_node(self) -> Optional['exec.SqlNode']:
+        from .sql_node import SqlNode
+        if isinstance(self, SqlNode):
+            return self
+        if self.input is not None:
+            return self.input.get_sql_node()
+        return None
+
+    def set_limit(self, limit: int) -> None:
+        """Default implementation propagates to input"""
+        if self.input is not None:
+            self.input.set_limit(limit)

pixeltable/exec/expr_eval_node.py

@@ -5,10 +5,11 @@ import warnings
 from dataclasses import dataclass
 from typing import Iterable, List, Optional
 
-from tqdm import tqdm, TqdmWarning
+from tqdm import TqdmWarning, tqdm
 
-import pixeltable.exprs as exprs
+from pixeltable import exprs
 from pixeltable.func import CallableFunction
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
@@ -21,7 +22,7 @@ class ExprEvalNode(ExecNode):
     @dataclass
     class Cohort:
         """List of exprs that form an evaluation context and contain calls to at most one external function"""
-        exprs: List[exprs.Expr]
+        exprs_: List[exprs.Expr]
         batched_fn: Optional[CallableFunction]
         segment_ctxs: List['exprs.RowBuilder.EvalCtx']
         target_slot_idxs: List[int]
@@ -37,7 +38,7 @@ class ExprEvalNode(ExecNode):
         # we're only materializing exprs that are not already in the input
         self.target_exprs = [e for e in output_exprs if e.slot_idx not in input_slot_idxs]
         self.pbar: Optional[tqdm] = None
-        self.cohorts: List[List[ExprEvalNode.Cohort]] = []
+        self.cohorts: List[ExprEvalNode.Cohort] = []
         self._create_cohorts()
 
     def __next__(self) -> DataRowBatch:
@@ -87,6 +88,8 @@ class ExprEvalNode(ExecNode):
         for e in all_exprs:
             if not self._is_batched_fn_call(e):
                 continue
+            assert isinstance(e, exprs.FunctionCall)
+            assert isinstance(e.fn, CallableFunction)
             if current_batched_fn is None or current_batched_fn != e.fn:
                 # create a new cohort
                 cohorts.append([])
@@ -95,8 +98,8 @@ class ExprEvalNode(ExecNode):
 
         # expand the cohorts to include all exprs that are in the same evaluation context as the external calls;
         # cohorts are evaluated in order, so we can exclude the target slots from preceding cohorts and input slots
-        exclude = set([e.slot_idx for e in self.input_exprs])
-        all_target_slot_idxs = set([e.slot_idx for e in self.target_exprs])
+        exclude = set(e.slot_idx for e in self.input_exprs)
+        all_target_slot_idxs = set(e.slot_idx for e in self.target_exprs)
         target_slot_idxs: List[List[int]] = []  # the ones materialized by each cohort
         for i in range(len(cohorts)):
             cohorts[i] = self.row_builder.get_dependencies(
@@ -105,7 +108,7 @@ class ExprEvalNode(ExecNode):
                 [e.slot_idx for e in cohorts[i] if e.slot_idx in all_target_slot_idxs])
             exclude.update(target_slot_idxs[-1])
 
-        all_cohort_slot_idxs = set([e.slot_idx for cohort in cohorts for e in cohort])
+        all_cohort_slot_idxs = set(e.slot_idx for cohort in cohorts for e in cohort)
         remaining_slot_idxs = set(all_target_slot_idxs) - all_cohort_slot_idxs
         if len(remaining_slot_idxs) > 0:
             cohorts.append(self.row_builder.get_dependencies(
@@ -163,9 +166,10 @@ class ExprEvalNode(ExecNode):
                             rows[row_idx], segment_ctx, self.ctx.profile, ignore_errors=self.ctx.ignore_errors)
                 else:
                     fn_call = segment_ctx.exprs[0]
+                    assert isinstance(fn_call, exprs.FunctionCall)
                     # make a batched external function call
-                    arg_batches = [[] for _ in range(len(fn_call.args))]
-                    kwarg_batches = {k: [] for k in fn_call.kwargs.keys()}
+                    arg_batches: list[list[exprs.Expr]] = [[] for _ in range(len(fn_call.args))]
+                    kwarg_batches: dict[str, list[exprs.Expr]] = {k: [] for k in fn_call.kwargs.keys()}
 
                     valid_batch_idxs: List[int] = []  # rows with exceptions are not valid
                     for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
@@ -175,12 +179,15 @@ class ExprEvalNode(ExecNode):
                             continue
                         valid_batch_idxs.append(row_idx)
                         args, kwargs = fn_call._make_args(row)
-                        [arg_batches[i].append(args[i]) for i in range(len(args))]
-                        [kwarg_batches[k].append(kwargs[k]) for k in kwargs.keys()]
+                        for i in range(len(args)):
+                            arg_batches[i].append(args[i])
+                        for k in kwargs.keys():
+                            kwarg_batches[k].append(kwargs[k])
                     num_valid_batch_rows = len(valid_batch_idxs)
 
                     if ext_batch_size is None:
                         # we need to choose a batch size based on the args
+                        assert isinstance(fn_call.fn, CallableFunction)
                         sample_args = [arg_batches[i][0] for i in range(len(arg_batches))]
                         ext_batch_size = fn_call.fn.get_batch_size(*sample_args)
 
@@ -200,6 +207,7 @@ class ExprEvalNode(ExecNode):
                             for k in kwarg_batches.keys()
                         }
                         start_ts = time.perf_counter()
+                        assert isinstance(fn_call.fn, CallableFunction)
                         result_batch = fn_call.fn.exec_batch(*call_args, **call_kwargs)
                         self.ctx.profile.eval_time[fn_call.slot_idx] += time.perf_counter() - start_ts
                         self.ctx.profile.eval_count[fn_call.slot_idx] += num_ext_batch_rows

pixeltable/exec/in_memory_data_node.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Optional
+from typing import Any, Iterator, Optional
 
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
@@ -18,19 +18,26 @@ class InMemoryDataNode(ExecNode):
     - with the values provided in the input rows
     - if an input row doesn't provide a value, sets the slot to the column default
     """
+    tbl: catalog.TableVersion
+    input_rows: list[dict[str, Any]]
+    start_row_id: int
+    output_rows: Optional[DataRowBatch]
+
+    # output_exprs is declared in the superclass, but we redeclare it here with a more specific type
+    output_exprs: list[exprs.ColumnRef]
+
     def __init__(
         self, tbl: catalog.TableVersion, rows: list[dict[str, Any]],
         row_builder: exprs.RowBuilder, start_row_id: int,
     ):
-        # we materialize all output slots
-        output_exprs = [e for e in row_builder.get_output_exprs() if isinstance(e, exprs.ColumnRef)]
+        # we materialize the input slots
+        output_exprs = list(row_builder.input_exprs)
         super().__init__(row_builder, output_exprs, [], None)
         assert tbl.is_insertable()
         self.tbl = tbl
         self.input_rows = rows
         self.start_row_id = start_row_id
-        self.has_returned_data = False
-        self.output_rows: Optional[DataRowBatch] = None
+        self.output_rows = None
 
     def _open(self) -> None:
         """Create row batch and populate with self.input_rows"""
@@ -67,12 +74,8 @@ class InMemoryDataNode(ExecNode):
                 assert col_info is not None
                 self.output_rows[row_idx][col_info.slot_idx] = None
 
-        self.output_rows.set_row_ids([self.start_row_id + i for i in range(len(self.output_rows))])
         self.ctx.num_rows = len(self.output_rows)
 
-    def __next__(self) -> DataRowBatch:
-        if self.has_returned_data:
-            raise StopIteration
-        self.has_returned_data = True
+    def __iter__(self) -> Iterator[DataRowBatch]:
         _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_rows)} output_rows')
-        return self.output_rows
+        yield self.output_rows