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

pixeltable/catalog/insertable_table.py

@@ -2,19 +2,25 @@ from __future__ import annotations
 
 import enum
 import logging
-from typing import TYPE_CHECKING, Any, Literal, Optional, overload
+from typing import TYPE_CHECKING, Any, Literal, Sequence, cast, overload
 from uuid import UUID
 
+import pydantic
+import pydantic_core
+
 import pixeltable as pxt
 from pixeltable import exceptions as excs, type_system as ts
 from pixeltable.env import Env
 from pixeltable.utils.filecache import FileCache
+from pixeltable.utils.pydantic import is_json_convertible
 
-from .globals import MediaValidation, UpdateStatus
+from .globals import MediaValidation
 from .table import Table
-from .table_version import TableVersion
+from .table_version import TableVersion, TableVersionMd
 from .table_version_handle import TableVersionHandle
 from .table_version_path import TableVersionPath
+from .tbl_ops import CreateStoreTableOp, TableOp
+from .update_status import UpdateStatus
 
 if TYPE_CHECKING:
     from pixeltable import exprs
@@ -51,72 +57,63 @@ class InsertableTable(Table):
     def __init__(self, dir_id: UUID, tbl_version: TableVersionHandle):
         tbl_version_path = TableVersionPath(tbl_version)
         super().__init__(tbl_version.id, dir_id, tbl_version.get().name, tbl_version_path)
+        self._tbl_version = tbl_version
 
-    @classmethod
-    def _display_name(cls) -> str:
+    def _display_name(self) -> str:
+        assert not self._tbl_version_path.is_replica()
         return 'table'
 
     @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,
         media_validation: MediaValidation,
-    ) -> InsertableTable:
+        create_default_idxs: bool,
+    ) -> tuple[TableVersionMd, list[TableOp]]:
         columns = cls._create_columns(schema)
         cls._verify_schema(columns)
         column_names = [col.name for col in columns]
         for pk_col in primary_key:
             if pk_col not in column_names:
-                raise excs.Error(f'Primary key column {pk_col} not found in table schema')
+                raise excs.Error(f'Primary key column {pk_col!r} not found in table schema.')
             col = columns[column_names.index(pk_col)]
             if col.col_type.nullable:
-                raise excs.Error(f'Primary key column {pk_col} cannot be nullable')
+                raise excs.Error(f'Primary key column {pk_col!r} cannot be nullable.')
             col.is_pk = True
 
