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

pixeltable/exprs/row_builder.py

@@ -1,21 +1,25 @@
 from __future__ import annotations
 
+import dataclasses
 import sys
 import time
-from dataclasses import dataclass
-from typing import Any, Iterable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Sequence, TypeVar
 from uuid import UUID
 
+import numpy as np
 import sqlalchemy as sql
 
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.func as func
-import pixeltable.utils as utils
+from pixeltable import catalog, exceptions as excs, exprs, utils
+from pixeltable.env import Env
+from pixeltable.utils.misc import non_none_dict_factory
+
 from .data_row import DataRow
-from .expr import Expr
+from .expr import Expr, ExprScope
 from .expr_set import ExprSet
 
+if TYPE_CHECKING:
+    from .column_ref import ColumnRef
+
 
 class ExecProfile:
     def __init__(self, row_builder: RowBuilder):
@@ -30,14 +34,16 @@ class ExecProfile:
             per_call_time = self.eval_time[i] / self.eval_count[i]
             calls_per_row = self.eval_count[i] / num_rows
             multiple_str = f'({calls_per_row}x)' if calls_per_row > 1 else ''
-            print(f'{self.row_builder.unique_exprs[i]}: {utils.print_perf_counter_delta(per_call_time)} {multiple_str}')
+            Env.get().console_logger.info(
+                f'{self.row_builder.unique_exprs[i]}: {utils.print_perf_counter_delta(per_call_time)} {multiple_str}'
+            )
 
 
-@dataclass
-class ColumnSlotIdx:
+class ColumnSlotIdx(NamedTuple):
     """Info for how to locate materialized column in DataRow
     TODO: can this be integrated into RowBuilder directly?
     """
+
     col: catalog.Column
     slot_idx: int
 
@@ -47,7 +53,14 @@ 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
     next_slot_idx: int
     input_expr_slot_idxs: set[int]
@@ -59,61 +72,84 @@ 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
+    for_view_load: bool  # True if this RowBuilder represents a view load
+
+    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]
+    unstored_iter_outputs: dict[UUID, list['ColumnRef']]
 
     # transitive dependents for the purpose of exception propagation: an exception for slot i is propagated to
     # _exc_dependents[i]
     # (list of set of slot_idxs, indexed by slot_idx)
     _exc_dependents: list[set[int]]
 
+    # dependents[i] = direct dependents of expr with slot idx i; dependents[i, j] == True: expr j depends on expr i
+    dependents: np.ndarray  # of bool
+    transitive_dependents: np.ndarray  # of bool
+    # dependencies[i] = direct dependencies of expr with slot idx i; transpose of dependents
+    dependencies: np.ndarray  # of bool
+    # num_dependencies[i] = number of direct dependencies of expr with slot idx i
+    num_dependencies: np.ndarray  # of int
+
     # records the output_expr that a subexpr belongs to
     # (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"""
+
         slot_idxs: list[int]  # slot idxs of exprs needed to evaluate target exprs; does not contain duplicates
         exprs: list[Expr]  # exprs corresponding to slot_idxs
         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]
+        self,
+        output_exprs: Sequence[Expr],
+        columns: Sequence[catalog.Column],
+        input_exprs: Iterable[Expr],
+        tbl: catalog.TableVersion | None = None,
+        for_view_load: bool = False,
     ):
-        """
-        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?
-        """
+        from .column_property_ref import ColumnPropertyRef
+        from .column_ref import ColumnRef
+
         self.unique_exprs: ExprSet[Expr] = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
 
-        # record input and output exprs; make copies to avoid reusing execution state
+        # record input exprs; make copies to avoid reusing execution state
         unique_input_exprs = [self._record_unique_expr(e.copy(), recursive=False) for e in input_exprs]
+
         self.input_expr_slot_idxs = {e.slot_idx for e in unique_input_exprs}
 
         resolve_cols = set(columns)
-        self.output_exprs = ExprSet([
+        self.output_exprs = ExprSet(
             self._record_unique_expr(e.copy().resolve_computed_cols(resolve_cols=resolve_cols), recursive=True)
             for e in output_exprs
-        ])
+        )
 
         # 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:
         #   * for write-validated columns, we need to create validating ColumnRefs
         #   * 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.for_view_load = for_view_load
+        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:
             expr: Expr
             if col.is_computed:
@@ -125,7 +161,8 @@ class RowBuilder:
             else:
                 # record a ColumnRef so that references to this column resolve to the same slot idx
                 perform_validation = (
-                    None if not col.col_type.is_media_type()
+                    None
+                    if not col.col_type.is_media_type()
                     else col.media_validation == catalog.MediaValidation.ON_WRITE
                 )
                 expr = ColumnRef(col, perform_validation=perform_validation)
