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

pixeltable/__version__.py

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

pixeltable/catalog/__init__.py

@@ -1,7 +1,7 @@
 from .catalog import Catalog
 from .column import Column
 from .dir import Dir
-from .globals import UpdateStatus, is_valid_identifier, is_valid_path
+from .globals import UpdateStatus, is_valid_identifier, is_valid_path, MediaValidation
 from .insertable_table import InsertableTable
 from .named_function import NamedFunction
 from .path import Path

pixeltable/catalog/column.py

@@ -8,24 +8,43 @@ import sqlalchemy as sql
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable import exprs
-
-from .globals import is_valid_identifier
+from .globals import is_valid_identifier, MediaValidation
 
 if TYPE_CHECKING:
     from .table_version import TableVersion
 
 _logger = logging.getLogger('pixeltable')
 
+
 class Column:
     """Representation of a column in the schema of a Table/DataFrame.
 
     A Column contains all the metadata necessary for executing queries and updates against a particular version of a
     table/view.
     """
+    name: str
+    id: Optional[int]
+    col_type: ts.ColumnType
+    stored: bool
+    is_pk: bool
+    _media_validation: Optional[MediaValidation]  # if not set, TableVersion.media_validation applies
+    schema_version_add: Optional[int]
+    schema_version_drop: Optional[int]
+    _records_errors: Optional[bool]
+    sa_col: Optional[sql.schema.Column]
+    sa_col_type: Optional[sql.sqltypes.TypeEngine]
+    sa_errormsg_col: Optional[sql.schema.Column]
+    sa_errortype_col: Optional[sql.schema.Column]
+    compute_func: Optional[Callable]
+    _value_expr: Optional[exprs.Expr]
+    value_expr_dict: Optional[dict[str, Any]]
+    dependent_cols: set[Column]
+    tbl: Optional[TableVersion]
+
     def __init__(
             self, name: Optional[str], col_type: Optional[ts.ColumnType] = None,
             computed_with: Optional[Union[exprs.Expr, Callable]] = None,
-            is_pk: bool = False, stored: bool = True,
+            is_pk: bool = False, stored: bool = True, media_validation: Optional[MediaValidation] = None,
             col_id: Optional[int] = None, schema_version_add: Optional[int] = None,
             schema_version_drop: Optional[int] = None, sa_col_type: Optional[sql.sqltypes.TypeEngine] = None,
             records_errors: Optional[bool] = None, value_expr_dict: Optional[dict[str, Any]] = None,
@@ -61,8 +80,8 @@ class Column:
         if col_type is None and computed_with is None:
             raise excs.Error(f'Column `{name}`: col_type is required if computed_with is not specified')
 
-        self._value_expr: Optional[exprs.Expr] = None
-        self.compute_func: Optional[Callable] = None
+        self._value_expr = None
+        self.compute_func = None
         self.value_expr_dict = value_expr_dict
         if computed_with is not None:
             value_expr = exprs.Expr.from_object(computed_with)
@@ -86,24 +105,24 @@ class Column:
         assert self.col_type is not None
 
         self.stored = stored
-        self.dependent_cols: set[Column] = set()  # cols with value_exprs that reference us; set by TableVersion
+        self.dependent_cols = set()  # cols with value_exprs that reference us; set by TableVersion
         self.id = col_id
         self.is_pk = is_pk
+        self._media_validation = media_validation
         self.schema_version_add = schema_version_add
         self.schema_version_drop = schema_version_drop
 
         self._records_errors = records_errors
 
         # column in the stored table for the values of this Column
-        self.sa_col: Optional[sql.schema.Column] = None
+        self.sa_col = None
         self.sa_col_type = sa_col_type
 
         # computed cols also have storage columns for the exception string and type
-        self.sa_errormsg_col: Optional[sql.schema.Column] = None
-        self.sa_errortype_col: Optional[sql.schema.Column] = None
+        self.sa_errormsg_col = None
+        self.sa_errortype_col = None
 
-        from .table_version import TableVersion
-        self.tbl: Optional[TableVersion] = None  # set by owning TableVersion
+        self.tbl = None  # set by owning TableVersion
 
     @property
     def value_expr(self) -> Optional[exprs.Expr]:
@@ -160,6 +179,13 @@ class Column:
         assert self.tbl is not None
         return f'{self.tbl.name}.{self.name}'
 
+    @property
+    def media_validation(self) -> MediaValidation:
+        if self._media_validation is not None:
+            return self._media_validation
+        assert self.tbl is not None
+        return self.tbl.media_validation
+
     def source(self) -> None:
         """
         If this is a computed col and the top-level expr is a function call, print the source, if possible.

pixeltable/catalog/globals.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
 import dataclasses
+import enum
 import itertools
 import logging
 from typing import Optional
 
+import pixeltable.exceptions as excs
+
 _logger = logging.getLogger('pixeltable')
 
 # name of the position column in a component view
@@ -34,6 +38,20 @@ class UpdateStatus:
         self.cols_with_excs = list(dict.fromkeys(self.cols_with_excs + other.cols_with_excs))
         return self
 
+
+class MediaValidation(enum.Enum):
+    ON_READ = 0
+    ON_WRITE = 1
+
+    @classmethod
+    def validated(cls, name: str, error_prefix: str) -> MediaValidation:
+        try:
+            return cls[name.upper()]
+        except KeyError:
+            val_strs = ', '.join(f'{s.lower()!r}' for s in cls.__members__.keys())
+            raise excs.Error(f'{error_prefix} must be one of: [{val_strs}]')
+
+
 def is_valid_identifier(name: str) -> bool:
     return name.isidentifier() and not name.startswith('_')
 

pixeltable/catalog/insertable_table.py

@@ -13,7 +13,7 @@ from pixeltable.env import Env
 from pixeltable.utils.filecache import FileCache
 
 from .catalog import Catalog
-from .globals import UpdateStatus
+from .globals import UpdateStatus, MediaValidation
 from .table import Table
 from .table_version import TableVersion
 from .table_version_path import TableVersionPath
@@ -35,8 +35,8 @@ class InsertableTable(Table):
     # MODULE-LOCAL, NOT PUBLIC
     @classmethod
     def _create(
-            cls, dir_id: UUID, name: str, schema: dict[str, ts.ColumnType], df: Optional[pxt.DataFrame], primary_key: List[str],
-            num_retained_versions: int, comment: str
+        cls, dir_id: UUID, name: str, schema: dict[str, ts.ColumnType], df: Optional[pxt.DataFrame],
+        primary_key: List[str], num_retained_versions: int, comment: str, media_validation: MediaValidation
     ) -> InsertableTable:
         columns = cls._create_columns(schema)
         cls._verify_schema(columns)
@@ -50,7 +50,9 @@ class InsertableTable(Table):
             col.is_pk = True
 
         with orm.Session(Env.get().engine, future=True) as session:
-            _, tbl_version = TableVersion.create(session, dir_id, name, columns, num_retained_versions, comment)
+            _, tbl_version = TableVersion.create(
+                session, dir_id, name, columns, num_retained_versions=num_retained_versions, comment=comment,
+                media_validation=media_validation)
             tbl = cls(dir_id, tbl_version)
             # TODO We need to commit before doing the insertion, in order to avoid a primary key (version) collision
             #   when the table metadata gets updated. Once we have a notion of user-defined transactions in

pixeltable/catalog/table.py

@@ -24,7 +24,7 @@ import pixeltable.type_system as ts
 from pixeltable.utils.filecache import FileCache
 
 from .column import Column
-from .globals import _ROWID_COLUMN_NAME, UpdateStatus, is_system_column_name, is_valid_identifier
+from .globals import _ROWID_COLUMN_NAME, UpdateStatus, is_system_column_name, is_valid_identifier, MediaValidation
 from .schema_object import SchemaObject
 from .table_version import TableVersion
 from .table_version_path import TableVersionPath
@@ -91,6 +91,7 @@ class Table(SchemaObject):
                     'num_retained_versions': 10,
                     'is_view': False,
                     'is_snapshot': False,
+                    'media_validation': 'on_write',
                 }
                 ```
         """
@@ -101,6 +102,7 @@ class Table(SchemaObject):
         md['schema_version'] = self._tbl_version.schema_version
         md['comment'] = self._comment
         md['num_retained_versions'] = self._num_retained_versions
+        md['media_validation'] = self._media_validation.name.lower()
         return md
 
     @property
@@ -244,6 +246,10 @@ class Table(SchemaObject):
     def _num_retained_versions(self):
         return self._tbl_version.num_retained_versions
 
+    @property
+    def _media_validation(self) -> MediaValidation:
+        return self._tbl_version.media_validation
+
     def _description(self) -> pd.DataFrame:
         cols = self._tbl_version_path.columns()
         df = pd.DataFrame({
@@ -422,7 +428,7 @@ class Table(SchemaObject):
         (on account of containing Python Callables or Exprs).
         """
         assert isinstance(spec, dict)
-        valid_keys = {'type', 'value', 'stored'}
+        valid_keys = {'type', 'value', 'stored', 'media_validation'}
         has_type = False
         for k in spec.keys():
             if k not in valid_keys:
@@ -449,6 +455,9 @@ class Table(SchemaObject):
                 if 'type' in spec:
                     raise excs.Error(f'Column {name}: "type" is redundant if value is a Pixeltable expression')
 
+        if 'media_validation' in spec:
+            _ = catalog.MediaValidation.validated(spec['media_validation'], f'Column {name}: media_validation')
+
         if 'stored' in spec and not isinstance(spec['stored'], bool):
             raise excs.Error(f'Column {name}: "stored" must be a bool, got {spec["stored"]}')
         if not has_type:
@@ -462,6 +471,7 @@ class Table(SchemaObject):
             col_type: Optional[ts.ColumnType] = None
             value_expr: Optional[exprs.Expr] = None
             primary_key: Optional[bool] = None
+            media_validation: Optional[catalog.MediaValidation] = None
             stored = True
 
             if isinstance(spec, (ts.ColumnType, type, _GenericAlias)):
@@ -484,9 +494,15 @@ class Table(SchemaObject):
                     value_expr = value_expr.copy()
                 stored = spec.get('stored', True)
                 primary_key = spec.get('primary_key')
+                media_validation_str = spec.get('media_validation')
+                media_validation = (
+                    catalog.MediaValidation[media_validation_str.upper()] if media_validation_str is not None
+                    else None
+                )
 
             column = Column(
-                name, col_type=col_type, computed_with=value_expr, stored=stored, is_pk=primary_key)
+                name, col_type=col_type, computed_with=value_expr, stored=stored, is_pk=primary_key,
+                media_validation=media_validation)
             columns.append(column)
         return columns
 

pixeltable/catalog/table_version.py

@@ -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.
 
@@ -710,20 +724,22 @@ 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())
+                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=print_stats, rowids=rowids())
+            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, *,
-        rowids: Optional[Iterator[int]] = None, print_stats: bool = False,
+        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, rowids=rowids)
+            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
@@ -1203,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)
@@ -1211,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/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,9 +49,10 @@ class View(Table):
 
     @classmethod
     def _create(
-            cls, dir_id: UUID, name: str, base: TableVersionPath, additional_columns: Dict[str, Any],
+            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,
-            iterator_cls: Optional[Type[ComponentIterator]], iterator_args: Optional[Dict]
+            media_validation: MediaValidation,
+            iterator_cls: Optional[type[ComponentIterator]], iterator_args: Optional[dict]
     ) -> View:
         columns = cls._create_columns(additional_columns)
         cls._verify_schema(columns)
@@ -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

@@ -371,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:
@@ -399,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

pixeltable/env.py

@@ -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, SqlAggregationNode, SqlNode

pixeltable/exec/aggregation_node.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Iterable, Optional, Any, Iterator
+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
 
@@ -28,13 +29,15 @@ class AggregationNode(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=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)
 

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

@@ -49,7 +49,7 @@ class DataRowBatch:
     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)

pixeltable/exec/exec_node.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import abc
-from typing import Iterable, Optional, List, TYPE_CHECKING, Iterator
+from typing import TYPE_CHECKING, Iterable, Iterator, List, Optional
 
 import pixeltable.exprs as exprs
+
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext