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

pixeltable/exec/sql_node.py

@@ -1,31 +1,34 @@
+import datetime
 import logging
 import warnings
 from decimal import Decimal
-from typing import Iterable, Iterator, NamedTuple, Optional, TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, AsyncIterator, Iterable, NamedTuple, Sequence
 from uuid import UUID
 
 import sqlalchemy as sql
 
-import pixeltable.catalog as catalog
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exprs
+from pixeltable.env import Env
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
 if TYPE_CHECKING:
     import pixeltable.plan
+    from pixeltable.plan import SampleClause
 
 _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
@@ -53,56 +56,91 @@ def combine_order_by_clauses(clauses: Iterable[OrderByClause]) -> Optional[Order
 
 
 def print_order_by_clause(clause: OrderByClause) -> str:
-    return ', '.join([
+    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 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
+    sql_select_list_exprs: exprs.ExprSet
+    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], row_builder: exprs.RowBuilder,
-            select_list: Iterable[exprs.Expr], sql_elements: exprs.SqlElementCache, set_pk: bool = False
+        self,
+        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
-        assert all(not isinstance(e, exprs.Literal) for e in select_list)  # we're never asked to materialize literals
+        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
+        # unstored iter columns: we also need to retrieve whatever is needed to materialize the
+        # iter args and stored outputs
         for iter_arg in row_builder.unstored_iter_args.values():
             sql_subexprs = iter_arg.subexprs(filter=self.sql_elements.contains, traverse_matches=False)
-            for e in sql_subexprs:
-                self.select_list.add(e)
+            self.select_list.update(sql_subexprs)
+        # We query for unstored outputs only if we're not loading a view; when we're loading a view, we are populating
+        # those columns, so we need to keep them out of the select list. This isn't a problem, because view loads never
+        # need to call set_pos().
+        # TODO: This is necessary because create_view_load_plan passes stored output columns to `RowBuilder` via the
+        #     `columns` parameter (even though they don't appear in `output_exprs`). This causes them to be recorded as
+        #     expressions in `RowBuilder`, which creates a conflict if we add them here. If `RowBuilder` is restructured
+        #     to keep them out of `unique_exprs`, then this conditional can be removed.
+        if not row_builder.for_view_load:
+            for outputs in row_builder.unstored_iter_outputs.values():
+                self.select_list.update(outputs)
         super().__init__(row_builder, self.select_list, [], None)  # we materialize self.select_list
 
         if tbl is not None:
@@ -115,9 +153,13 @@ class SqlNode(ExecNode):
         if set_pk:
             # we also need to retrieve the pk columns
             assert tbl is not None
-            self.num_pk_cols = len(tbl.tbl_version.store_tbl.pk_columns())
+            self.num_pk_cols = len(tbl.tbl_version.get().store_tbl.pk_columns())
+            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
@@ -128,14 +170,38 @@ class SqlNode(ExecNode):
         self.where_clause_element = None
         self.order_by_clause = []
 
-    def _create_stmt(self) -> sql.Select:
-        """Create Select from local state"""
+        if self.tbl is not None:
+            tv = self.tbl.tbl_version._tbl_version
+            if tv is not None:
+                assert tv.is_validated
 
-        assert self.sql_elements.contains_all(self.select_list)
-        sql_select_list = [self.sql_elements.get(e) for e in self.select_list]
+    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
-            sql_select_list += self.tbl.tbl_version.store_tbl.pk_columns()
+            assert self.tbl.tbl_version.get().is_validated
+            return self.tbl.tbl_version.get().store_tbl.pk_columns()
+        return []
+
+    def _init_exec_state(self) -> None:
+        assert self.sql_elements.contains_all(self.select_list)
+        self.sql_select_list_exprs = exprs.ExprSet(self.select_list)
+        self.cellmd_item_idxs = exprs.ExprDict((ref, self.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: self.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: self.sql_select_list_exprs.add(cellmd_ref) for cellmd_ref in column_cellmd_refs
+        }
+
+    def _create_stmt(self) -> sql.Select:
+        """Create Select from local state"""
+        self._init_exec_state()
+        sql_select_list = [self.sql_elements.get(e) for e in self.sql_select_list_exprs] + self._pk_col_items()
         stmt = sql.select(*sql_select_list)
 
         where_clause_element = (
@@ -161,9 +227,10 @@ 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) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
+    def to_cte(self, keep_pk: bool = False) -> tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]] | None:
         """
-        Returns a CTE that materializes the output of this node plus a mapping from select list expr to output column
+        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
 
         Returns:
             (CTE, dict from Expr to output column)
@@ -171,11 +238,11 @@ class SqlNode(ExecNode):
         if self.py_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:
+            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()
-            assert len(self.cte.c) == len(self.select_list)
-        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))
+        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:
@@ -186,8 +253,11 @@ class SqlNode(ExecNode):
 
     @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: 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:
@@ -205,31 +275,35 @@ class SqlNode(ExecNode):
             exact_version_only = set()
         candidates = tbl.get_tbl_versions()
         assert len(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)
+        joined_tbls: list[catalog.TableVersionHandle] = [candidates[0]]
+        for t in candidates[1:]:
+            if t.id in refd_tbl_ids:
+                joined_tbls.append(t)
 
         first = True
-        prev_tbl: catalog.TableVersion
-        for tbl in joined_tbls[::-1]:
+        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)}')
             if first:
-                stmt = stmt.select_from(tbl.store_tbl.sa_tbl)
+                stmt = stmt.select_from(tv.store_tbl.sa_tbl)
                 first = False
             else:
-                # join tbl to prev_tbl on prev_tbl's rowid cols
-                prev_tbl_rowid_cols = prev_tbl.store_tbl.rowid_columns()
-                tbl_rowid_cols = tbl.store_tbl.rowid_columns()
-                rowid_clauses = \
-                    [c1 == c2 for c1, c2 in zip(prev_tbl_rowid_cols, tbl_rowid_cols[:len(prev_tbl_rowid_cols)])]
-                stmt = stmt.join(tbl.store_tbl.sa_tbl, sql.and_(*rowid_clauses))
-            if tbl.id in exact_version_only:
-                stmt = stmt.where(tbl.store_tbl.v_min_col == tbl.version)
+                # join tv to prev_tv on prev_tv's rowid cols
+                prev_tbl_rowid_cols = prev_tv.store_tbl.rowid_columns()
+                tbl_rowid_cols = tv.store_tbl.rowid_columns()
+                rowid_clauses = [
+                    c1 == c2 for c1, c2 in zip(prev_tbl_rowid_cols, tbl_rowid_cols[: len(prev_tbl_rowid_cols)])
+                ]
+                stmt = stmt.join(tv.store_tbl.sa_tbl, sql.and_(*rowid_clauses))
+
+            if t.id in exact_version_only:
+                stmt = stmt.where(tv.store_tbl.v_min_col == tv.version)
             else:
-                stmt = stmt \
-                    .where(tbl.store_tbl.v_min_col <= tbl.version) \
-                    .where(tbl.store_tbl.v_max_col > tbl.version)
-            prev_tbl = tbl
+                stmt = stmt.where(tv.store_tbl.sa_tbl.c.v_min <= tv.version)
+                stmt = stmt.where(tv.store_tbl.sa_tbl.c.v_max > tv.version)
+            prev_tv = tv
+
         return stmt
 
     def set_where(self, where_clause: exprs.Expr) -> None:
@@ -255,18 +329,18 @@ class SqlNode(ExecNode):
         self.limit = limit
 
     def _log_explain(self, stmt: sql.Select) -> None:
+        conn = Env.get().conn
         try:
             # don't set dialect=Env.get().engine.dialect: x % y turns into x %% y, which results in a syntax error
             stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
-            explain_result = self.ctx.conn.execute(sql.text(f'EXPLAIN {stmt_str}'))
+            explain_result = 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')
+            _logger.warning(f'EXPLAIN failed with error: {e}')
 
-    def __iter__(self) -> Iterator[DataRowBatch]:
+    async def __aiter__(self) -> AsyncIterator[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:
@@ -274,35 +348,65 @@ class SqlNode(ExecNode):
                 stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
                 _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
             except Exception:
-                pass
+                # log something if we can't log the compiled stmt
+                _logger.debug(f'SqlLookupNode proto-stmt:\n{stmt}')
             self._log_explain(stmt)
 
-            result_cursor = self.ctx.conn.execute(stmt)
-            for warning in w:
+            conn = Env.get().conn
+            result_cursor = conn.execute(stmt)
+            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:]))
+                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]
 
@@ -324,7 +428,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')
@@ -340,22 +444,35 @@ 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.TableVersion]
+
+    exact_version_only: list[catalog.TableVersionHandle]
 
     def __init__(
-        self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
+        self,
+        tbl: catalog.TableVersionPath,
+        row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
-        set_pk: bool = False, exact_version_only: Optional[list[catalog.TableVersion]] = None
+        columns: list[catalog.Column],
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
+        set_pk: bool = False,
+        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 = []
@@ -367,27 +484,41 @@ class SqlScanNode(SqlNode):
         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})
+            self.tbl, stmt, refd_tbl_ids, exact_version_only={t.id for t in self.exact_version_only}
+        )
         return stmt
 
 
 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__(
-        self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
-        select_list: Iterable[exprs.Expr], sa_key_cols: list[sql.Column], key_vals: list[tuple],
+        self,
+        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)
 
@@ -401,30 +532,33 @@ 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]]
+    group_by_items: list[exprs.Expr] | None
+    input_cte: sql.CTE | None
 
     def __init__(
-        self, row_builder: exprs.RowBuilder,
+        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
-        """
-        _, input_col_map = input.to_cte()
+        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:
-        stmt = super()._create_stmt()
+        stmt = super()._create_stmt().select_from(self.input_cte)
         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)
