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

pixeltable/exprs/column_ref.py

@@ -1,20 +1,29 @@
 from __future__ import annotations
 
-import copy
-from typing import Any, Optional, Sequence
+import warnings
+from typing import TYPE_CHECKING, Any, Sequence, cast
 from uuid import UUID
 
+import PIL.Image
 import sqlalchemy as sql
 
-import pixeltable as pxt
-from pixeltable import catalog, exceptions as excs, iterators as iters
+import pixeltable.catalog as catalog
+import pixeltable.exceptions as excs
+import pixeltable.iterators as iters
+import pixeltable.type_system as ts
+from pixeltable.catalog.table_version import TableVersionKey
 
 from ..utils.description_helper import DescriptionHelper
+from ..utils.filecache import FileCache
 from .data_row import DataRow
 from .expr import Expr
+from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
+if TYPE_CHECKING:
+    from pixeltable._query import Query, ResultSet
+
 
 class ColumnRef(Expr):
     """A reference to a table column
@@ -41,37 +50,38 @@ class ColumnRef(Expr):
     insert them into the EvalCtxs as needed
     """
 
-    col: catalog.Column
-    reference_tbl: Optional[catalog.TableVersionPath]
+    col: catalog.Column  # TODO: merge with col_handle
+    col_handle: catalog.ColumnHandle
+    reference_tbl: catalog.TableVersionPath | None
     is_unstored_iter_col: bool
-    iter_arg_ctx: Optional[RowBuilder.EvalCtx]
-    base_rowid_len: int
-    base_rowid: Sequence[Optional[Any]]
-    iterator: Optional[iters.ComponentIterator]
-    pos_idx: Optional[int]
-    id: int
     perform_validation: bool  # if True, performs media validation
+    iter_arg_ctx: RowBuilder.EvalCtx | None
+    iter_outputs: list[ColumnRef] | None
+    base_rowid_len: int  # number of rowid columns in the base table
+
+    # execution state
+    base_rowid: Sequence[Any | None]
+    iterator: iters.ComponentIterator | None
+    pos_idx: int
 
     def __init__(
         self,
         col: catalog.Column,
-        reference_tbl: Optional[catalog.TableVersionPath] = None,
-        perform_validation: Optional[bool] = None,
+        reference_tbl: catalog.TableVersionPath | None = None,
+        perform_validation: bool | None = None,
     ):
         super().__init__(col.col_type)
-        assert col.tbl is not None
         self.col = col
         self.reference_tbl = reference_tbl
-        self.is_unstored_iter_col = (
-            col.tbl.get().is_component_view and col.tbl.get().is_iterator_column(col) and not col.is_stored
-        )
+        self.col_handle = col.handle
+
+        self.is_unstored_iter_col = col.is_iterator_col and not col.is_stored
         self.iter_arg_ctx = None
-        # number of rowid columns in the base table
-        self.base_rowid_len = col.tbl.get().base.get().num_rowid_columns() if self.is_unstored_iter_col else 0
-        self.base_rowid = [None] * self.base_rowid_len
+        self.iter_outputs = None
+        self.base_rowid_len = 0
+        self.base_rowid = []
         self.iterator = None
-        # index of the position column in the view's primary key; don't try to reference tbl.store_tbl here
-        self.pos_idx = col.tbl.get().num_rowid_columns() - 1 if self.is_unstored_iter_col else None
+        self.pos_idx = 0
 
         self.perform_validation = False
         if col.col_type.is_media_type():
@@ -90,22 +100,27 @@ class ColumnRef(Expr):
             self.components = [non_validating_col_ref]
         self.id = self._create_id()
 
-    def set_iter_arg_ctx(self, iter_arg_ctx: RowBuilder.EvalCtx) -> None:
+    def set_iter_arg_ctx(self, iter_arg_ctx: RowBuilder.EvalCtx, iter_outputs: list[ColumnRef]) -> None:
         self.iter_arg_ctx = iter_arg_ctx
+        self.iter_outputs = iter_outputs
+        # If this is an unstored iterator column, then the iterator outputs may be needed in order to properly set the
+        # iterator position. Therefore, we need to add them as components in order to ensure they're marked as
+        # eval dependencies.
+        self.components.extend(iter_outputs)
         assert len(self.iter_arg_ctx.target_slot_idxs) == 1  # a single inline dict
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [
             *super()._id_attrs(),
-            ('tbl_id', self.col.tbl.id),
+            ('tbl_id', self.col.tbl_handle.id),
             ('col_id', self.col.id),
             ('perform_validation', self.perform_validation),
         ]
 
     # override
     def _retarget(self, tbl_versions: dict[UUID, catalog.TableVersion]) -> ColumnRef:
