pixeltable 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

pixeltable/exec/sql_node.py

@@ -1,73 +1,175 @@
-from typing import List, Optional, Tuple, Iterable, Set
-from uuid import UUID
 import logging
 import warnings
+from decimal import Decimal
+from typing import Optional, Iterable, Iterator, NamedTuple
+from uuid import UUID
 
 import sqlalchemy as sql
 
+import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
-import pixeltable.exprs as exprs
-import pixeltable.catalog as catalog
-
 
 _logger = logging.getLogger('pixeltable')
 
+
+class OrderByItem(NamedTuple):
+    expr: exprs.Expr
+    asc: Optional[bool]
+
+
+OrderByClause = list[OrderByItem]
+
+
+def combine_order_by_clauses(clauses: Iterable[OrderByClause]) -> Optional[OrderByClause]:
+    """Returns a clause that's compatible with 'clauses', or None if that doesn't exist.
+    Two clauses are compatible if for each of their respective items c1[i] and c2[i]
+    a) the exprs are identical and
+    b) the asc values are identical or at least one is None (None serves as a wildcard)
+    """
+    result: OrderByClause = []
+    for clause in clauses:
+        combined: OrderByClause = []
+        for item1, item2 in zip(result, clause):
+            if item1.expr.id != item2.expr.id:
+                return None
+            if item1.asc is not None and item2.asc is not None and item1.asc != item2.asc:
+                return None
+            asc = item1.asc if item1.asc is not None else item2.asc
+            combined.append(OrderByItem(item1.expr, asc))
+
+        # add remaining ordering of the longer list
+        prefix_len = min(len(result), len(clause))
+        if len(result) > prefix_len:
+            combined.extend(result[prefix_len:])
+        elif len(clause) > prefix_len:
+            combined.extend(clause[prefix_len:])
+        result = combined
+    return result
+
+
+def print_order_by_clause(clause: OrderByClause) -> str:
+    return ', '.join([
+        f'({item.expr}{", asc=True" if item.asc is True else ""}{", asc=False" if item.asc is False else ""})'
+        for item in clause
+    ])
+
+
 class SqlNode(ExecNode):
-    """Materializes data from the store via a Select stmt."""
+    """
+    Materializes data from the store via a Select stmt.
+    This only provides the select list. The subclasses are responsible for the From clause and any additional clauses.
+    """
+
+    tbl: Optional[catalog.TableVersionPath]
+    select_list: exprs.ExprSet
+    set_pk: bool
+    num_pk_cols: int
+    filter: Optional[exprs.Expr]
+    filter_eval_ctx: Optional[exprs.RowBuilder.EvalCtx]
+    cte: Optional[sql.CTE]
+    sql_elements: exprs.SqlElementCache
+    limit: Optional[int]
+    order_by_clause: OrderByClause
 
     def __init__(
-            self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
+            self, tbl: Optional[catalog.TableVersionPath], row_builder: exprs.RowBuilder,
             select_list: Iterable[exprs.Expr], sql_elements: exprs.SqlElementCache, set_pk: bool = False
     ):
         """
-        Initialize self.stmt with expressions derived from select_list.
-
-        This only provides the select list. The subclass is responsible for the From clause and any additional clauses.
+        If row_builder contains references to unstored iter columns, expands the select list to include their
+        SQL-materializable subexpressions.
 
         Args:
             select_list: output of the query
             set_pk: if True, sets the primary for each DataRow
         """
         # create Select stmt
+        self.sql_elements = sql_elements
         self.tbl = tbl
-        target = tbl.tbl_version  # the stored table we're scanning
-        self.sql_exprs = exprs.ExprSet(select_list)
+        self.select_list = exprs.ExprSet(select_list)
         # unstored iter columns: we also need to retrieve whatever is needed to materialize the iter args
         for iter_arg in row_builder.unstored_iter_args.values():
-            sql_subexprs = iter_arg.subexprs(filter=sql_elements.contains, traverse_matches=False)
+            sql_subexprs = iter_arg.subexprs(filter=self.sql_elements.contains, traverse_matches=False)
             for e in sql_subexprs:
-                self.sql_exprs.add(e)
-        super().__init__(row_builder, self.sql_exprs, [], None)  # we materialize self.sql_exprs
+                self.select_list.add(e)
+        super().__init__(row_builder, self.select_list, [], None)  # we materialize self.select_list
 
-        # change rowid refs against a base table to rowid refs against the target table, so that we minimize
-        # the number of tables that need to be joined to the target table
-        for rowid_ref in [e for e in self.sql_exprs if isinstance(e, exprs.RowidRef)]:
-            rowid_ref.set_tbl(tbl)
+        if tbl is not None:
+            # minimize the number of tables that need to be joined to the target table
+            self.retarget_rowid_refs(tbl, self.select_list)
 
-        sql_select_list = [sql_elements.get(e) for e in self.sql_exprs]
-        assert len(sql_select_list) == len(self.sql_exprs)
-        assert all(e is not None for e in sql_select_list)
+        assert self.sql_elements.contains(self.select_list)
         self.set_pk = set_pk
         self.num_pk_cols = 0
         if set_pk:
             # we also need to retrieve the pk columns
-            pk_columns = target.store_tbl.pk_columns()
-            self.num_pk_cols = len(pk_columns)
-            sql_select_list += pk_columns
-
-        self.stmt = sql.select(*sql_select_list)
+            assert tbl is not None
+            self.num_pk_cols = len(tbl.tbl_version.store_tbl.pk_columns())
 
         # additional state
-        self.result_cursor: Optional[sql.engine.CursorResult] = None
+        self.result_cursor = None
         # the filter is provided by the subclass
