pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exceptions.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 from types import TracebackType
 from typing import TYPE_CHECKING, Any
 
@@ -10,8 +9,13 @@ class Error(Exception):
     pass
 
 
-@dataclass
 class ExprEvalError(Exception):
+    """
+    Used during query execution to signal expr evaluation failures.
+
+    NOT A USER-FACING EXCEPTION. All ExprEvalError instances need to be converted into Error instances.
+    """
+
     expr: 'exprs.Expr'
     expr_msg: str
     exc: Exception
@@ -19,6 +23,26 @@ class ExprEvalError(Exception):
     input_vals: list[Any]
     row_num: int
 
+    def __init__(
+        self,
+        expr: 'exprs.Expr',
+        expr_msg: str,
+        exc: Exception,
+        exc_tb: TracebackType,
+        input_vals: list[Any],
+        row_num: int,
+    ) -> None:
+        exct = type(exc)
+        super().__init__(
+            f'Expression evaluation failed with an error of type `{exct.__module__}.{exct.__qualname__}`:\n{expr}'
+        )
+        self.expr = expr
+        self.expr_msg = expr_msg
+        self.exc = exc
+        self.exc_tb = exc_tb
+        self.input_vals = input_vals
+        self.row_num = row_num
+
 
 class PixeltableWarning(Warning):
     pass

pixeltable/exec/__init__.py

@@ -1,10 +1,15 @@
+# ruff: noqa: F401
+
 from .aggregation_node import AggregationNode
 from .cache_prefetch_node import CachePrefetchNode
+from .cell_materialization_node import CellMaterializationNode
+from .cell_reconstruction_node import CellReconstructionNode
 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 .expr_eval import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
+from .object_store_save_node import ObjectStoreSaveNode
 from .row_update_node import RowUpdateNode
-from .sql_node import SqlLookupNode, SqlScanNode, SqlAggregationNode, SqlNode, SqlJoinNode
+from .sql_node import SqlAggregationNode, SqlJoinNode, SqlLookupNode, SqlNode, SqlSampleNode, SqlScanNode

pixeltable/exec/aggregation_node.py

@@ -2,32 +2,38 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Any, Iterable, Iterator, Optional, cast
+from typing import Any, AsyncIterator, Iterable, cast
 
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exceptions as excs, exprs
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
 _logger = logging.getLogger('pixeltable')
 
+
 class AggregationNode(ExecNode):
     """
     In-memory aggregation for UDAs.
 
     At the moment, this returns all results in a single DataRowBatch.
     """
-    group_by: Optional[list[exprs.Expr]]
+
+    group_by: list[exprs.Expr] | None
     input_exprs: list[exprs.Expr]
     agg_fn_eval_ctx: exprs.RowBuilder.EvalCtx
     agg_fn_calls: list[exprs.FunctionCall]
     output_batch: DataRowBatch
+    limit: int | None
 
     def __init__(
-            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: Optional[list[exprs.Expr]],
-            agg_fn_calls: list[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
+        self,
+        tbl: catalog.TableVersionHandle,
+        row_builder: exprs.RowBuilder,
+        group_by: list[exprs.Expr] | None,
+        agg_fn_calls: list[exprs.FunctionCall],
+        input_exprs: Iterable[exprs.Expr],
+        input: ExecNode,
     ):
         output_exprs: list[exprs.Expr] = [] if group_by is None else list(group_by)
         output_exprs.extend(agg_fn_calls)
@@ -39,51 +45,63 @@ class AggregationNode(ExecNode):
         # we need to make sure to refer to the same exprs that RowBuilder.eval() will use
         self.agg_fn_calls = [cast(exprs.FunctionCall, e) for e in self.agg_fn_eval_ctx.target_exprs]
         # create output_batch here, rather than in __iter__(), so we don't need to remember tbl and row_builder
-        self.output_batch = DataRowBatch(tbl, row_builder, 0)
+        self.output_batch = DataRowBatch(row_builder)
+        self.limit = None
+
+    def set_limit(self, limit: int) -> None:
+        # we can't propagate the limit to our input
+        self.limit = limit
 
     def _reset_agg_state(self, row_num: int) -> None:
         for fn_call in self.agg_fn_calls:
             try:
                 fn_call.reset_agg()
-            except Exception as e:
+            except Exception as exc:
                 _, _, exc_tb = sys.exc_info()
                 expr_msg = f'init() function of the aggregate {fn_call}'
-                raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, [], row_num)
+                raise excs.ExprEvalError(fn_call, expr_msg, exc, exc_tb, [], row_num) from exc
 
     def _update_agg_state(self, row: exprs.DataRow, row_num: int) -> None:
         for fn_call in self.agg_fn_calls:
             try:
                 fn_call.update(row)
