pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/exec/cell_reconstruction_node.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import io
+import logging
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+import numpy as np
+import PIL.Image
+
+import pixeltable.type_system as ts
+from pixeltable import exprs
+from pixeltable.utils import parse_local_file_path
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+from .globals import INLINED_OBJECT_MD_KEY, InlinedObjectMd
+
+_logger = logging.getLogger('pixeltable')
+
+
+def json_has_inlined_objs(element: Any) -> bool:
+    """Returns True if element contains inlined objects produced by CellMaterializationNode."""
+    if isinstance(element, list):
+        return any(json_has_inlined_objs(v) for v in element)
+    if isinstance(element, dict):
+        if INLINED_OBJECT_MD_KEY in element:
+            return True
+        return any(json_has_inlined_objs(v) for v in element.values())
+    return False
+
+
+def reconstruct_json(element: Any, urls: list[str], file_handles: dict[Path, io.BufferedReader]) -> Any:
+    """Recursively reconstructs inlined objects in a json structure."""
+    if isinstance(element, list):
+        return [reconstruct_json(v, urls, file_handles) for v in element]
+    if isinstance(element, dict):
+        if INLINED_OBJECT_MD_KEY in element:
+            obj_md = InlinedObjectMd.from_dict(element[INLINED_OBJECT_MD_KEY])
+            url = urls[obj_md.url_idx]
+            local_path = parse_local_file_path(url)
+            if local_path not in file_handles:
+                file_handles[local_path] = open(local_path, 'rb')  # noqa: SIM115
+            fp = file_handles[local_path]
+
+            if obj_md.type == ts.ColumnType.Type.ARRAY.name:
+                fp.seek(obj_md.array_md.start)
+                ar = load_array(
+                    fp, obj_md.array_md.start, obj_md.array_md.end, obj_md.array_md.is_bool, obj_md.array_md.shape
+                )
+                return ar
+            else:
+                fp.seek(obj_md.img_start)
+                bytesio = io.BytesIO(fp.read(obj_md.img_end - obj_md.img_start))
+                img = PIL.Image.open(bytesio)
+                img.load()
+                assert fp.tell() == obj_md.img_end, f'{fp.tell()} != {obj_md.img_end} / {obj_md.img_start}'
+                return img
+        else:
+            return {k: reconstruct_json(v, urls, file_handles) for k, v in element.items()}
+    return element
+
+
+def load_array(
+    fh: io.BufferedReader, start: int, end: int, is_bool_array: bool, shape: tuple[int, ...] | None
+) -> np.ndarray:
+    """Loads an array from a section of a file."""
+    fh.seek(start)
+    ar = np.load(fh, allow_pickle=False)
+    assert fh.tell() == end
+    if is_bool_array:
+        assert shape is not None
+        ar = np.unpackbits(ar, count=np.prod(shape)).reshape(shape).astype(bool)
+    return ar
+
+
+class CellReconstructionNode(ExecNode):
+    """
+    Reconstruction of stored json and array cells that were produced by CellMaterializationNode.
+    """
+
+    json_refs: list[exprs.ColumnRef]
+    array_refs: list[exprs.ColumnRef]
+    file_handles: dict[Path, io.BufferedReader]  # key: file path
+
+    def __init__(
+        self,
+        json_refs: list[exprs.ColumnRef],
+        array_refs: list[exprs.ColumnRef],
+        row_builder: exprs.RowBuilder,
+        input: ExecNode | None = None,
+    ):
+        super().__init__(row_builder, [], [], input)
+        self.json_refs = json_refs
+        self.array_refs = array_refs
+        self.file_handles = {}
+
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        async for batch in self.input:
+            for row in batch:
+                for col_ref in self.json_refs:
+                    val = row[col_ref.slot_idx]
+                    if val is None:
+                        continue
+                    cell_md = row.slot_md.get(col_ref.slot_idx)
+                    if cell_md is None or cell_md.file_urls is None or not json_has_inlined_objs(row[col_ref.slot_idx]):
+                        continue
+                    row[col_ref.slot_idx] = reconstruct_json(val, cell_md.file_urls, self.file_handles)
+
+                for col_ref in self.array_refs:
+                    cell_md = row.slot_md.get(col_ref.slot_idx)
+                    if cell_md is not None and cell_md.array_md is not None:
+                        assert row[col_ref.slot_idx] is None
+                        assert cell_md.file_urls is not None and len(cell_md.file_urls) == 1
+                        row[col_ref.slot_idx] = self._reconstruct_array(cell_md)
+                    else:
+                        assert row[col_ref.slot_idx] is None or isinstance(row[col_ref.slot_idx], np.ndarray)
+
+            yield batch
+
+    def close(self) -> None:
+        for fp in self.file_handles.values():
+            fp.close()
+
+    def _reconstruct_array(self, cell_md: exprs.CellMd) -> np.ndarray:
+        assert cell_md.array_md is not None
+        local_path = parse_local_file_path(cell_md.file_urls[0])
+        assert local_path is not None
+        if local_path not in self.file_handles:
+            self.file_handles[local_path] = open(str(local_path), 'rb')  # noqa: SIM115
+        fp = self.file_handles[local_path]
+        ar = load_array(
+            fp, cell_md.array_md.start, cell_md.array_md.end, bool(cell_md.array_md.is_bool), cell_md.array_md.shape
+        )
+        return ar

