pixeltable 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

pixeltable/exec/component_iteration_node.py

@@ -1,10 +1,12 @@
-from typing import Generator, Optional
+import inspect
+from typing import Iterator, Optional
 
-from .data_row_batch import DataRowBatch
-from .exec_node import ExecNode
 import pixeltable.catalog as catalog
-import pixeltable.exprs as exprs
 import pixeltable.exceptions as excs
+import pixeltable.exprs as exprs
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
 
 
 class ComponentIterationNode(ExecNode):
@@ -12,7 +14,7 @@ class ComponentIterationNode(ExecNode):
 
     Returns row batches of OUTPUT_BATCH_SIZE size.
     """
-    OUTPUT_BATCH_SIZE = 1024
+    __OUTPUT_BATCH_SIZE = 1024
 
     def __init__(self, view: catalog.TableVersion, input: ExecNode):
         assert view.is_component_view()
@@ -23,57 +25,76 @@ class ComponentIterationNode(ExecNode):
         self.iterator_args = iterator_args[0]
         assert isinstance(self.iterator_args, exprs.InlineDict)
         self.iterator_args_ctx = self.row_builder.create_eval_ctx([self.iterator_args])
-        self.iterator_output_schema, self.unstored_column_names = \
+        self.iterator_output_schema, self.unstored_column_names = (
             self.view.iterator_cls.output_schema(**self.iterator_args.to_kwargs())
+        )
         self.iterator_output_fields = list(self.iterator_output_schema.keys())
-        self.iterator_output_cols = \
-            {field_name: self.view.cols_by_name[field_name] for field_name in self.iterator_output_fields}
+        self.iterator_output_cols = {
+            field_name: self.view.cols_by_name[field_name] for field_name in self.iterator_output_fields
+        }
         # referenced iterator output fields
         self.refd_output_slot_idxs = {
             e.col.name: e.slot_idx for e in self.row_builder.unique_exprs
             if isinstance(e, exprs.ColumnRef) and e.col.name in self.iterator_output_fields
         }
-        self._output: Optional[Generator[DataRowBatch, None, None]] = None
+        self.__output: Optional[Iterator[DataRowBatch]] = None
 
-    def _output_batches(self) -> Generator[DataRowBatch, None, None]:
+    def __output_batches(self) -> Iterator[DataRowBatch]:
         output_batch = DataRowBatch(self.view, self.row_builder)
         for input_batch in self.input:
             for input_row in input_batch:
                 self.row_builder.eval(input_row, self.iterator_args_ctx)
                 iterator_args = input_row[self.iterator_args.slot_idx]
-                iterator = self.view.iterator_cls(**iterator_args)
-                for pos, component_dict in enumerate(iterator):
-                    output_row = output_batch.add_row()
-                    input_row.copy(output_row)
-                    # we're expanding the input and need to add the iterator position to the pk
-                    pk = output_row.pk[:-1] + (pos,) + output_row.pk[-1:]
-                    output_row.set_pk(pk)
-
-                    # verify and copy component_dict fields to their respective slots in output_row
-                    for field_name, field_val in component_dict.items():
-                        if field_name not in self.iterator_output_fields:
-                            raise excs.Error(
-                                f'Invalid field name {field_name} in output of {self.view.iterator_cls.__name__}')
-                        if field_name not in self.refd_output_slot_idxs:
-                            # we can ignore this
-                            continue
-                        output_col = self.iterator_output_cols[field_name]
-                        output_col.col_type.validate_literal(field_val)
-                        output_row[self.refd_output_slot_idxs[field_name]] = field_val
-                    if len(component_dict) != len(self.iterator_output_fields):
-                        missing_fields = set(self.refd_output_slot_idxs.keys()) - set(component_dict.keys())
-                        raise excs.Error(
-                            f'Invalid output of {self.view.iterator_cls.__name__}: '
-                            f'missing fields {", ".join(missing_fields)}')
-
-                    if len(output_batch) == self.OUTPUT_BATCH_SIZE:
-                        yield output_batch
-                        output_batch = DataRowBatch(self.view, self.row_builder)
+                assert isinstance(iterator_args, dict)
+                # We need to ensure that all of the required (non-nullable) parameters of the iterator are
+                # specified and are not null. If any of them are null, then we skip this row (i.e., we emit 0
+                # output rows for this input row).
+                if self.__non_nullable_args_specified(iterator_args):
+                    iterator = self.view.iterator_cls(**iterator_args)
+                    for pos, component_dict in enumerate(iterator):
+                        output_row = output_batch.add_row()
+                        input_row.copy(output_row)
+                        # we're expanding the input and need to add the iterator position to the pk
+                        self.__populate_output_row(output_row, pos, component_dict)
+                        if len(output_batch) == self.__OUTPUT_BATCH_SIZE:
+                            yield output_batch
+                            output_batch = DataRowBatch(self.view, self.row_builder)
 
         if len(output_batch) > 0:
             yield output_batch
 
+    def __non_nullable_args_specified(self, iterator_args: dict) -> bool:
+        """
+        Returns true if all non-nullable iterator arguments are not `None`.
+        """
+        input_schema = self.view.iterator_cls.input_schema()
+        for arg_name, arg_value in iterator_args.items():
+            col_type = input_schema[arg_name]
+            if arg_value is None and not col_type.nullable:
+                return False
+        return True
+
+    def __populate_output_row(self, output_row: exprs.DataRow, pos: int, component_dict: dict) -> None:
+        pk = output_row.pk[:-1] + (pos,) + output_row.pk[-1:]
+        output_row.set_pk(pk)
+        # verify and copy component_dict fields to their respective slots in output_row
+        for field_name, field_val in component_dict.items():
+            if field_name not in self.iterator_output_fields:
+                raise excs.Error(
+                    f'Invalid field name {field_name} in output of {self.view.iterator_cls.__name__}')
+            if field_name not in self.refd_output_slot_idxs:
+                # we can ignore this
+                continue
+            output_col = self.iterator_output_cols[field_name]
+            output_col.col_type.validate_literal(field_val)
+            output_row[self.refd_output_slot_idxs[field_name]] = field_val
+        if len(component_dict) != len(self.iterator_output_fields):
+            missing_fields = set(self.refd_output_slot_idxs.keys()) - set(component_dict.keys())
+            raise excs.Error(
+                f'Invalid output of {self.view.iterator_cls.__name__}: '
+                f'missing fields {", ".join(missing_fields)}')
+
     def __next__(self) -> DataRowBatch:
-        if self._output is None:
-            self._output = self._output_batches()
-        return next(self._output)
+        if self.__output is None:
+            self.__output = self.__output_batches()
+        return next(self.__output)

pixeltable/exec/data_row_batch.py

@@ -14,6 +14,13 @@ class DataRowBatch:
 
     Contains the metadata needed to initialize DataRows.
     """
