pixeltable 0.4.13__py3-none-any.whl → 0.4.15__py3-none-any.whl

pixeltable/catalog/table_version.py

@@ -18,9 +18,8 @@ from pixeltable import exprs, index
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
 from pixeltable.metadata import schema
-from pixeltable.utils.exception_handler import run_cleanup_on_exception
 from pixeltable.utils.filecache import FileCache
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.object_stores import ObjectOps
 
 from .tbl_ops import TableOp
 
@@ -327,7 +326,7 @@ class TableVersion:
             from .table_version_path import TableVersionPath
 
             # clear out any remaining media files from an aborted previous attempt
-            MediaStore.get().delete(self.id)
+            self.delete_media()
             view_path = TableVersionPath.from_dict(op.load_view_op.view_path)
             plan, _ = Planner.create_view_load_plan(view_path)
             _, row_counts = self.store_tbl.insert_rows(plan, v_min=self.version)
@@ -356,14 +355,23 @@ class TableVersion:
         cat = pxt.catalog.Catalog.get()
         # We're creating a new TableVersion replica, so we should never have seen this particular
         # TableVersion instance before.
-        assert tbl_version.effective_version is not None
-        assert (tbl_version.id, tbl_version.effective_version) not in cat._tbl_versions
+        # Actually this isn't true, because we might be re-creating a dropped replica.
+        # TODO: Understand why old TableVersions are kept around even for a dropped table.
+        # assert tbl_version.effective_version is not None
+        # assert (tbl_version.id, tbl_version.effective_version) not in cat._tbl_versions
         cat._tbl_versions[tbl_version.id, tbl_version.effective_version] = tbl_version
         tbl_version.init()
         tbl_version.store_tbl.create()
         tbl_version.store_tbl.ensure_columns_exist(col for col in tbl_version.cols if col.is_stored)
         return tbl_version
 
+    def delete_media(self, tbl_version: Optional[int] = None) -> None:
+        # Assemble a set of column destinations and delete objects from all of them
+        # None is a valid column destination which refers to the default object location
+        destinations = {col.destination for col in self.cols if col.is_stored}
+        for dest in destinations:
+            ObjectOps.delete(dest, self.id, tbl_version=tbl_version)
+
     def drop(self) -> None:
         # if self.is_view and self.is_mutable:
         #     # update mutable_views
@@ -374,7 +382,7 @@ class TableVersion:
         #     if self.base.get().is_mutable:
         #         self.base.get().mutable_views.remove(TableVersionHandle.create(self))
 
-        MediaStore.get().delete(self.id)
+        self.delete_media()
         FileCache.get().clear(tbl_id=self.id)
         self.store_tbl.drop()
 
@@ -595,18 +603,7 @@ class TableVersion:
         idx_info = self.IndexInfo(id=idx_id, name=idx_name, idx=idx, col=col, val_col=val_col, undo_col=undo_col)
         self._tbl_md.index_md[idx_id] = idx_md
         self.idxs_by_name[idx_name] = idx_info
-        try:
-            idx.create_index(self._store_idx_name(idx_id), val_col)
-        finally:
-
-            def cleanup_index() -> None:
-                """Delete the newly added in-memory index structure"""
-                del self.idxs_by_name[idx_name]
-                del self._tbl_md.index_md[idx_id]
-                self.next_idx_id = idx_id
-
-            # Run cleanup only if there has been an exception; otherwise, skip cleanup.
-            run_cleanup_on_exception(cleanup_index)
+        idx.create_index(self._store_idx_name(idx_id), val_col)
 
     def _add_index(self, col: Column, idx_name: Optional[str], idx: index.IndexBase) -> UpdateStatus:
         val_col, undo_vol = self._create_index_columns(idx)
@@ -741,21 +738,6 @@ class TableVersion:
                     num_excs += excs_per_col
                 computed_values += plan.ctx.num_computed_exprs * row_count
             finally:
-                # Ensure cleanup occurs if an exception or keyboard interruption happens during `load_column()`.
-                def cleanup_on_error() -> None:
-                    """Delete columns that are added as part of current add_columns operation and re-initialize
-                    the sqlalchemy schema"""
-                    self.cols = [col for col in self.cols if col not in cols_to_add]
-                    for col in cols_to_add:
-                        # remove columns that we already added
-                        if col.id in self.cols_by_id:
-                            del self.cols_by_id[col.id]
-                        if col.name is not None and col.name in self.cols_by_name:
-                            del self.cols_by_name[col.name]
-                    self.store_tbl.create_sa_tbl()
-
-                # Run cleanup only if there has been an exception; otherwise, skip cleanup.
-                run_cleanup_on_exception(cleanup_on_error)
                 plan.close()
 
         pxt.catalog.Catalog.get().record_column_dependencies(self)
