pixeltable 0.3.0__tar.gz → 0.3.2__tar.gz

{pixeltable-0.3.0 → pixeltable-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pixeltable
-Version: 0.3.0
+Version: 0.3.2
 Summary: AI Data Infrastructure: Declarative, Multimodal, and Incremental
 Home-page: https://pixeltable.com/
 License: Apache-2.0
@@ -27,11 +27,14 @@ Requires-Dist: av (>=10.0.0)
 Requires-Dist: beautifulsoup4 (>=4.0.0,<5.0.0)
 Requires-Dist: cloudpickle (>=2.2.1,<3.0.0)
 Requires-Dist: ftfy (>=6.2.0,<7.0.0)
+Requires-Dist: httpcore (>=1.0.3)
+Requires-Dist: httpx (>=0.27)
 Requires-Dist: jinja2 (>=3.1.3,<4.0.0)
 Requires-Dist: jmespath (>=1.0.1,<2.0.0)
 Requires-Dist: jsonschema (>=4.1.0)
 Requires-Dist: lxml (>=5.0)
 Requires-Dist: more-itertools (>=10.2,<11.0)
+Requires-Dist: nest_asyncio (>=1.5)
 Requires-Dist: numpy (>=1.25,<2.0)
 Requires-Dist: pandas (>=2.0,<3.0)
 Requires-Dist: pgvector (>=0.2.1,<0.3.0)

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/__version__.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced during build.
-__version__ = "0.3.0"
-__version_tuple__ = (0, 3, 0)
+__version__ = "0.3.2"
+__version_tuple__ = (0, 3, 2)

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/catalog/insertable_table.py

@@ -68,7 +68,7 @@ class InsertableTable(Table):
             cat.tbls[tbl._id] = tbl
 
             _logger.info(f'Created table `{name}`, id={tbl_version.id}')
-            print(f'Created table `{name}`.')
+            Env.get().console_logger.info(f'Created table `{name}`.')
             return tbl
 
     def get_metadata(self) -> dict[str, Any]:
@@ -134,7 +134,7 @@ class InsertableTable(Table):
             f'Inserted {status.num_rows} row{"" if status.num_rows == 1 else "s"} '
             f'with {status.num_excs} error{"" if status.num_excs == 1 else "s"}{cols_with_excs_str}.'
         )
-        print(msg)
+        Env.get().console_logger.info(msg)
         _logger.info(f'InsertableTable {self._name}: {msg}')
         FileCache.get().emit_eviction_warnings()
         return status
@@ -164,7 +164,7 @@ class InsertableTable(Table):
                     row[col_name] = checked_val
                 except TypeError as e:
                     msg = str(e)