pixeltable/exec/component_iteration_node.py

@@ -40,7 +40,7 @@ class ComponentIterationNode(ExecNode):
         }
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
-        output_batch = DataRowBatch(self.view, self.row_builder)
+        output_batch = DataRowBatch(self.row_builder)
         async for input_batch in self.input:
             for input_row in input_batch:
                 self.row_builder.eval(input_row, self.iterator_args_ctx)
@@ -52,13 +52,14 @@ class ComponentIterationNode(ExecNode):
                 if self.__non_nullable_args_specified(iterator_args):
                     iterator = self.view.get().iterator_cls(**iterator_args)
                     for pos, component_dict in enumerate(iterator):
-                        output_row = output_batch.add_row()
+                        output_row = self.row_builder.make_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)
+                        output_batch.add_row(output_row)
                         if len(output_batch) == self.__OUTPUT_BATCH_SIZE:
                             yield output_batch
-                            output_batch = DataRowBatch(self.view, self.row_builder)
+                            output_batch = DataRowBatch(self.row_builder)
 
         if len(output_batch) > 0:
             yield output_batch

pixeltable/exec/data_row_batch.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import logging
-from typing import Iterator, Optional
+from typing import Iterator
 
-from pixeltable import catalog, exprs
-from pixeltable.utils.media_store import MediaStore
+from pixeltable import exprs
 
 _logger = logging.getLogger('pixeltable')
 
@@ -13,53 +12,20 @@ class DataRowBatch:
     """Set of DataRows, indexed by rowid.
 
     Contains the metadata needed to initialize DataRows.
+
+    Requires either num_rows or rows to be specified, but not both.
     """
 
-    tbl: Optional[catalog.TableVersionHandle]
     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.TableVersionHandle],
-        row_builder: exprs.RowBuilder,
-        num_rows: Optional[int] = None,
-        rows: Optional[list[exprs.DataRow]] = None,
-    ):
-        """
-        Requires either num_rows or rows to be specified, but not both.
-        """
-        assert num_rows is None or rows is None
-        self.tbl = tbl
+    def __init__(self, row_builder: exprs.RowBuilder, rows: list[exprs.DataRow] | None = None):
         self.row_builder = row_builder
-        self.img_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_image_type()]
-        # non-image media slots
-        self.media_slot_idxs = [
-            e.slot_idx
-            for e in row_builder.unique_exprs
-            if e.col_type.is_media_type() and not e.col_type.is_image_type()
-        ]
-        self.array_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_array_type()]
-        if rows is not None:
-            self.rows = rows
-        else:
-            if num_rows is None:
-                num_rows = 0
-            self.rows = [
-                exprs.DataRow(
-                    row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
-                )
-                for _ in range(num_rows)
-            ]
+        self.rows = [] if rows is None else rows
 
-    def add_row(self, row: Optional[exprs.DataRow] = None) -> exprs.DataRow:
+    def add_row(self, row: exprs.DataRow | None) -> exprs.DataRow:
         if row is None:
-            row = exprs.DataRow(
-                self.row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
-            )
+            row = self.row_builder.make_row()
         self.rows.append(row)
         return row
 
@@ -72,28 +38,5 @@ class DataRowBatch:
     def __getitem__(self, index: int) -> exprs.DataRow:
         return self.rows[index]
 