@@ -1236,7 +1218,7 @@ class TableVersion:
             )
 
         # delete newly-added data
-        MediaStore.get().delete(self.id, tbl_version=self.version)
+        self.delete_media(tbl_version=self.version)
         conn.execute(sql.delete(self.store_tbl.sa_tbl).where(self.store_tbl.sa_tbl.c.v_min == self.version))
 
         # revert new deletions

pixeltable/catalog/update_status.py

@@ -57,27 +57,35 @@ class UpdateStatus:
     """
 
     updated_cols: list[str] = field(default_factory=list)
+    """Columns that were updated."""
     cols_with_excs: list[str] = field(default_factory=list)
+    """Columns that encountered exceptions."""
 
     # stats for the rows affected by the operation
     row_count_stats: RowCountStats = field(default_factory=RowCountStats)
+    """Row count statistics for rows affected by this operation."""
 
     # stats for changes cascaded to other tables
     cascade_row_count_stats: RowCountStats = field(default_factory=RowCountStats)
+    """Row count statistics for changes cascaded to other tables."""
 
     # stats for the rows affected by the operation in an external store
     ext_row_count_stats: RowCountStats = field(default_factory=RowCountStats)
+    """Row count statistics for rows affected in an external store."""
 
     @property
     def num_rows(self) -> int:
+        """Total number of rows affected (including cascaded changes)."""
         return self.row_count_stats.num_rows + self.cascade_row_count_stats.num_rows
 
     @property
     def num_excs(self) -> int:
+        """Total number of exceptions encountered (including cascaded changes)."""
         return self.row_count_stats.num_excs + self.cascade_row_count_stats.num_excs
 
     @property
     def num_computed_values(self) -> int:
+        """Total number of computed values affected (including cascaded changes)."""
         return self.row_count_stats.computed_values + self.cascade_row_count_stats.computed_values
 
     def insert_to_update(self) -> 'UpdateStatus':
@@ -164,16 +172,20 @@ class UpdateStatus:
 
     @property
     def external_rows_updated(self) -> int:
+        """Number of rows updated in an external store."""
         return self.ext_row_count_stats.upd_rows
 
     @property
     def external_rows_created(self) -> int:
+        """Number of rows created in an external store."""
         return self.ext_row_count_stats.ins_rows
 
     @property
     def external_rows_deleted(self) -> int:
+        """Number of rows deleted from an external store."""
         return self.ext_row_count_stats.del_rows
 
     @property
     def ext_num_rows(self) -> int:
+        """Total number of rows affected in an external store."""
         return self.ext_row_count_stats.num_rows

pixeltable/catalog/view.py

@@ -47,17 +47,13 @@ class View(Table):
             self._tbl_version = tbl_version_path.tbl_version
 
     def _display_name(self) -> str:
-        name: str
-        if self._tbl_version_path.is_snapshot():
-            name = 'snapshot'
-        elif self._tbl_version_path.is_view():
-            name = 'view'
-        else:
-            assert self._tbl_version_path.is_replica()
-            name = 'table'
         if self._tbl_version_path.is_replica():
-            name = f'{name}-replica'
-        return name
+            return 'replica'
+        if self._tbl_version_path.is_snapshot():
+            return 'snapshot'
+        if self._tbl_version_path.is_view():
+            return 'view'
+        return 'table'
 
     @classmethod
     def select_list_to_additional_columns(cls, select_list: list[tuple[exprs.Expr, Optional[str]]]) -> dict[str, dict]:
@@ -270,12 +266,12 @@ class View(Table):
             # Update name and path with version qualifiers.
             md['name'] = f'{self._name}:{self._tbl_version_path.version()}'
             md['path'] = f'{self._path()}:{self._tbl_version_path.version()}'
-        base_tbl = self._get_base_table()
-        if base_tbl is None:
-            md['base'] = None
-        else:
+        base_tbl_id = self._base_tbl_id
+        if base_tbl_id is not None:
+            base_tbl = self._get_base_table()
+            base_path = '<anonymous base table>' if base_tbl is None else base_tbl._path()
             base_version = self._effective_base_versions[0]
-            md['base'] = base_tbl._path() if base_version is None else f'{base_tbl._path()}:{base_version}'
+            md['base'] = base_path if base_version is None else f'{base_path}:{base_version}'
         return md
 
     def insert(
@@ -294,17 +290,21 @@ class View(Table):
     def delete(self, where: Optional[exprs.Expr] = None) -> UpdateStatus:
         raise excs.Error(f'{self._display_str()}: Cannot delete from a {self._display_name()}.')
 
-    def _get_base_table(self) -> Optional['Table']:
+    @property
+    def _base_tbl_id(self) -> Optional[UUID]:
         if self._tbl_version_path.tbl_id != self._id:
             # _tbl_version_path represents a different schema object from this one. This can only happen if this is a
             # named pure snapshot.
-            base_id = self._tbl_version_path.tbl_id
-        elif self._tbl_version_path.base is None:
+            return self._tbl_version_path.tbl_id
+        if self._tbl_version_path.base is None:
             return None
-        else:
-            base_id = self._tbl_version_path.base.tbl_id
-        with catalog.Catalog.get().begin_xact(tbl_id=base_id, for_write=False):
-            return catalog.Catalog.get().get_table_by_id(base_id)
+        return self._tbl_version_path.base.tbl_id
+
+    def _get_base_table(self) -> Optional['Table']:
+        """Returns None if there is no base table, or if the base table is hidden."""
+        base_tbl_id = self._base_tbl_id
+        with catalog.Catalog.get().begin_xact(tbl_id=base_tbl_id, for_write=False):
+            return catalog.Catalog.get().get_table_by_id(base_tbl_id)
 
     @property
     def _effective_base_versions(self) -> list[Optional[int]]:

pixeltable/config.py

@@ -161,6 +161,8 @@ KNOWN_CONFIG_OPTIONS = {
         'hide_warnings': 'Hide warnings from the console',
         'verbosity': 'Verbosity level for console output',
         'api_key': 'API key for Pixeltable cloud',
+        'r2_profile': 'AWS config profile name used to access R2 storage',
+        's3_profile': 'AWS config profile name used to access S3 storage',
     },
     'anthropic': {'api_key': 'Anthropic API key'},
     'bedrock': {'api_key': 'AWS Bedrock API key'},

pixeltable/dataframe.py

@@ -456,6 +456,7 @@ class DataFrame:
 
     @property
     def schema(self) -> dict[str, ColumnType]:
+        """Column names and types in this DataFrame."""
         return self._schema
 
     def bind(self, args: dict[str, Any]) -> DataFrame:
@@ -1276,10 +1277,11 @@ class DataFrame:
 
         # TODO: Reconcile these with Table.__check_mutable()
         assert len(self._from_clause.tbls) == 1
-        if self._first_tbl.is_snapshot():
-            raise excs.Error(f'Cannot use `{op_name}` on a snapshot.')
+        # First check if it's a replica, since every replica handle is also a snapshot
         if self._first_tbl.is_replica():
             raise excs.Error(f'Cannot use `{op_name}` on a replica.')
+        if self._first_tbl.is_snapshot():
+            raise excs.Error(f'Cannot use `{op_name}` on a snapshot.')
 
     def _validate_mutable_op_sequence(self, op_name: str, allow_select: bool) -> None:
         """Tests whether the sequence of operations on this DataFrame is valid for a mutation operation."""

pixeltable/env.py

@@ -28,6 +28,7 @@ import nest_asyncio  # type: ignore[import-untyped]
 import pixeltable_pgserver
 import sqlalchemy as sql
 from pillow_heif import register_heif_opener  # type: ignore[import-untyped]
+from sqlalchemy import orm
 from tenacity import retry, stop_after_attempt, wait_exponential_jitter
 from tqdm import TqdmWarning
 
@@ -36,6 +37,7 @@ from pixeltable.config import Config
 from pixeltable.utils.console_output import ConsoleLogger, ConsoleMessageFilter, ConsoleOutputHandler, map_level
 from pixeltable.utils.dbms import CockroachDbms, Dbms, PostgresqlDbms
 from pixeltable.utils.http_server import make_server
+from pixeltable.utils.object_stores import ObjectPath, StorageObjectAddress
 
 if TYPE_CHECKING:
     import spacy
@@ -58,7 +60,8 @@ class Env:
     _log_fmt_str = '%(asctime)s %(levelname)s %(name)s %(filename)s:%(lineno)d: %(message)s'
 
     _media_dir: Optional[Path]
-    _file_cache_dir: Optional[Path]  # cached media files with external URL
+    _object_soa: Optional[StorageObjectAddress]
+    _file_cache_dir: Optional[Path]  # cached object files with external URL
     _dataset_cache_dir: Optional[Path]  # cached datasets (eg, pytorch or COCO)
     _log_dir: Optional[Path]  # log files
     _tmp_dir: Optional[Path]  # any tmp files
@@ -88,7 +91,7 @@ class Env:
 
     _resource_pool_info: dict[str, Any]
     _current_conn: Optional[sql.Connection]
-    _current_session: Optional[sql.orm.Session]
+    _current_session: Optional[orm.Session]
     _current_isolation_level: Optional[Literal['REPEATABLE_READ', 'SERIALIZABLE']]
     _dbms: Optional[Dbms]
     _event_loop: Optional[asyncio.AbstractEventLoop]  # event loop for ExecNode
@@ -120,7 +123,8 @@ class Env:
         assert self._instance is None, 'Env is a singleton; use Env.get() to access the instance'
 
         self._media_dir = None  # computed media files
-        self._file_cache_dir = None  # cached media files with external URL
+        self._object_soa = None  # computed object files in StorageObjectAddress format
+        self._file_cache_dir = None  # cached object files with external URL
         self._dataset_cache_dir = None  # cached datasets (eg, pytorch or COCO)
         self._log_dir = None  # log files
         self._tmp_dir = None  # any tmp files
@@ -224,7 +228,7 @@ class Env:
         return self._current_conn
 
     @property
-    def session(self) -> Optional[sql.orm.Session]:
+    def session(self) -> Optional[orm.Session]:
         assert self._current_session is not None
         return self._current_session
 
@@ -258,7 +262,7 @@ class Env:
                 self._current_isolation_level = 'SERIALIZABLE'
                 with (
                     self.engine.connect().execution_options(isolation_level=self._current_isolation_level) as conn,
-                    sql.orm.Session(conn) as session,
+                    orm.Session(conn) as session,
                     conn.begin(),
                 ):
                     self._current_conn = conn
@@ -363,6 +367,7 @@ class Env:
 
         if not self._media_dir.exists():
             self._media_dir.mkdir()
+        self._object_soa = ObjectPath.parse_object_storage_addr(str(self._media_dir), may_contain_object_name=False)
         if not self._file_cache_dir.exists():
             self._file_cache_dir.mkdir()
         if not self._dataset_cache_dir.exists():
@@ -615,15 +620,17 @@ class Env:
         Args:
             - name: The name of the client
         """
