pixeltable 0.3.15__py3-none-any.whl → 0.4.0rc2__py3-none-any.whl

pixeltable/catalog/table_version_handle.py

@@ -41,8 +41,21 @@ class TableVersionHandle:
     def get(self) -> TableVersion:
         from .catalog import Catalog
 
-        if self._tbl_version is None:
-            self._tbl_version = Catalog.get().get_tbl_version(self.id, self.effective_version)
+        cat = Catalog.get()
+        if self._tbl_version is None or not self._tbl_version.is_validated:
+            if self.effective_version is not None and self._tbl_version is not None:
+                # this is a snapshot version; we need to make sure we refer to the instance cached
+                # in Catalog, in order to avoid mixing sa_tbl instances in the same transaction
+                # (which will lead to duplicates in the From clause generated in SqlNode.create_from_clause())
+                assert (self.id, self.effective_version) in cat._tbl_versions
+                self._tbl_version = cat._tbl_versions[self.id, self.effective_version]
+                self._tbl_version.is_validated = True
+            else:
+                self._tbl_version = Catalog.get().get_tbl_version(self.id, self.effective_version)
+        if self.effective_version is None:
+            # make sure we don't see a discarded instance of a live TableVersion
+            tvs = list(Catalog.get()._tbl_versions.values())
+            assert self._tbl_version in tvs
         return self._tbl_version
 
     def as_dict(self) -> dict:

pixeltable/catalog/table_version_path.py

@@ -4,9 +4,11 @@ import logging
 from typing import Optional
 from uuid import UUID
 
+from pixeltable.env import Env
 from pixeltable.metadata import schema
 
 from .column import Column
+from .table_version import TableVersion
 from .table_version_handle import TableVersionHandle
 
 _logger = logging.getLogger('pixeltable')
@@ -22,15 +24,28 @@ class TableVersionPath:
 
     TableVersionPath contains all metadata needed to execute queries and updates against a particular version of a
     table/view.
