pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/exprs/literal.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import datetime
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import sqlalchemy as sql
@@ -18,7 +18,7 @@ from .sql_element_cache import SqlElementCache
 class Literal(Expr):
     val: Any
 
-    def __init__(self, val: Any, col_type: Optional[ts.ColumnType] = None):
+    def __init__(self, val: Any, col_type: ts.ColumnType | None = None):
         if col_type is not None:
             val = col_type.create_literal(val)
         else:
@@ -42,7 +42,7 @@ class Literal(Expr):
         self.val = val
         self.id = self._create_id()
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         return 'Literal'
 
     def __str__(self) -> str:
@@ -69,7 +69,7 @@ class Literal(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('val', self.val)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         # Return a sql object so that constants can participate in SQL expressions
         return sql.sql.expression.literal(self.val, type_=self.col_type.to_sa_type())
 
@@ -97,7 +97,7 @@ class Literal(Expr):
         else:
             return {'val': self.val, **super()._as_dict()}
 
-    def as_literal(self) -> Optional[Literal]:
+    def as_literal(self) -> Literal | None:
         return self
 
     @classmethod

pixeltable/exprs/method_ref.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -55,7 +55,7 @@ class MethodRef(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('method_name', self.method_name)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/object_ref.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -43,7 +43,7 @@ class ObjectRef(Expr):
     def _equals(self, other: ObjectRef) -> bool:
         return self.id == other.id
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/row_builder.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
+import dataclasses
 import sys
 import time
-from dataclasses import dataclass
-from typing import Any, Iterable, Optional, Sequence
+from typing import Any, Iterable, NamedTuple, Sequence
 from uuid import UUID
 
 import numpy as np
+import sqlalchemy as sql
 
-from pixeltable import catalog, exceptions as excs, utils
+from pixeltable import catalog, exceptions as excs, exprs, utils
 from pixeltable.env import Env
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.misc import non_none_dict_factory
 
 from .data_row import DataRow
 from .expr import Expr, ExprScope
@@ -35,8 +36,7 @@ class ExecProfile:
             )
 
 
-@dataclass
-class ColumnSlotIdx:
+class ColumnSlotIdx(NamedTuple):
     """Info for how to locate materialized column in DataRow
     TODO: can this be integrated into RowBuilder directly?
     """
@@ -50,6 +50,12 @@ class RowBuilder:
 
     For ColumnRefs to unstored iterator columns:
     - in order for them to be executable, we also record the iterator args and pass them to the ColumnRef
+
+    Args:
+        output_exprs: list of Exprs to be evaluated
+        columns: list of columns to be materialized
+        input_exprs: list of Exprs that are excluded from evaluation (because they're already materialized)
+    TODO: enforce that output_exprs doesn't overlap with input_exprs?
     """
 
     unique_exprs: ExprSet
@@ -63,7 +69,8 @@ class RowBuilder:
 
     input_exprs: ExprSet
 
-    table_columns: list[ColumnSlotIdx]
+    tbl: catalog.TableVersion | None  # reference table of the RowBuilder; used to identify pk columns for writes
+    table_columns: dict[catalog.Column, int | None]  # value: slot idx, if the result of an expr
     default_eval_ctx: EvalCtx
     unstored_iter_args: dict[UUID, Expr]
 
@@ -84,7 +91,12 @@ class RowBuilder:
     # (a subexpr can be shared across multiple output exprs)
     output_expr_ids: list[set[int]]
 
-    @dataclass
+    img_slot_idxs: list[int]  # Indices of image slots
+    media_slot_idxs: list[int]  # Indices of non-image media slots
+    array_slot_idxs: list[int]  # Indices of array slots
+    json_slot_idxs: list[int]  # Indices of json slots
+
+    @dataclasses.dataclass
     class EvalCtx:
         """Context for evaluating a set of target exprs"""
 
@@ -93,14 +105,13 @@ class RowBuilder:
         target_slot_idxs: list[int]  # slot idxs of target exprs; might contain duplicates
         target_exprs: list[Expr]  # exprs corresponding to target_slot_idxs
 
-    def __init__(self, output_exprs: Sequence[Expr], columns: Sequence[catalog.Column], input_exprs: Iterable[Expr]):
-        """
-        Args:
-            output_exprs: list of Exprs to be evaluated
-            columns: list of columns to be materialized
-            input_exprs: list of Exprs that are excluded from evaluation (because they're already materialized)
-        TODO: enforce that output_exprs doesn't overlap with input_exprs?
-        """
+    def __init__(
+        self,
+        output_exprs: Sequence[Expr],
+        columns: Sequence[catalog.Column],
+        input_exprs: Iterable[Expr],
+        tbl: catalog.TableVersion | None = None,
+    ):
         self.unique_exprs: ExprSet[Expr] = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
 
@@ -117,7 +128,7 @@ class RowBuilder:
         )
 
         # if init(columns):
-        # - we are creating table rows and need to record columns for create_table_row()
+        # - we are creating table rows and need to record columns for create_store_table_row()
         # - output_exprs materialize those columns
         # - input_exprs are ColumnRefs of the non-computed columns (ie, what needs to be provided as input)
         # - media validation:
@@ -125,7 +136,8 @@ class RowBuilder:
         #   * further references to that column (eg, computed cols) need to resolve to the validating ColumnRef
         from .column_ref import ColumnRef
 
-        self.table_columns: list[ColumnSlotIdx] = []
+        self.tbl = tbl
+        self.table_columns = {}
         self.input_exprs = ExprSet()
         validating_colrefs: dict[Expr, Expr] = {}  # key: non-validating colref, value: corresp. validating colref
         for col in columns:
@@ -171,18 +183,18 @@ class RowBuilder:
         col_refs = [e for e in self.unique_exprs if isinstance(e, ColumnRef)]
 
         def refs_unstored_iter_col(col_ref: ColumnRef) -> bool:
-            tbl = col_ref.col.tbl
+            tbl = col_ref.col.get_tbl()
             return tbl.is_component_view and tbl.is_iterator_column(col_ref.col) and not col_ref.col.is_stored
 
         unstored_iter_col_refs = [col_ref for col_ref in col_refs if refs_unstored_iter_col(col_ref)]
-        component_views = [col_ref.col.tbl for col_ref in unstored_iter_col_refs]
+        component_views = [col_ref.col.get_tbl() for col_ref in unstored_iter_col_refs]
         unstored_iter_args = {view.id: view.iterator_args.copy() for view in component_views}
         self.unstored_iter_args = {
             id: self._record_unique_expr(arg, recursive=True) for id, arg in unstored_iter_args.items()
         }
 
         for col_ref in unstored_iter_col_refs:
-            iter_arg_ctx = self.create_eval_ctx([unstored_iter_args[col_ref.col.tbl.id]])
+            iter_arg_ctx = self.create_eval_ctx([unstored_iter_args[col_ref.col.get_tbl().id]])
             col_ref.set_iter_arg_ctx(iter_arg_ctx)
 
         # we guarantee that we can compute the expr DAG in a single front-to-back pass
@@ -201,7 +213,7 @@ class RowBuilder:
                 # this is input and therefore doesn't depend on other exprs
                 continue
             # error properties don't have exceptions themselves
-            if isinstance(expr, ColumnPropertyRef) and expr.is_error_prop():
+            if isinstance(expr, ColumnPropertyRef) and expr.is_cellmd_prop():
                 continue
             dependency_idxs = [d.slot_idx for d in expr.dependencies()]
             self.dependencies[expr.slot_idx, dependency_idxs] = True
@@ -227,13 +239,32 @@ class RowBuilder:
         for e in self.output_exprs:
             self._record_output_expr_id(e, e.slot_idx)
 
+        self.img_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_image_type()]
+        self.media_slot_idxs = [
+            e.slot_idx for e in self.unique_exprs if e.col_type.is_media_type() and not e.col_type.is_image_type()
+        ]
+        self.array_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_array_type()]
+        self.json_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_json_type()]
+
     def add_table_column(self, col: catalog.Column, slot_idx: int) -> None:
-        """Record a column that is part of the table row"""
-        self.table_columns.append(ColumnSlotIdx(col, slot_idx))
+        """Record an output column for which the value is produced via expr evaluation"""
+        assert self.tbl is not None
+        assert col.is_stored
+        self.table_columns[col] = slot_idx
 
-    def output_slot_idxs(self) -> list[ColumnSlotIdx]:
-        """Return ColumnSlotIdx for output columns"""
-        return self.table_columns
+    def add_table_columns(self, cols: list[catalog.Column]) -> None:
+        """Record output columns whose values are materialized into DataRow.cell_vals"""
+        for col in cols:
+            self.table_columns[col] = None
+
+    @property
+    def media_output_col_info(self) -> list[ColumnSlotIdx]:
+        """Return slot idxs for media output columns whose values are produced by expr evaluation"""
+        return [
+            ColumnSlotIdx(col, slot_idx)
+            for col, slot_idx in self.table_columns.items()
+            if col.col_type.is_media_type() and slot_idx is not None
+        ]
 
     @property
     def num_materialized(self) -> int:
@@ -277,7 +308,7 @@ class RowBuilder:
             self._record_output_expr_id(d, output_expr_id)
 
     def _compute_dependencies(
-        self, target_slot_idxs: list[int], excluded_slot_idxs: list[int], target_scope: Optional[ExprScope] = None
+        self, target_slot_idxs: list[int], excluded_slot_idxs: list[int], target_scope: ExprScope | None = None
     ) -> list[int]:
         """Compute exprs needed to materialize the given target slots, excluding 'excluded_slot_idxs'
 
@@ -331,7 +362,7 @@ class RowBuilder:
             self.__set_slot_idxs_aux(c)
 
     def get_dependencies(
-        self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None, limit_scope: bool = True
+        self, targets: Iterable[Expr], exclude: Iterable[Expr] | None = None, limit_scope: bool = True
     ) -> list[Expr]:
         """
         Return list of dependencies needed to evaluate the given target exprs (expressed as slot idxs).
@@ -349,7 +380,7 @@ class RowBuilder:
             return []
         # make sure we only refer to recorded exprs
         targets = [self.unique_exprs[e] for e in targets]
-        target_scope: Optional[ExprScope] = None
+        target_scope: ExprScope | None = None
         if limit_scope:
             # make sure all targets are from the same scope
             target_scopes = {e.scope() for e in targets}
@@ -367,7 +398,7 @@ class RowBuilder:
         return [self.unique_exprs[id] for id in result_ids]
 
     def create_eval_ctx(
-        self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None, limit_scope: bool = True
+        self, targets: Iterable[Expr], exclude: Iterable[Expr] | None = None, limit_scope: bool = True
     ) -> EvalCtx:
         """Return EvalCtx for targets"""
         targets = list(targets)
@@ -396,9 +427,9 @@ class RowBuilder:
         self,
         data_row: DataRow,
         ctx: EvalCtx,
-        profile: Optional[ExecProfile] = None,
+        profile: ExecProfile | None = None,
         ignore_errors: bool = False,
-        force_eval: Optional[ExprScope] = None,
+        force_eval: ExprScope | None = None,
     ) -> None:
         """
         Populates the slots in data_row given in ctx.
@@ -427,33 +458,76 @@ class RowBuilder:
                         expr, f'expression {expr}', data_row.get_exc(expr.slot_idx), exc_tb, input_vals, 0
                     ) from exc
 
-    def create_table_row(self, data_row: DataRow, exc_col_ids: set[int]) -> tuple[dict[str, Any], int]:
-        """Create a table row from the slots that have an output column assigned
+    def create_store_table_row(
+        self, data_row: DataRow, cols_with_excs: set[int] | None, pk: tuple[int, ...]
+    ) -> tuple[list[Any], int]:
+        """Create a store table row from the slots that have an output column assigned
 
-        Return tuple[dict that represents a stored row (can be passed to sql.insert()), # of exceptions]
+        Return tuple[list of row values in `self.table_columns` order, # of exceptions]
             This excludes system columns.
+            Row values are converted to their store type.
         """
+        from pixeltable.exprs.column_property_ref import ColumnPropertyRef
+
         num_excs = 0
-        table_row: dict[str, Any] = {}
-        for info in self.table_columns:
-            col, slot_idx = info.col, info.slot_idx
+        table_row: list[Any] = list(pk)
+        # Nulls in JSONB columns need to be stored as sql.sql.null(), otherwise it stores a json 'null'
+        for col, slot_idx in self.table_columns.items():
+            if col.id in data_row.cell_vals:
+                table_row.append(data_row.cell_vals[col.id])
+                if col.stores_cellmd:
+                    if data_row.cell_md[col.id] is None:
+                        table_row.append(sql.sql.null())
+                    else:
+                        # we want to minimize the size of the stored dict and use dict_factory to remove Nones
+                        md = dataclasses.asdict(data_row.cell_md[col.id], dict_factory=non_none_dict_factory)
+                        assert len(md) > 0
+                        table_row.append(md)
+                if slot_idx is not None and data_row.has_exc(slot_idx):
+                    num_excs += 1
+                    if cols_with_excs is not None:
+                        cols_with_excs.add(col.id)
+                continue
+
             if data_row.has_exc(slot_idx):
-                # exceptions get stored in the errortype/-msg columns
                 exc = data_row.get_exc(slot_idx)
                 num_excs += 1
-                exc_col_ids.add(col.id)
-                table_row[col.store_name()] = None
-                table_row[col.errortype_store_name()] = type(exc).__name__
-                table_row[col.errormsg_store_name()] = str(exc)
+                if cols_with_excs is not None:
+                    cols_with_excs.add(col.id)
+                table_row.append(sql.sql.null() if col.col_type.is_json_type() else None)
+                if col.stores_cellmd:
+                    # exceptions get stored in the errortype/-msg properties of the cellmd column
+                    table_row.append(ColumnPropertyRef.create_cellmd_exc(exc))
             else:
-                if col.col_type.is_image_type() and data_row.file_urls[slot_idx] is None:
-                    # we have yet to store this image
-                    filepath = str(MediaStore.prepare_media_path(col.tbl.id, col.id, col.tbl.version))
-                    data_row.flush_img(slot_idx, filepath)
-                val = data_row.get_stored_val(slot_idx, col.get_sa_col_type())
-                table_row[col.store_name()] = val
-                # we unfortunately need to set these, even if there are no errors
-                table_row[col.errortype_store_name()] = None
-                table_row[col.errormsg_store_name()] = None
+                val = data_row.get_stored_val(slot_idx, col.sa_col_type)
+                table_row.append(val)
+                if col.stores_cellmd:
+                    table_row.append(sql.sql.null())  # placeholder for cellmd column
 
         return table_row, num_excs
+
+    def store_column_names(self) -> list[str]:
+        """
+        Returns the list of store column names corresponding to the table_columns of this RowBuilder.
+        The second tuple element of the return value is a dictionary containing all media columns in the
+        table; it's the mapping {list_index: column}.
+        """
+        assert self.tbl is not None, self.table_columns
+        store_col_names: list[str] = [pk_col.name for pk_col in self.tbl.store_tbl.pk_columns()]
+
+        for col in self.table_columns:
+            store_col_names.append(col.store_name())
+            if col.stores_cellmd:
+                store_col_names.append(col.cellmd_store_name())
+
+        return store_col_names
+
+    def make_row(self) -> exprs.DataRow:
+        """Creates a new DataRow with the current row_builder's configuration."""
+        return exprs.DataRow(
+            size=self.num_materialized,
+            img_slot_idxs=self.img_slot_idxs,
+            media_slot_idxs=self.media_slot_idxs,
+            array_slot_idxs=self.array_slot_idxs,
+            json_slot_idxs=self.json_slot_idxs,
+        )

pixeltable/exprs/rowid_ref.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Optional, cast
+from typing import Any, cast
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -25,18 +25,18 @@ class RowidRef(Expr):
     (with and without a TableVersion).
     """
 
-    tbl: Optional[catalog.TableVersionHandle]
-    normalized_base: Optional[catalog.TableVersionHandle]
+    tbl: catalog.TableVersionHandle | None
+    normalized_base: catalog.TableVersionHandle | None
     tbl_id: UUID
     normalized_base_id: UUID
     rowid_component_idx: int
 
     def __init__(
         self,
-        tbl: Optional[catalog.TableVersionHandle],
+        tbl: catalog.TableVersionHandle | None,
         idx: int,
-        tbl_id: Optional[UUID] = None,
-        normalized_base_id: Optional[UUID] = None,
+        tbl_id: UUID | None = None,
+        normalized_base_id: UUID | None = None,
     ):
         super().__init__(ts.IntType(nullable=False))
         self.tbl = tbl
@@ -57,7 +57,7 @@ class RowidRef(Expr):
         self.rowid_component_idx = idx
         self.id = self._create_id()
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         return str(self)
 
     def _equals(self, other: RowidRef) -> bool:
@@ -98,17 +98,13 @@ class RowidRef(Expr):
         self.tbl = tbl.tbl_version
         self.tbl_id = self.tbl.id
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         tbl = self.tbl.get() if self.tbl is not None else catalog.Catalog.get().get_tbl_version(self.tbl_id, None)
         assert tbl.is_validated
         rowid_cols = tbl.store_tbl.rowid_columns()
         assert self.rowid_component_idx <= len(rowid_cols), (
             f'{self.rowid_component_idx} not consistent with {rowid_cols}'
         )
-        # _logger.debug(
-        #     f'RowidRef.sql_expr: tbl={tbl.id}{tbl.effective_version} sa_tbl={id(tbl.store_tbl.sa_tbl):x} '
-        #     f'tv={id(tbl):x}'
-        # )
         return rowid_cols[self.rowid_component_idx]
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,6 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import sqlalchemy as sql
 
@@ -12,34 +14,43 @@ from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
+if TYPE_CHECKING:
+    from pixeltable.catalog.table_version import TableVersion
+
 
 class SimilarityExpr(Expr):
-    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: Optional[str] = None):
+    """
+    A similarity expression against an embedding index.
+    """
+
+    idx_id: int
+    idx_name: str
+
+    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: str | None = None):
+        from pixeltable.index import EmbeddingIndex
+
         super().__init__(ts.FloatType())
         item_expr = Expr.from_object(item)
         if item_expr is None or not (item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()):
             raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not a {type(item)}')
-        assert item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()
 
         self.components = [col_ref, item_expr]
 
-        from pixeltable import index
-
         # determine index to use
-        idx_dict = col_ref.find_embedding_index(idx_name, 'similarity')
-        assert len(idx_dict) == 1
-        self.idx_info = next(iter(idx_dict.values()))
-        idx = self.idx_info.idx
-        assert isinstance(idx, index.EmbeddingIndex)
+        idx_info = col_ref.tbl.get().get_idx(col_ref.col, idx_name, EmbeddingIndex)
+        self.idx_id = idx_info.id
+        self.idx_name = idx_info.name
+        idx = idx_info.idx
+        assert isinstance(idx, EmbeddingIndex)
 
         if item_expr.col_type.is_string_type() and idx.string_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have a '
+                f'Embedding index {idx_info.name!r} on column {idx_info.col.name!r} does not have a '
                 f'string embedding and does not support string queries'
             )
         if item_expr.col_type.is_image_type() and idx.image_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have an '
+                f'Embedding index {idx_info.name!r} on column {idx_info.col.name!r} does not have an '
                 f'image embedding and does not support image queries'
             )
         self.id = self._create_id()
@@ -48,39 +59,53 @@ class SimilarityExpr(Expr):
         return f'{self.components[0]}.similarity({self.components[1]})'
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return [*super()._id_attrs(), ('idx_name', self.idx_info.name)]
+        return [*super()._id_attrs(), ('idx_id', self.idx_id)]
 
     def default_column_name(self) -> str:
         return 'similarity'
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        # TODO: validate that the index still exists
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
+
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
             raise excs.Error('similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
-        from pixeltable import index
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.similarity_clause(idx_info.val_col, item)
 
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
+    def as_order_by_clause(self, is_asc: bool) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
 
-    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ColumnElement]:
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
             raise excs.Error('similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
-        from pixeltable import index
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.order_by_clause(idx_info.val_col, item, is_asc)
+
+    def _resolve_idx(self) -> 'TableVersion.IndexInfo':
+        from pixeltable.index import EmbeddingIndex
 
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
+        # resolve idx_id
+        col_ref = self.components[0]
+        if self.idx_id not in col_ref.tbl.get().idxs:
+            raise excs.Error(f'Index {self.idx_name!r} not found')
+        idx_info = col_ref.tbl.get().idxs[self.idx_id]
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         raise excs.Error('similarity(): cannot be used in a computed column')
 
     def _as_dict(self) -> dict:
-        return {'idx_name': self.idx_info.name, **super()._as_dict()}
+        return {'idx_name': self.idx_name, **super()._as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> 'SimilarityExpr':
-        iname = d.get('idx_name')
+        idx_name = d.get('idx_name')
         assert len(components) == 2
         assert isinstance(components[0], ColumnRef)
-        return cls(components[0], components[1], idx_name=iname)
+        return cls(components[0], components[1], idx_name=idx_name)

pixeltable/exprs/sql_element_cache.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional
+from typing import Iterable
 
 import sqlalchemy as sql
 
@@ -9,9 +9,9 @@ from .expr_dict import ExprDict
 class SqlElementCache:
     """Cache of sql.ColumnElements for exprs"""
 
-    cache: dict[int, Optional[sql.ColumnElement]]  # key: Expr.id
+    cache: dict[int, sql.ColumnElement | None]  # key: Expr.id
 
-    def __init__(self, elements: Optional[ExprDict[sql.ColumnElement]] = None):
+    def __init__(self, elements: ExprDict[sql.ColumnElement] | None = None):
         self.cache = {}
         if elements is not None:
             for e, el in elements.items():
@@ -21,7 +21,7 @@ class SqlElementCache:
         for e, el in elements.items():
             self.cache[e.id] = el
 
-    def get(self, e: Expr) -> Optional[sql.ColumnElement]:
+    def get(self, e: Expr) -> sql.ColumnElement | None:
         """Returns the sql.ColumnElement for the given Expr, or None if Expr.to_sql() returns None."""
         try:
             return self.cache[e.id]

pixeltable/exprs/string_op.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -60,7 +60,7 @@ class StringOp(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('operator', self.operator.value)]
 
-    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
         left = sql_elements.get(self._op1)
         right = sql_elements.get(self._op2)
         if left is None or right is None:
@@ -68,7 +68,7 @@ class StringOp(Expr):
         if self.operator == StringOperator.CONCAT:
             return left.concat(right)
         if self.operator == StringOperator.REPEAT:
-            return sql.func.repeat(sql.cast(left, sql.String), sql.cast(right, sql.Integer))
+            return sql.func.repeat(left.cast(sql.String), right.cast(sql.Integer))
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
@@ -76,7 +76,7 @@ class StringOp(Expr):
         op2_val = data_row[self._op2.slot_idx]
         data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
 
-    def eval_nullable(self, op1_val: Union[str, None], op2_val: Union[int, str, None]) -> Union[str, None]:
+    def eval_nullable(self, op1_val: str | None, op2_val: int | str | None) -> str | None:
         """
         Return the result of evaluating the expression on two nullable int/float operands,
         None is interpreted as SQL NULL
@@ -85,7 +85,7 @@ class StringOp(Expr):
             return None
         return self.eval_non_null(op1_val, op2_val)
 
-    def eval_non_null(self, op1_val: str, op2_val: Union[int, str]) -> str:
+    def eval_non_null(self, op1_val: str, op2_val: int | str) -> str:
         """
         Return the result of evaluating the expression on two int/float operands
         """