-        target = tbl_versions[self.col.tbl.id]
-        assert self.col.id in target.cols_by_id
+        target = tbl_versions[self.col.tbl_handle.id]
+        assert self.col.id in target.cols_by_id, f'{target}: {self.col.id} not in {list(target.cols_by_id.keys())}'
         col = target.cols_by_id[self.col.id]
         return ColumnRef(col, self.reference_tbl)
 
@@ -113,12 +128,16 @@ class ColumnRef(Expr):
         from .column_property_ref import ColumnPropertyRef
 
         # resolve column properties
+        if name == ColumnPropertyRef.Property.CELLMD.name.lower():
+            # This is not user accessible, but used internally to store cell metadata
+            return super().__getattr__(name)
+
         if (
             name == ColumnPropertyRef.Property.ERRORTYPE.name.lower()
             or name == ColumnPropertyRef.Property.ERRORMSG.name.lower()
         ):
-            property_is_present = self.col.is_stored and (self.col.is_computed or self.col_type.is_media_type())
-            if not property_is_present:
+            is_valid = (self.col.is_computed or self.col.col_type.is_media_type()) and self.col.is_stored
+            if not is_valid:
                 raise excs.Error(f'{name} only valid for a stored computed or media column: {self}')
             return ColumnPropertyRef(self, ColumnPropertyRef.Property[name.upper()])
         if (
@@ -138,78 +157,156 @@ class ColumnRef(Expr):
 
         return super().__getattr__(name)
 
-    def find_embedding_index(
-        self, idx_name: Optional[str], method_name: str
-    ) -> dict[str, catalog.TableVersion.IndexInfo]:
-        """Return IndexInfo for a column, with an optional given name"""
-        from pixeltable import index
+    def recompute(self, *, cascade: bool = True, errors_only: bool = False) -> catalog.UpdateStatus:
+        cat = catalog.Catalog.get()
+        # lock_mutable_tree=True: we need to be able to see whether any transitive view has column dependents
+        with cat.begin_xact(tbl=self.reference_tbl, for_write=True, lock_mutable_tree=True):
+            tbl_version = self.col_handle.tbl_version.get()
+            if tbl_version.id != self.reference_tbl.tbl_id:
+                raise excs.Error('Cannot recompute column of a base.')
+            if tbl_version.is_snapshot:
+                raise excs.Error('Cannot recompute column of a snapshot.')
+            col_name = self.col_handle.get().name
+            status = tbl_version.recompute_columns([col_name], errors_only=errors_only, cascade=cascade)
+            FileCache.get().emit_eviction_warnings()
+            return status
+
+    def similarity(
+        self,
+        item: Any = None,
+        *,
+        string: str | None = None,
+        image: PIL.Image.Image | None = None,
+        audio: str | None = None,
+        video: str | None = None,
+        idx: str | None = None,
+    ) -> Expr:
+        from .similarity_expr import SimilarityExpr
 
-        # determine index to use
-        idx_info_dict = self.col.get_idx_info(self.reference_tbl)
+        if item is not None:
+            warnings.warn(
+                'Use of similarity() without specifying an explicit modality is deprecated -- '
+                'since version 0.5.7. Please use one of the following instead:\n'
+                '  .similarity(string=...)\n'
+                '  .similarity(image=...)\n'
+                '  .similarity(audio=...)\n'
+                '  .similarity(video=...)',
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        arg_count = (string is not None) + (image is not None) + (audio is not None) + (video is not None)
+
+        if item is not None and arg_count != 0:
+            raise excs.Error('similarity(): `item` is deprecated and cannot be used together with modality arguments')
+
+        if arg_count > 1:
+            raise excs.Error('similarity(): expected exactly one of string=..., image=..., audio=..., video=...')
+
+        expr: Expr
+
+        if item is not None:
+            if isinstance(item, Expr):  # This can happen when using similarity() with @query
+                if not (item.col_type.is_string_type() or item.col_type.is_image_type()):
+                    raise excs.Error(f'similarity(): expected `String` or `Image`; got `{item.col_type}`')
+                expr = item
+            else:
+                if not isinstance(item, (str, PIL.Image.Image)):
+                    raise excs.Error(f'similarity(): expected `str` or `PIL.Image.Image`; got `{type(item).__name__}`')
+                expr = Expr.from_object(item)
+                assert expr.col_type.is_string_type() or expr.col_type.is_image_type()
+
+        if string is not None:
+            if isinstance(string, Expr):
+                if not string.col_type.is_string_type():
+                    raise excs.Error(f'similarity(string=...): expected `String`; got `{expr.col_type}`')
+                expr = string
+            else:
+                if not isinstance(string, str):
+                    raise excs.Error(f'similarity(string=...): expected `str`; got `{type(string).__name__}`')
+                expr = Expr.from_object(string)
+                assert expr.col_type.is_string_type()
+
+        if image is not None:
+            if isinstance(image, Expr):
+                if not image.col_type.is_image_type():
+                    raise excs.Error(f'similarity(image=...): expected `Image`; got `{image.col_type}`')
+                expr = image
+            else:
+                if not isinstance(image, PIL.Image.Image):
+                    raise excs.Error(f'similarity(image=...): expected `PIL.Image.Image`; got `{type(image).__name__}`')
+                expr = Expr.from_object(image)
+                assert expr.col_type.is_image_type()
+
+        if audio is not None:
+            if isinstance(audio, Expr):
+                if not audio.col_type.is_audio_type():
+                    raise excs.Error(f'similarity(audio=...): expected `Audio`; got `{audio.col_type}`')
+                expr = audio
+            else:
+                if not isinstance(audio, str):
+                    raise excs.Error(
+                        f'similarity(audio=...): expected `str` (path to audio file); got `{type(audio).__name__}`'
+                    )
+                expr = Literal(audio, ts.AudioType())
+
+        if video is not None:
+            if isinstance(video, Expr):
+                if not video.col_type.is_video_type():
+                    raise excs.Error(f'similarity(video=...): expected `Video`; got `{video.col_type}`')
+                expr = video
+            else:
+                if not isinstance(video, str):
+                    raise excs.Error(
+                        f'similarity(video=...): expected `str` (path to video file); got `{type(video).__name__}`'
+                    )
+                expr = Literal(video, ts.VideoType())
 
-        embedding_idx_info = {
-            info: value for info, value in idx_info_dict.items() if isinstance(value.idx, index.EmbeddingIndex)
-        }
-        if len(embedding_idx_info) == 0:
-            raise excs.Error(f'No indices found for {method_name!r} on column {self.col.name!r}')
-        if idx_name is not None and idx_name not in embedding_idx_info:
-            raise excs.Error(f'Index {idx_name!r} not found for {method_name!r} on column {self.col.name!r}')
-        if len(embedding_idx_info) > 1:
-            if idx_name is None:
-                raise excs.Error(
-                    f'Column {self.col.name!r} has multiple indices; use the index name to disambiguate: '
-                    f'`{method_name}(..., idx=<index_name>)`'
-                )
-            idx_info = {idx_name: embedding_idx_info[idx_name]}
-        else:
-            idx_info = embedding_idx_info
-        return idx_info
+        return SimilarityExpr(self, expr, idx_name=idx)
 
-    def similarity(self, item: Any, *, idx: Optional[str] = None) -> Expr:
-        from .similarity_expr import SimilarityExpr
+    def embedding(self, *, idx: str | None = None) -> ColumnRef:
+        from pixeltable.index import EmbeddingIndex
 
-        return SimilarityExpr(self, item, idx_name=idx)
+        idx_info = self.tbl.get().get_idx(self.col, idx, EmbeddingIndex)
+        return ColumnRef(idx_info.val_col)
 
-    def embedding(self, *, idx: Optional[str] = None) -> ColumnRef:
-        idx_info = self.find_embedding_index(idx, 'embedding')
-        assert len(idx_info) == 1
-        col = copy.copy(next(iter(idx_info.values())).val_col)
-        col.name = f'{self.col.name}_embedding_{idx if idx is not None else ""}'
-        col.create_sa_cols()
-        return ColumnRef(col)
+    @property
+    def tbl(self) -> catalog.TableVersionHandle:
+        return self.reference_tbl.tbl_version if self.reference_tbl is not None else self.col.tbl_handle
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         return self.col.name if self.col is not None else None
 
     def _equals(self, other: ColumnRef) -> bool:
         return self.col == other.col and self.perform_validation == other.perform_validation
 
-    def _df(self) -> 'pxt.dataframe.DataFrame':
-        from pixeltable import plan
+    def select(self) -> 'Query':
+        import pixeltable.plan as plan
+        from pixeltable._query import Query
 
         if self.reference_tbl is None:
             # No reference table; use the current version of the table to which the column belongs
-            tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
+            tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl_handle.id)
             return tbl.select(self)
         else:
-            # Explicit reference table; construct a DataFrame directly from it
-            return pxt.DataFrame(plan.FromClause([self.reference_tbl])).select(self)
+            # Explicit reference table; construct a Query directly from it
+            return Query(plan.FromClause([self.reference_tbl])).select(self)
 
-    def show(self, *args: Any, **kwargs: Any) -> 'pxt.dataframe.DataFrameResultSet':
-        return self._df().show(*args, **kwargs)
+    def show(self, *args: Any, **kwargs: Any) -> 'ResultSet':
+        return self.select().show(*args, **kwargs)
 
-    def head(self, *args: Any, **kwargs: Any) -> 'pxt.dataframe.DataFrameResultSet':
-        return self._df().head(*args, **kwargs)
+    def head(self, *args: Any, **kwargs: Any) -> 'ResultSet':
+        return self.select().head(*args, **kwargs)
 
-    def tail(self, *args: Any, **kwargs: Any) -> 'pxt.dataframe.DataFrameResultSet':
-        return self._df().tail(*args, **kwargs)
+    def tail(self, *args: Any, **kwargs: Any) -> 'ResultSet':
+        return self.select().tail(*args, **kwargs)
 
     def count(self) -> int:
-        return self._df().count()
+        return self.select().count()
 
-    def distinct(self) -> 'pxt.dataframe.DataFrame':
+    def distinct(self) -> 'Query':
         """Return distinct values in this column."""
-        return self._df().distinct()
+        return self.select().distinct()
 
     def __str__(self) -> str:
         if self.col.name is None:
@@ -224,17 +321,32 @@ class ColumnRef(Expr):
         return self._descriptors().to_html()
 
     def _descriptors(self) -> DescriptionHelper:
-        tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
+        with catalog.Catalog.get().begin_xact():
+            tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl_handle.id)
         helper = DescriptionHelper()
