pixeltable 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/exec/sql_node.py

@@ -1,73 +1,176 @@
-from typing import List, Optional, Tuple, Iterable, Set
-from uuid import UUID
 import logging
 import warnings
+from decimal import Decimal
+from typing import Iterable, Iterator, NamedTuple, Optional
+from uuid import UUID
 
 import sqlalchemy as sql
 
-from .data_row_batch import DataRowBatch
-from .exec_node import ExecNode
-import pixeltable.exprs as exprs
 import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
 
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
 
 _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_all(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_all(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.ColumnElement] = []
+        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.all_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:
@@ -80,17 +183,18 @@ class SqlNode(ExecNode):
         """
         # we need to include at least the root
         if refd_tbl_ids is None:
-            refd_tbl_ids = {}
+            refd_tbl_ids = set()
         if exact_version_only is None:
-            exact_version_only = {}
+            exact_version_only = set()
         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)
 
         first = True
+        prev_tbl: catalog.TableVersion
         for tbl in joined_tbls[::-1]:
             if first:
                 stmt = stmt.select_from(tbl.store_tbl.sa_tbl)
@@ -111,72 +215,94 @@ 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
+    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
+        with warnings.catch_warnings(record=True) as w:
+            stmt = self._create_stmt()
             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
-
-            if needs_row:
-                output_row = output_batch.add_row()
+                # log stmt, if possible
+                stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
+                _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
+            except Exception:
+                pass
+            self._log_explain(stmt)
+
+            result_cursor = self.ctx.conn.execute(stmt)
+            for warning in w:
+                pass
+
+        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
+
+        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 +315,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 +336,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.all_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 +366,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 +380,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.all_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/arithmetic_expr.py

@@ -6,6 +6,7 @@ import sqlalchemy as sql
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+
 from .data_row import DataRow
 from .expr import Expr
 from .globals import ArithmeticOperator
@@ -86,6 +87,7 @@ class ArithmeticExpr(Expr):
                 return sql.sql.expression.cast(sql.func.floor(left / right), sql.Integer)
             if self.col_type.is_float_type():
                 return sql.sql.expression.cast(sql.func.floor(left / right), sql.Float)
+        assert False
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         op1_val = data_row[self._op1.slot_idx]
@@ -121,7 +123,7 @@ class ArithmeticExpr(Expr):
         return {'operator': self.operator.value, **super()._as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, components: list[Expr]) -> ArithmeticExpr:
         assert 'operator' in d
         assert len(components) == 2
         return cls(ArithmeticOperator(d['operator']), components[0], components[1])

pixeltable/exprs/array_slice.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional, Union
 
 import sqlalchemy as sql
 
@@ -15,7 +15,7 @@ class ArraySlice(Expr):
     """
     Slice operation on an array, eg, t.array_col[:, 1:2].
     """
-    def __init__(self, arr: Expr, index: Tuple):
+    def __init__(self, arr: Expr, index: tuple[Union[int, slice], ...]):
         assert arr.col_type.is_array_type()
         # determine result type
         super().__init__(arr.col_type)
@@ -24,7 +24,7 @@ class ArraySlice(Expr):
         self.id = self._create_id()
 
     def __str__(self) -> str:
-        index_strs: List[str] = []
+        index_strs: list[str] = []
         for el in self.index:
             if isinstance(el, int):
                 index_strs.append(str(el))
@@ -39,7 +39,7 @@ class ArraySlice(Expr):
     def _equals(self, other: ArraySlice) -> bool:
         return self.index == other.index
 
-    def _id_attrs(self) -> List[Tuple[str, Any]]:
+    def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('index', self.index)]
 
     def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
@@ -49,8 +49,8 @@ class ArraySlice(Expr):
         val = data_row[self._array.slot_idx]
         data_row[self.slot_idx] = val[self.index]
 
-    def _as_dict(self) -> Dict:
-        index = []
+    def _as_dict(self) -> dict:
+        index: list[Any] = []
         for el in self.index:
             if isinstance(el, slice):
                 index.append([el.start, el.stop, el.step])
@@ -59,7 +59,7 @@ class ArraySlice(Expr):
         return {'index': index, **super()._as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, components: list[Expr]) -> ArraySlice:
         assert 'index' in d
         index = []
         for el in d['index']: