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

pixeltable/exec/sql_node.py

@@ -1,7 +1,8 @@
+import datetime
 import logging
 import warnings
 from decimal import Decimal
-from typing import TYPE_CHECKING, AsyncIterator, Iterable, NamedTuple, Optional, Sequence
+from typing import TYPE_CHECKING, AsyncIterator, Iterable, NamedTuple, Sequence
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -21,13 +22,13 @@ _logger = logging.getLogger('pixeltable')
 
 class OrderByItem(NamedTuple):
     expr: exprs.Expr
-    asc: Optional[bool]
+    asc: bool | None
 
 
 OrderByClause = list[OrderByItem]
 
 
-def combine_order_by_clauses(clauses: Iterable[OrderByClause]) -> Optional[OrderByClause]:
+def combine_order_by_clauses(clauses: Iterable[OrderByClause]) -> OrderByClause | None:
     """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
@@ -65,49 +66,66 @@ def print_order_by_clause(clause: OrderByClause) -> str:
 
 class SqlNode(ExecNode):
     """
-    Materializes data from the store via an SQL statement.
+    Materializes data from the store via a SQL statement.
     This only provides the select list. The subclasses are responsible for the From clause and any additional clauses.
     The pk columns are not included in the select list.
     If set_pk is True, they are added to the end of the result set when creating the SQL statement
     so they can always be referenced as cols[-num_pk_cols:] in the result set.
     The pk_columns consist of the rowid columns of the target table followed by the version number.
+
+    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
     """
 
-    tbl: Optional[catalog.TableVersionPath]
+    tbl: catalog.TableVersionPath | None
     select_list: exprs.ExprSet
+    columns: list[catalog.Column]  # for which columns to populate DataRow.cell_vals/cell_md
+    cell_md_refs: list[exprs.ColumnPropertyRef]  # of ColumnRefs which also need DataRow.slot_cellmd for evaluation
     set_pk: bool
     num_pk_cols: int
-    py_filter: Optional[exprs.Expr]  # a predicate that can only be run in Python
-    py_filter_eval_ctx: Optional[exprs.RowBuilder.EvalCtx]
-    cte: Optional[sql.CTE]
+    py_filter: exprs.Expr | None  # a predicate that can only be run in Python
+    py_filter_eval_ctx: exprs.RowBuilder.EvalCtx | None
+    cte: sql.CTE | None
     sql_elements: exprs.SqlElementCache
 
+    # execution state
+    cellmd_item_idxs: exprs.ExprDict[int]  # cellmd expr -> idx in sql select list
+    column_item_idxs: dict[catalog.Column, int]  # column -> idx in sql select list
+    column_cellmd_item_idxs: dict[catalog.Column, int]  # column -> idx in sql select list
+    result_cursor: sql.engine.CursorResult | None
+
     # where_clause/-_element: allow subclass to set one or the other (but not both)
-    where_clause: Optional[exprs.Expr]
-    where_clause_element: Optional[sql.ColumnElement]
+    where_clause: exprs.Expr | None
+    where_clause_element: sql.ColumnElement | None
 
     order_by_clause: OrderByClause
-    limit: Optional[int]
+    limit: int | None
 
     def __init__(
         self,
-        tbl: Optional[catalog.TableVersionPath],
+        tbl: catalog.TableVersionPath | None,
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
         sql_elements: exprs.SqlElementCache,
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
         set_pk: bool = False,
     ):
-        """
-        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
+        self.columns = columns
+        if cell_md_col_refs is not None:
+            assert all(ref.col.stores_cellmd for ref in cell_md_col_refs)
+            self.cell_md_refs = [
+                exprs.ColumnPropertyRef(ref, exprs.ColumnPropertyRef.Property.CELLMD) for ref in cell_md_col_refs
+            ]
+        else:
+            self.cell_md_refs = []
         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():
@@ -130,6 +148,9 @@ class SqlNode(ExecNode):
             assert self.num_pk_cols > 1
 
         # additional state