-        self.filter: Optional[exprs.Expr] = None
-        self.filter_eval_ctx: Optional[exprs.EvalContext] = None
+        self.filter = None
+        self.filter_eval_ctx = None
+        self.cte = None
+        self.limit = None
+        self.order_by_clause = []
+
+    def _create_stmt(self) -> sql.Select:
+        """Create Select from local state"""
+
+        assert self.sql_elements.contains(self.select_list)
+        sql_select_list = [self.sql_elements.get(e) for e in self.select_list]
+        if self.set_pk:
+            sql_select_list += self.tbl.tbl_version.store_tbl.pk_columns()
+        stmt = sql.select(*sql_select_list)
+
+        order_by_clause: list[sql.ClauseElement] = []
+        for e, asc in self.order_by_clause:
+            if isinstance(e, exprs.SimilarityExpr):
+                order_by_clause.append(e.as_order_by_clause(asc))
+            else:
+                order_by_clause.append(self.sql_elements.get(e).desc() if asc is False else self.sql_elements.get(e))
+        stmt = stmt.order_by(*order_by_clause)
+
+        if self.filter is None and self.limit is not None:
+            # if we don't have a Python filter, we can apply the limit to stmt
+            stmt = stmt.limit(self.limit)
+
+        return stmt
+
+    def _ordering_tbl_ids(self) -> set[UUID]:
+        return exprs.Expr.list_tbl_ids(e for e, _ in self.order_by_clause)
+
+    def to_cte(self) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
+        """
+        Returns a CTE that materializes the output of this node plus a mapping from select list expr to output column
+
+        Returns:
+            (CTE, dict from Expr to output column)
+        """
+        if self.filter is not None:
+            # the filter needs to run in Python
+            return None
+        self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
+        if self.cte is None:
+            self.cte = self._create_stmt().cte()
+            assert len(self.cte.c) == len(self.select_list)
+        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))
+
+    @classmethod
+    def retarget_rowid_refs(cls, target: catalog.TableVersionPath, expr_seq: Iterable[exprs.Expr]) -> None:
+        """Change rowid refs to point to target"""
+        for e in expr_seq:
+            if isinstance(e, exprs.RowidRef):
+                e.set_tbl(target)
 
     @classmethod
     def create_from_clause(
-            cls, tbl: catalog.TableVersionPath, stmt: sql.Select, refd_tbl_ids: Optional[Set[UUID]] = None,
-            exact_version_only: Optional[Set[UUID]] = None
+            cls, tbl: catalog.TableVersionPath, stmt: sql.Select, refd_tbl_ids: Optional[set[UUID]] = None,
+            exact_version_only: Optional[set[UUID]] = None
     ) -> sql.Select:
         """Add From clause to stmt for tables/views referenced by materialized_exprs
         Args:
@@ -85,7 +187,7 @@ class SqlNode(ExecNode):
             exact_version_only = {}
         candidates = tbl.get_tbl_versions()
         assert len(candidates) > 0
-        joined_tbls: List[catalog.TableVersion] = [candidates[0]]
+        joined_tbls: list[catalog.TableVersion] = [candidates[0]]
         for tbl in candidates[1:]:
             if tbl.id in refd_tbl_ids:
                 joined_tbls.append(tbl)
@@ -111,72 +213,97 @@ class SqlNode(ExecNode):
             prev_tbl = tbl
         return stmt
 
-    def _log_explain(self, conn: sql.engine.Connection) -> None:
+    def add_order_by(self, ordering: OrderByClause) -> None:
+        """Add Order By clause to stmt"""
+        if self.tbl is not None:
+            # change rowid refs against a base table to rowid refs against the target table, so that we minimize
+            # the number of tables that need to be joined to the target table
+            self.retarget_rowid_refs(self.tbl, [e for e, _ in ordering])
+        combined = combine_order_by_clauses([self.order_by_clause, ordering])
+        assert combined is not None
+        self.order_by_clause = combined
+
+    def set_limit(self, limit: int) -> None:
+        self.limit = limit
+
+    def _log_explain(self, stmt: sql.Select) -> None:
         try:
             # don't set dialect=Env.get().engine.dialect: x % y turns into x %% y, which results in a syntax error
-            stmt_str = str(self.stmt.compile(compile_kwargs={'literal_binds': True}))
+            stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
             explain_result = self.ctx.conn.execute(sql.text(f'EXPLAIN {stmt_str}'))
             explain_str = '\n'.join([str(row) for row in explain_result])
             _logger.debug(f'SqlScanNode explain:\n{explain_str}')
         except Exception as e:
             _logger.warning(f'EXPLAIN failed')
 
-    def __next__(self) -> DataRowBatch:
-        if self.result_cursor is None:
-            # run the query; do this here rather than in _open(), exceptions are only expected during iteration
-            assert self.ctx.conn is not None
-            try:
-                self._log_explain(self.ctx.conn)
-                with warnings.catch_warnings(record=True) as w:
-                    self.result_cursor = self.ctx.conn.execute(self.stmt)
-                    for warning in w:
-                        pass
-                self.has_more_rows = True
-            except Exception as e:
-                self.has_more_rows = False
-                raise e
-
-        if not self.has_more_rows:
-            raise StopIteration
-
-        output_batch = DataRowBatch(self.tbl.tbl_version, self.row_builder)
-        needs_row = True
-        while self.ctx.batch_size == 0 or len(output_batch) < self.ctx.batch_size:
-            try:
-                sql_row = next(self.result_cursor)
-            except StopIteration:
-                self.has_more_rows = False
-                break
+    def __iter__(self) -> Iterator[DataRowBatch]:
+        # run the query; do this here rather than in _open(), exceptions are only expected during iteration
+        assert self.ctx.conn is not None
+        try:
+            with warnings.catch_warnings(record=True) as w:
+                stmt = self._create_stmt()
+                try:
+                    # log stmt, if possible
+                    stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
+                    _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
+                except Exception as e:
+                    pass
+                self._log_explain(stmt)
+
+                result_cursor = self.ctx.conn.execute(stmt)
+                for warning in w:
+                    pass
+        except Exception as e:
+            raise e
+
+        tbl_version = self.tbl.tbl_version if self.tbl is not None else None
+        output_batch = DataRowBatch(tbl_version, self.row_builder)
+        output_row: Optional[exprs.DataRow] = None
+        num_rows_returned = 0
 
-            if needs_row:
-                output_row = output_batch.add_row()
+        for sql_row in result_cursor:
+            output_row = output_batch.add_row(output_row)
+
+            # populate output_row
             if self.num_pk_cols > 0:
                 output_row.set_pk(tuple(sql_row[-self.num_pk_cols:]))
             # copy the output of the SQL query into the output row
-            for i, e in enumerate(self.sql_exprs):
+            for i, e in enumerate(self.select_list):
                 slot_idx = e.slot_idx
-                output_row[slot_idx] = sql_row[i]
+                # certain numerical operations can produce Decimals (eg, SUM(<int column>)); we need to convert them
+                if isinstance(sql_row[i], Decimal):
+                    if e.col_type.is_int_type():
+                        output_row[slot_idx] = int(sql_row[i])
+                    elif e.col_type.is_float_type():
+                        output_row[slot_idx] = float(sql_row[i])
+                    else:
+                        raise RuntimeError(f'Unexpected Decimal value for {e}')
+                else:
+                    output_row[slot_idx] = sql_row[i]
+
             if self.filter is not None:
+                # evaluate filter
                 self.row_builder.eval(output_row, self.filter_eval_ctx, profile=self.ctx.profile)
-                if output_row[self.filter.slot_idx]:
-                    needs_row = True
-                    if self.limit > 0 and len(output_batch) >= self.limit:
-                        self.has_more_rows = False
-                        break
-                else:
-                    # we re-use this row for the next sql row if it didn't pass the filter
-                    needs_row = False
-                    output_row.clear()
+            if self.filter is not None and not output_row[self.filter.slot_idx]:
+                # we re-use this row for the next sql row since it didn't pass the filter
+                output_row = output_batch.pop_row()
+                output_row.clear()
+            else:
+                # reset output_row in order to add new one
+                output_row = None
+                num_rows_returned += 1
 
-        if not needs_row:
-            # the last row didn't pass the filter
-            assert self.filter is not None
-            output_batch.pop_row()
+            if self.limit is not None and num_rows_returned == self.limit:
+                break
+
+            if self.ctx.batch_size > 0 and len(output_batch) == self.ctx.batch_size:
+                _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
+                yield output_batch
+                output_batch = DataRowBatch(tbl_version, self.row_builder)
 
-        _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
-        if len(output_batch) == 0:
-            raise StopIteration
-        return output_batch
+        if len(output_batch) > 0:
+            _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
+            yield output_batch
 
     def _close(self) -> None:
         if self.result_cursor is not None:
@@ -189,12 +316,14 @@ class SqlScanNode(SqlNode):
 
     Supports filtering and ordering.
     """