-                    raise excs.Error(f'Error in column {col.name}: {msg[0].lower() + msg[1:]}\nRow: {row}')
+                    raise excs.Error(f'Error in column {col.name}: {msg[0].lower() + msg[1:]}\nRow: {row}') from e
 
     def delete(self, where: Optional['pxt.exprs.Expr'] = None) -> UpdateStatus:
         """Delete rows in this table.

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/catalog/table.py

@@ -1345,7 +1345,7 @@ class Table(SchemaObject):
             raise excs.Error(f'Table `{self._name}` already has an external store with that name: {store.name}')
         _logger.info(f'Linking external store `{store.name}` to table `{self._name}`')
         self._tbl_version.link_external_store(store)
-        print(f'Linked external store `{store.name}` to table `{self._name}`.')
+        env.Env.get().console_logger.info(f'Linked external store `{store.name}` to table `{self._name}`.')
 
     def unlink_external_stores(
             self,
@@ -1381,7 +1381,7 @@ class Table(SchemaObject):
 
         for store in stores:
             self._tbl_version.unlink_external_store(store, delete_external_data=delete_external_data)
-            print(f'Unlinked external store from table `{self._name}`: {store}')
+            env.Env.get().console_logger.info(f'Unlinked external store from table `{self._name}`: {store}')
 
     def sync(
             self,

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/catalog/table_version.py

@@ -532,7 +532,7 @@ class TableVersion:
             f'Added {status.num_rows} column value{"" if status.num_rows == 1 else "s"} '
             f'with {status.num_excs} error{"" if status.num_excs == 1 else "s"}.'
         )
-        print(msg)
+        Env.get().console_logger.info(msg)
         _logger.info(f'Columns {[col.name for col in cols]}: {msg}')
         return status
 
@@ -734,7 +734,8 @@ class TableVersion:
         if conn is None:
             with Env.get().engine.begin() as conn:
                 return self._insert(
-                    plan, conn, time.time(), print_stats=print_stats, rowids=rowids(), abort_on_exc=fail_on_exception)
+                    plan, conn, time.time(), print_stats=print_stats, rowids=rowids(),
+                    abort_on_exc=fail_on_exception)
         else:
             return self._insert(
                 plan, conn, time.time(), print_stats=print_stats, rowids=rowids(), abort_on_exc=fail_on_exception)

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/catalog/view.py

@@ -156,7 +156,7 @@ class View(Table):
                 plan, num_values_per_row = Planner.create_view_load_plan(view._tbl_version_path)
                 num_rows, num_excs, cols_with_excs = tbl_version.store_tbl.insert_rows(
                     plan, session.connection(), v_min=tbl_version.version)
-                print(f'Created view `{name}` with {num_rows} rows, {num_excs} exceptions.')
+                Env.get().console_logger.info(f'Created view `{name}` with {num_rows} rows, {num_excs} exceptions.')
 
             session.commit()
             cat = Catalog.get()

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/dataframe.py

@@ -8,18 +8,15 @@ import json
 import logging
 import traceback
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Hashable, Iterator, Optional, Sequence, Union, Literal
+from typing import TYPE_CHECKING, Any, Callable, Hashable, Iterator, Optional, Sequence, Union, AsyncIterator, NoReturn
 
+import numpy as np
 import pandas as pd
-import pandas.io.formats.style
 import sqlalchemy as sql
 
-import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
 import pixeltable.type_system as ts
-from pixeltable import exec
-from pixeltable import plan
+from pixeltable import catalog, exec, exprs, plan
 from pixeltable.catalog import is_valid_identifier
 from pixeltable.catalog.globals import UpdateStatus
 from pixeltable.env import Env
@@ -29,6 +26,7 @@ from pixeltable.utils.formatter import Formatter
 
 if TYPE_CHECKING:
     import torch
+    import torch.utils.data
 
 __all__ = ['DataFrame']
 
@@ -268,6 +266,20 @@ class DataFrame:
         else:
             yield from exec_plan(conn)
 
+    async def _aexec(self, conn: sql.engine.Connection) -> AsyncIterator[exprs.DataRow]:
+        """Run the query and return rows as a generator.
+        This function must not modify the state of the DataFrame, otherwise it breaks dataset caching.
+        """
+        plan = self._create_query_plan()
+        plan.ctx.set_conn(conn)
+        plan.open()
+        try:
+            async for row_batch in plan:
+                for row in row_batch:
+                    yield row
+        finally:
+            plan.close()
+
     def _create_query_plan(self) -> exec.ExecNode:
         # construct a group-by clause if we're grouping by a table
         group_by_clause: Optional[list[exprs.Expr]] = None
@@ -392,26 +404,29 @@ class DataFrame:
             group_by_clause=group_by_clause, grouping_tbl=self.grouping_tbl,
             order_by_clause=order_by_clause, limit=self.limit_val)
 
+    def _raise_expr_eval_err(self, e: excs.ExprEvalError) -> NoReturn:
+        msg = f'In row {e.row_num} the {e.expr_msg} encountered exception ' f'{type(e.exc).__name__}:\n{str(e.exc)}'
+        if len(e.input_vals) > 0:
+            input_msgs = [
+                f"'{d}' = {d.col_type.print_value(e.input_vals[i])}" for i, d in enumerate(e.expr.dependencies())
+            ]
+            msg += f'\nwith {", ".join(input_msgs)}'
+        assert e.exc_tb is not None
+        stack_trace = traceback.format_tb(e.exc_tb)
+        if len(stack_trace) > 2:
+            # append a stack trace if the exception happened in user code
+            # (frame 0 is ExprEvaluator and frame 1 is some expr's eval()
+            nl = '\n'
+            # [-1:0:-1]: leave out entry 0 and reverse order, so that the most recent frame is at the top
+            msg += f'\nStack:\n{nl.join(stack_trace[-1:1:-1])}'
+        raise excs.Error(msg)
+
     def _output_row_iterator(self, conn: Optional[sql.engine.Connection] = None) -> Iterator[list]:
         try:
             for data_row in self._exec(conn):
                 yield [data_row[e.slot_idx] for e in self._select_list_exprs]
         except excs.ExprEvalError as e:
-            msg = f'In row {e.row_num} the {e.expr_msg} encountered exception ' f'{type(e.exc).__name__}:\n{str(e.exc)}'
-            if len(e.input_vals) > 0:
-                input_msgs = [
-                    f"'{d}' = {d.col_type.print_value(e.input_vals[i])}" for i, d in enumerate(e.expr.dependencies())
-                ]
-                msg += f'\nwith {", ".join(input_msgs)}'
-            assert e.exc_tb is not None
-            stack_trace = traceback.format_tb(e.exc_tb)
-            if len(stack_trace) > 2:
-                # append a stack trace if the exception happened in user code
-                # (frame 0 is ExprEvaluator and frame 1 is some expr's eval()
-                nl = '\n'
-                # [-1:0:-1]: leave out entry 0 and reverse order, so that the most recent frame is at the top
-                msg += f'\nStack:\n{nl.join(stack_trace[-1:1:-1])}'
-            raise excs.Error(msg)
+            self._raise_expr_eval_err(e)
         except sql.exc.DBAPIError as e:
             raise excs.Error(f'Error during SQL execution:\n{e}')
 
@@ -421,6 +436,18 @@ class DataFrame:
     def _collect(self, conn: Optional[sql.engine.Connection] = None) -> DataFrameResultSet:
         return DataFrameResultSet(list(self._output_row_iterator(conn)), self.schema)
 
+    async def _acollect(self, conn: sql.engine.Connection) -> DataFrameResultSet:
+        try:
+            result = [
+                [row[e.slot_idx] for e in self._select_list_exprs]
+                async for row in self._aexec(conn)
+            ]
+            return DataFrameResultSet(result, self.schema)
+        except excs.ExprEvalError as e:
+            self._raise_expr_eval_err(e)
+        except sql.exc.DBAPIError as e:
+            raise excs.Error(f'Error during SQL execution:\n{e}')
+
     def count(self) -> int:
         """Return the number of rows in the DataFrame.
 
