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

pixeltable/catalog/insertable_table.py

@@ -1,172 +1,309 @@
 from __future__ import annotations
 
+import enum
 import logging
-from typing import Any, Iterable, Literal, Optional, overload
+from typing import TYPE_CHECKING, Any, Literal, Sequence, cast, overload
 from uuid import UUID
 
-import sqlalchemy.orm as orm
+import pydantic
+import pydantic_core
 
 import pixeltable as pxt
-import pixeltable.type_system as ts
-from pixeltable import exceptions as excs
+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 .catalog import Catalog
-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
+    from pixeltable.globals import TableDataSource
+    from pixeltable.io.table_data_conduit import TableDataConduit
 
 _logger = logging.getLogger('pixeltable')
 
 
+class OnErrorParameter(enum.Enum):
+    """Supported values for the on_error parameter"""
+
+    ABORT = 'abort'
+    IGNORE = 'ignore'
+
+    @classmethod
+    def is_valid(cls, v: Any) -> bool:
+        if isinstance(v, str):
+            return v.lower() in [c.value for c in cls]
+        return False
+
+    @classmethod
+    def fail_on_exception(cls, v: Any) -> bool:
+        if not cls.is_valid(v):
+            raise ValueError(f'Invalid value for on_error: {v}')
+        if isinstance(v, str):
+            return v.lower() != cls.IGNORE.value
+        return True
+
+
 class InsertableTable(Table):
     """A `Table` that allows inserting and deleting rows."""
 
-    def __init__(self, dir_id: UUID, tbl_version: TableVersion):
+    def __init__(self, dir_id: UUID, tbl_version: TableVersionHandle):
         tbl_version_path = TableVersionPath(tbl_version)
-        super().__init__(tbl_version.id, dir_id, tbl_version.name, tbl_version_path)
+        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'
 
-    # 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, media_validation: MediaValidation
-    ) -> InsertableTable:
+        cls,
+        name: str,
+        schema: dict[str, ts.ColumnType],
+        primary_key: list[str],
+        num_retained_versions: int,
+        comment: str,
+        media_validation: MediaValidation,
+        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
 
-        with orm.Session(Env.get().engine, future=True) as session:
-            _, 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, we can wrap the create/insert in a transaction to avoid this.
-            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, conn=session.connection(), fail_on_exception=True)
-            session.commit()
-            cat = Catalog.get()
-            cat.tbl_dependents[tbl._id] = []
-            cat.tbls[tbl._id] = tbl
-
-            _logger.info(f'Created table `{name}`, id={tbl_version.id}')
-            print(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
+        md = TableVersion.create_initial_md(
+            name,
+            columns,
+            num_retained_versions,
+            comment,
+            media_validation,
+            create_default_idxs=create_default_idxs,
+            view_md=None,
+        )
+
+        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,
-        rows: Iterable[dict[str, Any]],
+        source: TableDataSource | None = 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,
-        on_error: Literal['abort', 'ignore'] = 'abort'
+        **kwargs: Any,
     ) -> UpdateStatus: ...
 
     @overload
     def insert(
-        self,
-        *,
-        print_stats: bool = False,
-        on_error: Literal['abort', 'ignore'] = 'abort',
-        **kwargs: Any
+        self, /, *, on_error: Literal['abort', 'ignore'] = 'abort', print_stats: bool = False, **kwargs: Any
     ) -> UpdateStatus: ...
 
