pixeltable 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/catalog/table_version.py

@@ -6,7 +6,7 @@ import inspect
 import logging
 import time
 import uuid
-from typing import TYPE_CHECKING, Any, Iterable, Optional
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -26,7 +26,7 @@ from pixeltable.utils.media_store import MediaStore
 
 from ..func.globals import resolve_symbol
 from .column import Column
-from .globals import _POS_COLUMN_NAME, _ROWID_COLUMN_NAME, UpdateStatus, is_valid_identifier
+from .globals import _POS_COLUMN_NAME, _ROWID_COLUMN_NAME, UpdateStatus, is_valid_identifier, MediaValidation
 
 if TYPE_CHECKING:
     from pixeltable import exec, store
@@ -53,6 +53,7 @@ class TableVersion:
     name: str
     version: int
     comment: str
+    media_validation: MediaValidation
     num_retained_versions: int
     schema_version: int
     view_md: Optional[schema.ViewMd]
@@ -109,6 +110,7 @@ class TableVersion:
         self.view_md = tbl_md.view_md  # save this as-is, it's needed for _create_md()
         is_view = tbl_md.view_md is not None
         self.is_snapshot = (is_view and tbl_md.view_md.is_snapshot) or bool(is_snapshot)
+        self.media_validation = MediaValidation[schema_version_md.media_validation.upper()]
         # a mutable TableVersion doesn't have a static version
         self.effective_version = self.version if self.is_snapshot else None
 
@@ -182,7 +184,7 @@ class TableVersion:
     @classmethod
     def create(
             cls, session: orm.Session, dir_id: UUID, name: str, cols: list[Column], num_retained_versions: int,
-            comment: str, base_path: Optional[pxt.catalog.TableVersionPath] = None,
+            comment: str, media_validation: MediaValidation, base_path: Optional[pxt.catalog.TableVersionPath] = None,
             view_md: Optional[schema.ViewMd] = None
     ) -> tuple[UUID, Optional[TableVersion]]:
         # assign ids
@@ -214,11 +216,17 @@ class TableVersion:
             tbl_id=tbl_record.id, version=0, md=dataclasses.asdict(table_version_md))
 
         # create schema.TableSchemaVersion
-        schema_col_md = {col.id: schema.SchemaColumn(pos=pos, name=col.name) for pos, col in enumerate(cols)}
+        schema_col_md: dict[int, schema.SchemaColumn] = {}
+        for pos, col in enumerate(cols):
+            md = schema.SchemaColumn(
+                pos=pos, name=col.name,
+                media_validation=col._media_validation.name.lower() if col._media_validation is not None else None)
+            schema_col_md[col.id] = md
 
         schema_version_md = schema.TableSchemaVersionMd(
             schema_version=0, preceding_schema_version=None, columns=schema_col_md,
-            num_retained_versions=num_retained_versions, comment=comment)
+            num_retained_versions=num_retained_versions, comment=comment,
+            media_validation=media_validation.name.lower())
         schema_version_record = schema.TableSchemaVersion(
             tbl_id=tbl_record.id, schema_version=0, md=dataclasses.asdict(schema_version_md))
 
@@ -285,10 +293,15 @@ class TableVersion:
         self.cols_by_name = {}
         self.cols_by_id = {}
         for col_md in tbl_md.column_md.values():
-            col_name = schema_version_md.columns[col_md.id].name if col_md.id in schema_version_md.columns else None
+            schema_col_md = schema_version_md.columns[col_md.id] if col_md.id in schema_version_md.columns else None
+            col_name = schema_col_md.name if schema_col_md is not None else None
+            media_val = (
+                MediaValidation[schema_col_md.media_validation.upper()]
+                if schema_col_md is not None and schema_col_md.media_validation is not None else None
+            )
             col = Column(
                 col_id=col_md.id, name=col_name, col_type=ts.ColumnType.from_dict(col_md.col_type),
-                is_pk=col_md.is_pk, stored=col_md.stored,
+                is_pk=col_md.is_pk, stored=col_md.stored, media_validation=media_val,
                 schema_version_add=col_md.schema_version_add, schema_version_drop=col_md.schema_version_drop,
                 value_expr_dict=col_md.value_expr)
             col.tbl = self