-            except Exception as e:
+            except Exception as exc:
                 _, _, exc_tb = sys.exc_info()
                 expr_msg = f'update() function of the aggregate {fn_call}'
                 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)
+                raise excs.ExprEvalError(fn_call, expr_msg, exc, exc_tb, input_vals, row_num) from exc
 
-    def __iter__(self) -> Iterator[DataRowBatch]:
-        prev_row: Optional[exprs.DataRow] = None
-        current_group: Optional[list[Any]] = None  # the values of the group-by exprs
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        prev_row: exprs.DataRow | None = None
+        current_group: list[Any] | None = None  # the values of the group-by exprs
         num_input_rows = 0
-        for row_batch in self.input:
+        num_output_rows = 0
+        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
+
                 if current_group is None:
                     current_group = group
                     self._reset_agg_state(0)
+
                 if group != current_group:
                     # we're entering a new group, emit a row for the previous one
                     self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
                     self.output_batch.add_row(prev_row)
+                    num_output_rows += 1
+                    if self.limit is not None and num_output_rows == self.limit:
+                        yield self.output_batch
+                        return
                     current_group = group
                     self._reset_agg_state(0)
                 self._update_agg_state(row, 0)
                 prev_row = row
-        # emit the last group
-        self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
-        self.output_batch.add_row(prev_row)
 
-        self.output_batch.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        if prev_row is not None:
+            # emit the last group
+            self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
+            self.output_batch.add_row(prev_row)
+
         _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(self.output_batch.rows)} rows')
         yield self.output_batch
-

pixeltable/exec/cache_prefetch_node.py

@@ -9,14 +9,12 @@ import urllib.request
 from collections import deque
 from concurrent import futures
 from pathlib import Path
-from typing import Optional, Any, Iterator
+from typing import AsyncIterator, Iterator
 from uuid import UUID
 
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
-from pixeltable import catalog
+from pixeltable import exceptions as excs, exprs
 from pixeltable.utils.filecache import FileCache
+from pixeltable.utils.object_stores import ObjectOps
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -28,49 +26,45 @@ class CachePrefetchNode(ExecNode):
     """Brings files with external URLs into the cache
 
     TODO:
-    - adapting the number of download threads at runtime to maximize throughput
+    - Process a row at a time and limit the number of in-flight rows to control memory usage
+    - Create asyncio.Tasks to consume our input in order to increase concurrency.
     """
+
+    QUEUE_DEPTH_HIGH_WATER = 50  # target number of in-flight requests
+    QUEUE_DEPTH_LOW_WATER = 20  # target number of in-flight requests
     BATCH_SIZE = 16
-    NUM_EXECUTOR_THREADS = 16
+    MAX_WORKERS = 15
 
     retain_input_order: bool  # if True, return rows in the exact order they were received
     file_col_info: list[exprs.ColumnSlotIdx]
-    boto_client: Optional[Any]
-    boto_client_lock: threading.Lock
 
     # execution state
-    batch_tbl_version: Optional[catalog.TableVersion]  # needed to construct output batches
     num_returned_rows: int
 
     # ready_rows: rows that are ready to be returned, ordered by row idx;
     # the implied row idx of ready_rows[0] is num_returned_rows