-        _, tbl_version = TableVersion.create(
-            dir_id,
+        md = TableVersion.create_initial_md(
             name,
             columns,
-            num_retained_versions=num_retained_versions,
-            comment=comment,
-            media_validation=media_validation,
+            num_retained_versions,
+            comment,
+            media_validation,
+            create_default_idxs=create_default_idxs,
+            view_md=None,
         )
-        tbl = cls(dir_id, TableVersionHandle.create(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, we can wrap the create/insert in a transaction to avoid this.
-        session = Env.get().session
-        session.commit()
-        if df is not None:
-            # A DataFrame was provided, so insert its contents into the table
-            # (using the same DB session as the table creation)
-            tbl_version.insert(None, df, fail_on_exception=True)
-        session.commit()
-
-        _logger.info(f'Created table `{name}`, id={tbl_version.id}')
-        Env.get().console_logger.info(f'Created table `{name}`.')
-        return tbl
-
-    def get_metadata(self) -> dict[str, Any]:
-        md = super().get_metadata()
-        md['is_view'] = False
-        md['is_snapshot'] = False
-        return md
+
+        ops = [
+            TableOp(
+                tbl_id=md.tbl_md.tbl_id,
+                op_sn=0,
+                num_ops=1,
+                needs_xact=False,
+                create_store_table_op=CreateStoreTableOp(),
+            )
+        ]
+        return md, ops
 
     @overload
     def insert(
         self,
-        source: Optional[TableDataSource] = None,
+        source: TableDataSource | None = None,
         /,
         *,
-        source_format: Optional[Literal['csv', 'excel', 'parquet', 'json']] = None,
-        schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
+        source_format: Literal['csv', 'excel', 'parquet', 'json'] | None = None,
+        schema_overrides: dict[str, ts.ColumnType] | None = None,
         on_error: Literal['abort', 'ignore'] = 'abort',
         print_stats: bool = False,
         **kwargs: Any,
@@ -129,54 +126,73 @@ class InsertableTable(Table):
 
     def insert(
         self,
-        source: Optional[TableDataSource] = None,
+        source: TableDataSource | None = None,
         /,
         *,
-        source_format: Optional[Literal['csv', 'excel', 'parquet', 'json']] = None,
-        schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
+        source_format: Literal['csv', 'excel', 'parquet', 'json'] | None = None,
+        schema_overrides: dict[str, ts.ColumnType] | None = None,
         on_error: Literal['abort', 'ignore'] = 'abort',
         print_stats: bool = False,
         **kwargs: Any,
     ) -> UpdateStatus:
+        from pixeltable.catalog import Catalog
         from pixeltable.io.table_data_conduit import UnkTableDataConduit
 
-        table = self
-        if source is None:
-            source = [kwargs]
-            kwargs = None
+        if source is not None and isinstance(source, Sequence) and len(source) == 0:
+            raise excs.Error('Cannot insert an empty sequence.')
+        fail_on_exception = OnErrorParameter.fail_on_exception(on_error)
 
-        tds = UnkTableDataConduit(
-            source, source_format=source_format, src_schema_overrides=schema_overrides, extra_fields=kwargs
-        )
-        data_source = tds.specialize()
-        if data_source.source_column_map is None:
-            data_source.src_pk = []
+        with Catalog.get().begin_xact(tbl=self._tbl_version_path, for_write=True, lock_mutable_tree=True):
+            table = self
 
-        assert isinstance(table, Table)
-        data_source.add_table_info(table)
-        data_source.prepare_for_insert_into_table()
+            # TODO: unify with TableDataConduit
+            if source is not None and isinstance(source, Sequence) and isinstance(source[0], pydantic.BaseModel):
+                status = self._insert_pydantic(
+                    cast(Sequence[pydantic.BaseModel], source),  # needed for mypy
+                    print_stats=print_stats,
+                    fail_on_exception=fail_on_exception,
+                )
+                Env.get().console_logger.info(status.insert_msg)
+                FileCache.get().emit_eviction_warnings()
+                return status
 
-        fail_on_exception = OnErrorParameter.fail_on_exception(on_error)
-        return table.insert_table_data_source(
-            data_source=data_source, fail_on_exception=fail_on_exception, print_stats=print_stats
-        )
+            if source is None:
+                source = [kwargs]
+                kwargs = None
+
+            tds = UnkTableDataConduit(
+                source, source_format=source_format, src_schema_overrides=schema_overrides, extra_fields=kwargs
+            )
+            data_source = tds.specialize()
+            if data_source.source_column_map is None:
+                data_source.src_pk = []
+
+            assert isinstance(table, Table)
+            data_source.add_table_info(table)
+            data_source.prepare_for_insert_into_table()
+
+            return table.insert_table_data_source(
+                data_source=data_source, fail_on_exception=fail_on_exception, print_stats=print_stats
+            )
 
     def insert_table_data_source(
         self, data_source: TableDataConduit, fail_on_exception: bool, print_stats: bool = False
     ) -> pxt.UpdateStatus:
         """Insert row batches into this table from a `TableDataConduit`."""
-        from pixeltable.io.table_data_conduit import DFTableDataConduit
+        from pixeltable.catalog import Catalog
+        from pixeltable.io.table_data_conduit import QueryTableDataConduit
 
-        status = pxt.UpdateStatus()
-        with Env.get().begin_xact():
-            if isinstance(data_source, DFTableDataConduit):
+        with Catalog.get().begin_xact(tbl=self._tbl_version_path, for_write=True, lock_mutable_tree=True):
+            if isinstance(data_source, QueryTableDataConduit):
+                status = pxt.UpdateStatus()
                 status += self._tbl_version.get().insert(
-                    rows=None, df=data_source.pxt_df, print_stats=print_stats, fail_on_exception=fail_on_exception
+                    rows=None, query=data_source.pxt_query, print_stats=print_stats, fail_on_exception=fail_on_exception
                 )
             else:
+                status = pxt.UpdateStatus()
                 for row_batch in data_source.valid_row_batch():
                     status += self._tbl_version.get().insert(
-                        rows=row_batch, df=None, print_stats=print_stats, fail_on_exception=fail_on_exception
+                        rows=row_batch, query=None, print_stats=print_stats, fail_on_exception=fail_on_exception
                     )
 
         Env.get().console_logger.info(status.insert_msg)
@@ -184,34 +200,110 @@ class InsertableTable(Table):
         FileCache.get().emit_eviction_warnings()
         return status
 
-    def _validate_input_rows(self, rows: list[dict[str, Any]]) -> None:
-        """Verify that the input rows match the table schema"""
-        valid_col_names = set(self._schema.keys())
-        reqd_col_names = set(self._tbl_version_path.tbl_version.get().get_required_col_names())
-        computed_col_names = set(self._tbl_version_path.tbl_version.get().get_computed_col_names())
-        for row in rows:
-            assert isinstance(row, dict)
-            col_names = set(row.keys())
-            if len(reqd_col_names - col_names) > 0:
-                raise excs.Error(f'Missing required column(s) ({", ".join(reqd_col_names - col_names)}) in row {row}')
-
-            for col_name, val in row.items():
-                if col_name not in valid_col_names:
-                    raise excs.Error(f'Unknown column name {col_name} in row {row}')
-                if col_name in computed_col_names:
-                    raise excs.Error(f'Value for computed column {col_name} in row {row}')
-
-                # validate data
-                col = self._tbl_version_path.get_column(col_name)
-                try:
-                    # basic sanity checks here
-                    checked_val = col.col_type.create_literal(val)
-                    row[col_name] = checked_val
-                except TypeError as e:
-                    msg = str(e)
-                    raise excs.Error(f'Error in column {col.name}: {msg[0].lower() + msg[1:]}\nRow: {row}') from e
-
-    def delete(self, where: Optional['exprs.Expr'] = None) -> UpdateStatus:
+    def _insert_pydantic(
+        self, rows: Sequence[pydantic.BaseModel], print_stats: bool = False, fail_on_exception: bool = True
+    ) -> UpdateStatus:
+        model_class = type(rows[0])
+        self._validate_pydantic_model(model_class)
+        # convert rows one-by-one in order to be able to print meaningful error messages
+        pxt_rows: list[dict[str, Any]] = []
+        for i, row in enumerate(rows):
+            try:
+                pxt_rows.append(row.model_dump(mode='json'))
+            except pydantic_core.PydanticSerializationError as e:
+                raise excs.Error(f'Row {i}: error serializing pydantic model to JSON:\n{e}') from e
+
+        # explicitly check that all required columns are present and non-None in the rows,
+        # because we ignore nullability when validating the pydantic model
+        reqd_col_names = [col.name for col in self._tbl_version_path.columns() if col.is_required_for_insert]
+        for i, pxt_row in enumerate(pxt_rows):
+            if type(rows[i]) is not model_class:
+                raise excs.Error(
+                    f'Expected an instance of `{model_class.__name__}`; got `{type(rows[i]).__name__}` (in row {i})'
+                )
+            for col_name in reqd_col_names:
+                if pxt_row.get(col_name) is None:
+                    raise excs.Error(f'Missing required column {col_name!r} in row {i}')
+
+        status = self._tbl_version.get().insert(
+            rows=pxt_rows, query=None, print_stats=print_stats, fail_on_exception=fail_on_exception
+        )
+        return status
+
+    def _validate_pydantic_model(self, model: type[pydantic.BaseModel]) -> None:
+        """
+        Check if a Pydantic model is compatible with this table for insert operations.
+
+        A model is compatible if:
+        - All required table columns have corresponding model fields with compatible types
+        - Model does not define fields for computed columns
+        - Model field types are compatible with table column types
+        """
+        assert isinstance(model, type) and issubclass(model, pydantic.BaseModel)
+
+        schema = self._get_schema()
+        required_cols = set(self._tbl_version.get().get_required_col_names())
+        computed_cols = set(self._tbl_version.get().get_computed_col_names())
+        model_fields = model.model_fields
+        model_field_names = set(model_fields.keys())
+
+        missing_required = required_cols - model_field_names
+        if missing_required:
+            raise excs.Error(
+                f'Pydantic model `{model.__name__}` is missing required columns: ' + ', '.join(missing_required)
+            )
+
+        computed_in_model = computed_cols & model_field_names
+        if computed_in_model:
+            raise excs.Error(
+                f'Pydantic model `{model.__name__}` has fields for computed columns: ' + ', '.join(computed_in_model)
+            )
+
+        # validate type compatibility
+        common_fields = model_field_names & set(schema.keys())
+        if len(common_fields) == 0:
+            raise excs.Error(
+                f'Pydantic model `{model.__name__}` has no fields that map to columns in table {self._name!r}'
+            )
+        for field_name in common_fields:
+            pxt_col_type = schema[field_name]
+            model_field = model_fields[field_name]
+            model_type = model_field.annotation
+
+            # we ignore nullability: we want to accept optional model fields for required table columns, as long as
+            # the model instances provide a non-null value
+            # allow_enum=True: model_dump(mode='json') converts enums to their values
+            inferred_pxt_type = ts.ColumnType.from_python_type(model_type, infer_pydantic_json=True)
+            if inferred_pxt_type is None:
+                raise excs.Error(
+                    f'Pydantic model `{model.__name__}`: cannot infer Pixeltable type for column {field_name!r}'
+                )
+
+            if pxt_col_type.is_media_type():
+                # media types require file paths, either as str or Path
+                if not inferred_pxt_type.is_string_type():
+                    raise excs.Error(
+                        f'Column {field_name!r} requires a `str` or `Path` field in `{model.__name__}`, but it is '
+                        f'`{model_type.__name__}`'
+                    )
+            else:
+                if not pxt_col_type.is_supertype_of(inferred_pxt_type, ignore_nullable=True):
+                    raise excs.Error(
+                        f'Pydantic model `{model.__name__}` has incompatible type `{model_type.__name__}` '
+                        f'for column {field_name!r} (of Pixeltable type `{pxt_col_type}`)'
+                    )
+
+                if (
+                    isinstance(model_type, type)
+                    and issubclass(model_type, pydantic.BaseModel)
+                    and not is_json_convertible(model_type)
+                ):
+                    raise excs.Error(
+                        f'Pydantic model `{model.__name__}` has field {field_name!r} with nested model '
+                        f'`{model_type.__name__}`, which is not JSON-convertible'
+                    )
+
+    def delete(self, where: 'exprs.Expr' | None = None) -> UpdateStatus:
         """Delete rows in this table.
 
         Args:
@@ -226,16 +318,17 @@ class InsertableTable(Table):
 
             >>> tbl.delete(tbl.a > 5)
         """
-        with Env.get().begin_xact():
+        from pixeltable.catalog import Catalog
+
+        with Catalog.get().begin_xact(tbl=self._tbl_version_path, for_write=True, lock_mutable_tree=True):
             return self._tbl_version.get().delete(where=where)
 
-    @property
-    def _base_table(self) -> Optional['Table']:
+    def _get_base_table(self) -> 'Table' | None:
         return None
 
     @property
-    def _effective_base_versions(self) -> list[Optional[int]]:
+    def _effective_base_versions(self) -> list[int | None]:
         return []
 
     def _table_descriptor(self) -> str:
-        return f'Table {self._path!r}'
+        return self._display_str()

pixeltable/catalog/path.py

@@ -1,20 +1,55 @@
 from __future__ import annotations
 
 import logging
-from typing import Iterator
+from typing import NamedTuple
 
 from pixeltable import exceptions as excs
 
-from .globals import is_valid_path
+from .globals import is_valid_identifier
 
 _logger = logging.getLogger('pixeltable')
 
 
-class Path:
-    def __init__(self, path: str, empty_is_valid: bool = False, allow_system_paths: bool = False):
-        if not is_valid_path(path, empty_is_valid, allow_system_paths):
-            raise excs.Error(f"Invalid path format: '{path}'")
-        self.components = path.split('.')
+class Path(NamedTuple):
+    components: list[str]
+    version: int | None = None
+
+    @classmethod
+    def parse(
+        cls,
+        path: str,
+        allow_empty_path: bool = False,
+        allow_system_path: bool = False,
+        allow_versioned_path: bool = False,
+    ) -> Path:
+        components: list[str]
+        version: int | None
+        if ':' in path:
+            parts = path.split(':')
+            if len(parts) != 2:
+                raise excs.Error(f'Invalid path: {path}')
+            try:
+                components = parts[0].split('.')
+                version = int(parts[1])
+            except ValueError:
+                raise excs.Error(f'Invalid path: {path}') from None
+        else:
+            components = path.split('.')
+            version = None
+
+        if components == [''] and not allow_empty_path:
+            raise excs.Error(f'Invalid path: {path}')
+
+        if components != [''] and not all(
+            is_valid_identifier(c, allow_system_identifiers=allow_system_path, allow_hyphens=True) for c in components
+        ):
+            raise excs.Error(f'Invalid path: {path}')
+
+        if version is not None and not allow_versioned_path:
+            raise excs.Error(f'Versioned path not allowed here: {path}')
+
+        assert len(components) > 0
+        return Path(components, version)
 
     @property
     def len(self) -> int:
@@ -22,7 +57,6 @@ class Path:
 
     @property
     def name(self) -> str:
-        assert len(self.components) > 0
         return self.components[-1]
 
     @property
@@ -36,18 +70,15 @@ class Path:
     @property
     def parent(self) -> Path:
         if len(self.components) == 1:
-            if self.is_root:
-                return self
-            else:
-                return Path('', empty_is_valid=True, allow_system_paths=True)
+            return ROOT_PATH  # Includes the case of the root path, which is its own parent.
         else:
-            return Path('.'.join(self.components[:-1]), allow_system_paths=True)
+            return Path(self.components[:-1])
 
     def append(self, name: str) -> Path:
         if self.is_root:
-            return Path(name, allow_system_paths=True)
+            return Path([name])
         else:
-            return Path(f'{self}.{name}', allow_system_paths=True)
+            return Path([*self.components, name])
 
     def is_ancestor(self, other: Path, is_parent: bool = False) -> bool:
         """
@@ -60,22 +91,31 @@ class Path:
         is_prefix = self.components == other.components[: self.len]
         return is_prefix and (self.len == (other.len - 1) or not is_parent)
 
-    def ancestors(self) -> Iterator[Path]:
+    def ancestors(self) -> list[Path]:
         """
-        Return all ancestors of this path in top-down order including root.
+        Return all proper ancestors of this path in top-down order including root.
         If this path is for the root directory, which has no parent, then None is returned.
         """
         if self.is_root:
-            return
+            return []
         else:
-            for i in range(0, len(self.components)):
-                yield Path('.'.join(self.components[0:i]), empty_is_valid=True)
+            return [Path(self.components[:i]) if i > 0 else ROOT_PATH for i in range(len(self.components))]
 
     def __repr__(self) -> str:
         return repr(str(self))
 
     def __str__(self) -> str:
-        return '.'.join(self.components)
+        base = '.'.join(self.components)
+        if self.version is not None:
+            return f'{base}:{self.version}'
+        else:
+            return base
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Path) and str(self) == str(other)
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
 
-    def __lt__(self, other: Path) -> bool:
-        return str(self) < str(other)
+ROOT_PATH = Path([''])

pixeltable/catalog/schema_object.py

@@ -1,9 +1,7 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING
 from uuid import UUID
 
-from pixeltable.env import Env
-
 if TYPE_CHECKING:
     from pixeltable import catalog
 
@@ -16,45 +14,43 @@ class SchemaObject:
 
     _id: UUID
     _name: str
-    _dir_id: Optional[UUID]
+    _dir_id: UUID | None
 
-    def __init__(self, obj_id: UUID, name: str, dir_id: Optional[UUID]):
+    def __init__(self, obj_id: UUID, name: str, dir_id: UUID | None):
         # make these private so they don't collide with column names (id and name are fairly common)
+        assert dir_id is None or isinstance(dir_id, UUID), type(dir_id)
         self._id = obj_id
         self._name = name
         self._dir_id = dir_id
 
-    def _parent(self) -> Optional['catalog.Dir']:
+    def _parent(self) -> 'catalog.Dir | None':
         """Returns the parent directory of this schema object."""
         from .catalog import Catalog
 
-        with Env.get().begin_xact():
+        with Catalog.get().begin_xact(for_write=False):
             if self._dir_id is None:
                 return None
             return Catalog.get().get_dir(self._dir_id)
 
-    @property
     def _path(self) -> str:
         """Returns the path to this schema object."""
         from .catalog import Catalog
 
         assert self._dir_id is not None
-        with Env.get().begin_xact():
+        with Catalog.get().begin_xact(for_write=False):
             path = Catalog.get().get_dir_path(self._dir_id)
             return str(path.append(self._name))
 
-    def get_metadata(self) -> dict[str, Any]:
-        """Returns metadata associated with this schema object."""
-        return {'name': self._name, 'path': self._path}
-
-    @classmethod
     @abstractmethod
-    def _display_name(cls) -> str:
+    def _display_name(self) -> str:
         """
         Return name displayed in error messages.
         """
         pass
 
+    def _display_str(self) -> str:
+        return f'{self._display_name()} {self._path()!r}'
+
     def _move(self, new_name: str, new_dir_id: UUID) -> None:
         """Subclasses need to override this to make the change persistent"""
         self._name = new_name