-    def flush_imgs(
-        self,
-        idx_range: Optional[slice] = None,
-        stored_img_info: Optional[list[exprs.ColumnSlotIdx]] = None,
-        flushed_slot_idxs: Optional[list[int]] = None,
-    ) -> None:
-        """Flushes images in the given range of rows."""
-        assert self.tbl is not None
-        if stored_img_info is None:
-            stored_img_info = []
-        if flushed_slot_idxs is None:
-            flushed_slot_idxs = []
-        if len(stored_img_info) == 0 and len(flushed_slot_idxs) == 0:
-            return
-        if idx_range is None:
-            idx_range = slice(0, len(self.rows))
-        for row in self.rows[idx_range]:
-            for info in stored_img_info:
-                filepath = str(MediaStore.prepare_media_path(self.tbl.id, info.col.id, self.tbl.get().version))
-                row.flush_img(info.slot_idx, filepath)
-            for slot_idx in flushed_slot_idxs:
-                row.flush_img(slot_idx)
-
     def __iter__(self) -> Iterator[exprs.DataRow]:
         return iter(self.rows)

pixeltable/exec/exec_context.py

@@ -1,4 +1,4 @@
-from typing import Optional
+import random
 
 import sqlalchemy as sql
 
@@ -8,13 +8,24 @@ from pixeltable import exprs
 class ExecContext:
     """Class for execution runtime constants"""
 
+    row_builder: exprs.RowBuilder
+    profile: exprs.ExecProfile
+    show_pbar: bool
+    batch_size: int
+    num_rows: int | None
+    conn: sql.engine.Connection | None
+    pk_clause: list[sql.ClauseElement] | None
+    num_computed_exprs: int
+    ignore_errors: bool
+    random_seed: int  # general-purpose source of randomness with execution scope
+
     def __init__(
         self,
         row_builder: exprs.RowBuilder,
         *,
         show_pbar: bool = False,
         batch_size: int = 0,
-        pk_clause: Optional[list[sql.ClauseElement]] = None,
+        pk_clause: list[sql.ClauseElement] | None = None,
         num_computed_exprs: int = 0,
         ignore_errors: bool = False,
     ):
@@ -23,8 +34,9 @@ class ExecContext:
         self.row_builder = row_builder
         self.profile = exprs.ExecProfile(row_builder)
         # num_rows is used to compute the total number of computed cells used for the progress bar
-        self.num_rows: Optional[int] = None
-        self.conn: Optional[sql.engine.Connection] = None  # if present, use this to execute SQL queries
+        self.num_rows = None
+        self.conn = None  # if present, use this to execute SQL queries
         self.pk_clause = pk_clause
         self.num_computed_exprs = num_computed_exprs
         self.ignore_errors = ignore_errors
+        self.random_seed = random.randint(0, 1 << 63)

pixeltable/exec/exec_node.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 import abc
-import asyncio
 import logging
-from typing import AsyncIterator, Iterable, Iterator, Optional, TypeVar
+from typing import AsyncIterator, Iterable, Iterator, TypeVar
 
 from pixeltable import exprs
+from pixeltable.env import Env
 
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
@@ -18,17 +18,16 @@ class ExecNode(abc.ABC):
 
     output_exprs: Iterable[exprs.Expr]
     row_builder: exprs.RowBuilder
-    input: Optional[ExecNode]
+    input: ExecNode | None
     flushed_img_slots: list[int]  # idxs of image slots of our output_exprs dependencies
-    stored_img_cols: list[exprs.ColumnSlotIdx]
-    ctx: Optional[ExecContext]
+    ctx: ExecContext | None
 
     def __init__(
         self,
         row_builder: exprs.RowBuilder,
         output_exprs: Iterable[exprs.Expr],
         input_exprs: Iterable[exprs.Expr],
-        input: Optional[ExecNode] = None,
+        input: ExecNode | None = None,
     ):
         assert all(expr.is_valid for expr in output_exprs)
         self.output_exprs = output_exprs