+    where_clause: Optional[exprs.Expr]
+    exact_version_only: list[catalog.TableVersion]
+
     def __init__(
             self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
             select_list: Iterable[exprs.Expr],
             where_clause: Optional[exprs.Expr] = None, filter: Optional[exprs.Expr] = None,
-            order_by_items: Optional[List[Tuple[exprs.Expr, bool]]] = None,
-            limit: int = 0, set_pk: bool = False, exact_version_only: Optional[List[catalog.TableVersion]] = None
+            set_pk: bool = False, exact_version_only: Optional[list[catalog.TableVersion]] = None
     ):
         """
         Args:
@@ -208,52 +337,29 @@ class SqlScanNode(SqlNode):
         sql_elements = exprs.SqlElementCache()
         super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=set_pk)
         # create Select stmt
-        if order_by_items is None:
-            order_by_items = []
         if exact_version_only is None:
             exact_version_only = []
         target = tbl.tbl_version  # the stored table we're scanning
         self.filter = filter
         self.filter_eval_ctx = \
             row_builder.create_eval_ctx([filter], exclude=select_list) if filter is not None else None
-        self.limit = limit
 
-        where_clause_tbl_ids = where_clause.tbl_ids() if where_clause is not None else set()
-        refd_tbl_ids = exprs.Expr.list_tbl_ids(self.sql_exprs) | where_clause_tbl_ids
-        self.stmt = self.create_from_clause(
-            tbl, self.stmt, refd_tbl_ids, exact_version_only={t.id for t in exact_version_only})
-
-        # change rowid refs against a base table to rowid refs against the target table, so that we minimize
-        # the number of tables that need to be joined to the target table
-        for rowid_ref in [e for e, _ in order_by_items if isinstance(e, exprs.RowidRef)]:
-            rowid_ref.set_tbl(tbl)
-        order_by_clause: List[sql.ClauseElement] = []
-        for e, asc in order_by_items:
-            if isinstance(e, exprs.SimilarityExpr):
-                order_by_clause.append(e.as_order_by_clause(asc))
-            else:
-                order_by_clause.append(sql_elements.get(e).desc() if not asc else sql_elements.get(e))
+        self.where_clause = where_clause
+        self.exact_version_only = exact_version_only
 
-        if where_clause is not None:
-            sql_where_clause = sql_elements.get(where_clause)
+    def _create_stmt(self) -> sql.Select:
+        stmt = super()._create_stmt()
+        where_clause_tbl_ids = self.where_clause.tbl_ids() if self.where_clause is not None else set()
+        refd_tbl_ids = exprs.Expr.list_tbl_ids(self.select_list) | where_clause_tbl_ids | self._ordering_tbl_ids()
+        stmt = self.create_from_clause(
+            self.tbl, stmt, refd_tbl_ids, exact_version_only={t.id for t in self.exact_version_only})
+
+        if self.where_clause is not None:
+            sql_where_clause = self.sql_elements.get(self.where_clause)
             assert sql_where_clause is not None
-            self.stmt = self.stmt.where(sql_where_clause)
-        if len(order_by_clause) > 0:
-            self.stmt = self.stmt.order_by(*order_by_clause)
-        elif target.id in row_builder.unstored_iter_args:
-            # we are referencing unstored iter columns from this view and try to order by our primary key,
-            # which ensures that iterators will see monotonically increasing pos values
-            self.stmt = self.stmt.order_by(*self.tbl.store_tbl.rowid_columns())
-        if limit != 0 and self.filter is None:
-            # if we need to do post-SQL filtering, we can't use LIMIT
-            self.stmt = self.stmt.limit(limit)
+            stmt = stmt.where(sql_where_clause)
 
-        try:
-            # log stmt, if possible
-            stmt_str = str(self.stmt.compile(compile_kwargs={'literal_binds': True}))
-            _logger.debug(f'SqlScanNode stmt:\n{stmt_str}')
-        except Exception as e:
-            pass
+        return stmt
 
 
 class SqlLookupNode(SqlNode):
@@ -261,8 +367,7 @@ class SqlLookupNode(SqlNode):
     Materializes data from the store via a Select stmt with a WHERE clause that matches a list of key values
     """
 