@@ -436,12 +570,16 @@ class SqlJoinNode(SqlNode):
     """
     Materializes data from the store via a Select ... From ... that contains joins
     """
+
     input_ctes: list[sql.CTE]
     join_clauses: list['pixeltable.plan.JoinClause']
 
     def __init__(
-        self, row_builder: exprs.RowBuilder,
-        inputs: Sequence[SqlNode], join_clauses: list['pixeltable.plan.JoinClause'], select_list: Iterable[exprs.Expr]
+        self,
+        row_builder: exprs.RowBuilder,
+        inputs: Sequence[SqlNode],
+        join_clauses: list['pixeltable.plan.JoinClause'],
+        select_list: Iterable[exprs.Expr],
     ):
         assert len(inputs) > 1
         assert len(inputs) == len(join_clauses) + 1
@@ -452,20 +590,181 @@ 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
+
         stmt = super()._create_stmt()
         stmt = stmt.select_from(self.input_ctes[0])
         for i in range(len(self.join_clauses)):
             join_clause = self.join_clauses[i]
             on_clause = (
-                self.sql_elements.get(join_clause.join_predicate) if join_clause.join_type != plan.JoinType.CROSS
+                self.sql_elements.get(join_clause.join_predicate)
+                if join_clause.join_type != plan.JoinType.CROSS
                 else sql.sql.expression.literal(True)
             )
