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

pixeltable/io/external_store.py

@@ -3,17 +3,13 @@ from __future__ import annotations
 import abc
 import itertools
 import logging
-import time
-from dataclasses import dataclass
-from typing import Any, Optional
-from uuid import UUID
+from typing import Any
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-from pixeltable import Table, Column
-import sqlalchemy as sql
-
-from pixeltable.catalog import TableVersion
+from pixeltable import Column, Table
+from pixeltable.catalog import ColumnHandle, TableVersion
+from pixeltable.catalog.update_status import UpdateStatus
 
 _logger = logging.getLogger('pixeltable')
 
@@ -25,6 +21,8 @@ class ExternalStore(abc.ABC):
     and stateful external stores.
     """
 
+    __name: str
+
     def __init__(self, name: str) -> None:
         self.__name = name
 
@@ -33,25 +31,21 @@ class ExternalStore(abc.ABC):
         return self.__name
 
     @abc.abstractmethod
-    def link(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
-        """
-        Called by `TableVersion.link()` to implement store-specific logic.
-        """
+    def link(self, tbl_version: TableVersion) -> None:
+        """Creates store-specific metadata needed to implement sync()."""
 
     @abc.abstractmethod
-    def unlink(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
-        """
-        Called by `TableVersion.unlink()` to implement store-specific logic.
-        """
+    def unlink(self, tbl_version: TableVersion) -> None:
+        """Removes store-specific metadata created in link()."""
 
     @abc.abstractmethod
-    def get_local_columns(self) -> list[Column]:
+    def get_local_columns(self) -> list[ColumnHandle]:
         """
         Gets a list of all local (Pixeltable) columns that are associated with this external store.
         """
 
     @abc.abstractmethod
-    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> UpdateStatus:
         """
         Called by `Table.sync()` to implement store-specific synchronization logic.
         """
@@ -70,9 +64,12 @@ class Project(ExternalStore, abc.ABC):
     additional capabilities specific to such projects.
     """
 
-    stored_proxies: dict[Column, Column]
+    _col_mapping: dict[ColumnHandle, str]  # col -> external col name
+    stored_proxies: dict[ColumnHandle, ColumnHandle]  # original col -> proxy col
 
-    def __init__(self, name: str, col_mapping: dict[Column, str], stored_proxies: Optional[dict[Column, Column]]):
+    def __init__(
+        self, name: str, col_mapping: dict[ColumnHandle, str], stored_proxies: dict[ColumnHandle, ColumnHandle] | None
+    ):
         super().__init__(name)
         self._col_mapping = col_mapping
 
@@ -87,68 +84,63 @@ class Project(ExternalStore, abc.ABC):
         # Note from aaron-siegel: This methodology is inefficient in the case where a table has many views with a high
         # proportion of overlapping rows, all proxying the same base column.
         if stored_proxies is None:
-            self.stored_proxies: dict[Column, Column] = {}
+            self.stored_proxies: dict[ColumnHandle, ColumnHandle] = {}
         else:
             self.stored_proxies = stored_proxies
 
-    def get_local_columns(self) -> list[Column]:
+    def get_local_columns(self) -> list[ColumnHandle]:
         return list(self.col_mapping.keys())
 