@@ -40,43 +39,19 @@ class ExecNode(abc.ABC):
         self.flushed_img_slots = [
             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 = []
-        self.ctx = None  # all nodes of a tree share the same context
+        self.ctx = input.ctx if input is not None else None
 
     def set_ctx(self, ctx: ExecContext) -> None:
         self.ctx = ctx
         if self.input is not None:
             self.input.set_ctx(ctx)
 
-    def set_stored_img_cols(self, stored_img_cols: list[exprs.ColumnSlotIdx]) -> None:
-        self.stored_img_cols = stored_img_cols
-        # propagate batch size to the source
-        if self.input is not None:
-            self.input.set_stored_img_cols(stored_img_cols)
-
     @abc.abstractmethod
     def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         pass
 
     def __iter__(self) -> Iterator[DataRowBatch]:
-        running_loop: Optional[asyncio.AbstractEventLoop] = None
-        loop: asyncio.AbstractEventLoop
-        try:
-            # check if we are already in an event loop (eg, Jupyter's); if so, patch it to allow
-            # multiple run_until_complete()
-            running_loop = asyncio.get_running_loop()
-            import nest_asyncio  # type: ignore[import-untyped]
-
-            nest_asyncio.apply()
-            loop = running_loop
-            _logger.debug('Patched running loop')
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        if _logger.isEnabledFor(logging.DEBUG):
-            loop.set_debug(True)
-
+        loop = Env.get().event_loop
         aiter = self.__aiter__()
         try:
             while True:
@@ -84,9 +59,11 @@ class ExecNode(abc.ABC):
                 yield batch
         except StopAsyncIteration:
             pass
-        finally:
-            if loop != running_loop:
-                loop.close()
+        # TODO:
+        #  - we seem to have some tasks that aren't accounted for by ExprEvalNode and don't get cancelled by the time
+        #    we end up here
+        # - however, blindly cancelling all pending tasks doesn't work when running in a jupyter environment, which
+        #   creates tasks on its own
 
     def open(self) -> None:
         """Bottom-up initialization of nodes for execution. Must be called before __next__."""
@@ -108,7 +85,7 @@ class ExecNode(abc.ABC):
 
     T = TypeVar('T', bound='ExecNode')
 
-    def get_node(self, node_class: type[T]) -> Optional[T]:
+    def get_node(self, node_class: type[T]) -> T | None:
         if isinstance(self, node_class):
             return self
         if self.input is not None:

pixeltable/exec/expr_eval/evaluators.py

@@ -5,7 +5,7 @@ import datetime
 import itertools
 import logging
 import sys
-from typing import Any, Callable, Iterator, Optional, cast
+from typing import Any, Callable, Iterator, cast
 
 from pixeltable import exprs, func
 
@@ -64,11 +64,11 @@ class FnCallEvaluator(Evaluator):
 
     fn_call: exprs.FunctionCall
     fn: func.CallableFunction
-    scalar_py_fn: Optional[Callable]  # only set for non-batching CallableFunctions
+    scalar_py_fn: Callable | None  # only set for non-batching CallableFunctions
 
     # only set if fn.is_batched
-    call_args_queue: Optional[asyncio.Queue[FnCallArgs]]  # FnCallArgs waiting for execution
-    batch_size: Optional[int]
+    call_args_queue: asyncio.Queue[FnCallArgs] | None  # FnCallArgs waiting for execution
+    batch_size: int | None
 
     def __init__(self, fn_call: exprs.FunctionCall, dispatcher: Dispatcher, exec_ctx: ExecCtx):
         super().__init__(dispatcher, exec_ctx)
@@ -160,8 +160,8 @@ class FnCallEvaluator(Evaluator):
 
     def _create_batch_call_args(self, call_args: list[FnCallArgs]) -> FnCallArgs:
         """Roll call_args into a single batched FnCallArgs"""
-        batch_args: list[list[Optional[Any]]] = [[None] * len(call_args) for _ in range(len(self.fn_call.arg_idxs))]
-        batch_kwargs: dict[str, list[Optional[Any]]] = {k: [None] * len(call_args) for k in self.fn_call.kwarg_idxs}
+        batch_args: list[list[Any | None]] = [[None] * len(call_args) for _ in range(len(self.fn_call.arg_idxs))]
+        batch_kwargs: dict[str, list[Any | None]] = {k: [None] * len(call_args) for k in self.fn_call.kwarg_idxs}
         assert isinstance(self.fn, func.CallableFunction)
         for i, item in enumerate(call_args):
             for j in range(len(item.args)):
@@ -311,6 +311,7 @@ class JsonMapperDispatcher(Evaluator):
                     img_slot_idxs=[],
                     media_slot_idxs=[],
                     array_slot_idxs=[],
+                    json_slot_idxs=[],
                     parent_row=row,
                     parent_slot_idx=self.e.slot_idx,
                 )

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -4,7 +4,7 @@ import asyncio
 import logging
 import traceback
 from types import TracebackType