+        self.cellmd_item_idxs = exprs.ExprDict()
+        self.column_item_idxs = {}
+        self.column_cellmd_item_idxs = {}
         self.result_cursor = None
         # the filter is provided by the subclass
         self.py_filter = None
@@ -145,10 +166,9 @@ class SqlNode(ExecNode):
             if tv is not None:
                 assert tv.is_validated
 
-    def _create_pk_cols(self) -> list[sql.Column]:
-        """Create a list of pk columns"""
-        # we need to retrieve the pk columns
+    def _pk_col_items(self) -> list[sql.Column]:
         if self.set_pk:
+            # we need to retrieve the pk columns
             assert self.tbl is not None
             assert self.tbl.tbl_version.get().is_validated
             return self.tbl.tbl_version.get().store_tbl.pk_columns()
@@ -158,7 +178,19 @@ class SqlNode(ExecNode):
         """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] + self._create_pk_cols()
+        sql_select_list_exprs = exprs.ExprSet(self.select_list)
+        self.cellmd_item_idxs = exprs.ExprDict((ref, sql_select_list_exprs.add(ref)) for ref in self.cell_md_refs)
+        column_refs = [exprs.ColumnRef(col) for col in self.columns]
+        self.column_item_idxs = {col_ref.col: sql_select_list_exprs.add(col_ref) for col_ref in column_refs}
+        column_cellmd_refs = [
+            exprs.ColumnPropertyRef(col_ref, exprs.ColumnPropertyRef.Property.CELLMD)
+            for col_ref in column_refs
+            if col_ref.col.stores_cellmd
+        ]
+        self.column_cellmd_item_idxs = {
+            cellmd_ref.col_ref.col: sql_select_list_exprs.add(cellmd_ref) for cellmd_ref in column_cellmd_refs
+        }
+        sql_select_list = [self.sql_elements.get(e) for e in sql_select_list_exprs] + self._pk_col_items()
         stmt = sql.select(*sql_select_list)
 
         where_clause_element = (
@@ -184,7 +216,7 @@ class SqlNode(ExecNode):
     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, keep_pk: bool = False) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
+    def to_cte(self, keep_pk: bool = False) -> tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]] | None:
         """
         Creates a CTE that materializes the output of this node plus a mapping from select list expr to output column.
         keep_pk: if True, the PK columns are included in the CTE Select statement
@@ -199,9 +231,7 @@ class SqlNode(ExecNode):
             if not keep_pk:
                 self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
             self.cte = self._create_stmt().cte()
-        pk_count = self.num_pk_cols if self.set_pk else 0
-        assert len(self.select_list) + pk_count == len(self.cte.c)
-        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))  # skip pk cols
+        return self.cte, exprs.ExprDict(zip(list(self.select_list) + self.cell_md_refs, self.cte.c))  # skip pk cols
 
     @classmethod
     def retarget_rowid_refs(cls, target: catalog.TableVersionPath, expr_seq: Iterable[exprs.Expr]) -> None:
@@ -215,8 +245,8 @@ class SqlNode(ExecNode):
         cls,
         tbl: catalog.TableVersionPath,
         stmt: sql.Select,