-        cl = _registered_clients[name]
-        if cl.client_obj is not None:
-            return cl.client_obj  # Already initialized
-
-        # Construct a client, retrieving each parameter from config.
+        # Return the existing client if it has already been constructed
+        with _registered_clients_lock:
+            cl = _registered_clients[name]
+            if cl.client_obj is not None:
+                return cl.client_obj  # Already initialized
 
+        # Retrieve parameters required to construct the requested client.
         init_kwargs: dict[str, Any] = {}
         for param in cl.params.values():
             # Determine the type of the parameter for proper config parsing.
+            pname = param.name
             t = param.annotation
             # Deference Optional[T]
             if typing.get_origin(t) in (typing.Union, types.UnionType):
@@ -633,27 +640,31 @@ class Env:
                 elif args[1] is type(None):
                     t = args[0]
             assert isinstance(t, type), t
-            arg: Any = Config.get().get_value(param.name, t, section=name)
+            arg: Any = Config.get().get_value(pname, t, section=name)
             if arg is not None:
-                init_kwargs[param.name] = arg
+                init_kwargs[pname] = arg
             elif param.default is inspect.Parameter.empty:
                 raise excs.Error(
-                    f'`{name}` client not initialized: parameter `{param.name}` is not configured.\n'
-                    f'To fix this, specify the `{name.upper()}_{param.name.upper()}` environment variable, '
-                    f'or put `{param.name.lower()}` in the `{name.lower()}` section of $PIXELTABLE_HOME/config.toml.'
+                    f'`{name}` client not initialized: parameter `{pname}` is not configured.\n'
+                    f'To fix this, specify the `{name.upper()}_{pname.upper()}` environment variable, '
+                    f'or put `{pname.lower()}` in the `{name.lower()}` section of $PIXELTABLE_HOME/config.toml.'
                 )
 
