pixeltable 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

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.1"
+__version_tuple__ = (0, 3, 1)

pixeltable/catalog/table_version.py

@@ -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/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/env.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from abc import abstractmethod
 import datetime
 import glob
 import http.server
@@ -15,9 +16,9 @@ import sys
 import threading
 import uuid
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Type
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 import pixeltable_pgserver
@@ -33,6 +34,11 @@ if TYPE_CHECKING:
     import spacy
 
 
+_logger = logging.getLogger('pixeltable')
+
+T = TypeVar('T')
+
+
 class Env:
     """
     Store for runtime globals.
@@ -70,6 +76,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:
@@ -121,6 +129,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
@@ -609,6 +619,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 +706,6 @@ class Config:
     """
     __config: dict[str, Any]
 
-    T = TypeVar('T')
-
     @classmethod
     def from_file(cls, path: Path) -> Config:
         """
@@ -767,3 +785,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/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/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/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/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/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:

pixeltable/exec/exec_node.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional, TypeVar
+import asyncio
+import logging
+import sys
+from typing import Iterable, Iterator, Optional, TypeVar, AsyncIterator
 
 import pixeltable.exprs as exprs
-
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
 
+_logger = logging.getLogger('pixeltable')
 
 class ExecNode(abc.ABC):
     """Base class of all execution nodes"""
@@ -17,7 +20,6 @@ class ExecNode(abc.ABC):
     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],
@@ -34,7 +36,6 @@ class ExecNode(abc.ABC):
         ]
         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
@@ -47,15 +48,34 @@ class ExecNode(abc.ABC):
         if self.input is not None:
             self.input.set_stored_img_cols(stored_img_cols)
 
-    # TODO: make this an abstractmethod when __next__() is removed
-    def __iter__(self) -> Iterator[DataRowBatch]:
-        return self
+    @abc.abstractmethod
+    def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        pass
 
-    # TODO: remove this and switch every subclass over to implementing __iter__
-    def __next__(self) -> DataRowBatch:
-        if self.__iter is None:
-            self.__iter = iter(self)
-        return next(self.__iter)
+    def __iter__(self) -> Iterator[DataRowBatch]:
+        try:
+            # check if we are already in an event loop (eg, Jupyter's); if so, patch it to allow nested event loops
+            _ = asyncio.get_event_loop()
+            import nest_asyncio  # type: ignore
+            nest_asyncio.apply()
+        except RuntimeError:
+            pass
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        if 'pytest' in sys.modules:
+            loop.set_debug(True)
+
+        aiter = self.__aiter__()
+        try:
+            while True:
+                batch: DataRowBatch = loop.run_until_complete(aiter.__anext__())
+                yield batch
+        except StopAsyncIteration:
+            pass
+        finally:
+            loop.close()
 
     def open(self) -> None:
         """Bottom-up initialization of nodes for execution. Must be called before __next__."""

pixeltable/exec/expr_eval/__init__.py

@@ -0,0 +1 @@
+from .expr_eval_node import ExprEvalNode