-from typing import AsyncIterator, Iterable, Optional, Union
+from typing import AsyncIterator, Iterable
 
 import numpy as np
 
@@ -49,17 +49,17 @@ class ExprEvalNode(ExecNode):
     # execution state
     tasks: set[asyncio.Task]  # collects all running tasks to prevent them from getting gc'd
     exc_event: asyncio.Event  # set if an exception needs to be propagated
-    error: Optional[Union[excs.Error, excs.ExprEvalError]]  # exception that needs to be propagated
+    error: Exception | None  # exception that needs to be propagated
     completed_rows: asyncio.Queue[exprs.DataRow]  # rows that have completed evaluation
     completed_event: asyncio.Event  # set when completed_rows is non-empty
     input_iter: AsyncIterator[DataRowBatch]
-    current_input_batch: Optional[DataRowBatch]  # batch from which we're currently consuming rows
+    current_input_batch: DataRowBatch | None  # batch from which we're currently consuming rows
     input_row_idx: int  # next row to consume from current_input_batch
-    next_input_batch: Optional[DataRowBatch]  # read-ahead input batch
+    next_input_batch: DataRowBatch | None  # read-ahead input batch
     avail_input_rows: int  # total number across both current_/next_input_batch
     input_complete: bool  # True if we've received all input batches
     num_in_flight: int  # number of dispatched rows that haven't completed
-    row_pos_map: Optional[dict[int, int]]  # id(row) -> position of row in input; only set if maintain_input_order
+    row_pos_map: dict[int, int] | None  # id(row) -> position of row in input; only set if maintain_input_order
     output_buffer: RowBuffer  # holds rows that are ready to be returned, in order
 
     # debugging
@@ -133,10 +133,10 @@ class ExprEvalNode(ExecNode):
         except StopAsyncIteration:
             self.input_complete = True
             _logger.debug(f'finished input: #input_rows={self.num_input_rows}, #avail={self.avail_input_rows}')
-        except excs.Error as err:
-            self.error = err
+        # make sure to pass DBAPIError through, so the transaction handling logic sees it
+        except Exception as exc:
+            self.error = exc
             self.exc_event.set()
-        # TODO: should we also handle Exception here and create an excs.Error from it?
 
     @property
     def total_buffered(self) -> int:
@@ -217,9 +217,10 @@ class ExprEvalNode(ExecNode):
 
         row: exprs.DataRow
         exc_event_aw = asyncio.create_task(self.exc_event.wait(), name='exc_event.wait()')
-        input_batch_aw: Optional[asyncio.Task] = None
-        completed_aw: Optional[asyncio.Task] = None
+        input_batch_aw: asyncio.Task | None = None
+        completed_aw: asyncio.Task | None = None
         closed_evaluators = False  # True after calling Evaluator.close()
+        exprs.Expr.prepare_list(self.exec_ctx.all_exprs)
 
         try:
             while True:
@@ -240,7 +241,7 @@ class ExprEvalNode(ExecNode):
                         # make sure we top up our in-flight rows before yielding
                         self._dispatch_input_rows()
                         self._log_state(f'yielding {len(batch_rows)} rows')
-                        yield DataRowBatch(tbl=None, row_builder=self.row_builder, rows=batch_rows)
+                        yield DataRowBatch(row_builder=self.row_builder, rows=batch_rows)
                         # at this point, we may have more completed rows
 
                 assert self.completed_rows.empty()  # all completed rows should be sitting in output_buffer
@@ -254,7 +255,7 @@ class ExprEvalNode(ExecNode):
                         batch_rows = self.output_buffer.get_rows(self.output_buffer.num_ready)
                         self.num_output_rows += len(batch_rows)
                         self._log_state(f'yielding {len(batch_rows)} rows')
-                        yield DataRowBatch(tbl=None, row_builder=self.row_builder, rows=batch_rows)
+                        yield DataRowBatch(row_builder=self.row_builder, rows=batch_rows)
 
                     assert self.output_buffer.num_rows == 0
                     return
@@ -306,6 +307,9 @@ class ExprEvalNode(ExecNode):
                     task.cancel()
             _ = await asyncio.gather(*active_tasks, return_exceptions=True)
 