@@ -153,22 +190,39 @@ class RowBuilder:
         #   because that would cause them to be evaluated for every single row
         # - the separate eval ctx allows the ColumnRef to materialize the iterator args only when the underlying
         #   iterated object changes
+
         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
-            return tbl.is_component_view() and tbl.is_iterator_column(col_ref.col) and not col_ref.col.is_stored
+            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}
+
+        # the *stored* output columns of the unstored iterators
+        self.unstored_iter_outputs = {
+            view.id: [
+                self._record_unique_expr(ColumnRef(col), recursive=True)
+                for col in view.iterator_columns()
+                if col.is_stored
+            ]
+            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()
+            id: self._record_unique_expr(args, recursive=True) for id, args in unstored_iter_args.items()
         }
 
+        unstored_iter_col_refs = [
+            self._record_unique_expr(col_ref, recursive=True) for col_ref in unstored_iter_col_refs
+        ]
+
         for col_ref in unstored_iter_col_refs:
-            iter_arg_ctx = self.create_eval_ctx([unstored_iter_args[col_ref.col.tbl.id]])
-            col_ref.set_iter_arg_ctx(iter_arg_ctx)
+            iter_arg_ctx = self.create_eval_ctx([self.unstored_iter_args[col_ref.col.get_tbl().id]])
+            iter_outputs = self.unstored_iter_outputs[col_ref.col.get_tbl().id]
+            col_ref.set_iter_arg_ctx(iter_arg_ctx, iter_outputs)
 
         # we guarantee that we can compute the expr DAG in a single front-to-back pass
         for i, expr in enumerate(self.unique_exprs):
@@ -176,19 +230,31 @@ class RowBuilder:
 
         # determine transitive dependencies for the purpose of exception propagation
         # (list of set of slot_idxs, indexed by slot_idx)
+        # self.dependents = np.zeros((self.num_materialized, self.num_materialized), dtype=bool)
+        self.dependencies = np.zeros((self.num_materialized, self.num_materialized), dtype=bool)
         exc_dependencies: list[set[int]] = [set() for _ in range(self.num_materialized)]
-        from .column_property_ref import ColumnPropertyRef
+
         for expr in self.unique_exprs:
             if expr.slot_idx in self.input_expr_slot_idxs:
                 # 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
             for d in expr.dependencies():
                 exc_dependencies[expr.slot_idx].add(d.slot_idx)
                 exc_dependencies[expr.slot_idx].update(exc_dependencies[d.slot_idx])
 
+        self.num_dependencies = np.sum(self.dependencies, axis=1)
+        self.dependents = self.dependencies.T
+        self.transitive_dependents = np.zeros((self.num_materialized, self.num_materialized), dtype=bool)
+        for i in reversed(range(self.num_materialized)):
+            self.transitive_dependents[i] = self.dependents[i] | np.any(
+                self.transitive_dependents[self.dependents[i]], axis=0
+            )
+
         self._exc_dependents = [set() for _ in range(self.num_materialized)]
         for expr in self.unique_exprs:
             assert expr.slot_idx is not None
@@ -199,19 +265,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
 
-    def set_conn(self, conn: sql.engine.Connection) -> None:
-        from .function_call import FunctionCall
-        for expr in self.unique_exprs:
-            if isinstance(expr, FunctionCall) and isinstance(expr.fn, func.QueryTemplateFunction):
-                expr.fn.set_conn(conn)
+    @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:
@@ -226,7 +305,9 @@ class RowBuilder:
         self.next_slot_idx += 1
         return result
 
-    def _record_unique_expr(self, expr: Expr, recursive: bool) -> Expr:
+    T = TypeVar('T', bound=Expr)
+
+    def _record_unique_expr(self, expr: T, recursive: bool) -> T:
         """Records the expr if it's not a duplicate and assigns a slot idx to expr and its components"
         Returns:
             the unique expr
@@ -254,8 +335,14 @@ class RowBuilder:
         for d in e.dependencies():
             self._record_output_expr_id(d, output_expr_id)
 
-    def _compute_dependencies(self, target_slot_idxs: list[int], excluded_slot_idxs: list[int]) -> list[int]:
-        """Compute exprs needed to materialize the given target slots, excluding 'excluded_slot_idxs'"""
+    def _compute_dependencies(
+        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'
+
+        If target_scope != None, stops transitive dependency resolution when leaving target_scope (ie, includes
+        immediate dependents that aren't in target_scope, but doesn't resolve those).
+        """
         dependencies: list[set[int]] = [set() for _ in range(self.num_materialized)]  # indexed by slot_idx
         # doing this front-to-back ensures that we capture transitive dependencies
         max_target_slot_idx = max(target_slot_idxs)
