pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__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, TableVersionCompleteMd
 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[TableVersionCompleteMd, 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,11 +126,11 @@ 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,
@@ -141,8 +138,24 @@ class InsertableTable(Table):
         from pixeltable.catalog import Catalog
         from pixeltable.io.table_data_conduit import UnkTableDataConduit
 
-        with Catalog.get().begin_xact(tbl_id=self._id, for_write=True):
+        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)
+
+        with Catalog.get().begin_xact(tbl=self._tbl_version_path, for_write=True, lock_mutable_tree=True):
             table = self
+
+            # 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
+
             if source is None:
                 source = [kwargs]
                 kwargs = None
@@ -158,7 +171,6 @@ class InsertableTable(Table):
             data_source.add_table_info(table)
             data_source.prepare_for_insert_into_table()
 
-            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
             )
@@ -170,13 +182,14 @@ class InsertableTable(Table):
         from pixeltable.catalog import Catalog
         from pixeltable.io.table_data_conduit import DFTableDataConduit
 
-        status = pxt.UpdateStatus()
-        with Catalog.get().begin_xact(tbl_id=self._id, for_write=True):
+        with Catalog.get().begin_xact(tbl=self._tbl_version_path, for_write=True, lock_mutable_tree=True):
             if isinstance(data_source, DFTableDataConduit):
+                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
                 )
             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
@@ -187,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, df=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:
@@ -231,16 +320,15 @@ class InsertableTable(Table):
         """
         from pixeltable.catalog import Catalog
 
-        with Catalog.get().begin_xact(tbl_id=self._id, for_write=True):
+        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,53 @@
 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_path) 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 +55,6 @@ class Path:
 
     @property
     def name(self) -> str:
-        assert len(self.components) > 0
         return self.components[-1]
 
     @property
@@ -36,18 +68,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 +89,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,5 +1,5 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING
 from uuid import UUID
 
 if TYPE_CHECKING:
@@ -14,15 +14,16 @@ 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
 
@@ -40,18 +41,16 @@ class SchemaObject:
             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