-    stmt: sql.Select
-    where_clause: sql.ColumnElement[bool]
+    where_clause: sql.ColumnElement
 
     def __init__(
             self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
@@ -276,21 +381,45 @@ class SqlLookupNode(SqlNode):
         """
         sql_elements = exprs.SqlElementCache()
         super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=True)
-        target = tbl.tbl_version  # the stored table we're scanning
-        refd_tbl_ids = exprs.Expr.list_tbl_ids(self.sql_exprs)
-        self.stmt = self.create_from_clause(tbl, self.stmt, refd_tbl_ids)
         # Where clause: (key-col-1, key-col-2, ...) IN ((val-1, val-2, ...), ...)
         self.where_clause = sql.tuple_(*sa_key_cols).in_(key_vals)
-        self.stmt = self.stmt.where(self.where_clause)
 
-        if target.id in row_builder.unstored_iter_args:
-            # we are referencing unstored iter columns from this view and try to order by our primary key,
-            # which ensures that iterators will see monotonically increasing pos values
-            self.stmt = self.stmt.order_by(*self.tbl.store_tbl.rowid_columns())
+    def _create_stmt(self) -> sql.Select:
+        stmt = super()._create_stmt()
+        refd_tbl_ids = exprs.Expr.list_tbl_ids(self.select_list) | self._ordering_tbl_ids()
+        stmt = self.create_from_clause(self.tbl, stmt, refd_tbl_ids)
+        stmt = stmt.where(self.where_clause)
+        return stmt
 
-        try:
-            # log stmt, if possible
-            stmt_str = str(self.stmt.compile(compile_kwargs={'literal_binds': True}))
-            _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
-        except Exception as e:
-            pass
+class SqlAggregationNode(SqlNode):
+    """
+    Materializes data from the store via a Select stmt with a WHERE clause that matches a list of key values
+    """
+
+    group_by_items: Optional[list[exprs.Expr]]
+
+    def __init__(
+            self, row_builder: exprs.RowBuilder,
+            input: SqlNode,
+            select_list: Iterable[exprs.Expr],
+            group_by_items: Optional[list[exprs.Expr]] = None,
+            limit: Optional[int] = None, exact_version_only: Optional[list[catalog.TableVersion]] = None
+    ):
+        """
+        Args:
+            select_list: can contain calls to AggregateFunctions
+            group_by_items: list of expressions to group by
+            limit: max number of rows to return: None = no limit
+        """
+        _, input_col_map = input.to_cte()
+        sql_elements = exprs.SqlElementCache(input_col_map)
+        super().__init__(None, row_builder, select_list, sql_elements)
+        self.group_by_items = group_by_items
+
+    def _create_stmt(self) -> sql.Select:
+        stmt = super()._create_stmt()
+        if self.group_by_items is not None:
+            sql_group_by_items = [self.sql_elements.get(e) for e in self.group_by_items]
+            assert all(e is not None for e in sql_group_by_items)
+            stmt = stmt.group_by(*sql_group_by_items)
+        return stmt

pixeltable/exprs/__init__.py

@@ -6,6 +6,7 @@ from .comparison import Comparison
 from .compound_predicate import CompoundPredicate
 from .data_row import DataRow
 from .expr import Expr
+from .expr_dict import ExprDict
 from .expr_set import ExprSet
 from .function_call import FunctionCall
 from .in_predicate import InPredicate

pixeltable/exprs/data_row.py

@@ -33,29 +33,40 @@ class DataRow:
     - ImageType: PIL.Image.Image
     - VideoType: local path if available, otherwise url
     """
-    def __init__(self, size: int, img_slot_idxs: List[int], media_slot_idxs: List[int], array_slot_idxs: List[int]):
-        self.vals: List[Any] = [None] * size  # either cell values or exceptions
-        self.has_val = [False] * size
-        self.excs: List[Optional[Exception]] = [None] * size
 
-        # control structures that are shared across all DataRows in a batch
-        self.img_slot_idxs = img_slot_idxs
-        self.media_slot_idxs = media_slot_idxs  # all media types aside from image
-        self.array_slot_idxs = array_slot_idxs
+    vals: list[Any]
+    has_val: list[bool]
+    excs: list[Optional[Exception]]
 
-        # the primary key of a store row is a sequence of ints (the number is different for table vs view)
-        self.pk: Optional[Tuple[int, ...]] = None
+    # control structures that are shared across all DataRows in a batch
+    img_slot_idxs: list[int]
+    media_slot_idxs: list[int]
+    array_slot_idxs: list[int]
 
-        # file_urls:
-        # - stored url of file for media in vals[i]
-        # - None if vals[i] is not media type
-        # - not None if file_paths[i] is not None
-        self.file_urls: List[Optional[str]] = [None] * size
+    # the primary key of a store row is a sequence of ints (the number is different for table vs view)
+    pk: Optional[tuple[int, ...]]
 
-        # file_paths:
-        # - local path of media file in vals[i]; points to the file cache if file_urls[i] is remote
-        # - None if vals[i] is not a media type or if there is no local file yet for file_urls[i]
-        self.file_paths: List[Optional[str]] = [None] * size
+    # file_urls:
+    # - stored url of file for media in vals[i]
+    # - None if vals[i] is not media type
+    # - not None if file_paths[i] is not None
+    file_urls: list[Optional[str]]
+
+    # file_paths:
+    # - local path of media file in vals[i]; points to the file cache if file_urls[i] is remote
+    # - None if vals[i] is not a media type or if there is no local file yet for file_urls[i]
+    file_paths: list[Optional[str]]
+
+    def __init__(self, size: int, img_slot_idxs: List[int], media_slot_idxs: List[int], array_slot_idxs: List[int]):
+        self.vals = [None] * size
+        self.has_val = [False] * size
+        self.excs = [None] * size
+        self.img_slot_idxs = img_slot_idxs
+        self.media_slot_idxs = media_slot_idxs
+        self.array_slot_idxs = array_slot_idxs
+        self.pk = None
+        self.file_urls = [None] * size
+        self.file_paths = [None] * size
 
     def clear(self) -> None:
         size = len(self.vals)

pixeltable/exprs/expr.py

@@ -7,11 +7,11 @@ import inspect
 import json
 import sys
 import typing
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, TypeVar, Union, overload, Iterable
 from uuid import UUID
 
 import sqlalchemy as sql
-from typing_extensions import Self
+from typing_extensions import _AnnotatedAlias, Self
 
 import pixeltable
 import pixeltable.catalog as catalog
@@ -281,9 +281,10 @@ class Expr(abc.ABC):
         """
         Iterate over all subexprs, including self.
         """
-        is_match = filter is None or filter(self)
-        if expr_class is not None:
-            is_match = is_match and isinstance(self, expr_class)
+        is_match = isinstance(self, expr_class) if expr_class is not None else True
+        # apply filter after checking for expr_class
+        if filter is not None and is_match:
+            is_match = filter(self)
         if not is_match or traverse_matches:
             for c in self.components:
                 yield from c.subexprs(expr_class=expr_class, filter=filter, traverse_matches=traverse_matches)
@@ -292,7 +293,7 @@ class Expr(abc.ABC):
 
     @overload
     def list_subexprs(
-        expr_list: list[Expr], *, filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
+        expr_list: Iterable[Expr], *, filter: Optional[Callable[[Expr], bool]] = None, traverse_matches: bool = True
     ) -> Iterator[Expr]: ...
 
     @overload
@@ -312,13 +313,11 @@ class Expr(abc.ABC):
 
     def _contains(self, cls: Optional[type[Expr]] = None, filter: Optional[Callable[[Expr], bool]] = None) -> bool:
         """
-        Returns True if any subexpr is an instance of cls.
+        Returns True if any subexpr is an instance of cls and/or matches filter.
         """
-        assert (cls is not None) != (filter is not None)  # need one of them
-        if cls is not None:
-            filter = lambda e: isinstance(e, cls)
+        assert cls is not None or filter is not None
         try:
-            _ = next(self.subexprs(filter=filter, traverse_matches=False))
+            _ = next(self.subexprs(expr_class=cls, filter=filter, traverse_matches=False))
             return True
         except StopIteration:
             return False
@@ -457,11 +456,13 @@ class Expr(abc.ABC):
         else:
             return InPredicate(self, value_set_literal=value_set)
 
-    def astype(self, new_type: ts.ColumnType) -> 'pixeltable.exprs.TypeCast':
+    def astype(self, new_type: Union[ts.ColumnType, type, _AnnotatedAlias]) -> 'pixeltable.exprs.TypeCast':
         from pixeltable.exprs import TypeCast
-        return TypeCast(self, new_type)
+        return TypeCast(self, ts.ColumnType.normalize_type(new_type))
 
-    def apply(self, fn: Callable, *, col_type: Optional[ts.ColumnType] = None) -> 'pixeltable.exprs.FunctionCall':
+    def apply(self, fn: Callable, *, col_type: Union[ts.ColumnType, type, _AnnotatedAlias, None] = None) -> 'pixeltable.exprs.FunctionCall':
+        if col_type is not None:
+            col_type = ts.ColumnType.normalize_type(col_type)
         function = self._make_applicator_function(fn, col_type)
         # Return a `FunctionCall` obtained by passing this `Expr` to the new `function`.
         return function(self)