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

pixeltable/exprs/expr.py

@@ -7,7 +7,7 @@ import inspect
 import json
 import sys
 import typing
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, Optional, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, TypeVar, overload
 from uuid import UUID
 
 import numpy as np
@@ -29,7 +29,7 @@ class ExprScope:
     parent is None: outermost scope
     """
 
-    def __init__(self, parent: Optional[ExprScope]):
+    def __init__(self, parent: ExprScope | None):
         self.parent = parent
 
     def is_contained_in(self, other: ExprScope) -> bool:
@@ -47,7 +47,7 @@ class Expr(abc.ABC):
     """
     Rules for using state in subclasses:
     - all state except for components and slot_idx is shared between copies of an Expr
-    - slot_idx is set during analysis (DataFrame.show())
+    - slot_idx is set during analysis (Query.show())
     - during eval(), components can only be accessed via self.components; any Exprs outside of that won't
       have slot_idx set
     """
@@ -61,13 +61,13 @@ class Expr(abc.ABC):
     # - set by the subclass's __init__()
     # - produced by _create_id()
     # - not expected to survive a serialize()/deserialize() roundtrip
-    id: Optional[int]
+    id: int | None
 
     # index of the expr's value in the data row:
     # - set for all materialized exprs
     # - None: not executable
     # - not set for subexprs that don't need to be materialized because the parent can be materialized via SQL
-    slot_idx: Optional[int]
+    slot_idx: int | None
 
     T = TypeVar('T', bound='Expr')
 
@@ -103,7 +103,7 @@ class Expr(abc.ABC):
         assert not has_rel_path, self._expr_tree()
         assert not self._has_relative_path(), self._expr_tree()
 
-    def _bind_rel_paths(self, mapper: Optional['exprs.JsonMapperDispatch'] = None) -> None:
+    def _bind_rel_paths(self, mapper: 'exprs.JsonMapperDispatch' | None = None) -> None:
         for c in self.components:
             c._bind_rel_paths(mapper)
 
@@ -118,7 +118,7 @@ class Expr(abc.ABC):
         for c in self.components:
             c._expr_tree_r(indent + 2, buf)
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         """
         Returns:
             None if this expression lacks a default name,
@@ -127,7 +127,7 @@ class Expr(abc.ABC):
         return None
 
     @property
-    def validation_error(self) -> Optional[str]:
+    def validation_error(self) -> str | None:
         """
         Subclasses can override this to indicate that validation has failed after a catalog load.
 
@@ -205,12 +205,12 @@ class Expr(abc.ABC):
         return result
 
     @classmethod
-    def copy_list(cls, expr_list: Optional[list[Expr]]) -> Optional[list[Expr]]:
+    def copy_list(cls, expr_list: list[Expr] | None) -> list[Expr] | None:
         if expr_list is None:
             return None
         return [e.copy() for e in expr_list]
 
-    def __deepcopy__(self, memo: Optional[dict[int, Any]] = None) -> Expr:
+    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Expr:
         # we don't need to create an actual deep copy because all state other than execution state is read-only
         if memo is None:
             memo = {}
@@ -241,7 +241,7 @@ class Expr(abc.ABC):
         for i in range(len(expr_list)):
             expr_list[i] = expr_list[i].substitute(spec)
 
-    def resolve_computed_cols(self, resolve_cols: Optional[set[catalog.Column]] = None) -> Expr:
+    def resolve_computed_cols(self, resolve_cols: set[catalog.Column] | None = None) -> Expr:
         """
         Recursively replace ColRefs to unstored computed columns with their value exprs.
         Also replaces references to stored computed columns in resolve_cols.
@@ -276,6 +276,13 @@ class Expr(abc.ABC):
         tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
         return self._retarget(tbl_versions)
 
+    @classmethod
+    def retarget_list(cls, expr_list: list[Expr], tbl: catalog.TableVersionPath) -> None:
+        """Retarget ColumnRefs in expr_list to the specific TableVersions in tbl."""
+        tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
+        for i in range(len(expr_list)):
+            expr_list[i] = expr_list[i]._retarget(tbl_versions)
+
     def _retarget(self, tbl_versions: dict[UUID, catalog.TableVersion]) -> Self:
         for i in range(len(self.components)):
             self.components[i] = self.components[i]._retarget(tbl_versions)