-        helper.append(f'Column\n{self.col.name!r}\n(of table {tbl._path!r})')
+        helper.append(f'Column\n{self.col.name!r}\n(of table {tbl._path()!r})')
         helper.append(tbl._col_descriptor([self.col.name]))
         idxs = tbl._index_descriptor([self.col.name])
         if len(idxs) > 0:
             helper.append(idxs)
         return helper
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        return None if self.perform_validation else self.col.sa_col
+    def prepare(self) -> None:
+        from pixeltable import store
+
+        if not self.is_unstored_iter_col:
+            return
+        col = self.col_handle.get()
+        self.base_rowid_len = col.get_tbl().base.get().num_rowid_columns()
+        self.base_rowid = [None] * self.base_rowid_len
+        assert isinstance(col.get_tbl().store_tbl, store.StoreComponentView)
+        self.pos_idx = cast(store.StoreComponentView, col.get_tbl().store_tbl).pos_col_idx
+
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
+        if self.perform_validation:
+            return None
+        self.col = self.col_handle.get()
+        return self.col.sa_col
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         if self.perform_validation:
@@ -273,31 +385,43 @@ class ColumnRef(Expr):
 
         # if this is a new base row, we need to instantiate a new iterator
         if self.base_rowid != data_row.pk[: self.base_rowid_len]:
+            assert self.iter_arg_ctx is not None
             row_builder.eval(data_row, self.iter_arg_ctx)
             iterator_args = data_row[self.iter_arg_ctx.target_slot_idxs[0]]
-            self.iterator = self.col.tbl.get().iterator_cls(**iterator_args)
+            self.iterator = self.col.get_tbl().iterator_cls(**iterator_args)
             self.base_rowid = data_row.pk[: self.base_rowid_len]
-        self.iterator.set_pos(data_row.pk[self.pos_idx])
+        stored_outputs = {col_ref.col.name: data_row[col_ref.slot_idx] for col_ref in self.iter_outputs}
+        assert all(name is not None for name in stored_outputs)
+        self.iterator.set_pos(data_row.pk[self.pos_idx], **stored_outputs)
         res = next(self.iterator)
         data_row[self.slot_idx] = res[self.col.name]
 
     def _as_dict(self) -> dict:
-        tbl = self.col.tbl
-        tbl_version = tbl.get().version if tbl.get().is_snapshot else None
+        tbl_handle = self.col.tbl_handle
         # we omit self.components, even if this is a validating ColumnRef, because init() will recreate the
         # non-validating component ColumnRef
+        assert tbl_handle.anchor_tbl_id is None  # TODO: support anchor_tbl_id for view-over-replica
         return {
-            'tbl_id': str(tbl.id),
-            'tbl_version': tbl_version,
+            'tbl_id': str(tbl_handle.id),
+            'tbl_version': tbl_handle.effective_version,
             'col_id': self.col.id,
             'reference_tbl': self.reference_tbl.as_dict() if self.reference_tbl is not None else None,
             'perform_validation': self.perform_validation,
         }
 
+    @classmethod
+    def get_column_id(cls, d: dict) -> catalog.QColumnId:
+        tbl_id, col_id = UUID(d['tbl_id']), d['col_id']
+        return catalog.QColumnId(tbl_id, col_id)
+
     @classmethod
     def get_column(cls, d: dict) -> catalog.Column:
         tbl_id, version, col_id = UUID(d['tbl_id']), d['tbl_version'], d['col_id']
-        tbl_version = catalog.Catalog.get().get_tbl_version(tbl_id, version)
+        # validate_initialized=False: this gets called as part of TableVersion.init()
+        # TODO: When we have views on replicas, we will need to store anchor_tbl_id in metadata as well.
+        tbl_version = catalog.Catalog.get().get_tbl_version(
+            TableVersionKey(tbl_id, version, None), validate_initialized=False
+        )
         # don't use tbl_version.cols_by_id here, this might be a snapshot reference to a column that was then dropped
         col = next(col for col in tbl_version.cols if col.id == col_id)
         return col

pixeltable/exprs/comparison.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -69,8 +69,8 @@ class Comparison(Expr):
     def _op2(self) -> Expr:
         return self.components[1]
 
-    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
-        from pixeltable import index
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
+        import pixeltable.index as index
 
         if str(self._op1.col_type.to_sa_type()) != str(self._op2.col_type.to_sa_type()):
             # Comparing columns of different SQL types (e.g., string vs. json); this can only be done in Python
@@ -81,9 +81,9 @@ class Comparison(Expr):
         if self.is_search_arg_comparison:
             # reference the index value column if there is an index and this is not a snapshot
             # (indices don't apply to snapshots)
-            tbl = self._op1.col.tbl.get()
+            tbl = self._op1.col.get_tbl()
             idx_info = [
-                info for info in self._op1.col.get_idx_info().values() if isinstance(info.idx, index.BtreeIndex)
+                info for info in tbl.idxs_by_col.get(self._op1.col.qid, []) if isinstance(info.idx, index.BtreeIndex)
             ]
             if len(idx_info) > 0 and not tbl.is_snapshot:
                 # there shouldn't be multiple B-tree indices on a column

pixeltable/exprs/compound_predicate.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import operator
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 import sqlalchemy as sql
 
@@ -36,7 +36,8 @@ class CompoundPredicate(Expr):
         return f' {self.operator} '.join([f'({e})' for e in self.components])
 
     @classmethod
-    def make_conjunction(cls, operands: list[Expr]) -> Optional[Expr]:
+    def make_conjunction(cls, operands: list[Expr | None]) -> Expr | None:
+        operands = [e for e in operands if e is not None]
         if len(operands) == 0:
             return None
         if len(operands) == 1:
@@ -60,14 +61,14 @@ class CompoundPredicate(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('operator', self.operator.value)]
 
-    def split_conjuncts(self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Optional[Expr]]:
+    def split_conjuncts(self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Expr | None]:
         if self.operator in (LogicalOperator.OR, LogicalOperator.NOT):
             return super().split_conjuncts(condition)
         matches = [op for op in self.components if condition(op)]
         non_matches = [op for op in self.components if not condition(op)]
         return (matches, self.make_conjunction(non_matches))
 
-    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
         sql_exprs = [sql_elements.get(op) for op in self.components]
         if any(e is None for e in sql_exprs):
             return None