+
+    TableVersionPath supplies metadata needed for query construction (eg, column names), for which it uses a
+    cached TableVersion instance.
+    - when running inside a transaction, this instance is guaranteed to be validated
+    - when running outside a transaction, we use an unvalidated instance in order to avoid repeated validation
+      on every metadata-related method call (the instance won't stay validated, because TableVersionHandle.get()
+      runs a local transaction, at the end of which the instance is again invalidated)
+    - supplying metadata from an unvalidated instance is okay, because it needs to get revalidated anyway when a
+      query actually runs (at which point there is a transaction context) - there is no guarantee that in between
+      constructing a DataFrame and executing it, the underlying table schema hasn't changed (eg, a concurrent process
+      could have dropped a column referenced in the query).
     """
 
     tbl_version: TableVersionHandle
     base: Optional[TableVersionPath]
+    _cached_tbl_version: Optional[TableVersion]
 
     def __init__(self, tbl_version: TableVersionHandle, base: Optional[TableVersionPath] = None):
         assert tbl_version is not None
         self.tbl_version = tbl_version
         self.base = base
+        self._cached_tbl_version = None
 
     @classmethod
     def from_md(cls, path: schema.TableVersionPath) -> TableVersionPath:
@@ -47,17 +62,40 @@ class TableVersionPath:
             result.extend(self.base.as_md())
         return result
 
+    def refresh_cached_md(self) -> None:
+        from pixeltable.catalog import Catalog
+
+        if Env.get().in_xact:
+            # when we're running inside a transaction, we need to make sure to supply current metadata;
+            # mixing stale metadata with current metadata leads to query construction failures
+            # (multiple sqlalchemy Table instances for the same underlying table create corrupted From clauses)
+            if self._cached_tbl_version is not None and self._cached_tbl_version.is_validated:
+                # nothing to refresh
+                return
+        elif self._cached_tbl_version is not None:
+            return
+
+        with Catalog.get().begin_xact(for_write=False):
+            self._cached_tbl_version = self.tbl_version.get()
+
+    def clear_cached_md(self) -> None:
+        self._cached_tbl_version = None
+        if self.base is not None:
+            self.base.clear_cached_md()
+
     def tbl_id(self) -> UUID:
         """Return the id of the table/view that this path represents"""
         return self.tbl_version.id
 
     def version(self) -> int:
         """Return the version of the table/view that this path represents"""
-        return self.tbl_version.get().version
+        self.refresh_cached_md()
+        return self._cached_tbl_version.version
 
     def tbl_name(self) -> str:
         """Return the name of the table/view that this path represents"""
-        return self.tbl_version.get().name
+        self.refresh_cached_md()
+        return self._cached_tbl_version.name
 
     def path_len(self) -> int:
         """Return the length of the path"""
@@ -65,18 +103,22 @@ class TableVersionPath:
 
     def is_snapshot(self) -> bool:
         """Return True if this is a path of snapshot versions"""
-        if not self.tbl_version.get().is_snapshot:
+        self.refresh_cached_md()
+        if not self._cached_tbl_version.is_snapshot:
             return False
         return self.base.is_snapshot() if self.base is not None else True
 
     def is_view(self) -> bool:
-        return self.tbl_version.get().is_view
+        self.refresh_cached_md()
+        return self._cached_tbl_version.is_view
 
     def is_component_view(self) -> bool:
-        return self.tbl_version.get().is_component_view
+        self.refresh_cached_md()
+        return self._cached_tbl_version.is_component_view
 
     def is_insertable(self) -> bool:
-        return self.tbl_version.get().is_insertable()
+        self.refresh_cached_md()
+        return self._cached_tbl_version.is_insertable
 
     def get_tbl_versions(self) -> list[TableVersionHandle]:
         """Return all tbl versions"""
@@ -100,11 +142,12 @@ class TableVersionPath:
 
     def columns(self) -> list[Column]:
         """Return all user columns visible in this tbl version path, including columns from bases"""
-        result = list(self.tbl_version.get().cols_by_name.values())
-        if self.base is not None and self.tbl_version.get().include_base_columns:
+        self.refresh_cached_md()
+        result = list(self._cached_tbl_version.cols_by_name.values())
+        if self.base is not None and self._cached_tbl_version.include_base_columns:
             base_cols = self.base.columns()
             # we only include base columns that don't conflict with one of our column names
-            result.extend(c for c in base_cols if c.name not in self.tbl_version.get().cols_by_name)
+            result.extend(c for c in base_cols if c.name not in self._cached_tbl_version.cols_by_name)
         return result
 
     def cols_by_name(self) -> dict[str, Column]:
@@ -119,19 +162,21 @@ class TableVersionPath:
 
     def get_column(self, name: str, include_bases: Optional[bool] = None) -> Optional[Column]:
         """Return the column with the given name, or None if not found"""
-        col = self.tbl_version.get().cols_by_name.get(name)
+        self.refresh_cached_md()
+        col = self._cached_tbl_version.cols_by_name.get(name)
         if col is not None:
             return col
-        elif self.base is not None and (include_bases or self.tbl_version.get().include_base_columns):
+        elif self.base is not None and (include_bases or self._cached_tbl_version.include_base_columns):
             return self.base.get_column(name)
         else:
             return None
 
     def get_column_by_id(self, tbl_id: UUID, col_id: int) -> Optional[Column]:
         """Return the column for the given tbl/col id"""
+        self.refresh_cached_md()
         if self.tbl_version.id == tbl_id:
-            assert col_id in self.tbl_version.get().cols_by_id
-            return self.tbl_version.get().cols_by_id[col_id]
+            assert col_id in self._cached_tbl_version.cols_by_id
+            return self._cached_tbl_version.cols_by_id[col_id]
         elif self.base is not None:
             return self.base.get_column_by_id(tbl_id, col_id)
         else:
@@ -139,11 +184,12 @@ class TableVersionPath:
 
     def has_column(self, col: Column, include_bases: bool = True) -> bool:
         """Return True if this table has the given column."""
+        self.refresh_cached_md()
         assert col.tbl is not None
         if (
             col.tbl.id == self.tbl_version.id
             and col.tbl.effective_version == self.tbl_version.effective_version
-            and col.id in self.tbl_version.get().cols_by_id
+            and col.id in self._cached_tbl_version.cols_by_id
         ):
             # the column is visible in this table version
             return True

pixeltable/catalog/view.py

@@ -12,6 +12,10 @@ from pixeltable import catalog, exprs, func
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
 
+if TYPE_CHECKING:
+    from pixeltable.plan import SampleClause
+
+
 from .column import Column
 from .globals import _POS_COLUMN_NAME, MediaValidation, UpdateStatus
 from .table import Table
@@ -66,6 +70,7 @@ class View(Table):
         select_list: Optional[list[tuple[exprs.Expr, Optional[str]]]],
         additional_columns: dict[str, Any],
         predicate: Optional['exprs.Expr'],
+        sample_clause: Optional['SampleClause'],
         is_snapshot: bool,
         num_retained_versions: int,
         comment: str,
@@ -73,6 +78,8 @@ class View(Table):
         iterator_cls: Optional[type[ComponentIterator]],
         iterator_args: Optional[dict],
     ) -> View:
+        from pixeltable.plan import SampleClause
+
         # Convert select_list to more additional_columns if present
         include_base_columns: bool = select_list is None
         select_list_columns: List[Column] = []
@@ -84,12 +91,23 @@ class View(Table):
         columns = select_list_columns + columns_from_additional_columns
         cls._verify_schema(columns)
 
-        # verify that filter can be evaluated in the context of the base
+        # verify that filters can be evaluated in the context of the base
         if predicate is not None:
             if not predicate.is_bound_by([base]):
                 raise excs.Error(f'Filter cannot be computed in the context of the base {base.tbl_name()}')
             # create a copy that we can modify and store
             predicate = predicate.copy()
+        if sample_clause is not None:
+            # make sure that the sample clause can be computed in the context of the base
+            if sample_clause.stratify_exprs is not None and not all(
+                stratify_expr.is_bound_by([base]) for stratify_expr in sample_clause.stratify_exprs
+            ):
+                raise excs.Error(f'Sample clause cannot be computed in the context of the base {base.tbl_name()}')
+            # create a copy that we can modify and store
+            sc = sample_clause
+            sample_clause = SampleClause(
+                sc.version, sc.n, sc.n_per_stratum, sc.fraction, sc.seed, sc.stratify_exprs.copy()
+            )
 
         # same for value exprs
         for col in columns:
@@ -160,6 +178,8 @@ class View(Table):
         # if this is a snapshot, we need to retarget all exprs to the snapshot tbl versions
         if is_snapshot:
             predicate = predicate.retarget(base_version_path) if predicate is not None else None
+            if sample_clause is not None:
+                exprs.Expr.retarget_list(sample_clause.stratify_exprs, base_version_path)
             iterator_args_expr = (
                 iterator_args_expr.retarget(base_version_path) if iterator_args_expr is not None else None
             )
@@ -171,6 +191,7 @@ class View(Table):
             is_snapshot=is_snapshot,
             include_base_columns=include_base_columns,
             predicate=predicate.as_dict() if predicate is not None else None,
+            sample_clause=sample_clause.as_dict() if sample_clause is not None else None,
             base_versions=base_version_path.as_md(),
             iterator_class_fqn=iterator_class_fqn,
             iterator_args=iterator_args_expr.as_dict() if iterator_args_expr is not None else None,
@@ -204,8 +225,17 @@ class View(Table):
 
             from pixeltable.plan import Planner
 
-            plan, _ = Planner.create_view_load_plan(view._tbl_version_path)
-            num_rows, num_excs, _ = tbl_version.store_tbl.insert_rows(plan, v_min=tbl_version.version)
+            try:
+                plan, _ = Planner.create_view_load_plan(view._tbl_version_path)
+                num_rows, num_excs, _ = tbl_version.store_tbl.insert_rows(plan, v_min=tbl_version.version)
+            except:
+                # we need to remove the orphaned TableVersion instance
+                del catalog.Catalog.get()._tbl_versions[tbl_version.id, tbl_version.effective_version]
+                base_tbl_version = base.tbl_version.get()
+                if tbl_version.effective_version is None and not base_tbl_version.is_snapshot:
+                    # also remove tbl_version from the base
+                    base_tbl_version.mutable_views.remove(TableVersionHandle.create(tbl_version))
+                raise
             Env.get().console_logger.info(f'Created view `{name}` with {num_rows} rows, {num_excs} exceptions.')
 
         session.commit()
@@ -285,16 +315,18 @@ class View(Table):
 
     def _table_descriptor(self) -> str:
         display_name = 'Snapshot' if self._snapshot_only else 'View'
-        result = [f'{display_name} {self._path!r}']
+        result = [f'{display_name} {self._path()!r}']
         bases_descrs: list[str] = []
         for base, effective_version in zip(self._base_tables, self._effective_base_versions):
             if effective_version is None:
-                bases_descrs.append(f'{base._path!r}')
+                bases_descrs.append(f'{base._path()!r}')
             else:
-                base_descr = f'{base._path}:{effective_version}'
+                base_descr = f'{base._path()}:{effective_version}'
                 bases_descrs.append(f'{base_descr!r}')
         result.append(f' (of {", ".join(bases_descrs)})')
 
         if self._tbl_version.get().predicate is not None:
             result.append(f'\nWhere: {self._tbl_version.get().predicate!s}')
+        if self._tbl_version.get().sample_clause is not None:
+            result.append(f'\nSample: {self._tbl_version.get().sample_clause!s}')
         return ''.join(result)