@@ -268,7 +355,11 @@ class RowBuilder:
             if expr.slot_idx in self.input_expr_slot_idxs:
                 # this is input and therefore doesn't depend on other exprs
                 continue
+            if target_scope is not None and expr.scope() != target_scope:
+                # don't resolve dependencies outside of target_scope
+                continue
             for d in expr.dependencies():
+                assert d.slot_idx is not None, f'{expr}, {d}'
                 if d.slot_idx in excluded_slot_idxs:
                     continue
                 dependencies[expr.slot_idx].add(d.slot_idx)
@@ -298,10 +389,15 @@ class RowBuilder:
         for c in e.components:
             self.__set_slot_idxs_aux(c)
 
-    def get_dependencies(self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None) -> list[Expr]:
+    def get_dependencies(
+        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).
         The exprs given in 'exclude' are excluded.
+        If limit_scope == True, only returns dependencies in the same scope and immediate (ie, not transitive)
+        dependencies from enclosing scopes.
+
         Returns:
             list of Exprs from unique_exprs (= with slot_idx set)
         """
@@ -312,38 +408,56 @@ class RowBuilder:
             return []
         # make sure we only refer to recorded exprs
         targets = [self.unique_exprs[e] for e in targets]
+        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}
+            assert len(target_scopes) == 1
+            target_scope = target_scopes.pop()
         exclude = [self.unique_exprs[e] for e in exclude]
         target_slot_idxs = [e.slot_idx for e in targets]
         excluded_slot_idxs = [e.slot_idx for e in exclude]
-        all_dependencies = set(self._compute_dependencies(target_slot_idxs, excluded_slot_idxs))
+        all_dependencies = set(
+            self._compute_dependencies(target_slot_idxs, excluded_slot_idxs, target_scope=target_scope)
+        )
         all_dependencies.update(target_slot_idxs)
         result_ids = list(all_dependencies)
         result_ids.sort()
         return [self.unique_exprs[id] for id in result_ids]
 
-    def create_eval_ctx(self, targets: Iterable[Expr], exclude: Optional[Iterable[Expr]] = None) -> EvalCtx:
+    def create_eval_ctx(
+        self, targets: Iterable[Expr], exclude: Iterable[Expr] | None = None, limit_scope: bool = True
+    ) -> EvalCtx:
         """Return EvalCtx for targets"""
         targets = list(targets)
         if exclude is None:
             exclude = []
         if len(targets) == 0:
             return self.EvalCtx([], [], [], [])
-        dependencies = self.get_dependencies(targets, exclude)
+        dependencies = self.get_dependencies(targets, exclude, limit_scope=limit_scope)
         targets = [self.unique_exprs[e] for e in targets]
         target_slot_idxs = [e.slot_idx for e in targets]
         ctx_slot_idxs = [e.slot_idx for e in dependencies]
         return self.EvalCtx(
-            slot_idxs=ctx_slot_idxs, exprs=[self.unique_exprs[slot_idx] for slot_idx in ctx_slot_idxs],
-            target_slot_idxs=target_slot_idxs, target_exprs=targets)
+            slot_idxs=ctx_slot_idxs,
+            exprs=[self.unique_exprs[slot_idx] for slot_idx in ctx_slot_idxs],
+            target_slot_idxs=target_slot_idxs,
+            target_exprs=targets,
+        )
 
     def set_exc(self, data_row: DataRow, slot_idx: int, exc: Exception) -> None:
         """Record an exception in data_row and propagate it to dependents"""
         data_row.set_exc(slot_idx, exc)
-        for slot_idx in self._exc_dependents[slot_idx]:
-            data_row.set_exc(slot_idx, exc)
+        for idx in self._exc_dependents[slot_idx]:
+            data_row.set_exc(idx, exc)
 
     def eval(
-            self, data_row: DataRow, ctx: EvalCtx, profile: Optional[ExecProfile] = None, ignore_errors: bool = False
+        self,
+        data_row: DataRow,
+        ctx: EvalCtx,
+        profile: ExecProfile | None = None,
+        ignore_errors: bool = False,
+        force_eval: ExprScope | None = None,
     ) -> None:
         """
         Populates the slots in data_row given in ctx.
@@ -351,10 +465,11 @@ class RowBuilder:
         and omits any of that expr's dependents's eval().
         profile: if present, populated with execution time of each expr.eval() call; indexed by expr.slot_idx
         ignore_errors: if False, raises ExprEvalError if any expr.eval() raises an exception
+        force_eval: forces exprs in the specified scope to be reevaluated, even if they already have a value
         """
         for expr in ctx.exprs:
             assert expr.slot_idx >= 0
-            if data_row.has_val[expr.slot_idx] or data_row.has_exc(expr.slot_idx):
+            if expr.scope() != force_eval and (data_row.has_val[expr.slot_idx] or data_row.has_exc(expr.slot_idx)):
                 continue
             try:
                 start_time = time.perf_counter()
@@ -368,31 +483,79 @@ class RowBuilder:
                 if not ignore_errors:
                     input_vals = [data_row[d.slot_idx] for d in expr.dependencies()]
                     raise excs.ExprEvalError(
-                        expr, f'expression {expr}', data_row.get_exc(expr.slot_idx), exc_tb, input_vals, 0)
+                        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:
-                val = data_row.get_stored_val(slot_idx, col.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,18 +1,21 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+import logging
+from typing import Any, cast
 from uuid import UUID
 
 import sqlalchemy as sql
 
-import pixeltable.catalog as catalog
-import pixeltable.type_system as ts
+from pixeltable import catalog, type_system as ts
+from pixeltable.catalog.table_version import TableVersionKey
 
 from .data_row import DataRow
 from .expr import Expr
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
+_logger = logging.getLogger('pixeltable')
+
 
 class RowidRef(Expr):
     """A reference to a part of a table rowid
@@ -22,9 +25,20 @@ class RowidRef(Expr):
     _from_dict()/init() is called, which is why this class effectively has two separate paths for construction
     (with and without a TableVersion).
     """
+
+    tbl: catalog.TableVersionHandle | None
+    normalized_base: catalog.TableVersionHandle | None
+    tbl_id: UUID
+    normalized_base_id: UUID
+    rowid_component_idx: int
+
     def __init__(
-            self, tbl: catalog.TableVersion, idx: int,
-            tbl_id: Optional[UUID] = None, normalized_base_id: Optional[UUID] = None):
+        self,
+        tbl: catalog.TableVersionHandle | None,
+        idx: int,
+        tbl_id: UUID | None = None,
+        normalized_base_id: UUID | None = None,
+    ):
         super().__init__(ts.IntType(nullable=False))
         self.tbl = tbl
         if tbl is not None:
@@ -32,8 +46,8 @@ class RowidRef(Expr):
             # (which has the same values as all its descendent views)
             normalized_base = tbl
             # don't try to reference tbl.store_tbl here
-            while normalized_base.base is not None and normalized_base.base.num_rowid_columns() > idx:
-                normalized_base = normalized_base.base
+            while normalized_base.get().base is not None and normalized_base.get().base.get().num_rowid_columns() > idx:
+                normalized_base = normalized_base.get().base
             self.normalized_base = normalized_base
         else:
             self.normalized_base = None
@@ -44,21 +58,35 @@ 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:
-        return self.normalized_base_id == other.normalized_base_id \
+        return (
+            self.normalized_base_id == other.normalized_base_id
             and self.rowid_component_idx == other.rowid_component_idx
+        )
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() +\
-            [('normalized_base_id', self.normalized_base_id), ('idx', self.rowid_component_idx)]
+        return [
+            *super()._id_attrs(),
+            ('normalized_base_id', self.normalized_base_id),
+            ('idx', self.rowid_component_idx),
+        ]
 
     def __repr__(self) -> str:
         # check if this is the pos column of a component view
-        tbl = self.tbl if self.tbl is not None else catalog.Catalog.get().tbl_versions[(self.tbl_id, None)]
-        if tbl.is_component_view() and self.rowid_component_idx == tbl.store_tbl.pos_col_idx:  # type: ignore[attr-defined]
+        from pixeltable import store
+
+        tbl = (
+            self.tbl.get()
+            if self.tbl is not None
+            else catalog.Catalog.get().get_tbl_version(TableVersionKey(self.tbl_id, None, None))
+        )
+        if (
+            tbl.is_component_view
+            and self.rowid_component_idx == cast(store.StoreComponentView, tbl.store_tbl).pos_col_idx
+        ):
             return catalog.globals._POS_COLUMN_NAME
         return ''
 
@@ -75,15 +103,25 @@ class RowidRef(Expr):
         self.tbl = tbl.tbl_version
         self.tbl_id = self.tbl.id
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        tbl = self.tbl if self.tbl is not None else catalog.Catalog.get().tbl_versions[(self.tbl_id, None)]
+    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(TableVersionKey(self.tbl_id, None, 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}'
+        )
         return rowid_cols[self.rowid_component_idx]
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         data_row[self.slot_idx] = data_row.pk[self.rowid_component_idx]
 
     def _as_dict(self) -> dict:
+        # TODO: Serialize the full TableVersionHandle, not just the UUID
+        assert self.tbl is None or self.tbl.anchor_tbl_id is None  # TODO: support anchor_tbl_id for view-over-replica
         return {
             'tbl_id': str(self.tbl_id),
             'normalized_base_id': str(self.normalized_base_id),