-        refd_tbl_ids: Optional[set[UUID]] = None,
-        exact_version_only: Optional[set[UUID]] = None,
+        refd_tbl_ids: set[UUID] | None = None,
+        exact_version_only: set[UUID] | None = None,
     ) -> sql.Select:
         """Add From clause to stmt for tables/views referenced by materialized_exprs
         Args:
@@ -240,7 +270,7 @@ class SqlNode(ExecNode):
                 joined_tbls.append(t)
 
         first = True
-        prev_tv: Optional[catalog.TableVersion] = None
+        prev_tv: catalog.TableVersion | None = None
         for t in joined_tbls[::-1]:
             tv = t.get()
             # _logger.debug(f'create_from_clause: tbl_id={tv.id} {id(tv.store_tbl.sa_tbl)}')
@@ -308,8 +338,7 @@ class SqlNode(ExecNode):
                 _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
             except Exception:
                 # log something if we can't log the compiled stmt
-                stmt_str = repr(stmt)
-                _logger.debug(f'SqlLookupNode proto-stmt:\n{stmt_str}')
+                _logger.debug(f'SqlLookupNode proto-stmt:\n{stmt}')
             self._log_explain(stmt)
 
             conn = Env.get().conn
@@ -317,28 +346,56 @@ class SqlNode(ExecNode):
             for _ 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
+        output_batch = DataRowBatch(self.row_builder)
+        output_row: exprs.DataRow | None = None
         num_rows_returned = 0
+        is_using_cockroachdb = Env.get().is_using_cockroachdb
+        tzinfo = Env.get().default_time_zone
 
         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 :]))
+
+            # column copies
+            for col, item_idx in self.column_item_idxs.items():
+                output_row.cell_vals[col.id] = sql_row[item_idx]
+            for col, item_idx in self.column_cellmd_item_idxs.items():
+                cell_md_dict = sql_row[item_idx]
+                output_row.cell_md[col.id] = exprs.CellMd(**cell_md_dict) if cell_md_dict is not None else None
+
+            # populate DataRow.slot_cellmd, where requested
+            for cellmd_ref, item_idx in self.cellmd_item_idxs.items():
+                cell_md_dict = sql_row[item_idx]
+                output_row.slot_md[cellmd_ref.col_ref.slot_idx] = (
+                    exprs.CellMd.from_dict(cell_md_dict) if cell_md_dict is not None else None
+                )
+
             # copy the output of the SQL query into the output row
             for i, e in enumerate(self.select_list):
                 slot_idx = e.slot_idx
-                # certain numerical operations can produce Decimals (eg, SUM(<int column>)); we need to convert them
                 if isinstance(sql_row[i], Decimal):
+                    # certain numerical operations can produce Decimals (eg, SUM(<int column>)); we need to convert them
                     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}')
+                elif is_using_cockroachdb and isinstance(sql_row[i], datetime.datetime):
+                    # Ensure that the datetime is timezone-aware and in the session time zone
+                    # cockroachDB returns timestamps in the session time zone, with numeric offset,
+                    # convert to the session time zone with the requested tzinfo for DST handling
+                    if e.col_type.is_timestamp_type():
+                        if isinstance(sql_row[i].tzinfo, datetime.timezone):
+                            output_row[slot_idx] = sql_row[i].astimezone(tz=tzinfo)
+                        else:
+                            output_row[slot_idx] = sql_row[i]
+                    else:
+                        raise RuntimeError(f'Unexpected datetime value for {e}')
                 else:
                     output_row[slot_idx] = sql_row[i]
 
@@ -360,7 +417,7 @@ class SqlNode(ExecNode):
             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)
+                output_batch = DataRowBatch(self.row_builder)
 
         if len(output_batch) > 0:
             _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
@@ -376,6 +433,11 @@ class SqlScanNode(SqlNode):
     Materializes data from the store via a Select stmt.
 
     Supports filtering and ordering.
+
+    Args:
+        select_list: output of the query
+        set_pk: if True, sets the primary for each DataRow
+        exact_version_only: tables for which we only want to see rows created at the current version
     """
 
     exact_version_only: list[catalog.TableVersionHandle]
@@ -385,17 +447,21 @@ class SqlScanNode(SqlNode):
         tbl: catalog.TableVersionPath,
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
         set_pk: bool = False,
-        exact_version_only: Optional[list[catalog.TableVersionHandle]] = None,
+        exact_version_only: list[catalog.TableVersionHandle] | None = None,
     ):
-        """
-        Args:
-            select_list: output of the query
-            set_pk: if True, sets the primary for each DataRow
-            exact_version_only: tables for which we only want to see rows created at the current version
-        """
         sql_elements = exprs.SqlElementCache()