-    ready_rows: deque[Optional[exprs.DataRow]]
+    ready_rows: deque[exprs.DataRow | None]
 
     in_flight_rows: dict[int, CachePrefetchNode.RowState]  # rows with in-flight urls; id(row) -> RowState
     in_flight_requests: dict[futures.Future, str]  # in-flight requests for urls; future -> URL
     in_flight_urls: dict[str, list[tuple[exprs.DataRow, exprs.ColumnSlotIdx]]]  # URL -> [(row, info)]
     input_finished: bool
-    row_idx: Iterator[Optional[int]]
+    row_idx: Iterator[int | None]
 
     @dataclasses.dataclass
     class RowState:
         row: exprs.DataRow
-        idx: Optional[int]  # position in input stream; None if we don't retain input order
+        idx: int | None  # position in input stream; None if we don't retain input order
         num_missing: int  # number of missing URLs in this row
 
     def __init__(
-            self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode,
-            retain_input_order: bool = True):
+        self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True
+    ):
         # input_/output_exprs=[]: we don't have anything to evaluate
         super().__init__(input.row_builder, [], [], input)
         self.retain_input_order = retain_input_order
         self.file_col_info = file_col_info
 
-        # clients for specific services are constructed as needed, because it's time-consuming
-        self.boto_client = None
-        self.boto_client_lock = threading.Lock()
-
-        self.batch_tbl_version = None
         self.num_returned_rows = 0
         self.ready_rows = deque()
         self.in_flight_rows = {}
@@ -78,26 +72,44 @@ class CachePrefetchNode(ExecNode):
         self.in_flight_urls = {}
         self.input_finished = False
         self.row_idx = itertools.count() if retain_input_order else itertools.repeat(None)
+        assert self.QUEUE_DEPTH_HIGH_WATER > self.QUEUE_DEPTH_LOW_WATER
 
-    def __iter__(self) -> Iterator[DataRowBatch]:
-        input_iter = iter(self.input)
-        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)
-
-            while True:
-                # try to assemble a full batch of output rows
-                if not self.__has_ready_batch() and len(self.in_flight_requests) > 0:
-                    self.__wait_for_requests()
+    @property
+    def queued_work(self) -> int:
+        return len(self.in_flight_requests)
 
-                # 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)
+    async def get_input_batch(self, input_iter: AsyncIterator[DataRowBatch]) -> DataRowBatch | None:
+        """Get the next batch of input rows, or None if there are no more rows"""
+        try:
+            input_batch = await anext(input_iter)
+            if input_batch is None:
+                self.input_finished = True
+            return input_batch
+        except StopAsyncIteration:
+            self.input_finished = True
+            return None
 
-                if len(self.ready_rows) > 0:
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        input_iter = aiter(self.input)
+        with futures.ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
+            while True:
+                # Create work to fill the queue to the high water mark ... ?without overrunning the in-flight row limit.
+                while not self.input_finished and self.queued_work < self.QUEUE_DEPTH_HIGH_WATER:
+                    input_batch = await self.get_input_batch(input_iter)
+                    if input_batch is not None:
+                        self.__process_input_batch(input_batch, executor)
+
+                # Wait for enough completions to enable more queueing or if we're done
+                while self.queued_work > self.QUEUE_DEPTH_LOW_WATER or (self.input_finished and self.queued_work > 0):
+                    done, _ = futures.wait(self.in_flight_requests, return_when=futures.FIRST_COMPLETED)
+                    self.__process_completions(done, ignore_errors=self.ctx.ignore_errors)
+
+                # Emit results to meet batch size requirements or empty the in-flight row queue
+                if self.__has_ready_batch() or (
+                    len(self.ready_rows) > 0 and self.input_finished and self.queued_work == 0
+                ):
                     # create DataRowBatch from the first BATCH_SIZE ready rows
-                    batch = DataRowBatch(self.batch_tbl_version, self.row_builder)
+                    batch = DataRowBatch(self.row_builder)
                     rows = [self.ready_rows.popleft() for _ in range(min(self.BATCH_SIZE, len(self.ready_rows)))]
                     for row in rows:
                         assert row is not None
@@ -106,23 +118,16 @@ class CachePrefetchNode(ExecNode):
                     _logger.debug(f'returning {len(rows)} rows')
                     yield batch
 
-                if self.input_finished and self.__num_pending_rows() == 0:
+                if self.input_finished and self.queued_work == 0 and len(self.ready_rows) == 0:
                     return
 
-    def __num_pending_rows(self) -> int:
-        return len(self.in_flight_rows) + len(self.ready_rows)
-
     def __has_ready_batch(self) -> bool:
         """True if there are >= BATCH_SIZES entries in ready_rows and the first BATCH_SIZE ones are all non-None"""
         return (
             sum(int(row is not None) for row in itertools.islice(self.ready_rows, self.BATCH_SIZE)) == self.BATCH_SIZE
         )
 
-    def __ready_prefix_len(self) -> int:
-        """Length of the non-None prefix of ready_rows (= what we can return right now)"""
-        return sum(1 for _ in itertools.takewhile(lambda x: x is not None, self.ready_rows))
-
-    def __add_ready_row(self, row: exprs.DataRow, row_idx: Optional[int]) -> None:
+    def __add_ready_row(self, row: exprs.DataRow, row_idx: int | None) -> None:
         if row_idx is None:
             self.ready_rows.append(row)
         else:
@@ -132,46 +137,36 @@ class CachePrefetchNode(ExecNode):
                 self.ready_rows.extend([None] * (idx - len(self.ready_rows) + 1))
             self.ready_rows[idx] = row
 
-    def __wait_for_requests(self) -> None:
-        """Wait for in-flight requests to complete until we have a full batch of rows"""
+    def __process_completions(self, done: set[futures.Future], ignore_errors: bool) -> None:
         file_cache = FileCache.get()
-        _logger.debug(f'waiting for requests; ready_batch_size={self.__ready_prefix_len()}')
-        while not self.__has_ready_batch() and len(self.in_flight_requests) > 0:
-            done, _ = futures.wait(self.in_flight_requests, return_when=futures.FIRST_COMPLETED)
-            for f in done:
-                url = self.in_flight_requests.pop(f)
-                tmp_path, exc = f.result()
-                local_path: Optional[Path] = None
-                if tmp_path is not None:
-                    # register the file with the cache for the first column in which it's missing
-                    assert url in self.in_flight_urls
-                    _, info = self.in_flight_urls[url][0]
-                    local_path = file_cache.add(info.col.tbl.id, info.col.id, url, tmp_path)
-                    _logger.debug(f'cached {url} as {local_path}')
-
-                # add the local path/exception to the slots that reference the url
-                for row, info in self.in_flight_urls.pop(url):
-                    if exc is not None:
-                        self.row_builder.set_exc(row, info.slot_idx, exc)
-                    else:
-                        assert local_path is not None
-                        row.set_file_path(info.slot_idx, str(local_path))
-                    state = self.in_flight_rows[id(row)]
-                    state.num_missing -= 1
-                    if state.num_missing == 0:
-                        del self.in_flight_rows[id(row)]
-                        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:
-        assert not self.input_finished
-        input_batch = next(input, None)
-        if input_batch is None:
-            self.input_finished = True
-            return
-        if self.batch_tbl_version is None:
-            self.batch_tbl_version = input_batch.tbl
-
+        for f in done:
+            url = self.in_flight_requests.pop(f)
+            tmp_path, exc = f.result()
+            if exc is not None and not ignore_errors:
+                raise exc
+            local_path: Path | None = None
+            if tmp_path is not None:
+                # register the file with the cache for the first column in which it's missing
+                assert url in self.in_flight_urls
+                _, info = self.in_flight_urls[url][0]
+                local_path = file_cache.add(info.col.get_tbl().id, info.col.id, url, tmp_path)
+                _logger.debug(f'cached {url} as {local_path}')
+
+            # add the local path/exception to the slots that reference the url
+            for row, info in self.in_flight_urls.pop(url):
+                if exc is not None:
+                    self.row_builder.set_exc(row, info.slot_idx, exc)
+                else:
+                    assert local_path is not None
+                    row.set_file_path(info.slot_idx, str(local_path))
+                state = self.in_flight_rows[id(row)]
+                state.num_missing -= 1
+                if state.num_missing == 0:
+                    del self.in_flight_rows[id(row)]
+                    self.__add_ready_row(row, state.idx)
+
+    def __process_input_batch(self, input_batch: DataRowBatch, executor: futures.ThreadPoolExecutor) -> None:
+        """Process a batch of input rows, submitting URLs for download and adding ready rows to ready_rows"""
         file_cache = FileCache.get()
 
         # URLs from this input batch that aren't already in the file cache;