-        cl.client_obj = cl.init_fn(**init_kwargs)
-        self._logger.info(f'Initialized `{name}` client.')
-        return cl.client_obj
+        # Construct the requested client
+        with _registered_clients_lock:
+            if cl.client_obj is not None:
+                return cl.client_obj  # Already initialized
+            cl.client_obj = cl.init_fn(**init_kwargs)
+            self._logger.info(f'Initialized `{name}` client with parameters: {init_kwargs}.')
+            return cl.client_obj
 
     def _start_web_server(self) -> None:
         """
         The http server root is the file system root.
         eg: /home/media/foo.mp4 is located at http://127.0.0.1:{port}/home/media/foo.mp4
         On Windows, the server will translate paths like http://127.0.0.1:{port}/c:/media/foo.mp4
-        This arrangement enables serving media hosted within _home,
-        as well as external media inserted into pixeltable or produced by pixeltable.
+        This arrangement enables serving objects hosted within _home,
+        as well as external objects inserted into pixeltable or produced by pixeltable.
         The port is chosen dynamically to prevent conflicts.
         """
         # Port 0 means OS picks one for us.
@@ -713,10 +724,12 @@ class Env:
     def __register_packages(self) -> None:
         """Declare optional packages that are utilized by some parts of the code."""
         self.__register_package('anthropic')