-    def link(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+    def link(self, tbl_version: TableVersion) -> None:
         # All of the media columns being linked need to either be stored computed columns, or else have stored proxies.
         # This ensures that the media in those columns resides in the media store.
         # First determine which columns (if any) need stored proxies, but don't have one yet.
         stored_proxies_needed: list[Column] = []
-        for col in self.col_mapping.keys():
+        for col_handle in self.col_mapping:
+            col = col_handle.get()
             if col.col_type.is_media_type() and not (col.is_stored and col.is_computed):
                 # If this column is already proxied in some other Project, use the existing proxy to avoid
                 # duplication. Otherwise, we'll create a new one.
                 for store in tbl_version.external_stores.values():
-                    if isinstance(store, Project) and col in store.stored_proxies:
-                        self.stored_proxies[col] = store.stored_proxies[col]
+                    if isinstance(store, Project) and col_handle in store.stored_proxies:
+                        self.stored_proxies[col_handle] = store.stored_proxies[col_handle]
                         break
-                if col not in self.stored_proxies:
+                if col_handle not in self.stored_proxies:
                     # We didn't find it in an existing Project
                     stored_proxies_needed.append(col)
+
         if len(stored_proxies_needed) > 0:
             _logger.info(f'Creating stored proxies for columns: {[col.name for col in stored_proxies_needed]}')
-            # Create stored proxies for columns that need one. Increment the schema version
-            # accordingly.
-            tbl_version.version += 1
-            preceding_schema_version = tbl_version.schema_version
-            tbl_version.schema_version = tbl_version.version
-            proxy_cols = [self.create_stored_proxy(tbl_version, col) for col in stored_proxies_needed]
+            # Create stored proxies for columns that need one
+            proxy_cols = [self.create_stored_proxy(col) for col in stored_proxies_needed]
             # Add the columns; this will also update table metadata.
-            tbl_version._add_columns(proxy_cols, conn, print_stats=False, on_error='ignore')
-            # We don't need to retain `UpdateStatus` since the stored proxies are intended to be
-            # invisible to the user.
-            tbl_version._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
+            tbl_version.add_columns(proxy_cols, print_stats=False, on_error='ignore')
+            self.stored_proxies.update(
+                {col.handle: proxy_col.handle for col, proxy_col in zip(stored_proxies_needed, proxy_cols)}
+            )
 
-    def unlink(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+    def unlink(self, tbl_version: TableVersion) -> None:
         # Determine which stored proxies can be deleted. (A stored proxy can be deleted if it is not referenced by
         # any *other* external store for this table.)
-        deletions_needed: set[Column] = set(self.stored_proxies.values())
+        deletions_needed: set[ColumnHandle] = set(self.stored_proxies.values())
         for name, store in tbl_version.external_stores.items():
             if isinstance(store, Project) and name != self.name:
                 deletions_needed = deletions_needed.difference(set(store.stored_proxies.values()))
         if len(deletions_needed) > 0:
-            _logger.info(f'Removing stored proxies for columns: {[col.name for col in deletions_needed]}')
-            # Delete stored proxies that are no longer needed.
-            tbl_version.version += 1
-            preceding_schema_version = tbl_version.schema_version
-            tbl_version.schema_version = tbl_version.version
-            tbl_version._drop_columns(deletions_needed)
+            _logger.info(f'Removing stored proxies for columns: {[col.get().name for col in deletions_needed]}')
+            tbl_version._drop_columns(col.get() for col in deletions_needed)
             self.stored_proxies.clear()
-            tbl_version._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
 
-    def create_stored_proxy(self, tbl_version: TableVersion, col: Column) -> Column:
+    def create_stored_proxy(self, col: Column) -> Column:
         """
         Creates a proxy column for the specified column. The proxy column will be created in the specified
         `TableVersion`.
         """
         from pixeltable import exprs
 
-        assert col.col_type.is_media_type() and not (col.is_stored and col.is_computed) and col not in self.stored_proxies
+        assert (
+            col.col_type.is_media_type() and not (col.is_stored and col.is_computed) and col not in self.stored_proxies
+        )
         proxy_col = Column(
             name=None,
             # Force images in the proxy column to be materialized inside the media store, in a normalized format.
@@ -157,17 +149,11 @@ class Project(ExternalStore, abc.ABC):
             #   Once `destination` is implemented, it can be replaced with a simple `ColumnRef`.
             computed_with=exprs.ColumnRef(col).apply(lambda x: x, col_type=col.col_type),
             stored=True,
-            col_id=tbl_version.next_col_id,
-            sa_col_type=col.col_type.to_sa_type(),
-            schema_version_add=tbl_version.schema_version
         )
-        proxy_col.tbl = tbl_version
-        tbl_version.next_col_id += 1
-        self.stored_proxies[col] = proxy_col
         return proxy_col
 
     @property
-    def col_mapping(self) -> dict[Column, str]:
+    def col_mapping(self) -> dict[ColumnHandle, str]:
         return self._col_mapping
 
     @abc.abstractmethod
@@ -197,12 +183,12 @@ class Project(ExternalStore, abc.ABC):
 
     @classmethod
     def validate_columns(
-            cls,
-            table: Table,
-            export_cols: dict[str, ts.ColumnType],
-            import_cols: dict[str, ts.ColumnType],
-            col_mapping: Optional[dict[str, str]]
-    ) -> dict[Column, str]:
+        cls,
+        table: Table,
+        export_cols: dict[str, ts.ColumnType],
+        import_cols: dict[str, ts.ColumnType],
+        col_mapping: dict[str, str] | None,
+    ) -> dict[ColumnHandle, str]:
         """
         Verifies that the specified `col_mapping` is valid. In particular, checks that:
         (i) the keys of `col_mapping` are valid columns of the specified `Table`;
@@ -212,6 +198,7 @@ class Project(ExternalStore, abc.ABC):
             external (import or export) columns.
         If validation fails, an exception will be raised. If validation succeeds, a new mapping will be returned
         in which the Pixeltable column names are resolved to the corresponding `Column` objects.
+        TODO: return columns as names or qualified ids
         """
         from pixeltable import exprs
 
@@ -219,32 +206,34 @@ class Project(ExternalStore, abc.ABC):
         if col_mapping is None:
             col_mapping = {col: col for col in itertools.chain(export_cols.keys(), import_cols.keys())}
 
-        resolved_col_mapping: dict[Column, str] = {}
+        resolved_col_mapping: dict[ColumnHandle, str] = {}
 
         # Validate names
-        t_cols = set(table._schema.keys())
+        t_cols = set(table._get_schema().keys())
         for t_col, ext_col in col_mapping.items():
             if t_col not in t_cols:
                 if is_user_specified_col_mapping:
                     raise excs.Error(
-                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table._name}` '
+                        f'Column name {t_col!r} appears as a key in `col_mapping`, but {table._display_str()} '
                         'contains no such column.'
                     )
                 else:
                     raise excs.Error(
-                        f'Column `{t_col}` does not exist in Table `{table._name}`. Either add a column `{t_col}`, '
-                        f'or specify a `col_mapping` to associate a different column with the external field `{ext_col}`.'
+                        f'Column {t_col!r} does not exist in {table._display_str()}. Either add a column {t_col!r}, '
+                        f'or specify a `col_mapping` to associate a different column with '
+                        f'the external field {ext_col!r}.'
                     )
             if ext_col not in export_cols and ext_col not in import_cols:
                 raise excs.Error(
-                    f'Column name `{ext_col}` appears as a value in `col_mapping`, but the external store '
-                    f'configuration has no column `{ext_col}`.'
+                    f'Column name {ext_col!r} appears as a value in `col_mapping`, but the external store '
+                    f'configuration has no column {ext_col!r}.'
                 )
             col_ref = table[t_col]
             assert isinstance(col_ref, exprs.ColumnRef)
-            resolved_col_mapping[col_ref.col] = ext_col
+            resolved_col_mapping[col_ref.col.handle] = ext_col
+
         # Validate column specs
-        t_col_types = table._schema
+        t_col_types = table._get_schema()
         for t_col, ext_col in col_mapping.items():
             t_col_type = t_col_types[t_col]
             if ext_col in export_cols:
@@ -252,65 +241,34 @@ class Project(ExternalStore, abc.ABC):
                 ext_col_type = export_cols[ext_col]
                 if not ext_col_type.is_supertype_of(t_col_type, ignore_nullable=True):
                     raise excs.Error(
-                        f'Column `{t_col}` cannot be exported to external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
+                        f'Column {t_col!r} cannot be exported to external column {ext_col!r} '
+                        f'(incompatible types; expecting `{ext_col_type}`)'
                     )
             if ext_col in import_cols:
                 # Validate that the external column can be assigned to the table column
                 if table._tbl_version_path.get_column(t_col).is_computed:
                     raise excs.Error(
-                        f'Column `{t_col}` is a computed column, which cannot be populated from an external column'
+                        f'Column {t_col!r} is a computed column, which cannot be populated from an external column'
                     )
                 ext_col_type = import_cols[ext_col]
                 if not t_col_type.is_supertype_of(ext_col_type, ignore_nullable=True):
                     raise excs.Error(
-                        f'Column `{t_col}` cannot be imported from external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
+                        f'Column {t_col!r} cannot be imported from external column {ext_col!r} '
+                        f'(incompatible types; expecting `{ext_col_type}`)'
                     )
         return resolved_col_mapping
 
-    @classmethod
-    def _column_as_dict(cls, col: Column) -> dict[str, Any]:
-        return {'tbl_id': str(col.tbl.id), 'col_id': col.id}
-
-    @classmethod
-    def _column_from_dict(cls, d: dict[str, Any]) -> Column:
-        from pixeltable.catalog import Catalog
-
-        tbl_id = UUID(d['tbl_id'])
-        col_id = d['col_id']
-        return Catalog.get().tbl_versions[(tbl_id, None)].cols_by_id[col_id]
-
-
-@dataclass(frozen=True)
-class SyncStatus:
-    external_rows_created: int = 0
-    external_rows_deleted: int = 0
-    external_rows_updated: int = 0
-    pxt_rows_updated: int = 0
-    num_excs: int = 0
-
-    def combine(self, other: 'SyncStatus') -> 'SyncStatus':
-        return SyncStatus(
-            external_rows_created=self.external_rows_created + other.external_rows_created,
-            external_rows_deleted=self.external_rows_deleted + other.external_rows_deleted,
-            external_rows_updated=self.external_rows_updated + other.external_rows_updated,
-            pxt_rows_updated=self.pxt_rows_updated + other.pxt_rows_updated,
-            num_excs=self.num_excs + other.num_excs
-        )
-
-    @classmethod
-    def empty(cls) -> 'SyncStatus':
-        return SyncStatus(0, 0, 0, 0, 0)
-
 
 class MockProject(Project):
     """A project that cannot be synced, used mainly for testing."""
+
     def __init__(
-            self,
-            name: str,
-            export_cols: dict[str, ts.ColumnType],
-            import_cols: dict[str, ts.ColumnType],
-            col_mapping: dict[Column, str],
-            stored_proxies: Optional[dict[Column, Column]] = None
+        self,
+        name: str,
+        export_cols: dict[str, ts.ColumnType],
+        import_cols: dict[str, ts.ColumnType],
+        col_mapping: dict[ColumnHandle, str],
+        stored_proxies: dict[ColumnHandle, ColumnHandle] | None = None,
     ):
         super().__init__(name, col_mapping, stored_proxies)
         self.export_cols = export_cols
@@ -319,12 +277,12 @@ class MockProject(Project):
 
     @classmethod
     def create(
-            cls,
-            t: Table,
-            name: str,
-            export_cols: dict[str, ts.ColumnType],
-            import_cols: dict[str, ts.ColumnType],
-            col_mapping: Optional[dict[str, str]] = None
+        cls,
+        t: Table,
+        name: str,
+        export_cols: dict[str, ts.ColumnType],
+        import_cols: dict[str, ts.ColumnType],
+        col_mapping: dict[str, str] | None = None,
     ) -> 'MockProject':
         col_mapping = cls.validate_columns(t, export_cols, import_cols, col_mapping)
         return cls(name, export_cols, import_cols, col_mapping)
@@ -335,7 +293,7 @@ class MockProject(Project):
     def get_import_columns(self) -> dict[str, ts.ColumnType]:
         return self.import_cols
 
-    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> UpdateStatus:
         raise NotImplementedError()
 
     def delete(self) -> None:
@@ -350,8 +308,8 @@ class MockProject(Project):
             'name': self.name,
             'export_cols': {k: v.as_dict() for k, v in self.export_cols.items()},
             'import_cols': {k: v.as_dict() for k, v in self.import_cols.items()},
-            'col_mapping': [[self._column_as_dict(k), v] for k, v in self.col_mapping.items()],
-            'stored_proxies': [[self._column_as_dict(k), self._column_as_dict(v)] for k, v in self.stored_proxies.items()]
+            'col_mapping': [[k.as_dict(), v] for k, v in self.col_mapping.items()],
+            'stored_proxies': [[k.as_dict(), v.as_dict()] for k, v in self.stored_proxies.items()],
         }
 
     @classmethod
@@ -360,11 +318,11 @@ class MockProject(Project):
             md['name'],
             {k: ts.ColumnType.from_dict(v) for k, v in md['export_cols'].items()},
             {k: ts.ColumnType.from_dict(v) for k, v in md['import_cols'].items()},
-            {cls._column_from_dict(entry[0]): entry[1] for entry in md['col_mapping']},
-            {cls._column_from_dict(entry[0]): cls._column_from_dict(entry[1]) for entry in md['stored_proxies']}
+            {ColumnHandle.from_dict(entry[0]): entry[1] for entry in md['col_mapping']},
+            {ColumnHandle.from_dict(entry[0]): ColumnHandle.from_dict(entry[1]) for entry in md['stored_proxies']},
         )
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, MockProject):
             return False
         return self.name == other.name

pixeltable/io/fiftyone.py

@@ -1,5 +1,5 @@
 import os
-from typing import Iterator, Optional, Union
+from typing import Any, Iterator
 
 import fiftyone as fo  # type: ignore[import-untyped]
 import fiftyone.utils.data as foud  # type: ignore[import-untyped]
@@ -9,17 +9,18 @@ import puremagic
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 from pixeltable import exprs
-from pixeltable.env import Env
+from pixeltable.utils.local_store import TempStore
 
 
 class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     """
     Implementation of a FiftyOne `DatasetImporter` that reads image data from a Pixeltable table.
     """
+
     __image_format: str  # format to use for any exported images that are not already stored on disk
     __labels: dict[str, tuple[exprs.Expr, type[fo.Label]]]  # label_name -> (expr, label_cls)
     __image_idx: int  # index of the image expr in the select list
-    __localpath_idx: Optional[int]  # index of the image localpath in the select list, if present
+    __localpath_idx: int | None  # index of the image localpath in the select list, if present
     __row_iter: Iterator[list]  # iterator over the table rows, to be convered to FiftyOne samples
 
     def __init__(
@@ -27,19 +28,14 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         tbl: pxt.Table,
         image: exprs.Expr,
         image_format: str,
-        classifications: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
-        detections: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
-        dataset_dir: Optional[os.PathLike] = None,
+        classifications: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
+        detections: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
+        dataset_dir: os.PathLike | None = None,
         shuffle: bool = False,
-        seed: Union[int, float, str, bytes, bytearray, None] = None,
-        max_samples: Optional[int] = None,
+        seed: int | float | str | bytes | bytearray | None = None,
+        max_samples: int | None = None,
     ):
-        super().__init__(
-            dataset_dir=dataset_dir,
-            shuffle=shuffle,
-            seed=seed,
-            max_samples=max_samples
-        )
+        super().__init__(dataset_dir=dataset_dir, shuffle=shuffle, seed=seed, max_samples=max_samples)
 
         self.__image_format = image_format
 
@@ -54,19 +50,18 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             if isinstance(exprs_, dict):
                 for label_name, expr in exprs_.items():
                     if not label_name.isidentifier():
-                        raise excs.Error(f"Invalid label name: {label_name}")
+                        raise excs.Error(f'Invalid label name: {label_name}')
                     if label_name in self.__labels:
-                        raise excs.Error(f"Duplicate label name: {label_name}")
+                        raise excs.Error(f'Duplicate label name: {label_name}')
                     self.__labels[label_name] = (expr, label_cls)
 
         # Now add the remaining labels, assigning unused default names.
         for exprs_, label_cls, default_name in label_categories:
             if exprs_ is None or isinstance(exprs_, dict):
                 continue
-            if isinstance(exprs_, exprs.Expr):
-                exprs_ = [exprs_]
-            assert isinstance(exprs_, list)
-            for expr in exprs_:
+            exprs_list = [exprs_] if isinstance(exprs_, exprs.Expr) else exprs_
+            assert isinstance(exprs_list, list)
+            for expr in exprs_list:
                 if default_name not in self.__labels:
                     name = default_name
                 else:
@@ -92,10 +87,10 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         else:
             self.__localpath_idx = None
 
-        df = tbl.select(*selection)
-        self.__row_iter = df._output_row_iterator()
+        query = tbl.select(*selection)
+        self.__row_iter = query._output_row_iterator()
 
-    def __next__(self) -> tuple[str, Optional[fo.ImageMetadata], Optional[dict[str, fo.Label]]]:
+    def __next__(self) -> tuple[str, fo.ImageMetadata | None, dict[str, fo.Label] | None]:
         row = next(self.__row_iter)
         img = row[self.__image_idx]
         assert isinstance(img, PIL.Image.Image)
@@ -105,7 +100,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             assert isinstance(file, str)
         else:
             # Write the dynamically created image to a temp file
-            file = str(Env.get().create_tmp_path(f'.{self.__image_format}'))
+            file = TempStore.create_path(extension=f'.{self.__image_format}')
             img.save(file, format=self.__image_format)
 
         metadata = fo.ImageMetadata(
@@ -113,7 +108,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             mime_type=puremagic.from_file(file, mime=True),
             width=img.width,
             height=img.height,
-            filepath=file,
+            filepath=str(file),
             num_channels=len(img.getbands()),
         )
 
@@ -129,7 +124,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             elif label_cls is fo.Detections:
                 label = fo.Detections(detections=self.__as_fo_detections(label_data))
             else:
-                assert False
+                raise AssertionError()
             labels[label_name] = label
 
         return file, metadata, labels
@@ -137,13 +132,9 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     def __as_fo_classifications(self, data: list) -> list[fo.Classification]:
         if not isinstance(data, list) or any('label' not in entry for entry in data):
             raise excs.Error(
-                f'Invalid classifications data: {data}\n'
-                "(Expected a list of dicts, each containing a 'label' key)"
+                f"Invalid classifications data: {data}\n(Expected a list of dicts, each containing a 'label' key)"
             )
-        return [
-            fo.Classification(label=entry['label'], confidence=entry.get('confidence'))
-            for entry in data
-        ]
+        return [fo.Classification(label=entry['label'], confidence=entry.get('confidence')) for entry in data]
 
     def __as_fo_detections(self, data: list) -> list[fo.Detections]:
         if not isinstance(data, list) or any('label' not in entry or 'bounding_box' not in entry for entry in data):
@@ -174,5 +165,5 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     def get_dataset_info(self) -> dict:
         pass
 
-    def close(self, *args) -> None:
+    def close(self, *args: Any) -> None:
         pass