@@ -349,7 +362,8 @@ class TableVersion:
             self.store_tbl = StoreTable(self)
 
     def _update_md(
-            self, timestamp: float, conn: sql.engine.Connection, update_tbl_version: bool = True, preceding_schema_version: Optional[int] = None
+        self, timestamp: float, conn: sql.engine.Connection, update_tbl_version: bool = True,
+        preceding_schema_version: Optional[int] = None
     ) -> None:
         """Writes table metadata to the database.
 
@@ -453,7 +467,9 @@ class TableVersion:
         self.idxs_by_name[idx_name] = idx_info
 
         # add the columns and update the metadata
-        status = self._add_columns([val_col, undo_col], conn)
+        # TODO support on_error='abort' for indices; it's tricky because of the way metadata changes are entangled
+        # with the database operations
+        status = self._add_columns([val_col, undo_col], conn, print_stats=False, on_error='ignore')
         # now create the index structure
         idx.create_index(self._store_idx_name(idx_id), val_col, conn)
 
@@ -478,7 +494,7 @@ class TableVersion:
             self._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
             _logger.info(f'Dropped index {idx_md.name} on table {self.name}')
 
-    def add_column(self, col: Column, print_stats: bool = False) -> UpdateStatus:
+    def add_column(self, col: Column, print_stats: bool, on_error: Literal['abort', 'ignore']) -> UpdateStatus:
         """Adds a column to the table.
         """
         assert not self.is_snapshot
@@ -498,9 +514,8 @@ class TableVersion:
         preceding_schema_version = self.schema_version
         self.schema_version = self.version
         with Env.get().engine.begin() as conn:
-            status = self._add_columns([col], conn, print_stats=print_stats)
+            status = self._add_columns([col], conn, print_stats=print_stats, on_error=on_error)
             _ = self._add_default_index(col, conn)
-            # TODO: what to do about errors?
             self._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
         _logger.info(f'Added column {col.name} to table {self.name}, new version: {self.version}')
 
@@ -512,7 +527,13 @@ class TableVersion:
         _logger.info(f'Column {col.name}: {msg}')
         return status
 
-    def _add_columns(self, cols: Iterable[Column], conn: sql.engine.Connection, print_stats: bool = False) -> UpdateStatus:
+    def _add_columns(
+        self,
+        cols: Iterable[Column],
+        conn: sql.engine.Connection,
+        print_stats: bool,
+        on_error: Literal['abort', 'ignore']
+    ) -> UpdateStatus:
         """Add and populate columns within the current transaction"""
         cols = list(cols)
         row_count = self.store_tbl.count(conn=conn)
@@ -550,10 +571,14 @@ class TableVersion:
             try:
                 plan.ctx.set_conn(conn)
                 plan.open()
-                num_excs = self.store_tbl.load_column(col, plan, value_expr_slot_idx, conn)
+                try:
+                    num_excs = self.store_tbl.load_column(col, plan, value_expr_slot_idx, conn, on_error)
+                except sql.exc.DBAPIError as exc:
+                    # Wrap the DBAPIError in an excs.Error to unify processing in the subsequent except block
+                    raise excs.Error(f'SQL error during execution of computed column `{col.name}`:\n{exc}') from exc
                 if num_excs > 0:
                     cols_with_excs.append(col)
-            except sql.exc.DBAPIError as e:
+            except excs.Error as exc:
                 self.cols.pop()
                 for col in cols:
                     # remove columns that we already added
@@ -564,7 +589,7 @@ class TableVersion:
                     del self.cols_by_id[col.id]
                 # we need to re-initialize the sqlalchemy schema
                 self.store_tbl.create_sa_tbl()
-                raise excs.Error(f'Error during SQL execution:\n{e}')
+                raise exc
             finally:
                 plan.close()
 
@@ -689,21 +714,32 @@ class TableVersion:
             plan = Planner.create_insert_plan(self, rows, ignore_errors=not fail_on_exception)
         else:
             plan = Planner.create_df_insert_plan(self, df, ignore_errors=not fail_on_exception)
+
+        # this is a base table; we generate rowids during the insert
+        def rowids() -> Iterator[int]:
+            while True:
+                rowid = self.next_rowid
+                self.next_rowid += 1
+                yield rowid
+
         if conn is None:
             with Env.get().engine.begin() as conn:
-                return self._insert(plan, conn, time.time(), print_stats)
+                return self._insert(
+                    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)
+            return self._insert(
+                plan, conn, time.time(), print_stats=print_stats, rowids=rowids(), abort_on_exc=fail_on_exception)
 
     def _insert(
-        self, exec_plan: 'exec.ExecNode', conn: sql.engine.Connection, timestamp: float, print_stats: bool = False,
+        self, exec_plan: 'exec.ExecNode', conn: sql.engine.Connection, timestamp: float, *,
+        rowids: Optional[Iterator[int]] = None, print_stats: bool = False, abort_on_exc: bool = False
     ) -> UpdateStatus:
         """Insert rows produced by exec_plan and propagate to views"""
         # we're creating a new version
         self.version += 1
         result = UpdateStatus()
-        num_rows, num_excs, cols_with_excs = self.store_tbl.insert_rows(exec_plan, conn, v_min=self.version)
-        self.next_rowid = num_rows
+        num_rows, num_excs, cols_with_excs = self.store_tbl.insert_rows(
+            exec_plan, conn, v_min=self.version, rowids=rowids, abort_on_exc=abort_on_exc)
         result.num_rows = num_rows
         result.num_excs = num_excs
         result.num_computed_values += exec_plan.ctx.num_computed_exprs * num_rows
@@ -714,7 +750,7 @@ class TableVersion:
         for view in self.mutable_views:
             from pixeltable.plan import Planner
             plan, _ = Planner.create_view_load_plan(view.path, propagates_insert=True)
-            status = view._insert(plan, conn, timestamp, print_stats)
+            status = view._insert(plan, conn, timestamp, print_stats=print_stats)
             result.num_rows += status.num_rows
             result.num_excs += status.num_excs
             result.num_computed_values += status.num_computed_values
@@ -751,9 +787,7 @@ class TableVersion:
                 raise excs.Error(f'Filter {analysis_info.filter} not expressible in SQL')
 
         with Env.get().engine.begin() as conn:
-            plan, updated_cols, recomputed_cols = (
-                Planner.create_update_plan(self.path, update_spec, [], where, cascade)
-            )
+            plan, updated_cols, recomputed_cols = Planner.create_update_plan(self.path, update_spec, [], where, cascade)
             from pixeltable.exprs import SqlElementCache
             result = self.propagate_update(
                 plan, where.sql_expr(SqlElementCache()) if where is not None else None, recomputed_cols,
@@ -1185,7 +1219,8 @@ class TableVersion:
             name=self.name, current_version=self.version, current_schema_version=self.schema_version,
             next_col_id=self.next_col_id, next_idx_id=self.next_idx_id, next_row_id=self.next_rowid,
             column_md=self._create_column_md(self.cols), index_md=self.idx_md,
-            external_stores=self._create_stores_md(self.external_stores.values()), view_md=self.view_md)
+            external_stores=self._create_stores_md(self.external_stores.values()), view_md=self.view_md,
+        )
 
     def _create_version_md(self, timestamp: float) -> schema.TableVersionMd:
         return schema.TableVersionMd(created_at=timestamp, version=self.version, schema_version=self.schema_version)
@@ -1193,11 +1228,14 @@ class TableVersion:
     def _create_schema_version_md(self, preceding_schema_version: int) -> schema.TableSchemaVersionMd:
         column_md: dict[int, schema.SchemaColumn] = {}
         for pos, col in enumerate(self.cols_by_name.values()):
-            column_md[col.id] = schema.SchemaColumn(pos=pos, name=col.name)
+            column_md[col.id] = schema.SchemaColumn(
+                pos=pos, name=col.name,
+                media_validation=col._media_validation.name.lower() if col._media_validation is not None else None)
         # preceding_schema_version to be set by the caller
         return schema.TableSchemaVersionMd(
             schema_version=self.schema_version, preceding_schema_version=preceding_schema_version,
-            columns=column_md, num_retained_versions=self.num_retained_versions, comment=self.comment)
+            columns=column_md, num_retained_versions=self.num_retained_versions, comment=self.comment,
+            media_validation=self.media_validation.name.lower())
 
     def as_dict(self) -> dict:
         return {'id': str(self.id), 'effective_version': self.effective_version}

pixeltable/catalog/table_version_path.py

@@ -91,14 +91,6 @@ class TableVersionPath:
         col = self.tbl_version.cols_by_name[col_name]
         return ColumnRef(col)
 
-    def __getitem__(self, index: object) -> Union[exprs.ColumnRef, pxt.DataFrame]:
-        """Return a ColumnRef for the given column name, or a DataFrame for the given slice.
-        """
-        if isinstance(index, str):
-            # basically <tbl>.<colname>
-            return self.__getattr__(index)
-        return pxt.DataFrame(self).__getitem__(index)
-
     def columns(self) -> list[Column]:
         """Return all user columns visible in this tbl version path, including columns from bases"""
         result = list(self.tbl_version.cols_by_name.values())

pixeltable/catalog/view.py

@@ -2,24 +2,21 @@ from __future__ import annotations
 
 import inspect
 import logging
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Set, Type
+from typing import TYPE_CHECKING, Any, Iterable, Optional
 from uuid import UUID
 
 import sqlalchemy.orm as orm
 
-import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
-import pixeltable.func as func
 import pixeltable.metadata.schema as md_schema
+import pixeltable.type_system as ts
+from pixeltable import catalog, exprs, func
 from pixeltable.env import Env
-from pixeltable.exceptions import Error
 from pixeltable.iterators import ComponentIterator
-from pixeltable.type_system import IntType, InvalidType
 
 from .catalog import Catalog
 from .column import Column
-from .globals import _POS_COLUMN_NAME, UpdateStatus
+from .globals import _POS_COLUMN_NAME, UpdateStatus, MediaValidation
 from .table import Table
 from .table_version import TableVersion
 from .table_version_path import TableVersionPath
@@ -52,11 +49,12 @@ class View(Table):
 
     @classmethod
     def _create(
-            cls, dir_id: UUID, name: str, base: TableVersionPath, schema: Dict[str, Any],
-            predicate: 'pxt.exprs.Expr', is_snapshot: bool, num_retained_versions: int, comment: str,
-            iterator_cls: Optional[Type[ComponentIterator]], iterator_args: Optional[Dict]
+            cls, dir_id: UUID, name: str, base: TableVersionPath, additional_columns: dict[str, Any],
+            predicate: Optional['pxt.exprs.Expr'], is_snapshot: bool, num_retained_versions: int, comment: str,
+            media_validation: MediaValidation,
+            iterator_cls: Optional[type[ComponentIterator]], iterator_args: Optional[dict]
     ) -> View:
-        columns = cls._create_columns(schema)
+        columns = cls._create_columns(additional_columns)
         cls._verify_schema(columns)
 
         # verify that filter can be evaluated in the context of the base
@@ -92,17 +90,17 @@ class View(Table):
                     func.Parameter(param_name, param_type, kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
                     for param_name, param_type in iterator_cls.input_schema().items()
                 ]
-                sig = func.Signature(InvalidType(), params)
+                sig = func.Signature(ts.InvalidType(), params)
                 from pixeltable.exprs import FunctionCall
                 FunctionCall.normalize_args(iterator_cls.__name__, sig, bound_args)
             except TypeError as e:
-                raise Error(f'Cannot instantiate iterator with given arguments: {e}')
+                raise excs.Error(f'Cannot instantiate iterator with given arguments: {e}')
 
             # prepend pos and output_schema columns to cols:
             # a component view exposes the pos column of its rowid;
             # we create that column here, so it gets assigned a column id;
             # stored=False: it is not stored separately (it's already stored as part of the rowid)
-            iterator_cols = [Column(_POS_COLUMN_NAME, IntType(), stored=False)]
+            iterator_cols = [Column(_POS_COLUMN_NAME, ts.IntType(), stored=False)]
             output_dict, unstored_cols = iterator_cls.output_schema(**bound_args)
             iterator_cols.extend([
                 Column(col_name, col_type, stored=col_name not in unstored_cols)
@@ -112,12 +110,12 @@ class View(Table):
             iterator_col_names = {col.name for col in iterator_cols}
             for col in columns:
                 if col.name in iterator_col_names:
-                    raise Error(f'Duplicate name: column {col.name} is already present in the iterator output schema')
+                    raise excs.Error(f'Duplicate name: column {col.name} is already present in the iterator output schema')
             columns = iterator_cols + columns
 
         with orm.Session(Env.get().engine, future=True) as session:
             from pixeltable.exprs import InlineDict
-            iterator_args_expr = InlineDict(iterator_args) if iterator_args is not None else None
+            iterator_args_expr: exprs.Expr = InlineDict(iterator_args) if iterator_args is not None else None
             iterator_class_fqn = f'{iterator_cls.__module__}.{iterator_cls.__name__}' if iterator_cls is not None \
                 else None
             base_version_path = cls._get_snapshot_path(base) if is_snapshot else base
@@ -142,7 +140,8 @@ class View(Table):
                 iterator_args=iterator_args_expr.as_dict() if iterator_args_expr is not None else None)
 
             id, tbl_version = TableVersion.create(
-                session, dir_id, name, columns, num_retained_versions, comment, base_path=base_version_path, view_md=view_md)
+                session, dir_id, name, columns, num_retained_versions, comment, media_validation=media_validation,
+                base_path=base_version_path, view_md=view_md)
             if tbl_version is None:
                 # this is purely a snapshot: we use the base's tbl version path
                 view = cls(id, dir_id, name, base_version_path, base.tbl_id(), snapshot_only=True)
@@ -168,11 +167,11 @@ class View(Table):
 
     @classmethod
     def _verify_column(
-            cls, col: Column, existing_column_names: Set[str], existing_query_names: Optional[Set[str]] = None
+            cls, col: Column, existing_column_names: set[str], existing_query_names: Optional[set[str]] = None
     ) -> None:
         # make sure that columns are nullable or have a default
         if not col.col_type.nullable and not col.is_computed:
-            raise Error(f'Column {col.name}: non-computed columns in views must be nullable')
+            raise excs.Error(f'Column {col.name}: non-computed columns in views must be nullable')
         super()._verify_column(col, existing_column_names, existing_query_names)
 
     @classmethod

pixeltable/dataframe.py

@@ -8,7 +8,7 @@ import logging
 import mimetypes
 import traceback
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, Iterator, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, Iterator, List, Optional, Sequence, Set, Tuple, Union
 
 import pandas as pd
 import pandas.io.formats.style
@@ -97,8 +97,8 @@ class DataFrameResultSet:
             return self._rows[index[0]][col_idx]
         raise excs.Error(f'Bad index: {index}')
 
-    def __iter__(self) -> DataFrameResultSetIterator:
-        return DataFrameResultSetIterator(self)
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        return (self._row_to_dict(i) for i in range(len(self)))
 
     def __eq__(self, other):
         if not isinstance(other, DataFrameResultSet):
@@ -106,19 +106,6 @@ class DataFrameResultSet:
         return self.to_pandas().equals(other.to_pandas())
 
 
-class DataFrameResultSetIterator:
-    def __init__(self, result_set: DataFrameResultSet):
-        self._result_set = result_set
-        self._idx = 0
-
-    def __next__(self) -> Dict[str, Any]:
-        if self._idx >= len(self._result_set):
-            raise StopIteration
-        row = self._result_set._row_to_dict(self._idx)
-        self._idx += 1
-        return row
-
-
 # # TODO: remove this; it's only here as a reminder that we still need to call release() in the current implementation
 # class AnalysisInfo:
 #     def __init__(self, tbl: catalog.TableVersion):
@@ -296,7 +283,7 @@ class DataFrame:
 
     def _create_query_plan(self) -> exec.ExecNode:
         # construct a group-by clause if we're grouping by a table
-        group_by_clause: List[exprs.Expr] = []
+        group_by_clause: Optional[list[exprs.Expr]] = None
         if self.grouping_tbl is not None:
             assert self.group_by_clause is None
             num_rowid_cols = len(self.grouping_tbl.store_tbl.rowid_columns())
@@ -315,8 +302,8 @@ class DataFrame:
             where_clause=self.where_clause,
             group_by_clause=group_by_clause,
             order_by_clause=self.order_by_clause if self.order_by_clause is not None else [],
-            limit=self.limit_val if self.limit_val is not None else 0,
-        )  # limit_val == 0: no limit_val
+            limit=self.limit_val
+        )
 
 
     def show(self, n: int = 20) -> DataFrameResultSet:
@@ -384,15 +371,10 @@ class DataFrame:
             group_by_clause=group_by_clause, grouping_tbl=self.grouping_tbl,
             order_by_clause=order_by_clause, limit=self.limit_val)
 
-    def collect(self) -> DataFrameResultSet:
-        return self._collect()
-
-    def _collect(self, conn: Optional[sql.engine.Connection] = None) -> DataFrameResultSet:
+    def _output_row_iterator(self, conn: Optional[sql.engine.Connection] = None) -> Iterator[list]:
         try:
-            result_rows = []
             for data_row in self._exec(conn):
-                result_row = [data_row[e.slot_idx] for e in self._select_list_exprs]
-                result_rows.append(result_row)
+                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:
@@ -412,7 +394,11 @@ class DataFrame:
         except sql.exc.DBAPIError as e:
             raise excs.Error(f'Error during SQL execution:\n{e}')
 
-        return DataFrameResultSet(result_rows, self.schema)
+    def collect(self) -> DataFrameResultSet:
+        return self._collect()
+
+    def _collect(self, conn: Optional[sql.engine.Connection] = None) -> DataFrameResultSet:
+        return DataFrameResultSet(list(self._output_row_iterator(conn)), self.schema)
 
     def count(self) -> int:
         from pixeltable.plan import Planner
@@ -629,17 +615,15 @@ class DataFrame:
         if self.limit_val is not None:
             raise excs.Error(f'Cannot use `{op_name}` after `limit`')
 
-    def __getitem__(self, index: object) -> DataFrame:
+    def __getitem__(self, index: Union[exprs.Expr, Sequence[exprs.Expr]]) -> DataFrame:
         """
         Allowed:
         - [List[Expr]]/[Tuple[Expr]]: setting the select list
         - [Expr]: setting a single-col select list
         """
-        if isinstance(index, tuple):
-            index = list(index)
         if isinstance(index, exprs.Expr):
-            index = [index]
-        if isinstance(index, list):
+            return self.select(index)
+        if isinstance(index, Sequence):
             return self.select(*index)
         raise TypeError(f'Invalid index type: {type(index)}')
 

pixeltable/env.py

@@ -342,7 +342,7 @@ class Env:
 
         if create_db:
             from pixeltable.metadata import schema
-            schema.Base.metadata.create_all(self._sa_engine)
+            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}')
@@ -494,13 +494,18 @@ class Env:
         self.__register_package('anthropic')
         self.__register_package('boto3')
         self.__register_package('datasets')
+        self.__register_package('fiftyone')
         self.__register_package('fireworks', library_name='fireworks-ai')
+        self.__register_package('huggingface_hub', library_name='huggingface-hub')
         self.__register_package('label_studio_sdk', library_name='label-studio-sdk')
+        self.__register_package('llama_cpp', library_name='llama-cpp-python')
         self.__register_package('mistralai')
         self.__register_package('mistune')
+        self.__register_package('ollama')
         self.__register_package('openai')
         self.__register_package('openpyxl')
         self.__register_package('pyarrow')
+        self.__register_package('replicate')
         self.__register_package('sentence_transformers', library_name='sentence-transformers')
         self.__register_package('spacy')
         self.__register_package('tiktoken')

pixeltable/exec/__init__.py

@@ -6,6 +6,5 @@ from .exec_context import ExecContext
 from .exec_node import ExecNode
 from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
-from .media_validation_node import MediaValidationNode
 from .row_update_node import RowUpdateNode
-from .sql_node import SqlLookupNode, SqlScanNode
+from .sql_node import SqlLookupNode, SqlScanNode, SqlAggregationNode, SqlNode

pixeltable/exec/aggregation_node.py

@@ -2,28 +2,43 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Iterable, List, Optional, Any
+from typing import Any, Iterable, Iterator, Optional, cast
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.exprs as 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]]
+    input_exprs: list[exprs.Expr]
+    agg_fn_eval_ctx: exprs.RowBuilder.EvalCtx
+    agg_fn_calls: list[exprs.FunctionCall]
+    output_batch: DataRowBatch
+
     def __init__(
-            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: List[exprs.Expr],
-            agg_fn_calls: List[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
+            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
     ):
-        super().__init__(row_builder, group_by + agg_fn_calls, input_exprs, input)
+        output_exprs: list[exprs.Expr] = [] if group_by is None else list(group_by)
+        output_exprs.extend(agg_fn_calls)
+        super().__init__(row_builder, output_exprs, input_exprs, input)
         self.input = input
         self.group_by = group_by
         self.input_exprs = list(input_exprs)
-        self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=input_exprs)
+        self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=self.input_exprs)
         # we need to make sure to refer to the same exprs that RowBuilder.eval() will use
-        self.agg_fn_calls = self.agg_fn_eval_ctx.target_exprs
+        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)
 
     def _reset_agg_state(self, row_num: int) -> None:
@@ -45,17 +60,14 @@ 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 __next__(self) -> DataRowBatch:
-        if self.output_batch is None:
-            raise StopIteration
-
+    def __iter__(self) -> Iterator[DataRowBatch]:
         prev_row: Optional[exprs.DataRow] = None
-        current_group: Optional[List[Any]] = None  # the values of the group-by exprs
+        current_group: Optional[list[Any]] = None  # the values of the group-by exprs
         num_input_rows = 0
         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]
+                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)
@@ -71,9 +83,7 @@ class AggregationNode(ExecNode):
         self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
         self.output_batch.add_row(prev_row)
 
-        result = self.output_batch
-        result.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
-        self.output_batch = None
-        _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(result.rows)} rows')
-        return result
+        self.output_batch.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        _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

@@ -79,7 +79,7 @@ class CachePrefetchNode(ExecNode):
 
         return input_batch
 
-    def _fetch_url(self, row: exprs.DataRow, slot_idx: int) -> Optional[str]:
+    def _fetch_url(self, row: exprs.DataRow, slot_idx: int) -> Optional[Path]:
         """Fetches a remote URL into Env.tmp_dir and returns its path"""
         url = row.file_urls[slot_idx]
         parsed = urllib.parse.urlparse(url)

pixeltable/exec/data_row_batch.py

@@ -14,6 +14,13 @@ class DataRowBatch:
 
     Contains the metadata needed to initialize DataRows.
     """