+    tbl: Optional[catalog.TableVersion]
+    row_builder: exprs.RowBuilder
+    img_slot_idxs: list[int]
+    media_slot_idxs: list[int]  # non-image media slots
+    array_slot_idxs: list[int]
+    rows: list[exprs.DataRow]
+
     def __init__(self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, len: int = 0):
         self.tbl = tbl
         self.row_builder = row_builder
@@ -39,13 +46,6 @@ class DataRowBatch:
     def pop_row(self) -> exprs.DataRow:
         return self.rows.pop()
 
-    def set_row_ids(self, row_ids: List[int]) -> None:
-        """Sets pks for rows in batch"""
-        assert self.tbl is not None
-        assert len(row_ids) == len(self.rows)
-        for row, row_id in zip(self.rows, row_ids):
-            row.set_pk((row_id, self.tbl))
-
     def __len__(self) -> int:
         return len(self.rows)
 

pixeltable/exec/exec_node.py

@@ -1,13 +1,25 @@
 from __future__ import annotations
-from typing import Iterable, Optional, List
+
 import abc
+from typing import Iterable, Optional, List, TYPE_CHECKING, Iterator
 
+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 +33,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 +48,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 +76,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.func import CallableFunction
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 

pixeltable/exec/in_memory_data_node.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Iterator
 
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
@@ -18,6 +18,11 @@ 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]
+
     def __init__(
         self, tbl: catalog.TableVersion, rows: list[dict[str, Any]],
         row_builder: exprs.RowBuilder, start_row_id: int,
@@ -29,8 +34,7 @@ class InMemoryDataNode(ExecNode):
         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 +71,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