@@ -540,10 +567,10 @@ class DataFrame:
         for raw_expr, name in base_list:
             if isinstance(raw_expr, exprs.Expr):
                 select_list.append((raw_expr, name))
-            elif isinstance(raw_expr, dict):
-                select_list.append((exprs.InlineDict(raw_expr), name))
-            elif isinstance(raw_expr, list):
-                select_list.append((exprs.InlineList(raw_expr), name))
+            elif isinstance(raw_expr, (dict, list, tuple)):
+                select_list.append((exprs.Expr.from_object(raw_expr), name))
+            elif isinstance(raw_expr, np.ndarray):
+                select_list.append((exprs.Expr.from_array(raw_expr), name))
             else:
                 select_list.append((exprs.Literal(raw_expr), name))
             expr = select_list[-1][0]
@@ -1031,8 +1058,6 @@ class DataFrame:
         else:
             return write_coco_dataset(self, dest_path)
 
-    # TODO Factor this out into a separate module.
-    # The return type is unresolvable, but torch can't be imported since it's an optional dependency.
     def to_pytorch_dataset(self, image_format: str = 'pt') -> 'torch.utils.data.IterableDataset':
         """
         Convert the dataframe to a pytorch IterableDataset suitable for parallel loading

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/env.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from abc import abstractmethod
 import datetime
 import glob
 import http.server
@@ -15,8 +16,9 @@ import sys
 import threading
 import uuid
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
+from sys import stdout
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
@@ -27,12 +29,18 @@ from tqdm import TqdmWarning
 
 import pixeltable.exceptions as excs
 from pixeltable import metadata
+from pixeltable.utils.console_output import ConsoleLogger, ConsoleMessageFilter, ConsoleOutputHandler, map_level
 from pixeltable.utils.http_server import make_server
 
 if TYPE_CHECKING:
     import spacy
 
 
+_logger = logging.getLogger('pixeltable')
+
+T = TypeVar('T')
+
+
 class Env:
     """
     Store for runtime globals.