-            is_outer = join_clause.join_type == plan.JoinType.LEFT or join_clause.join_type == plan.JoinType.FULL_OUTER
+            is_outer = join_clause.join_type in (plan.JoinType.LEFT, plan.JoinType.FULL_OUTER)
             stmt = stmt.join(
-                self.input_ctes[i + 1], onclause=on_clause, isouter=is_outer,
-                full=join_clause == plan.JoinType.FULL_OUTER)
-        return stmt
+                self.input_ctes[i + 1],
+                onclause=on_clause,
+                isouter=is_outer,
+                full=join_clause == plan.JoinType.FULL_OUTER,
+            )
+        return stmt
+
+
+class SqlSampleNode(SqlNode):
+    """
+    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.
+    """
+
+    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],
+        sample_clause: 'SampleClause',
+        stratify_exprs: list[exprs.Expr],
+    ):
+        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)
+        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.sample_clause = sample_clause
+
+    @classmethod
+    def key_sql_expr(cls, seed: sql.ColumnElement, sql_cols: Iterable[sql.ColumnElement]) -> sql.ColumnElement:
+        """Construct expression which is the ordering key for rows to be sampled
+        General SQL form is:
+        - MD5(<seed::text> [ + '___' + <rowid_col_val>::text]+
+        """
+        sql_expr: sql.ColumnElement = seed.cast(sql.String)
+        for e in sql_cols:
+            # 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_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
+        # 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:
+        from pixeltable.plan import SampleClause
+
+        self._init_exec_state()
+
+        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"""
+
+        sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
+        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
+        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')
+
+        final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
+        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_key_sql(row_rank_cte)
+            return sql.select(*final_columns).order_by(row_rank_cte.c.rank, secondary_order).limit(n)
+
+    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)
+        # 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 = (
+            sql.select(
+                *sql_strata_exprs,
+                sql.func.ceil(fraction_samples * sql.func.count(1).cast(sql.Integer)).label('s_s_size'),
+            )
+            .select_from(self.input_cte)
+            .group_by(*sql_strata_exprs)
+            .cte('per_strata_count_cte')
+        )
+
+        # Build a CTE that ranks the rows within each stratum
+        # Include all columns from the input CTE dynamically
+        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 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 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)
+            .select_from(row_rank_cte)
+            .join(per_strata_count_cte, join_c)
+            .where(row_rank_cte.c.rank <= per_strata_count_cte.c.s_s_size)
+        )
+
+        return stmt