-        super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=set_pk)
+        super().__init__(
+            tbl,
+            row_builder,
+            select_list,
+            columns=columns,
+            sql_elements=sql_elements,
+            set_pk=set_pk,
+            cell_md_col_refs=cell_md_col_refs,
+        )
         # create Select stmt
         if exact_version_only is None:
             exact_version_only = []
@@ -415,6 +481,11 @@ class SqlScanNode(SqlNode):
 class SqlLookupNode(SqlNode):
     """
     Materializes data from the store via a Select stmt with a WHERE clause that matches a list of key values
+
+    Args:
+        select_list: output of the query
+        sa_key_cols: list of key columns in the store table
+        key_vals: list of key values to look up
     """
 
     def __init__(
@@ -422,17 +493,21 @@ class SqlLookupNode(SqlNode):
         tbl: catalog.TableVersionPath,
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
         sa_key_cols: list[sql.Column],
         key_vals: list[tuple],
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
     ):
-        """
-        Args:
-            select_list: output of the query
-            sa_key_cols: list of key columns in the store table
-            key_vals: list of key values to look up
-        """
         sql_elements = exprs.SqlElementCache()
-        super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=True)
+        super().__init__(
+            tbl,
+            row_builder,
+            select_list,
+            columns=columns,
+            sql_elements=sql_elements,
+            set_pk=True,
+            cell_md_col_refs=cell_md_col_refs,
+        )
         # Where clause: (key-col-1, key-col-2, ...) IN ((val-1, val-2, ...), ...)
         self.where_clause_element = sql.tuple_(*sa_key_cols).in_(key_vals)
 
@@ -446,29 +521,29 @@ class SqlLookupNode(SqlNode):
 class SqlAggregationNode(SqlNode):
     """
     Materializes data from the store via a Select stmt with a WHERE clause that matches a list of key values
+
+    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
     """
 
-    group_by_items: Optional[list[exprs.Expr]]
-    input_cte: Optional[sql.CTE]
+    group_by_items: list[exprs.Expr] | None
+    input_cte: sql.CTE | None
 
     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,
+        group_by_items: list[exprs.Expr] | None = None,
+        limit: int | None = None,
+        exact_version_only: list[catalog.TableVersion] | None = 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
-        """
+        assert len(input.cell_md_refs) == 0  # there's no aggregation over json or arrays in SQL
         self.input_cte, input_col_map = input.to_cte()
         sql_elements = exprs.SqlElementCache(input_col_map)
-        super().__init__(None, row_builder, select_list, sql_elements)
+        super().__init__(None, row_builder, select_list, columns=[], sql_elements=sql_elements)
         self.group_by_items = group_by_items
 
     def _create_stmt(self) -> sql.Select:
@@ -504,7 +579,10 @@ class SqlJoinNode(SqlNode):
             input_cte, input_col_map = input_node.to_cte()
             self.input_ctes.append(input_cte)
             sql_elements.extend(input_col_map)
-        super().__init__(None, row_builder, select_list, sql_elements)
+        cell_md_col_refs = [cell_md_ref.col_ref for input in inputs for cell_md_ref in input.cell_md_refs]
+        super().__init__(
+            None, row_builder, select_list, columns=[], sql_elements=sql_elements, cell_md_col_refs=cell_md_col_refs
+        )
 
     def _create_stmt(self) -> sql.Select:
         from pixeltable import plan
@@ -530,40 +608,46 @@ class SqlJoinNode(SqlNode):
 
 class SqlSampleNode(SqlNode):
     """