@@ -302,18 +309,18 @@ class Expr(abc.ABC):
 
     @overload
     def subexprs(
-        self, *, filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+        self, *, filter: Callable[[Expr], bool] | None = None, traverse_matches: bool = True
     ) -> Iterator[Expr]: ...
 
     @overload
     def subexprs(
-        self, expr_class: type[T], filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+        self, expr_class: type[T], filter: Callable[[Expr], bool] | None = None, traverse_matches: bool = True
     ) -> Iterator[T]: ...
 
     def subexprs(
         self,
-        expr_class: Optional[type[T]] = None,
-        filter: Optional[Callable[[Expr], bool]] = None,
+        expr_class: type[T] | None = None,
+        filter: Callable[[Expr], bool] | None = None,
         traverse_matches: bool = True,
     ) -> Iterator[T]:
         """
@@ -332,11 +339,7 @@ class Expr(abc.ABC):
     @overload
     @classmethod
     def list_subexprs(
-        cls,
-        expr_list: Iterable[Expr],
-        *,
-        filter: Optional[Callable[[Expr], bool]] = None,
-        traverse_matches: bool = True,
+        cls, expr_list: Iterable[Expr], *, filter: Callable[[Expr], bool] | None = None, traverse_matches: bool = True
     ) -> Iterator[Expr]: ...
 
     @overload
@@ -345,7 +348,7 @@ class Expr(abc.ABC):
         cls,
         expr_list: Iterable[Expr],
         expr_class: type[T],
-        filter: Optional[Callable[[Expr], bool]] = None,
+        filter: Callable[[Expr], bool] | None = None,
         traverse_matches: bool = True,
     ) -> Iterator[T]: ...
 
@@ -353,15 +356,24 @@ class Expr(abc.ABC):
     def list_subexprs(
         cls,
         expr_list: Iterable[Expr],
-        expr_class: Optional[type[T]] = None,
-        filter: Optional[Callable[[Expr], bool]] = None,
+        expr_class: type[T] | None = None,
+        filter: Callable[[Expr], bool] | None = None,
         traverse_matches: bool = True,
     ) -> Iterator[T]:
         """Produce subexprs for all exprs in list. Can contain duplicates."""
         for e in expr_list:
             yield from e.subexprs(expr_class=expr_class, filter=filter, traverse_matches=traverse_matches)
 
-    def _contains(self, cls: Optional[type[Expr]] = None, filter: Optional[Callable[[Expr], bool]] = None) -> bool:
+    @classmethod
+    def list_contains(
+        cls,
+        expr_list: Iterable[Expr],
+        expr_class: type[Expr] | None = None,
+        filter: Callable[[Expr], bool] | None = None,
+    ) -> bool:
+        return any(e._contains(expr_class, filter) for e in expr_list)
+
+    def _contains(self, cls: type[Expr] | None = None, filter: Callable[[Expr], bool] | None = None) -> bool:
         """
         Returns True if any subexpr is an instance of cls and/or matches filter.
         """
@@ -380,34 +392,36 @@ class Expr(abc.ABC):
         from .column_ref import ColumnRef
         from .rowid_ref import RowidRef
 
-        return {ref.col.tbl.id for ref in self.subexprs(ColumnRef)} | {ref.tbl.id for ref in self.subexprs(RowidRef)}
+        return {ref.col.get_tbl().id for ref in self.subexprs(ColumnRef)} | {
+            ref.tbl.id for ref in self.subexprs(RowidRef)
+        }
 
     @classmethod
     def all_tbl_ids(cls, exprs_: Iterable[Expr]) -> set[UUID]:
         return {tbl_id for e in exprs_ for tbl_id in e.tbl_ids()}
 
     @classmethod
-    def get_refd_columns(cls, expr_dict: dict[str, Any]) -> list[catalog.Column]:
+    def get_refd_column_ids(cls, expr_dict: dict[str, Any]) -> set[catalog.QColumnId]:
         """Return Columns referenced by expr_dict."""
-        result: list[catalog.Column] = []
+        result: set[catalog.QColumnId] = set()
         assert '_classname' in expr_dict
         from .column_ref import ColumnRef
 
         if expr_dict['_classname'] == 'ColumnRef':
-            result.append(ColumnRef.get_column(expr_dict))
+            result.add(ColumnRef.get_column_id(expr_dict))
         if 'components' in expr_dict:
             for component_dict in expr_dict['components']:
-                result.extend(cls.get_refd_columns(component_dict))
+                result.update(cls.get_refd_column_ids(component_dict))
         return result
 
-    def as_literal(self) -> Optional[Expr]:
+    def as_literal(self) -> Expr | None:
         """
         Return a Literal expression if this expression can be evaluated to a constant value, otherwise return None.
         """
         return None
 
     @classmethod
-    def from_array(cls, elements: Iterable) -> Optional[Expr]:
+    def from_array(cls, elements: Iterable) -> Expr | None:
         from .inline_expr import InlineArray
         from .literal import Literal
 
@@ -430,7 +444,7 @@ class Expr(abc.ABC):
             return self
 
     @classmethod
-    def from_object(cls, o: object) -> Optional[Expr]:
+    def from_object(cls, o: object) -> Expr | None:
         """
         Try to turn a literal object into an Expr.
         """
@@ -460,7 +474,7 @@ class Expr(abc.ABC):
                 return Literal(o, col_type=obj_type)
         return None
 
-    def sql_expr(self, sql_elements: 'exprs.SqlElementCache') -> Optional[sql.ColumnElement]:
+    def sql_expr(self, sql_elements: 'exprs.SqlElementCache') -> sql.ColumnElement | None:
         """
         If this expr can be materialized directly in SQL:
         - returns a ColumnElement
@@ -479,6 +493,18 @@ class Expr(abc.ABC):
         """
         pass
 
+    def prepare(self) -> None:
+        """
+        Create execution state. This is called before the first eval() call.
+        """
+        for c in self.components:
+            c.prepare()
+
+    @classmethod
+    def prepare_list(cls, expr_list: Iterable[Expr]) -> None:
+        for e in expr_list:
+            e.prepare()
+
     def release(self) -> None:
         """
         Allow Expr class to tear down execution state. This is called after the last eval() call.
@@ -487,7 +513,7 @@ class Expr(abc.ABC):
             c.release()
 
     @classmethod
-    def release_list(cls, expr_list: list[Expr]) -> None:
+    def release_list(cls, expr_list: Iterable[Expr]) -> None:
         for e in expr_list:
             e.release()
 
@@ -543,7 +569,7 @@ class Expr(abc.ABC):
         else:
             return InPredicate(self, value_set_literal=value_set)
 
-    def astype(self, new_type: Union[ts.ColumnType, type, _AnnotatedAlias]) -> 'exprs.TypeCast':
+    def astype(self, new_type: ts.ColumnType | type | _AnnotatedAlias) -> 'exprs.TypeCast':
         from pixeltable.exprs import TypeCast
 
         # Interpret the type argument the same way we would if given in a schema
@@ -555,7 +581,7 @@ class Expr(abc.ABC):
         return TypeCast(self, col_type)
 
     def apply(
-        self, fn: Callable, *, col_type: Union[ts.ColumnType, type, _AnnotatedAlias, None] = None
+        self, fn: Callable, *, col_type: ts.ColumnType | type | _AnnotatedAlias | None = None
     ) -> 'exprs.FunctionCall':
         if col_type is not None:
             col_type = ts.ColumnType.normalize_type(col_type)
@@ -639,7 +665,7 @@ class Expr(abc.ABC):
 
     def _make_comparison(self, op: ComparisonOperator, other: object) -> 'exprs.Comparison':
         """
-        other: Union[Expr, LiteralPythonTypes]
+        other: Expr | LiteralPythonTypes
         """
         # TODO: check for compatibility
         from .comparison import Comparison
@@ -654,7 +680,7 @@ class Expr(abc.ABC):
     def __neg__(self) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, -1)
 
-    def __add__(self, other: object) -> Union[exprs.ArithmeticExpr, exprs.StringOp]:
+    def __add__(self, other: object) -> exprs.ArithmeticExpr | exprs.StringOp:
         if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
             return self._make_string_expr(StringOperator.CONCAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.ADD, other)
@@ -662,7 +688,7 @@ class Expr(abc.ABC):
     def __sub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __mul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __mul__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
             return self._make_string_expr(StringOperator.REPEAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, other)
@@ -676,7 +702,7 @@ class Expr(abc.ABC):
     def __floordiv__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.FLOORDIV, other)
 
-    def __radd__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __radd__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
             return self._rmake_string_expr(StringOperator.CONCAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.ADD, other)
@@ -684,7 +710,7 @@ class Expr(abc.ABC):
     def __rsub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._rmake_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __rmul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __rmul__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
             return self._rmake_string_expr(StringOperator.REPEAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.MUL, other)
@@ -726,7 +752,7 @@ class Expr(abc.ABC):
 
     def _make_arithmetic_expr(self, op: ArithmeticOperator, other: object) -> 'exprs.ArithmeticExpr':
         """
-        other: Union[Expr, LiteralPythonTypes]
+        other: Expr | LiteralPythonTypes
         """
         # TODO: check for compatibility
         from .arithmetic_expr import ArithmeticExpr
@@ -775,7 +801,7 @@ class Expr(abc.ABC):
 
         return CompoundPredicate(LogicalOperator.NOT, [self])
 
-    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]:
         """
         Returns clauses of a conjunction that meet condition in the first element.
         The second element contains remaining clauses, rolled into a conjunction.
@@ -786,7 +812,7 @@ class Expr(abc.ABC):
         else:
             return [], self
 
-    def _make_applicator_function(self, fn: Callable, col_type: Optional[ts.ColumnType]) -> 'func.Function':
+    def _make_applicator_function(self, fn: Callable, col_type: ts.ColumnType | None) -> 'func.Function':
         """
         Creates a unary pixeltable `Function` that encapsulates a python `Callable`. The result type of
         the new `Function` is given by `col_type`, and its parameter type will be `self.col_type`.

pixeltable/exprs/expr_dict.py

@@ -1,4 +1,4 @@
-from typing import Generic, Iterable, Iterator, Optional, TypeVar
+from typing import Generic, Iterable, Iterator, TypeVar
 
 from .expr import Expr
 
@@ -14,7 +14,7 @@ class ExprDict(Generic[T]):
 
     _data: dict[int, tuple[Expr, T]]
 
-    def __init__(self, iterable: Optional[Iterable[tuple[Expr, T]]] = None):
+    def __init__(self, iterable: Iterable[tuple[Expr, T]] | None = None):
         self._data = {}
 
         if iterable is not None:
@@ -39,7 +39,7 @@ class ExprDict(Generic[T]):
     def __contains__(self, key: Expr) -> bool:
         return key.id in self._data
 
-    def get(self, key: Expr, default: Optional[T] = None) -> Optional[T]:
+    def get(self, key: Expr, default: T | None = None) -> T | None:
         item = self._data.get(key.id)
         return item[1] if item is not None else default
 

pixeltable/exprs/expr_set.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Generic, Iterable, Iterator, Optional, TypeVar
+from typing import Generic, Iterable, Iterator, TypeVar
 
 from .expr import Expr
 
@@ -9,26 +9,33 @@ T = TypeVar('T', bound='Expr')
 
 class ExprSet(Generic[T]):
     """
-    A set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by Expr.id.
+    An ordered set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by
+    Expr.id.
     """
 
     exprs: dict[int, T]  # key: Expr.id
+    expr_offsets: dict[int, int]  # key: Expr.id, value: offset into self.exprs.keys()
     exprs_by_idx: dict[int, T]  # key: slot_idx
 
-    def __init__(self, elements: Optional[Iterable[T]] = None):
+    def __init__(self, elements: Iterable[T] | None = None):
         self.exprs = {}
+        self.expr_offsets = {}
         self.exprs_by_idx = {}
         if elements is not None:
             for e in elements:
                 self.add(e)
 
-    def add(self, expr: T) -> None:
-        if expr.id in self.exprs:
-            return
+    def add(self, expr: T) -> int:
+        """Returns offset corresponding to iteration order"""
+        offset = self.expr_offsets.get(expr.id)
+        if offset is not None:
+            return offset
+        offset = len(self.exprs)
         self.exprs[expr.id] = expr
-        if expr.slot_idx is None:
-            return
-        self.exprs_by_idx[expr.slot_idx] = expr
+        self.expr_offsets[expr.id] = offset
+        if expr.slot_idx is not None:
+            self.exprs_by_idx[expr.slot_idx] = expr
+        return offset
 
     def update(self, *others: Iterable[T]) -> None:
         for other in others:
@@ -44,7 +51,7 @@ class ExprSet(Generic[T]):
     def __iter__(self) -> Iterator[T]:
         return iter(self.exprs.values())
 
-    def __getitem__(self, index: object) -> Optional[T]:
+    def __getitem__(self, index: object) -> T | None:
         """Indexed lookup by slot_idx or Expr.id."""
         assert isinstance(index, (int, Expr))
         if isinstance(index, int):