@@ -179,7 +174,7 @@ class CachePrefetchNode(ExecNode):
         # the time it takes to get the next batch together
         cache_misses: list[str] = []
 
-        url_pos: dict[str, int] = {}  # url -> row_idx; used for logging
+        url_pos: dict[str, int | None] = {}  # url -> row_idx; used for logging
         for row in input_batch:
             # identify missing local files in input batch, or fill in their paths if they're already cached
             num_missing = 0
@@ -218,8 +213,10 @@ class CachePrefetchNode(ExecNode):
             _logger.debug(f'submitted {url} for idx {url_pos[url]}')
             self.in_flight_requests[f] = url
 
-    def __fetch_url(self, url: str) -> tuple[Optional[Path], Optional[Exception]]:
-        """Fetches a remote URL into Env.tmp_dir and returns its path"""
+    def __fetch_url(self, url: str) -> tuple[Path | None, Exception | None]:
+        """Fetches a remote URL into the TempStore and returns its path"""
+        from pixeltable.utils.local_store import TempStore
+
         _logger.debug(f'fetching url={url} thread_name={threading.current_thread().name}')
         parsed = urllib.parse.urlparse(url)
         # Use len(parsed.scheme) > 1 here to ensure we're not being passed
@@ -227,36 +224,17 @@ class CachePrefetchNode(ExecNode):
         assert len(parsed.scheme) > 1 and parsed.scheme != 'file'
         # preserve the file extension, if there is one
         extension = ''
-        if parsed.path != '':
+        if parsed.path:
             p = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
             extension = p.suffix
-        tmp_path = env.Env.get().create_tmp_path(extension=extension)
+        tmp_path = TempStore.create_path(extension=extension)
         try:
             _logger.debug(f'Downloading {url} to {tmp_path}')
-            if parsed.scheme == 's3':
-                from pixeltable.utils.s3 import get_client
-                with self.boto_client_lock:
-                    if self.boto_client is None:
-                        config = {
-                            'max_pool_connections': self.NUM_EXECUTOR_THREADS + 4,  # +4: leave some headroom
-                            'connect_timeout': 5,
-                            'read_timeout': 30,
-                            'retries': {'max_attempts': 3, 'mode': 'adaptive'},
-                        }
-                        self.boto_client = get_client(**config)
-                self.boto_client.download_file(parsed.netloc, parsed.path.lstrip('/'), str(tmp_path))
-            elif parsed.scheme == 'http' or parsed.scheme == 'https':
-                with urllib.request.urlopen(url) as resp, open(tmp_path, 'wb') as f:
-                    data = resp.read()
-                    f.write(data)
-            else:
-                assert False, f'Unsupported URL scheme: {parsed.scheme}'
+            ObjectOps.copy_object_to_local_file(url, tmp_path)
             _logger.debug(f'Downloaded {url} to {tmp_path}')
             return tmp_path, None
         except Exception as e:
             # we want to add the file url to the exception message
             exc = excs.Error(f'Failed to download {url}: {e}')
             _logger.debug(f'Failed to download {url}: {e}', exc_info=e)
-            if not self.ctx.ignore_errors:
-                raise exc from None  # suppress original exception
             return None, exc