-    Returns rows from a stratified sample with N samples per strata.
+    Returns rows sampled from the input node.
+
+    Args:
+        input: SqlNode to sample from
+        select_list: can contain calls to AggregateFunctions
+        sample_clause: specifies the sampling method
+        stratify_exprs: Analyzer processed list of expressions to stratify by.
     """
 
-    stratify_exprs: Optional[list[exprs.Expr]]
-    n_samples: Optional[int]
-    fraction_samples: Optional[float]
-    seed: int
-    input_cte: Optional[sql.CTE]
+    input_cte: sql.CTE | None
     pk_count: int
+    stratify_exprs: list[exprs.Expr] | None
+    sample_clause: 'SampleClause'
 
     def __init__(
         self,
         row_builder: exprs.RowBuilder,
         input: SqlNode,
         select_list: Iterable[exprs.Expr],
-        stratify_exprs: Optional[list[exprs.Expr]] = None,
-        sample_clause: Optional['SampleClause'] = None,
+        sample_clause: 'SampleClause',
+        stratify_exprs: list[exprs.Expr],
     ):
-        """
-        Args:
-            select_list: can contain calls to AggregateFunctions
-            stratify_exprs: list of expressions to group by
-            n: number of samples per strata
-        """
+        assert isinstance(input, SqlNode)
         self.input_cte, input_col_map = input.to_cte(keep_pk=True)
         self.pk_count = input.num_pk_cols
         assert self.pk_count > 1
         sql_elements = exprs.SqlElementCache(input_col_map)
-        super().__init__(input.tbl, row_builder, select_list, sql_elements, set_pk=True)
+        assert sql_elements.contains_all(stratify_exprs)
+        cell_md_col_refs = [cell_md_ref.col_ref for cell_md_ref in input.cell_md_refs]
+        super().__init__(
+            input.tbl,
+            row_builder,
+            select_list,
+            columns=[],
+            sql_elements=sql_elements,
+            cell_md_col_refs=cell_md_col_refs,
+            set_pk=True,
+        )
         self.stratify_exprs = stratify_exprs
-        self.n_samples = sample_clause.n
-        self.n_per_stratum = sample_clause.n_per_stratum
-        self.fraction_samples = sample_clause.fraction
-        self.seed = sample_clause.seed if sample_clause.seed is not None else 0
+        self.sample_clause = sample_clause
 
     @classmethod
     def key_sql_expr(cls, seed: sql.ColumnElement, sql_cols: Iterable[sql.ColumnElement]) -> sql.ColumnElement:
@@ -571,27 +655,48 @@ class SqlSampleNode(SqlNode):
         General SQL form is:
         - MD5(<seed::text> [ + '___' + <rowid_col_val>::text]+
         """
-        sql_expr: sql.ColumnElement = sql.cast(seed, sql.Text)
+        sql_expr: sql.ColumnElement = seed.cast(sql.String)
         for e in sql_cols:
-            sql_expr = sql_expr + sql.literal_column("'___'") + sql.cast(e, sql.Text)
+            # Quotes are required below to guarantee that the string is properly presented in SQL
+            sql_expr = sql_expr + sql.literal_column("'___'", sql.Text) + e.cast(sql.String)
         sql_expr = sql.func.md5(sql_expr)
         return sql_expr
 
-    def _create_order_by(self, cte: sql.CTE) -> sql.ColumnElement:
+    def _create_key_sql(self, cte: sql.CTE) -> sql.ColumnElement:
         """Create an expression for randomly ordering rows with a given seed"""
         rowid_cols = [*cte.c[-self.pk_count : -1]]  # exclude the version column
         assert len(rowid_cols) > 0
-        return self.key_sql_expr(sql.literal_column(str(self.seed)), rowid_cols)
+        # If seed is not set in the sample clause, use the random seed given by the execution context
+        seed = self.sample_clause.seed if self.sample_clause.seed is not None else self.ctx.random_seed
+        return self.key_sql_expr(sql.literal_column(str(seed)), rowid_cols)
 
     def _create_stmt(self) -> sql.Select:
-        if self.fraction_samples is not None:
-            return self._create_stmt_fraction(self.fraction_samples)
-        return self._create_stmt_n(self.n_samples, self.n_per_stratum)
+        from pixeltable.plan import SampleClause
+
+        if self.sample_clause.fraction is not None:
+            if len(self.stratify_exprs) == 0:
+                # If non-stratified sampling, construct a where clause, order_by, and limit clauses
+                s_key = self._create_key_sql(self.input_cte)
+
+                # Construct a suitable where clause
+                fraction_md5 = SampleClause.fraction_to_md5_hex(self.sample_clause.fraction)
+                order_by = self._create_key_sql(self.input_cte)
+                return sql.select(*self.input_cte.c).where(s_key < fraction_md5).order_by(order_by)
+
+            return self._create_stmt_stratified_fraction(self.sample_clause.fraction)
+        else:
+            if len(self.stratify_exprs) == 0:
+                # No stratification, just return n samples from the input CTE
+                order_by = self._create_key_sql(self.input_cte)
+                return sql.select(*self.input_cte.c).order_by(order_by).limit(self.sample_clause.n)
+
+            return self._create_stmt_stratified_n(self.sample_clause.n, self.sample_clause.n_per_stratum)
+
+    def _create_stmt_stratified_n(self, n: int | None, n_per_stratum: int | None) -> sql.Select:
+        """Create a Select stmt that returns n samples across all strata or n_per_stratum samples per stratum"""
 
-    def _create_stmt_n(self, n: Optional[int], n_per_stratum: Optional[int]) -> sql.Select:
-        """Create a Select stmt that returns n samples across all strata"""
         sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
-        order_by = self._create_order_by(self.input_cte)
+        order_by = self._create_key_sql(self.input_cte)
 
         # Create a list of all columns plus the rank
         # Get all columns from the input CTE dynamically
@@ -605,15 +710,15 @@ class SqlSampleNode(SqlNode):
         if n_per_stratum is not None:
             return sql.select(*final_columns).filter(row_rank_cte.c.rank <= n_per_stratum)
         else:
-            secondary_order = self._create_order_by(row_rank_cte)
+            secondary_order = self._create_key_sql(row_rank_cte)
             return sql.select(*final_columns).order_by(row_rank_cte.c.rank, secondary_order).limit(n)
 
-    def _create_stmt_fraction(self, fraction_samples: float) -> sql.Select:
+    def _create_stmt_stratified_fraction(self, fraction_samples: float) -> sql.Select:
         """Create a Select stmt that returns a fraction of the rows per strata"""
 
         # Build the strata count CTE
         # Produces a table of the form:
-        #   ([stratify_exprs], s_s_size)
+        #   (*stratify_exprs, s_s_size)
         # where s_s_size is the number of samples to take from each stratum
         sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
         per_strata_count_cte = (
@@ -628,19 +733,19 @@ class SqlSampleNode(SqlNode):
 
         # Build a CTE that ranks the rows within each stratum
         # Include all columns from the input CTE dynamically
-        order_by = self._create_order_by(self.input_cte)
+        order_by = self._create_key_sql(self.input_cte)
         select_columns = [*self.input_cte.c]
         select_columns.append(
             sql.func.row_number().over(partition_by=sql_strata_exprs, order_by=order_by).label('rank')
         )
         row_rank_cte = sql.select(*select_columns).select_from(self.input_cte).cte('row_rank_cte')
 
-        # Build the join criterion dynamically to accommodate any number of group by columns
+        # Build the join criterion dynamically to accommodate any number of stratify_by expressions
         join_c = sql.true()
         for col in per_strata_count_cte.c[:-1]:
             join_c &= row_rank_cte.c[col.name].isnot_distinct_from(col)
 
-        # Join srcp with per_strata_count_cte to limit returns to the requested fraction of rows
+        # Join with per_strata_count_cte to limit returns to the requested fraction of rows
         final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
         stmt = (
             sql.select(*final_columns)

pixeltable/exprs/__init__.py

@@ -6,7 +6,7 @@ from .column_property_ref import ColumnPropertyRef
 from .column_ref import ColumnRef
 from .comparison import Comparison
 from .compound_predicate import CompoundPredicate
-from .data_row import DataRow
+from .data_row import ArrayMd, CellMd, DataRow
 from .expr import Expr
 from .expr_dict import ExprDict
 from .expr_set import ExprSet