+    tbl: Optional[catalog.TableVersion]
+    row_builder: exprs.RowBuilder
+    img_slot_idxs: list[int]
+    media_slot_idxs: list[int]  # non-image media slots
+    array_slot_idxs: list[int]
+    rows: list[exprs.DataRow]
+
     def __init__(self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, len: int = 0):
         self.tbl = tbl
         self.row_builder = row_builder
@@ -39,17 +46,10 @@ class DataRowBatch:
     def pop_row(self) -> exprs.DataRow:
         return self.rows.pop()
 
-    def set_row_ids(self, row_ids: List[int]) -> None:
-        """Sets pks for rows in batch"""
-        assert self.tbl is not None
-        assert len(row_ids) == len(self.rows)
-        for row, row_id in zip(self.rows, row_ids):
-            row.set_pk((row_id, self.tbl))
-
     def __len__(self) -> int:
         return len(self.rows)
 
-    def __getitem__(self, index: object) -> exprs.DataRow:
+    def __getitem__(self, index: int) -> exprs.DataRow:
         return self.rows[index]
 
     def flush_imgs(
@@ -74,21 +74,4 @@ class DataRowBatch:
                 row.flush_img(slot_idx)
 
     def __iter__(self) -> Iterator[exprs.DataRow]:
-        return DataRowBatchIterator(self)
-
-
-class DataRowBatchIterator:
-    """
-    Iterator over a DataRowBatch.
-    """
-    def __init__(self, batch: DataRowBatch):
-        self.row_batch = batch
-        self.index = 0
-
-    def __next__(self) -> exprs.DataRow:
-        if self.index >= len(self.row_batch.rows):
-            raise StopIteration
-        row = self.row_batch.rows[self.index]
-        self.index += 1
-        return row
-
+        return iter(self.rows)