+        self.__register_package('azure.storage.blob', library_name='azure-storage-blob')
         self.__register_package('boto3')
         self.__register_package('datasets')
         self.__register_package('fiftyone')
         self.__register_package('fireworks', library_name='fireworks-ai')
+        self.__register_package('google.cloud.storage', library_name='google-cloud-storage')
         self.__register_package('google.genai', library_name='google-genai')
         self.__register_package('groq')
         self.__register_package('huggingface_hub', library_name='huggingface-hub')
@@ -757,6 +770,10 @@ class Env:
             library_name=library_name or package_name,  # defaults to package_name unless specified otherwise
         )
 
+    def require_binary(self, binary_name: str) -> None:
+        if not shutil.which(binary_name):
+            raise excs.Error(f'{binary_name} is not installed or not in PATH. Please install it to use this feature.')
+
     def require_package(self, package_name: str, min_version: Optional[list[int]] = None) -> None:
         """
         Checks whether the specified optional package is available. If not, raises an exception
@@ -815,6 +832,12 @@ class Env:
         assert self._media_dir is not None
         return self._media_dir
 
+    @property
+    def object_soa(self) -> StorageObjectAddress:
+        assert self._media_dir is not None
+        assert self._object_soa is not None
+        return self._object_soa
+
     @property
     def file_cache_dir(self) -> Path:
         assert self._file_cache_dir is not None
@@ -947,11 +970,13 @@ def register_client(name: str) -> Callable:
     def decorator(fn: Callable) -> None:
         sig = inspect.signature(fn)
         params = dict(sig.parameters)
-        _registered_clients[name] = ApiClient(init_fn=fn, params=params)
+        with _registered_clients_lock:
+            _registered_clients[name] = ApiClient(init_fn=fn, params=params)
 
     return decorator
 
 
+_registered_clients_lock: threading.Lock = threading.Lock()
 _registered_clients: dict[str, ApiClient] = {}
 
 

pixeltable/exec/__init__.py

@@ -8,5 +8,6 @@ from .exec_context import ExecContext
 from .exec_node import ExecNode
 from .expr_eval import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
+from .object_store_save_node import ObjectStoreSaveNode
 from .row_update_node import RowUpdateNode
 from .sql_node import SqlAggregationNode, SqlJoinNode, SqlLookupNode, SqlNode, SqlSampleNode, SqlScanNode

pixeltable/exec/aggregation_node.py

@@ -103,6 +103,5 @@ class AggregationNode(ExecNode):
             self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
             self.output_batch.add_row(prev_row)
 
-        self.output_batch.flush_imgs(None, self.row_builder.stored_img_cols, self.flushed_img_slots)
         _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(self.output_batch.rows)} rows')
         yield self.output_batch