@@ -61,6 +69,7 @@ class Env:
     _httpd: Optional[http.server.HTTPServer]
     _http_address: Optional[str]
     _logger: logging.Logger
+    _console_logger: ConsoleLogger
     _default_log_level: int
     _logfilename: Optional[str]
     _log_to_stdout: bool
@@ -70,6 +79,8 @@ class Env:
     _stdout_handler: logging.StreamHandler
     _initialized: bool
 
+    _resource_pool_info: dict[str, Any]
+
     @classmethod
     def get(cls) -> Env:
         if cls._instance is None:
@@ -84,6 +95,8 @@ class Env:
         cls._instance = env
 
     def __init__(self):
+        assert self._instance is None, 'Env is a singleton; use Env.get() to access the instance'
+
         self._home = None
         self._media_dir = None  # computed media files
         self._file_cache_dir = None  # cached media files with external URL
@@ -121,6 +134,8 @@ class Env:
         self._stdout_handler.setFormatter(logging.Formatter(self._log_fmt_str))
         self._initialized = False
 
+        self._resource_pool_info = {}
+
     @property
     def config(self) -> Config:
         assert self._config is not None
@@ -221,6 +236,10 @@ class Env:
         else:
             return False
 
+    @property
+    def console_logger(self) -> ConsoleLogger:
+        return self._console_logger
+
     def _set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
         if self._initialized:
             return
@@ -278,6 +297,14 @@ class Env:
             warnings.simplefilter('ignore', category=UserWarning)
             warnings.simplefilter('ignore', category=FutureWarning)
 
+        # Set verbose level for user visible console messages
+        verbosity = map_level(self._config.get_int_value('verbosity'))
+        stdout_handler = ConsoleOutputHandler(stream=stdout)
+        stdout_handler.setLevel(verbosity)
+        stdout_handler.addFilter(ConsoleMessageFilter())
+        self._logger.addHandler(stdout_handler)
+        self._console_logger = ConsoleLogger(self._logger)
+
         # configure _logger to log to a file
         self._logfilename = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + '.log'
         fh = logging.FileHandler(self._log_dir / self._logfilename, mode='w')
@@ -351,7 +378,7 @@ class Env:
             schema.base_metadata.create_all(self._sa_engine)
             metadata.create_system_info(self._sa_engine)
 
-        print(f'Connected to Pixeltable database at: {self.db_url}')
+        self.console_logger.info(f'Connected to Pixeltable database at: {self.db_url}')
 
         # we now have a home directory and db; start other services
         self._set_up_runtime()
@@ -609,6 +636,16 @@ class Env:
     def create_tmp_path(self, extension: str = '') -> Path:
         return self._tmp_dir / f'{uuid.uuid4()}{extension}'
 
+
+    #def get_resource_pool_info(self, pool_id: str, pool_info_cls: Optional[Type[T]]) -> T:
+    def get_resource_pool_info(self, pool_id: str, make_pool_info: Optional[Callable[[], T]] = None) -> T:
+        """Returns the info object for the given id, creating it if necessary."""
+        info = self._resource_pool_info.get(pool_id)
+        if info is None and make_pool_info is not None:
+            info = make_pool_info()
+            self._resource_pool_info[pool_id] = info
+        return info
+
     @property
     def home(self) -> Path:
         assert self._home is not None
@@ -686,8 +723,6 @@ class Config:
     """
     __config: dict[str, Any]
 
-    T = TypeVar('T')
-
     @classmethod
     def from_file(cls, path: Path) -> Config:
         """