-    def insert(  # type: ignore[misc]
+    def insert(
         self,
-        rows: Optional[Iterable[dict[str, Any]]] = None,
+        source: TableDataSource | None = None,
         /,
         *,
-        print_stats: bool = False,
+        source_format: Literal['csv', 'excel', 'parquet', 'json'] | None = None,
+        schema_overrides: dict[str, ts.ColumnType] | None = None,
         on_error: Literal['abort', 'ignore'] = 'abort',
-        **kwargs: Any
+        print_stats: bool = False,
+        **kwargs: Any,
     ) -> UpdateStatus:
-        if rows is None:
-            rows = [kwargs]
-        else:
-            rows = list(rows)
-            if len(kwargs) > 0:
-                raise excs.Error('`kwargs` cannot be specified unless `rows is None`.')
-
-        fail_on_exception = on_error == 'abort'
-
-        if not isinstance(rows, list):
-            raise excs.Error('rows must be a list of dictionaries')
-        if len(rows) == 0:
-            raise excs.Error('rows must not be empty')
-        for row in rows:
-            if not isinstance(row, dict):
-                raise excs.Error('rows must be a list of dictionaries')
-        self._validate_input_rows(rows)
-        status = self._tbl_version.insert(rows, None, print_stats=print_stats, fail_on_exception=fail_on_exception)
-
-        if status.num_excs == 0:
-            cols_with_excs_str = ''
-        else:
-            cols_with_excs_str = \
-                f' across {len(status.cols_with_excs)} column{"" if len(status.cols_with_excs) == 1 else "s"}'
-            cols_with_excs_str += f' ({", ".join(status.cols_with_excs)})'
-        msg = (
-            f'Inserted {status.num_rows} row{"" if status.num_rows == 1 else "s"} '
-            f'with {status.num_excs} error{"" if status.num_excs == 1 else "s"}{cols_with_excs_str}.'
-        )
-        print(msg)
-        _logger.info(f'InsertableTable {self._name}: {msg}')
+        from pixeltable.catalog import Catalog
+        from pixeltable.io.table_data_conduit import UnkTableDataConduit
+
+        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
+
+            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.catalog import Catalog
+        from pixeltable.io.table_data_conduit import QueryTableDataConduit
+
+        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, 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, query=None, 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
 
-    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_required_col_names())
-        computed_col_names = set(self._tbl_version_path.tbl_version.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}')
-
-    def delete(self, where: Optional['pxt.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:
@@ -181,4 +318,17 @@ class InsertableTable(Table):
 
             >>> tbl.delete(tbl.a > 5)
         """
-        return self._tbl_version.delete(where=where)
+        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)
+
+    def _get_base_table(self) -> 'Table' | None:
+        return None
+
+    @property
+    def _effective_base_versions(self) -> list[int | None]:
+        return []
+
+    def _table_descriptor(self) -> str:
+        return self._display_str()

pixeltable/catalog/path.py

@@ -1,17 +1,55 @@
 from __future__ import annotations
 
 import logging
+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):
-        if not is_valid_path(path, empty_is_valid):
-            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:
@@ -19,28 +57,28 @@ class Path:
 
     @property
     def name(self) -> str:
-        assert len(self.components) > 0
         return self.components[-1]
 
     @property
     def is_root(self) -> bool:
-        return self.components[0] == ''
+        return not self.components[0]
+
+    @property
+    def is_system_path(self) -> bool:
+        return self.components[0].startswith('_')
 
     @property
     def parent(self) -> Path:
         if len(self.components) == 1:
-            if self.is_root:
-                return self
-            else:
-                return Path('', empty_is_valid=True)
+            return ROOT_PATH  # Includes the case of the root path, which is its own parent.
         else:
-            return Path('.'.join(self.components[:-1]))
+            return Path(self.components[:-1])
 
     def append(self, name: str) -> Path:
         if self.is_root:
-            return Path(name)
+            return Path([name])
         else:
-            return Path(f'{str(self)}.{name}')
+            return Path([*self.components, name])
 
     def is_ancestor(self, other: Path, is_parent: bool = False) -> bool:
         """
@@ -50,9 +88,34 @@ class Path:
             return False
         if self.is_root and (other.len == 1 or not is_parent):
             return True
-        is_prefix = self.components == other.components[:self.len]
+        is_prefix = self.components == other.components[: self.len]
         return is_prefix and (self.len == (other.len - 1) or not is_parent)
 
+    def ancestors(self) -> list[Path]:
+        """
+        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 []
+        else:
+            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))
+
 
+ROOT_PATH = Path([''])