+            # expr cleanup
+            exprs.Expr.release_list(self.exec_ctx.all_exprs)
+
     def dispatch_exc(
         self, rows: list[exprs.DataRow], slot_with_exc: int, exc_tb: TracebackType, exec_ctx: ExecCtx
     ) -> None:
@@ -390,6 +394,17 @@ class ExprEvalNode(ExecNode):
         # end the main loop if we had an unhandled exception
         try:
             t.result()
+        except KeyboardInterrupt:
+            # ExprEvalNode instances are long-running and reused across multiple operations.
+            # When a user interrupts an operation (Ctrl+C), the main evaluation loop properly
+            # handles the KeyboardInterrupt and terminates the current operation. However,
+            # background tasks spawned by evaluators may complete asynchronously after the
+            # operation has ended, and their done callbacks will fire during subsequent
+            # operations. These "phantom" KeyboardInterrupt exceptions from previous
+            # operations' background tasks should not interfere with new operations, so we
+            # absorb them here rather than propagating them via self.error/self.exc_event.
+            _logger.debug('Task completed with KeyboardInterrupt (user cancellation)')
+            pass
         except asyncio.CancelledError:
             pass
         except Exception as exc:

pixeltable/exec/expr_eval/globals.py

@@ -4,7 +4,7 @@ import abc
 import asyncio
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, Iterable, Optional, Protocol
+from typing import Any, Iterable, Protocol
 
 import numpy as np
 
@@ -18,11 +18,11 @@ class FnCallArgs:
     fn_call: exprs.FunctionCall
     rows: list[exprs.DataRow]
     # single call
-    args: Optional[list[Any]] = None
-    kwargs: Optional[dict[str, Any]] = None
+    args: list[Any] | None = None
+    kwargs: dict[str, Any] | None = None
     # batch call
-    batch_args: Optional[list[list[Optional[Any]]]] = None
-    batch_kwargs: Optional[dict[str, list[Optional[Any]]]] = None
+    batch_args: list[list[Any | None]] | None = None
+    batch_kwargs: dict[str, list[Any | None]] | None = None
 
     @property
     def pxt_fn(self) -> func.CallableFunction:
@@ -56,6 +56,7 @@ class Scheduler(abc.ABC):
         request: FnCallArgs
         num_retries: int
         exec_ctx: ExecCtx
+        retry_after: float | None = None  # time.monotonic()
 
         def __lt__(self, other: Scheduler.QueueItem) -> bool:
             # prioritize by number of retries (more retries = higher priority)
@@ -148,6 +149,7 @@ class ExecCtx:
     gc_targets: np.ndarray  # bool per slot; True if this is an intermediate expr (ie, not part of our output)
     eval_ctx: np.ndarray  # bool per slot; EvalCtx.slot_idxs as a mask
     literals: dict[int, Any]  # key: slot idx; value: literal value for this slot; used to pre-populate rows
+    all_exprs: list[exprs.Expr]  # all evaluated exprs; needed for cleanup
 
     def __init__(
         self,
@@ -164,6 +166,7 @@ class ExecCtx:
         self.gc_targets[[e.slot_idx for e in self.row_builder.output_exprs]] = False
 
         output_ctx = self.row_builder.create_eval_ctx(output_exprs, exclude=input_exprs)
+        self.all_exprs = output_ctx.exprs
         self.literals = {e.slot_idx: e.val for e in output_ctx.exprs if isinstance(e, exprs.Literal)}
         self.eval_ctx = np.zeros(self.row_builder.num_materialized, dtype=bool)
         non_literal_slot_idxs = [e.slot_idx for e in output_ctx.exprs if not isinstance(e, exprs.Literal)]

pixeltable/exec/expr_eval/row_buffer.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional
 
 import numpy as np
 
@@ -14,7 +13,7 @@ class RowBuffer:
     """Fixed-length circular buffer of DataRows; knows how to maintain input order"""
 
     size: int
-    row_pos_map: Optional[dict[int, int]]  # id(row) -> position of row in output; None if not maintaining order
+    row_pos_map: dict[int, int] | None  # id(row) -> position of row in output; None if not maintaining order
     num_rows: int  # number of rows in the buffer
     num_ready: int  # number of consecutive non-None rows at head
     buffer: np.ndarray  # of object