@@ -767,3 +802,73 @@ class PackageInfo:
     is_installed: bool
     library_name: str  # pypi library name (may be different from package name)
     version: Optional[list[int]] = None  # installed version, as a list of components (such as [3,0,2] for "3.0.2")
+
+
+TIME_FORMAT = '%H:%M.%S %f'
+
+
+@dataclass
+class RateLimitsInfo:
+    """
+    Abstract base class for resource pools made up of rate limits for different resources.
+
+    Rate limits and currently remaining resources are periodically reported via record().
+
+    Subclasses provide operational customization via:
+    - get_retry_delay()
+    - get_request_resources(self, ...) -> dict[str, int]
+    with parameters that are a subset of those of the udf that creates the subclass's instance
+    """
+
+    # get_request_resources:
+    # - Returns estimated resources needed for a specific request (ie, a single udf call) as a dict (key: resource name)
+    # - parameters are a subset of those of the udf
+    # - this is not a class method because the signature depends on the instantiating udf
+    get_request_resources: Callable[..., dict[str, int]]
+
+    resource_limits: dict[str, RateLimitInfo] = field(default_factory=dict)
+
+    def is_initialized(self) -> bool:
+        return len(self.resource_limits) > 0
+
+    def reset(self) -> None:
+        self.resource_limits.clear()
+
+    def record(self, **kwargs) -> None:
+        now = datetime.datetime.now(tz=datetime.timezone.utc)
+        if len(self.resource_limits) == 0:
+            self.resource_limits = {k: RateLimitInfo(k, now, *v) for k, v in kwargs.items() if v is not None}
+            # TODO: remove
+            for info in self.resource_limits.values():
+                _logger.debug(f'Init {info.resource} rate limit: rem={info.remaining} reset={info.reset_at.strftime(TIME_FORMAT)} delta={(info.reset_at - now).total_seconds()}')
+        else:
+            for k, v in kwargs.items():
+                if v is not None:
+                    self.resource_limits[k].update(now, *v)
+
+    @abstractmethod
+    def get_retry_delay(self, exc: Exception) -> Optional[float]:
+        """Returns number of seconds to wait before retry, or None if not retryable"""
+        pass
+
+
+@dataclass
+class RateLimitInfo:
+    """Container for rate limit-related information for a single resource."""
+    resource: str
+    recorded_at: datetime.datetime
+    limit: int
+    remaining: int
+    reset_at: datetime.datetime
+
+    def update(self, recorded_at: datetime.datetime, limit: int, remaining: int, reset_at: datetime.datetime) -> None:
+        # we always update everything, even though responses may come back out-of-order: we can't use reset_at to
+        # determine order, because it doesn't increase monotonically (the reeset duration shortens as output_tokens
+        # are freed up - going from max to actual)
+        self.recorded_at = recorded_at
+        self.limit = limit
+        self.remaining = remaining
+        reset_delta = reset_at - self.reset_at
+        self.reset_at = reset_at
+        # TODO: remove
+        _logger.debug(f'Update {self.resource} rate limit: rem={self.remaining} reset={self.reset_at.strftime(TIME_FORMAT)} reset_delta={reset_delta.total_seconds()} recorded_delta={(self.reset_at - recorded_at).total_seconds()}')

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/exec/__init__.py

@@ -4,7 +4,7 @@ from .component_iteration_node import ComponentIterationNode
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
 from .exec_node import ExecNode
-from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
 from .row_update_node import RowUpdateNode
 from .sql_node import SqlLookupNode, SqlScanNode, SqlAggregationNode, SqlNode, SqlJoinNode
+from .expr_eval import ExprEvalNode

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/exec/aggregation_node.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Any, Iterable, Iterator, Optional, cast
+from typing import Any, Iterable, Iterator, Optional, cast, AsyncIterator
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -60,11 +60,11 @@ class AggregationNode(ExecNode):
                 input_vals = [row[d.slot_idx] for d in fn_call.dependencies()]
                 raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, input_vals, row_num)
 
-    def __iter__(self) -> Iterator[DataRowBatch]:
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         prev_row: Optional[exprs.DataRow] = None
         current_group: Optional[list[Any]] = None  # the values of the group-by exprs
         num_input_rows = 0
-        for row_batch in self.input:
+        async for row_batch in self.input:
             num_input_rows += len(row_batch)
             for row in row_batch:
                 group = [row[e.slot_idx] for e in self.group_by] if self.group_by is not None else None

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/exec/cache_prefetch_node.py

@@ -9,7 +9,7 @@ import urllib.request
 from collections import deque
 from concurrent import futures
 from pathlib import Path
-from typing import Optional, Any, Iterator
+from typing import Optional, Any, Iterator, AsyncIterator
 from uuid import UUID
 
 import pixeltable.env as env
@@ -79,12 +79,12 @@ class CachePrefetchNode(ExecNode):
         self.input_finished = False
         self.row_idx = itertools.count() if retain_input_order else itertools.repeat(None)
 
-    def __iter__(self) -> Iterator[DataRowBatch]:
-        input_iter = iter(self.input)
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        input_iter = self.input.__aiter__()
         with futures.ThreadPoolExecutor(max_workers=self.NUM_EXECUTOR_THREADS) as executor:
             # we create enough in-flight requests to fill the first batch
             while not self.input_finished and self.__num_pending_rows() < self.BATCH_SIZE:
-                self.__submit_input_batch(input_iter, executor)
+                await self.__submit_input_batch(input_iter, executor)
 
             while True:
                 # try to assemble a full batch of output rows
@@ -93,7 +93,7 @@ class CachePrefetchNode(ExecNode):
 
                 # try to create enough in-flight requests to fill the next batch
                 while not self.input_finished and self.__num_pending_rows() < self.BATCH_SIZE:
-                    self.__submit_input_batch(input_iter, executor)
+                    await self.__submit_input_batch(input_iter, executor)
 
                 if len(self.ready_rows) > 0:
                     # create DataRowBatch from the first BATCH_SIZE ready rows
@@ -163,9 +163,15 @@ class CachePrefetchNode(ExecNode):
                         self.__add_ready_row(row, state.idx)
                         _logger.debug(f'row {state.idx} is ready (ready_batch_size={self.__ready_prefix_len()})')
 
-    def __submit_input_batch(self, input: Iterator[DataRowBatch], executor: futures.ThreadPoolExecutor) -> None:
+    async def __submit_input_batch(
+        self, input: AsyncIterator[DataRowBatch], executor: futures.ThreadPoolExecutor
+    ) -> None:
         assert not self.input_finished
-        input_batch = next(input, None)
+        input_batch: Optional[DataRowBatch]
+        try:
+            input_batch = await input.__anext__()
+        except StopAsyncIteration:
+            input_batch = None
         if input_batch is None:
             self.input_finished = True
             return

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/exec/component_iteration_node.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Iterator, Optional
+from typing import Iterator, Optional, AsyncIterator
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -37,11 +37,10 @@ class ComponentIterationNode(ExecNode):
             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[Iterator[DataRowBatch]] = None
 
-    def __output_batches(self) -> Iterator[DataRowBatch]:
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         output_batch = DataRowBatch(self.view, self.row_builder)
-        for input_batch in self.input:
+        async 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]
@@ -93,8 +92,3 @@ class ComponentIterationNode(ExecNode):
             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)

{pixeltable-0.3.0 → pixeltable-0.3.2}/pixeltable/exec/data_row_batch.py

@@ -21,7 +21,14 @@ class DataRowBatch:
     array_slot_idxs: list[int]
     rows: list[exprs.DataRow]
 
-    def __init__(self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, len: int = 0):
+    def __init__(
+        self, tbl: Optional[catalog.TableVersion], 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
         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()]
@@ -31,10 +38,15 @@ class DataRowBatch:
             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()]
-        self.rows = [
-            exprs.DataRow(row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
-            for _ in range(len)
-        ]
+        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)
+            ]
 
     def add_row(self, row: Optional[exprs.DataRow] = None) -> exprs.DataRow:
         if row is None: