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

pixeltable/plan.py

@@ -2,17 +2,17 @@ from __future__ import annotations
 
 import dataclasses
 import enum
-from typing import Any, Iterable, Optional, Sequence, Literal
+from textwrap import dedent
+from typing import Any, Iterable, Literal, Sequence, cast
 from uuid import UUID
 
+import pgvector.sqlalchemy  # type: ignore[import-untyped]
 import sqlalchemy as sql
 
 import pixeltable as pxt
-import pixeltable.exec as exec
-from pixeltable import catalog
-from pixeltable import exceptions as excs
-from pixeltable import exprs
-from pixeltable.exec.sql_node import OrderByItem, OrderByClause, combine_order_by_clauses, print_order_by_clause
+from pixeltable import catalog, exceptions as excs, exec, exprs
+from pixeltable.catalog import Column, TableVersionHandle
+from pixeltable.exec.sql_node import OrderByClause, OrderByItem, combine_order_by_clauses, print_order_by_clause
 
 
 def _is_agg_fn_call(e: exprs.Expr) -> bool:
@@ -20,7 +20,7 @@ def _is_agg_fn_call(e: exprs.Expr) -> bool:
 
 
 def _get_combined_ordering(
-        o1: list[tuple[exprs.Expr, bool]], o2: list[tuple[exprs.Expr, bool]]
+    o1: list[tuple[exprs.Expr, bool]], o2: list[tuple[exprs.Expr, bool]]
 ) -> list[tuple[exprs.Expr, bool]]:
     """Returns an ordering that's compatible with both o1 and o2, or an empty list if no such ordering exists"""
     result: list[tuple[exprs.Expr, bool]] = []
@@ -56,24 +56,103 @@ class JoinType(enum.Enum):
     def validated(cls, name: str, error_prefix: str) -> JoinType:
         try:
             return cls[name.upper()]
-        except KeyError:
-            val_strs = ', '.join(f'{s.lower()!r}' for s in cls.__members__.keys())
-            raise excs.Error(f'{error_prefix} must be one of: [{val_strs}]')
+        except KeyError as exc:
+            val_strs = ', '.join(f'{s.lower()!r}' for s in cls.__members__)
+            raise excs.Error(f'{error_prefix} must be one of: [{val_strs}]') from exc
 
 
 @dataclasses.dataclass
 class JoinClause:
     """Corresponds to a single 'JOIN ... ON (...)' clause in a SELECT statement; excludes the joined table."""
+
     join_type: JoinType
-    join_predicate: Optional[exprs.Expr]  # None for join_type == CROSS
+    join_predicate: exprs.Expr | None  # None for join_type == CROSS
 
 
 @dataclasses.dataclass
 class FromClause:
-    """Corresponds to the From-clause ('FROM <tbl> JOIN ... ON (...) JOIN ...') of a SELECT statement """
+    """Corresponds to the From-clause ('FROM <tbl> JOIN ... ON (...) JOIN ...') of a SELECT statement"""
+
     tbls: list[catalog.TableVersionPath]
     join_clauses: list[JoinClause] = dataclasses.field(default_factory=list)
 
+    @property
+    def _first_tbl(self) -> catalog.TableVersionPath:
+        assert len(self.tbls) == 1
+        return self.tbls[0]
+
+
+@dataclasses.dataclass
+class SampleClause:
+    """Defines a sampling clause for a table."""
+
+    version: int | None
+    n: int | None
+    n_per_stratum: int | None
+    fraction: float | None
+    seed: int | None
+    stratify_exprs: list[exprs.Expr] | None
+
+    # The version of the hashing algorithm used for ordering and fractional sampling.
+    CURRENT_VERSION = 1
+
+    def __post_init__(self) -> None:
+        # If no version was provided, provide the default version
+        if self.version is None:
+            self.version = self.CURRENT_VERSION
+
+    @property
+    def is_stratified(self) -> bool:
+        """Check if the sampling is stratified"""
+        return self.stratify_exprs is not None and len(self.stratify_exprs) > 0
+
+    @property
+    def is_repeatable(self) -> bool:
+        """Return true if the same rows will continue to be sampled if source rows are added or deleted."""
+        return not self.is_stratified and self.fraction is not None
+
+    def display_str(self, inline: bool = False) -> str:
+        return str(self)
+
+    def as_dict(self) -> dict:
+        """Return a dictionary representation of the object"""
+        d = dataclasses.asdict(self)
+        d['_classname'] = self.__class__.__name__
+        if self.is_stratified:
+            d['stratify_exprs'] = [e.as_dict() for e in self.stratify_exprs]
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict) -> SampleClause:
+        """Create a SampleClause from a dictionary representation"""
+        d_cleaned = {key: value for key, value in d.items() if key != '_classname'}
+        s = cls(**d_cleaned)
+        if s.is_stratified:
+            s.stratify_exprs = [exprs.Expr.from_dict(e) for e in d_cleaned.get('stratify_exprs', [])]
+        return s
+
+    def __repr__(self) -> str:
+        s = ','.join(e.display_str(inline=True) for e in self.stratify_exprs)
+        return (
+            f'sample_{self.version}(n={self.n}, n_per_stratum={self.n_per_stratum}, '
+            f'fraction={self.fraction}, seed={self.seed}, [{s}])'
+        )
+
+    @classmethod
+    def fraction_to_md5_hex(cls, fraction: float) -> str:
+        """Return the string representation of an approximation (to ~1e-9) of a fraction of the total space
+        of md5 hash values.
+        This is used for fractional sampling.
+        """
+        # Maximum count for the upper 32 bits of MD5: 2^32
+        max_md5_value = (2**32) - 1
+
+        # Calculate the fraction of this value
+        threshold_int = max_md5_value * int(1_000_000_000 * fraction) // 1_000_000_000
+
+        # Convert to hexadecimal string with padding
+        return format(threshold_int, '08x') + 'ffffffffffffffffffffffff'
+
 
 class Analyzer:
     """
@@ -83,26 +162,33 @@ class Analyzer:
     from_clause: FromClause
     all_exprs: list[exprs.Expr]  # union of all exprs, aside from sql_where_clause
     select_list: list[exprs.Expr]
-    group_by_clause: Optional[list[exprs.Expr]]  # None for non-aggregate queries; [] for agg query w/o grouping
+    group_by_clause: list[exprs.Expr] | None  # None for non-aggregate queries; [] for agg query w/o grouping
     grouping_exprs: list[exprs.Expr]  # [] for non-aggregate queries or agg query w/o grouping
     order_by_clause: OrderByClause
+    stratify_exprs: list[exprs.Expr]  # [] if no stratiifcation is required
+    sample_clause: SampleClause | None  # None if no sampling clause is present
 
     sql_elements: exprs.SqlElementCache
 
     # Where clause of the Select stmt of the SQL scan
-    sql_where_clause: Optional[exprs.Expr]
+    sql_where_clause: exprs.Expr | None
 
     # filter predicate applied to output rows of the SQL scan
-    filter: Optional[exprs.Expr]
+    filter: exprs.Expr | None
 
     agg_fn_calls: list[exprs.FunctionCall]  # grouping aggregation (ie, not window functions)
     window_fn_calls: list[exprs.FunctionCall]
     agg_order_by: list[exprs.Expr]
 
     def __init__(
-            self, from_clause: FromClause, select_list: Sequence[exprs.Expr],
-            where_clause: Optional[exprs.Expr] = None, group_by_clause: Optional[list[exprs.Expr]] = None,
-            order_by_clause: Optional[list[tuple[exprs.Expr, bool]]] = None):
+        self,
+        from_clause: FromClause,
+        select_list: Sequence[exprs.Expr],
+        where_clause: exprs.Expr | None = None,
+        group_by_clause: list[exprs.Expr] | None = None,
+        order_by_clause: list[tuple[exprs.Expr, bool]] | None = None,
+        sample_clause: SampleClause | None = None,
+    ):
         if order_by_clause is None:
             order_by_clause = []
         self.from_clause = from_clause
@@ -115,6 +201,11 @@ class Analyzer:
         self.group_by_clause = (
             [e.resolve_computed_cols() for e in group_by_clause] if group_by_clause is not None else None
         )
+        self.sample_clause = sample_clause
+        if self.sample_clause is not None and self.sample_clause.is_stratified:
+            self.stratify_exprs = [e.resolve_computed_cols() for e in sample_clause.stratify_exprs]
+        else:
+            self.stratify_exprs = []
         self.order_by_clause = [OrderByItem(e.resolve_computed_cols(), asc) for e, asc in order_by_clause]
 
         self.sql_where_clause = None
@@ -130,8 +221,11 @@ class Analyzer:
                 self.all_exprs.append(join_clause.join_predicate)
         if self.group_by_clause is not None:
             self.all_exprs.extend(self.group_by_clause)
+        self.all_exprs.extend(self.stratify_exprs)
         self.all_exprs.extend(e for e, _ in self.order_by_clause)
         if self.filter is not None:
+            if sample_clause is not None:
+                raise excs.Error(f'Filter {self.filter} not expressible in SQL')
             self.all_exprs.append(self.filter)
 
         self.agg_order_by = []
@@ -145,12 +239,17 @@ class Analyzer:
         candidates = self.select_list
         agg_fn_calls = exprs.ExprSet(
             exprs.Expr.list_subexprs(
-                candidates, expr_class=exprs.FunctionCall,
-                filter=lambda e: bool(e.is_agg_fn_call and not e.is_window_fn_call)))
+                candidates,
+                expr_class=exprs.FunctionCall,
+                filter=lambda e: bool(e.is_agg_fn_call and not e.is_window_fn_call),
+            )
+        )
         self.agg_fn_calls = list(agg_fn_calls)
         window_fn_calls = exprs.ExprSet(
             exprs.Expr.list_subexprs(
-                candidates, expr_class=exprs.FunctionCall, filter=lambda e: bool(e.is_window_fn_call)))
+                candidates, expr_class=exprs.FunctionCall, filter=lambda e: bool(e.is_window_fn_call)
+            )
+        )
         self.window_fn_calls = list(window_fn_calls)
         if len(self.agg_fn_calls) == 0:
             # nothing to do
@@ -164,19 +263,25 @@ class Analyzer:
         is_agg_output = [self._determine_agg_status(e, grouping_expr_ids)[0] for e in self.select_list]
         if is_agg_output.count(False) > 0:
             raise excs.Error(
-                f'Invalid non-aggregate expression in aggregate query: {self.select_list[is_agg_output.index(False)]}')
-
-        # check that filter doesn't contain aggregates
-        if self.filter is not None:
-            if any(_is_agg_fn_call(e) for e in self.filter.subexprs(expr_class=exprs.FunctionCall)):
-                raise excs.Error(f'Filter cannot contain aggregate functions: {self.filter}')
+                f'Invalid non-aggregate expression in aggregate query: {self.select_list[is_agg_output.index(False)]}'
+            )
+
+        # check that Where clause and filter doesn't contain aggregates
+        if self.sql_where_clause is not None and any(
+            _is_agg_fn_call(e) for e in self.sql_where_clause.subexprs(expr_class=exprs.FunctionCall)
+        ):
+            raise excs.Error(f'where() cannot contain aggregate functions: {self.sql_where_clause}')
+        if self.filter is not None and any(
+            _is_agg_fn_call(e) for e in self.filter.subexprs(expr_class=exprs.FunctionCall)
+        ):
+            raise excs.Error(f'where() cannot contain aggregate functions: {self.filter}')
 
         # check that grouping exprs don't contain aggregates and can be expressed as SQL (we perform sort-based
         # aggregation and rely on the SqlScanNode returning data in the correct order)
         for e in self.group_by_clause:
             if not self.sql_elements.contains(e):
                 raise excs.Error(f'Invalid grouping expression, needs to be expressible in SQL: {e}')
-            if e._contains(filter=lambda e: _is_agg_fn_call(e)):
+            if e._contains(filter=_is_agg_fn_call):
                 raise excs.Error(f'Grouping expression contains aggregate function: {e}')
 
     def _determine_agg_status(self, e: exprs.Expr, grouping_expr_ids: set[int]) -> tuple[bool, bool]:
@@ -194,14 +299,15 @@ class Analyzer:
             return True, False
         elif isinstance(e, exprs.Literal):
             return True, True
-        elif isinstance(e, exprs.ColumnRef) or isinstance(e, exprs.RowidRef):
+        elif isinstance(e, (exprs.ColumnRef, exprs.RowidRef)):
             # we already know that this isn't a grouping expr
             return False, True
         else:
             # an expression such as <grouping expr 1> + <grouping expr 2> can both be the output and input of agg
             assert len(e.components) > 0
             component_is_output, component_is_input = zip(
-                *[self._determine_agg_status(c, grouping_expr_ids) for c in e.components])
+                *[self._determine_agg_status(c, grouping_expr_ids) for c in e.components]
+            )
             is_output = component_is_output.count(True) == len(e.components)
             is_input = component_is_input.count(True) == len(e.components)
             if not is_output and not is_input:
@@ -224,13 +330,14 @@ class Analyzer:
         row_builder.set_slot_idxs(self.agg_fn_calls)
         row_builder.set_slot_idxs(self.agg_order_by)
 
-    def get_window_fn_ob_clause(self) -> Optional[OrderByClause]:
+    def get_window_fn_ob_clause(self) -> OrderByClause | None:
         clause: list[OrderByClause] = []
         for fn_call in self.window_fn_calls:
             # window functions require ordering by the group_by/order_by clauses
             group_by_exprs, order_by_exprs = fn_call.get_window_sort_exprs()
             clause.append(
-                [OrderByItem(e, None) for e in group_by_exprs] + [OrderByItem(e, True) for e in order_by_exprs])
+                [OrderByItem(e, None) for e in group_by_exprs] + [OrderByItem(e, True) for e in order_by_exprs]
+            )
         return combine_order_by_clauses(clause)
 
     def has_agg(self) -> bool:
@@ -239,103 +346,113 @@ class Analyzer:
 
 
 class Planner:
-    # TODO: create an exec.CountNode and change this to create_count_plan()
     @classmethod
-    def create_count_stmt(
-            cls, tbl: catalog.TableVersionPath, where_clause: Optional[exprs.Expr] = None
-    ) -> sql.Select:
-        stmt = sql.select(sql.func.count())
-        refd_tbl_ids: set[UUID] = set()
-        if where_clause is not None:
-            analyzer = cls.analyze(tbl, where_clause)
-            if analyzer.filter is not None:
-                raise excs.Error(f'Filter {analyzer.filter} not expressible in SQL')
-            clause_element = analyzer.sql_where_clause.sql_expr(analyzer.sql_elements)
-            assert clause_element is not None
-            stmt = stmt.where(clause_element)
-            refd_tbl_ids = where_clause.tbl_ids()
-        stmt = exec.SqlScanNode.create_from_clause(tbl, stmt, refd_tbl_ids)
-        return stmt
+    def create_count_stmt(cls, query: 'pxt.Query') -> sql.Select:
+        """Creates a SQL SELECT COUNT(*) statement for counting rows in a Query."""
+        # Create the query plan
+        plan = query._create_query_plan()
+        sql_node = plan.get_node(exec.SqlNode)
+        assert sql_node is not None
+        if sql_node.py_filter is not None:
+            raise excs.Error('count() cannot be used with Python-only filters. Use collect() instead.')
+        # Get the SQL statement from the SqlNode as a CTE
+        cte, _ = sql_node.to_cte(keep_pk=True)
+        count_stmt = sql.select(sql.func.count().label('all_count')).select_from(cte)
+        return count_stmt
 
     @classmethod
     def create_insert_plan(
         cls, tbl: catalog.TableVersion, rows: list[dict[str, Any]], ignore_errors: bool
     ) -> exec.ExecNode:
         """Creates a plan for TableVersion.insert()"""
-        assert not tbl.is_view()
+        assert not tbl.is_view
         # stored_cols: all cols we need to store, incl computed cols (and indices)
         stored_cols = [c for c in tbl.cols_by_id.values() if c.is_stored]
         assert len(stored_cols) > 0  # there needs to be something to store
-        row_builder = exprs.RowBuilder([], stored_cols, [])
+
+        cls.__check_valid_columns(tbl, stored_cols, 'inserted into')
+
+        row_builder = exprs.RowBuilder([], stored_cols, [], tbl)
 
         # create InMemoryDataNode for 'rows'
-        plan: exec.ExecNode = exec.InMemoryDataNode(tbl, rows, row_builder, tbl.next_rowid)
+        plan: exec.ExecNode = exec.InMemoryDataNode(tbl.handle, rows, row_builder, tbl.next_row_id)
 
-        media_input_col_info = [
-            exprs.ColumnSlotIdx(col_ref.col, col_ref.slot_idx)
-            for col_ref in row_builder.input_exprs
-            if isinstance(col_ref, exprs.ColumnRef) and col_ref.col_type.is_media_type()
-        ]
-        if len(media_input_col_info) > 0:
-            # prefetch external files for all input column refs
-            plan = exec.CachePrefetchNode(tbl.id, media_input_col_info, input=plan)
+        plan = cls._add_prefetch_node(tbl.id, row_builder.input_exprs, input_node=plan)
 
         computed_exprs = row_builder.output_exprs - row_builder.input_exprs
         if len(computed_exprs) > 0:
             # add an ExprEvalNode when there are exprs to compute
-            plan = exec.ExprEvalNode(row_builder, computed_exprs, plan.output_exprs, input=plan)
+            plan = exec.ExprEvalNode(
+                row_builder, computed_exprs, plan.output_exprs, input=plan, maintain_input_order=False
+            )
+        if any(c.col_type.supports_file_offloading() for c in stored_cols):
+            plan = exec.CellMaterializationNode(plan)
 
-        stored_col_info = row_builder.output_slot_idxs()
-        stored_img_col_info = [info for info in stored_col_info if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
         plan.set_ctx(
             exec.ExecContext(
-                row_builder, batch_size=0, show_pbar=True, num_computed_exprs=len(computed_exprs),
-                ignore_errors=ignore_errors))
+                row_builder,
+                batch_size=0,
+                show_pbar=True,
+                num_computed_exprs=len(computed_exprs),
+                ignore_errors=ignore_errors,
+            )
+        )
+        plan = cls._add_save_node(plan)
+
         return plan
 
     @classmethod
-    def create_df_insert_plan(
-        cls,
-        tbl: catalog.TableVersion,
-        df: 'pxt.DataFrame',
-        ignore_errors: bool
+    def rowid_columns(cls, target: TableVersionHandle, num_rowid_cols: int | None = None) -> list[exprs.Expr]:
+        """Return list of RowidRef for the given number of associated rowids"""
+        if num_rowid_cols is None:
+            num_rowid_cols = target.get().num_rowid_columns()
+        return [exprs.RowidRef(target, i) for i in range(num_rowid_cols)]
+
+    @classmethod
+    def create_query_insert_plan(
+        cls, tbl: catalog.TableVersion, query: 'pxt.Query', ignore_errors: bool
     ) -> exec.ExecNode:
-        assert not tbl.is_view()
-        plan = df._create_query_plan()  # ExecNode constructed by the DataFrame
+        assert not tbl.is_view
+        plan = query._create_query_plan()  # ExecNode constructed by the Query
 
         # Modify the plan RowBuilder to register the output columns
-        for col_name, expr in zip(df.schema.keys(), df._select_list_exprs):
+        needs_cell_materialization = False
+        for col_name, expr in zip(query.schema.keys(), query._select_list_exprs):
             assert col_name in tbl.cols_by_name
             col = tbl.cols_by_name[col_name]
             plan.row_builder.add_table_column(col, expr.slot_idx)
+            needs_cell_materialization = needs_cell_materialization or col.col_type.supports_file_offloading()
 
-        stored_col_info = plan.row_builder.output_slot_idxs()
-        stored_img_col_info = [info for info in stored_col_info if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
+        if needs_cell_materialization:
+            plan = exec.CellMaterializationNode(plan)
 
         plan.set_ctx(
             exec.ExecContext(
-                plan.row_builder, batch_size=0, show_pbar=True, num_computed_exprs=0,
-                ignore_errors=ignore_errors))
+                plan.row_builder, batch_size=0, show_pbar=True, num_computed_exprs=0, ignore_errors=ignore_errors
+            )
+        )
         plan.ctx.num_rows = 0  # Unknown
 
         return plan
 
     @classmethod
     def create_update_plan(
-            cls, tbl: catalog.TableVersionPath,
-            update_targets: dict[catalog.Column, exprs.Expr],
-            recompute_targets: list[catalog.Column],
-            where_clause: Optional[exprs.Expr], cascade: bool
+        cls,
+        tbl: catalog.TableVersionPath,
+        update_targets: dict[catalog.Column, exprs.Expr],
+        recompute_targets: list[catalog.Column],
+        where_clause: exprs.Expr | None,
+        cascade: bool,
     ) -> tuple[exec.ExecNode, list[str], list[catalog.Column]]:
         """Creates a plan to materialize updated rows.
+
         The plan:
         - retrieves rows that are visible at the current version of the table
         - materializes all stored columns and the update targets
         - if cascade is True, recomputes all computed columns that transitively depend on the updated columns
           and copies the values of all other stored columns
         - if cascade is False, copies all columns that aren't update targets from the original rows
+
         Returns:
             - root node of the plan
             - list of qualified column names that are getting updated
@@ -343,46 +460,178 @@ class Planner:
         """
         # retrieve all stored cols and all target exprs
         assert isinstance(tbl, catalog.TableVersionPath)
-        target = tbl.tbl_version  # the one we need to update
+        target = tbl.tbl_version.get()  # the one we need to update
         updated_cols = list(update_targets.keys())
+        recomputed_cols: set[Column]
         if len(recompute_targets) > 0:
-            recomputed_cols = set(recompute_targets)
+            assert len(update_targets) == 0
+            recomputed_cols = {*recompute_targets}
+            if cascade:
+                recomputed_cols |= target.get_dependent_columns(recomputed_cols)
         else:
             recomputed_cols = target.get_dependent_columns(updated_cols) if cascade else set()
-            # regardless of cascade, we need to update all indices on any updated column
-            idx_val_cols = target.get_idx_val_columns(updated_cols)
-            recomputed_cols.update(idx_val_cols)
-            # we only need to recompute stored columns (unstored ones are substituted away)
-            recomputed_cols = {c for c in recomputed_cols if c.is_stored}
-        recomputed_base_cols = {col for col in recomputed_cols if col.tbl == target}
+        # regardless of cascade, we need to update all indices on any updated/recomputed column
+        modified_base_cols = [c for c in set(updated_cols) | recomputed_cols if c.get_tbl().id == target.id]
+        idx_val_cols = target.get_idx_val_columns(modified_base_cols)
+        recomputed_cols.update(idx_val_cols)
+        # we only need to recompute stored columns (unstored ones are substituted away)
+        recomputed_cols = {c for c in recomputed_cols if c.is_stored}
+
+        cls.__check_valid_columns(tbl.tbl_version.get(), recomputed_cols, 'updated in')
+
+        # our query plan
+        # - evaluates the update targets and recomputed columns
+        # - copies all other stored columns
+        recomputed_base_cols = {col for col in recomputed_cols if col.get_tbl().id == tbl.tbl_version.id}
         copied_cols = [
-            col for col in target.cols_by_id.values()
-            if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
+            col
+            for col in target.cols_by_id.values()
+            if col.is_stored and col not in updated_cols and col not in recomputed_base_cols
         ]
-        select_list: list[exprs.Expr] = [exprs.ColumnRef(col) for col in copied_cols]
-        select_list.extend(update_targets.values())
+        select_list: list[exprs.Expr] = list(update_targets.values())
 
-        recomputed_exprs = \
-            [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols]
+        recomputed_exprs = [
+            c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols
+        ]
         # recomputed cols reference the new values of the updated cols
         spec: dict[exprs.Expr, exprs.Expr] = {exprs.ColumnRef(col): e for col, e in update_targets.items()}
         exprs.Expr.list_substitute(recomputed_exprs, spec)
         select_list.extend(recomputed_exprs)
 
         # we need to retrieve the PK columns of the existing rows
-        plan = cls.create_query_plan(FromClause(tbls=[tbl]), select_list, where_clause=where_clause, ignore_errors=True)
-        all_base_cols = copied_cols + updated_cols + list(recomputed_base_cols)  # same order as select_list
+        plan = cls.create_query_plan(
+            FromClause(tbls=[tbl]),
+            select_list=select_list,
+            columns=copied_cols,
+            where_clause=where_clause,
+            ignore_errors=True,
+        )
+        evaluated_cols = updated_cols + list(recomputed_base_cols)  # same order as select_list
         # update row builder with column information
-        for i, col in enumerate(all_base_cols):
+        plan.row_builder.add_table_columns(copied_cols)
+        for i, col in enumerate(evaluated_cols):
             plan.row_builder.add_table_column(col, select_list[i].slot_idx)
+        plan.ctx.num_computed_exprs = len(recomputed_exprs)
+
+        plan = cls._add_cell_materialization_node(plan)
+        plan = cls._add_save_node(plan)
+
         recomputed_user_cols = [c for c in recomputed_cols if c.name is not None]
-        return plan, [f'{c.tbl.name}.{c.name}' for c in updated_cols + recomputed_user_cols], recomputed_user_cols
+        return plan, [f'{c.get_tbl().name}.{c.name}' for c in updated_cols + recomputed_user_cols], recomputed_user_cols
+
+    @classmethod
+    def __check_valid_columns(
+        cls, tbl: catalog.TableVersion, cols: Iterable[Column], op_name: Literal['inserted into', 'updated in']
+    ) -> None:
+        for col in cols:
+            if col.value_expr is not None and not col.value_expr.is_valid:
+                raise excs.Error(
+                    dedent(
+                        f"""
+                        Data cannot be {op_name} the table {tbl.name!r},
+                        because the column {col.name!r} is currently invalid:
+                        {{validation_error}}
+                        """
+                    )
+                    .strip()
+                    .format(validation_error=col.value_expr.validation_error)
+                )
+
+    @classmethod
+    def _cell_md_col_refs(cls, expr_list: Iterable[exprs.Expr]) -> list[exprs.ColumnRef]:
+        """Return list of ColumnRefs that need their cellmd values for reconstruction"""
+        json_col_refs = list(
+            exprs.Expr.list_subexprs(
+                expr_list,
+                expr_class=exprs.ColumnRef,
+                filter=lambda e: cast(exprs.ColumnRef, e).col.col_type.is_json_type(),
+                traverse_matches=False,
+            )
+        )
+
+        def needs_reconstruction(e: exprs.Expr) -> bool:
+            assert isinstance(e, exprs.ColumnRef)
+            # Vector-typed array columns are used for vector indexes, and are stored in the db
+            return e.col.col_type.is_array_type() and not isinstance(e.col.sa_col_type, pgvector.sqlalchemy.Vector)
+
+        array_col_refs = list(
+            exprs.Expr.list_subexprs(
+                expr_list, expr_class=exprs.ColumnRef, filter=needs_reconstruction, traverse_matches=False
+            )
+        )
+
+        binary_col_refs = list(
+            exprs.Expr.list_subexprs(
+                expr_list,
+                expr_class=exprs.ColumnRef,
+                filter=lambda e: cast(exprs.ColumnRef, e).col.col_type.is_binary_type(),
+                traverse_matches=False,
+            )
+        )
+
+        return json_col_refs + array_col_refs + binary_col_refs
+
+    @classmethod
+    def _add_cell_materialization_node(cls, input: exec.ExecNode) -> exec.ExecNode:
+        # we need a CellMaterializationNode if any of the evaluated output columns are json or array-typed
+        has_target_cols = any(
+            col.col_type.supports_file_offloading()
+            for col, slot_idx in input.row_builder.table_columns.items()
+            if slot_idx is not None
+        )
+        if has_target_cols:
+            return exec.CellMaterializationNode(input)
+        else:
+            return input
+
+    @classmethod
+    def _add_cell_reconstruction_node(cls, expr_list: list[exprs.Expr], input: exec.ExecNode) -> exec.ExecNode:
+        """
+        Add a CellReconstructionNode, if required by any of the exprs in expr_list.
+
+        Cell reconstruction is required for
+        1) all json-typed ColumnRefs that are not used as part of a JsonPath (the latter does its own reconstruction)
+           or as part of a ColumnPropertyRef
+        2) all array-typed ColumnRefs that are not used as part of a ColumnPropertyRef
+        """
+
+        def json_filter(e: exprs.Expr) -> bool:
+            if isinstance(e, exprs.JsonPath):
+                return not e.is_relative_path() and isinstance(e.anchor, exprs.ColumnRef)
+            if isinstance(e, exprs.ColumnPropertyRef):
+                return e.col_ref.col.col_type.is_json_type()
+            return isinstance(e, exprs.ColumnRef) and e.col.col_type.is_json_type()
+
+        def array_filter(e: exprs.Expr) -> bool:
+            if isinstance(e, exprs.ColumnPropertyRef):
+                return e.col_ref.col.col_type.is_array_type()
+            if not isinstance(e, exprs.ColumnRef):
+                return False
+            # Vector-typed array columns are used for vector indexes, and are stored in the db
+            return e.col.col_type.is_array_type() and not isinstance(e.col.sa_col_type, pgvector.sqlalchemy.Vector)
+
+        def binary_filter(e: exprs.Expr) -> bool:
+            return isinstance(e, exprs.ColumnRef) and e.col.col_type.is_binary_type()
+
+        json_candidates = list(exprs.Expr.list_subexprs(expr_list, filter=json_filter, traverse_matches=False))
+        json_refs = [e for e in json_candidates if isinstance(e, exprs.ColumnRef)]
+        array_candidates = list(exprs.Expr.list_subexprs(expr_list, filter=array_filter, traverse_matches=False))
+        array_refs = [e for e in array_candidates if isinstance(e, exprs.ColumnRef)]
+        binary_refs = list(
+            exprs.Expr.list_subexprs(expr_list, exprs.ColumnRef, filter=binary_filter, traverse_matches=False)
+        )
+        if len(json_refs) > 0 or len(array_refs) > 0 or len(binary_refs) > 0:
+            return exec.CellReconstructionNode(json_refs, array_refs, binary_refs, input.row_builder, input=input)
+        else:
+            return input
 
     @classmethod
     def create_batch_update_plan(
-        cls, tbl: catalog.TableVersionPath,
-        batch: list[dict[catalog.Column, exprs.Expr]], rowids: list[tuple[int, ...]],
-        cascade: bool
+        cls,
+        tbl: catalog.TableVersionPath,
+        batch: list[dict[catalog.Column, exprs.Expr]],
+        rowids: list[tuple[int, ...]],
+        cascade: bool,
     ) -> tuple[exec.ExecNode, exec.RowUpdateNode, sql.ColumnElement[bool], list[catalog.Column], list[catalog.Column]]:
         """
         Returns:
@@ -393,9 +642,9 @@ class Planner:
         - list of user-visible columns that are being recomputed
         """
         assert isinstance(tbl, catalog.TableVersionPath)
-        target = tbl.tbl_version  # the one we need to update
-        sa_key_cols: list[sql.Column] = []
-        key_vals: list[tuple] = []
+        target = tbl.tbl_version.get()  # the one we need to update
+        sa_key_cols: list[sql.Column]
+        key_vals: list[tuple]
         if len(rowids) > 0:
             sa_key_cols = target.store_tbl.rowid_columns()
             key_vals = rowids
@@ -408,21 +657,23 @@ class Planner:
         updated_cols = batch[0].keys() - target.primary_key_columns()
         recomputed_cols = target.get_dependent_columns(updated_cols) if cascade else set()
         # regardless of cascade, we need to update all indices on any updated column
-        idx_val_cols = target.get_idx_val_columns(updated_cols)
+        modified_base_cols = [c for c in set(updated_cols) | recomputed_cols if c.get_tbl().id == target.id]
+        idx_val_cols = target.get_idx_val_columns(modified_base_cols)
         recomputed_cols.update(idx_val_cols)
         # we only need to recompute stored columns (unstored ones are substituted away)
         recomputed_cols = {c for c in recomputed_cols if c.is_stored}
-        recomputed_base_cols = {col for col in recomputed_cols if col.tbl == target}
+        recomputed_base_cols = {col for col in recomputed_cols if col.get_tbl().id == target.id}
         copied_cols = [
-            col for col in target.cols_by_id.values()
-            if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
+            col
+            for col in target.cols_by_id.values()
+            if col.is_stored and col not in updated_cols and col not in recomputed_base_cols
         ]
-        select_list: list[exprs.Expr] = [exprs.ColumnRef(col) for col in copied_cols]
-        select_list.extend(exprs.ColumnRef(col) for col in updated_cols)
+        select_list: list[exprs.Expr] = [exprs.ColumnRef(col) for col in updated_cols]
 
-        recomputed_exprs = \
-            [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols]
-        # the RowUpdateNode updates columns in-place, ie, in the original ColumnRef; no further sustitution is needed
+        recomputed_exprs = [
+            c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols
+        ]
+        # the RowUpdateNode updates columns in-place, ie, in the original ColumnRef; no further substitution is needed
         select_list.extend(recomputed_exprs)
 
         # ExecNode tree (from bottom to top):
@@ -430,36 +681,54 @@ class Planner:
         # - RowUpdateNode to update the retrieved rows
         # - ExprEvalNode to evaluate the remaining output exprs
         analyzer = Analyzer(FromClause(tbls=[tbl]), select_list)
-        sql_exprs = list(exprs.Expr.list_subexprs(
-            analyzer.all_exprs, filter=analyzer.sql_elements.contains, traverse_matches=False))
-        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], sql_exprs)
+        sql_exprs = list(
+            exprs.Expr.list_subexprs(analyzer.all_exprs, filter=analyzer.sql_elements.contains, traverse_matches=False)
+        )
+        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], sql_exprs, target)
         analyzer.finalize(row_builder)
-        sql_lookup_node = exec.SqlLookupNode(tbl, row_builder, sql_exprs, sa_key_cols, key_vals)
+
+        cell_md_col_refs = cls._cell_md_col_refs(sql_exprs)
+        sql_lookup_node = exec.SqlLookupNode(
+            tbl,
+            row_builder,
+            sql_exprs,
+            columns=copied_cols,
+            sa_key_cols=sa_key_cols,
+            key_vals=key_vals,
+            cell_md_col_refs=cell_md_col_refs,
+        )
         col_vals = [{col: row[col].val for col in updated_cols} for row in batch]
         row_update_node = exec.RowUpdateNode(tbl, key_vals, len(rowids) > 0, col_vals, row_builder, sql_lookup_node)
         plan: exec.ExecNode = row_update_node
         if not cls._is_contained_in(analyzer.select_list, sql_exprs):
             # we need an ExprEvalNode to evaluate the remaining output exprs
             plan = exec.ExprEvalNode(row_builder, analyzer.select_list, sql_exprs, input=plan)
+
         # update row builder with column information
-        all_base_cols = copied_cols + list(updated_cols) + list(recomputed_base_cols)  # same order as select_list
+        evaluated_cols = list(updated_cols) + list(recomputed_base_cols)  # same order as select_list
         row_builder.set_slot_idxs(select_list, remove_duplicates=False)
-        for i, col in enumerate(all_base_cols):
+        plan.row_builder.add_table_columns(copied_cols)
+        for i, col in enumerate(evaluated_cols):
             plan.row_builder.add_table_column(col, select_list[i].slot_idx)
-
-        ctx = exec.ExecContext(row_builder)
-        # we're returning everything to the user, so we might as well do it in a single batch
+        ctx = exec.ExecContext(row_builder, num_computed_exprs=len(recomputed_exprs))
+        # TODO: correct batch size?
         ctx.batch_size = 0
         plan.set_ctx(ctx)
+
+        plan = cls._add_cell_materialization_node(plan)
+        plan = cls._add_save_node(plan)
         recomputed_user_cols = [c for c in recomputed_cols if c.name is not None]
         return (
-            plan, row_update_node, sql_lookup_node.where_clause_element, list(updated_cols) + recomputed_user_cols,
-            recomputed_user_cols
+            plan,
+            row_update_node,
+            sql_lookup_node.where_clause_element,
+            list(updated_cols) + recomputed_user_cols,
+            recomputed_user_cols,
         )
 
     @classmethod
     def create_view_update_plan(
-            cls, view: catalog.TableVersionPath, recompute_targets: list[catalog.Column]
+        cls, view: catalog.TableVersionPath, recompute_targets: list[catalog.Column]
     ) -> exec.ExecNode:
         """Creates a plan to materialize updated rows for a view, given that the base table has been updated.
         The plan:
@@ -477,27 +746,33 @@ class Planner:
             - list of columns that are being recomputed
         """
         assert isinstance(view, catalog.TableVersionPath)
-        assert view.is_view()
-        target = view.tbl_version  # the one we need to update
+        assert view.is_view
+        target = view.tbl_version.get()  # the one we need to update
         # retrieve all stored cols and all target exprs
         recomputed_cols = set(recompute_targets.copy())
-        copied_cols = [col for col in target.cols_by_id.values() if col.is_stored and not col in recomputed_cols]
+        copied_cols = [col for col in target.cols_by_id.values() if col.is_stored and col not in recomputed_cols]
         select_list: list[exprs.Expr] = [exprs.ColumnRef(col) for col in copied_cols]
         # resolve recomputed exprs to stored columns in the base
-        recomputed_exprs = \
-            [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_cols) for c in recomputed_cols]
+        recomputed_exprs = [
+            c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_cols) for c in recomputed_cols
+        ]
         select_list.extend(recomputed_exprs)
 
         # we need to retrieve the PK columns of the existing rows
         plan = cls.create_query_plan(
-            FromClause(tbls=[view]), select_list, where_clause=target.predicate, ignore_errors=True,
-            exact_version_only=view.get_bases())
-        for i, col in enumerate(copied_cols + list(recomputed_cols)):  # same order as select_list
+            FromClause(tbls=[view]),
+            select_list,
+            where_clause=target.predicate,
+            ignore_errors=True,
+            exact_version_only=view.get_bases(),
+        )
+        plan.ctx.num_computed_exprs = len(recomputed_exprs)
+        materialized_cols = copied_cols + list(recomputed_cols)  # same order as select_list
+        for i, col in enumerate(materialized_cols):
             plan.row_builder.add_table_column(col, select_list[i].slot_idx)
-        # TODO: avoid duplication with view_load_plan() logic (where does this belong?)
-        stored_img_col_info = \
-            [info for info in plan.row_builder.output_slot_idxs() if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
+        plan = cls._add_cell_materialization_node(plan)
+        plan = cls._add_save_node(plan)
+
         return plan
 
     @classmethod
@@ -515,45 +790,61 @@ class Planner:
             - number of materialized values per row
         """
         assert isinstance(view, catalog.TableVersionPath)
-        assert view.is_view()
+        assert view.is_view
         # things we need to materialize as DataRows:
         # 1. stored computed cols
         # - iterator columns are effectively computed, just not with a value_expr
         # - we can ignore stored non-computed columns because they have a default value that is supplied directly by
         #   the store
-        target = view.tbl_version  # the one we need to populate
+        target = view.tbl_version.get()  # the one we need to populate
         stored_cols = [c for c in target.cols_by_id.values() if c.is_stored]
         # 2. for component views: iterator args
         iterator_args = [target.iterator_args] if target.iterator_args is not None else []
 
-        row_builder = exprs.RowBuilder(iterator_args, stored_cols, [])
+        from_clause = FromClause(tbls=[view.base])
+        base_analyzer = Analyzer(
+            from_clause, iterator_args, where_clause=target.predicate, sample_clause=target.sample_clause
+        )
+        row_builder = exprs.RowBuilder(base_analyzer.all_exprs, stored_cols, [], target, for_view_load=True)
 
+        # if we're propagating an insert, we only want to see those base rows that were created for the current version
         # execution plan:
         # 1. materialize exprs computed from the base that are needed for stored view columns
         # 2. if it's an iterator view, expand the base rows into component rows
         # 3. materialize stored view columns that haven't been produced by step 1
         base_output_exprs = [e for e in row_builder.default_eval_ctx.exprs if e.is_bound_by([view.base])]
         view_output_exprs = [
-            e for e in row_builder.default_eval_ctx.target_exprs
+            e
+            for e in row_builder.default_eval_ctx.target_exprs
             if e.is_bound_by([view]) and not e.is_bound_by([view.base])
         ]
-        # if we're propagating an insert, we only want to see those base rows that were created for the current version
-        base_analyzer = Analyzer(FromClause(tbls=[view.base]), base_output_exprs, where_clause=target.predicate)
+
+        # Create a new analyzer reflecting exactly what is required from the base table
+        base_analyzer = Analyzer(
+            from_clause, base_output_exprs, where_clause=target.predicate, sample_clause=target.sample_clause
+        )
         base_eval_ctx = row_builder.create_eval_ctx(base_analyzer.all_exprs)
         plan = cls._create_query_plan(
-            row_builder=row_builder, analyzer=base_analyzer, eval_ctx=base_eval_ctx, with_pk=True,
-            exact_version_only=view.get_bases() if propagates_insert else [])
+            row_builder=row_builder,
+            analyzer=base_analyzer,
+            eval_ctx=base_eval_ctx,
+            with_pk=True,
+            exact_version_only=view.get_bases() if propagates_insert else [],
+        )
         exec_ctx = plan.ctx
-        if target.is_component_view():
-            plan = exec.ComponentIterationNode(target, plan)
+        if target.is_component_view:
+            plan = exec.ComponentIterationNode(view.tbl_version, plan)
         if len(view_output_exprs) > 0:
             plan = exec.ExprEvalNode(
-                row_builder, output_exprs=view_output_exprs, input_exprs=base_output_exprs,input=plan)
+                row_builder, output_exprs=view_output_exprs, input_exprs=base_output_exprs, input=plan
+            )
 
-        stored_img_col_info = [info for info in row_builder.output_slot_idxs() if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
         exec_ctx.ignore_errors = True
         plan.set_ctx(exec_ctx)
+        if any(c.col_type.supports_file_offloading() for c in stored_cols):
+            plan = exec.CellMaterializationNode(plan)
+        plan = cls._add_save_node(plan)
+
         return plan, len(row_builder.default_eval_ctx.target_exprs)
 
     @classmethod
@@ -564,8 +855,8 @@ class Planner:
                 raise excs.Error(f'Join predicate {join_clause.join_predicate} not expressible in SQL')
 
     @classmethod
-    def _verify_ordering(cls, analyzer: Analyzer, verify_agg: bool) -> None:
-        """Verify that the various ordering requirements don't conflict"""
+    def _create_combined_ordering(cls, analyzer: Analyzer, verify_agg: bool) -> OrderByClause | None:
+        """Verify that the various ordering requirements don't conflict and return a combined ordering"""
         ob_clauses: list[OrderByClause] = [analyzer.order_by_clause.copy()]
 
         if verify_agg:
@@ -577,13 +868,15 @@ class Planner:
                 ob_clauses.append(ordering)
             for fn_call in analyzer.agg_fn_calls:
                 # agg functions with an ordering requirement are implicitly ascending
-                ordering = (
-                    [OrderByItem(e, None) for e in analyzer.group_by_clause]
-                    + [OrderByItem(e, True) for e in fn_call.get_agg_order_by()]
-                )
+                ordering = [OrderByItem(e, None) for e in analyzer.group_by_clause] + [
+                    OrderByItem(e, True) for e in fn_call.get_agg_order_by()
+                ]
                 ob_clauses.append(ordering)
-        if len(ob_clauses) <= 1:
-            return
+
+        if len(ob_clauses) == 0:
+            return None
+        elif len(ob_clauses) == 1:
+            return ob_clauses[0]
 
         combined_ordering = ob_clauses[0]
         for ordering in ob_clauses[1:]:
@@ -591,60 +884,101 @@ class Planner:
             if combined is None:
                 raise excs.Error(
                     f'Incompatible ordering requirements: '
-                    f'{print_order_by_clause(combined_ordering)} vs {print_order_by_clause(ordering)}')
+                    f'{print_order_by_clause(combined_ordering)} vs {print_order_by_clause(ordering)}'
+                )
             combined_ordering = combined
+        return combined_ordering
+
+    @classmethod
+    def _add_save_node(cls, input_node: exec.ExecNode) -> exec.ExecNode:
+        """Add an ObjectStoreSaveNode, if needed."""
+        media_col_info = input_node.row_builder.media_output_col_info
+        if len(media_col_info) == 0:
+            return input_node
+        else:
+            return exec.ObjectStoreSaveNode(media_col_info, input_node)
 
     @classmethod
     def _is_contained_in(cls, l1: Iterable[exprs.Expr], l2: Iterable[exprs.Expr]) -> bool:
         """Returns True if l1 is contained in l2"""
-        s1, s2 = set(e.id for e in l1), set(e.id for e in l2)
-        return s1 <= s2
+        return {e.id for e in l1} <= {e.id for e in l2}
 
     @classmethod
-    def _insert_prefetch_node(cls, tbl_id: UUID, row_builder: exprs.RowBuilder, input: exec.ExecNode) -> exec.ExecNode:
-        """Returns a CachePrefetchNode into the plan if needed, otherwise returns input"""
+    def _add_prefetch_node(
+        cls, tbl_id: UUID, expressions: Iterable[exprs.Expr], input_node: exec.ExecNode
+    ) -> exec.ExecNode:
+        """Add a CachePrefetch node, if needed."""
         # we prefetch external files for all media ColumnRefs, even those that aren't part of the dependencies
         # of output_exprs: if unstored iterator columns are present, we might need to materialize ColumnRefs that
         # aren't explicitly captured as dependencies
-        media_col_refs = [
-            e for e in list(row_builder.unique_exprs) if isinstance(e, exprs.ColumnRef) and e.col_type.is_media_type()
-        ]
+        media_col_refs = [e for e in expressions if isinstance(e, exprs.ColumnRef) and e.col_type.is_media_type()]
         if len(media_col_refs) == 0:
-            return input
+            return input_node
         # we need to prefetch external files for media column types
         file_col_info = [exprs.ColumnSlotIdx(e.col, e.slot_idx) for e in media_col_refs]
-        prefetch_node = exec.CachePrefetchNode(tbl_id, file_col_info, input)
+        prefetch_node = exec.CachePrefetchNode(tbl_id, file_col_info, input_node)
         return prefetch_node
 
     @classmethod
     def create_query_plan(
-            cls, from_clause: FromClause, select_list: Optional[list[exprs.Expr]] = None,
-            where_clause: Optional[exprs.Expr] = None, group_by_clause: Optional[list[exprs.Expr]] = None,
-            order_by_clause: Optional[list[tuple[exprs.Expr, bool]]] = None, limit: Optional[int] = None,
-            ignore_errors: bool = False, exact_version_only: Optional[list[catalog.TableVersion]] = None
+        cls,
+        from_clause: FromClause,
+        select_list: list[exprs.Expr] | None = None,
+        columns: list[catalog.Column] | None = None,
+        where_clause: exprs.Expr | None = None,
+        group_by_clause: list[exprs.Expr] | None = None,
+        order_by_clause: list[tuple[exprs.Expr, bool]] | None = None,
+        limit: exprs.Expr | None = None,
+        sample_clause: SampleClause | None = None,
+        ignore_errors: bool = False,
+        exact_version_only: list[catalog.TableVersionHandle] | None = None,
     ) -> exec.ExecNode:
-        """Return plan for executing a query.
+        """
+        Return plan for executing a query.
+
+        The plan:
+        - materializes the values of select_list exprs into their respective slots
+        - materializes cell values of 'columns' (and their cellmd, if applicable) into DataRow.cell_vals/cell_md
+
         Updates 'select_list' in place to make it executable.
         TODO: make exact_version_only a flag and use the versions from tbl
         """
         if select_list is None:
             select_list = []
+        if columns is None:
+            columns = []
         if order_by_clause is None:
             order_by_clause = []
         if exact_version_only is None:
             exact_version_only = []
+
         analyzer = Analyzer(
-            from_clause, select_list, where_clause=where_clause, group_by_clause=group_by_clause,
-            order_by_clause=order_by_clause)
-        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], [])
+            from_clause,
+            select_list,
+            where_clause=where_clause,
+            group_by_clause=group_by_clause,
+            order_by_clause=order_by_clause,
+            sample_clause=sample_clause,
+        )
+        # If the from_clause has a single table, we can use it as the context table for the RowBuilder.
+        # Otherwise there is no context table, but that's ok, because the context table is only needed for
+        # table mutations, which can't happen during a join.
+        context_tbl = from_clause.tbls[0].tbl_version.get() if len(from_clause.tbls) == 1 else None
+        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], [], context_tbl)
 
         analyzer.finalize(row_builder)
         # select_list: we need to materialize everything that's been collected
         # with_pk: for now, we always retrieve the PK, because we need it for the file cache
         eval_ctx = row_builder.create_eval_ctx(analyzer.select_list)
         plan = cls._create_query_plan(
-            row_builder=row_builder, analyzer=analyzer, eval_ctx=eval_ctx, limit=limit, with_pk=True,
-            exact_version_only=exact_version_only)
+            row_builder=row_builder,
+            analyzer=analyzer,
+            eval_ctx=eval_ctx,
+            columns=columns,
+            limit=limit,
+            with_pk=True,
+            exact_version_only=exact_version_only,
+        )
         plan.ctx.ignore_errors = ignore_errors
         select_list.clear()
         select_list.extend(analyzer.select_list)
@@ -652,9 +986,14 @@ class Planner:
 
     @classmethod
     def _create_query_plan(
-        cls, row_builder: exprs.RowBuilder, analyzer: Analyzer, eval_ctx: exprs.RowBuilder.EvalCtx,
-        limit: Optional[int] = None, with_pk: bool = False,
-        exact_version_only: Optional[list[catalog.TableVersion]] = None
+        cls,
+        row_builder: exprs.RowBuilder,
+        analyzer: Analyzer,
+        eval_ctx: exprs.RowBuilder.EvalCtx,
+        columns: list[catalog.Column] | None = None,
+        limit: exprs.Expr | None = None,
+        with_pk: bool = False,
+        exact_version_only: list[catalog.TableVersionHandle] | None = None,
     ) -> exec.ExecNode:
         """
         Create plan to materialize eval_ctx.
@@ -664,36 +1003,45 @@ class Planner:
                 in the context of that table version (eg, if 'tbl' is a view, 'plan_target' might be the base)
         TODO: make exact_version_only a flag and use the versions from tbl
         """
+        if columns is None:
+            columns = []
         if exact_version_only is None:
             exact_version_only = []
         sql_elements = analyzer.sql_elements
-        is_python_agg = (
-            not sql_elements.contains_all(analyzer.agg_fn_calls)
-            or not sql_elements.contains_all(analyzer.window_fn_calls)
+        is_python_agg = not sql_elements.contains_all(analyzer.agg_fn_calls) or not sql_elements.contains_all(
+            analyzer.window_fn_calls
         )
         ctx = exec.ExecContext(row_builder)
-        cls._verify_ordering(analyzer, verify_agg=is_python_agg)
+
+        combined_ordering = cls._create_combined_ordering(analyzer, verify_agg=is_python_agg)
         cls._verify_join_clauses(analyzer)
 
         # materialized with SQL table scans (ie, single-table SELECT statements):
         # - select list subexprs that aren't aggregates
         # - join clause subexprs
         # - subexprs of Where clause conjuncts that can't be run in SQL
-        # - all grouping exprs, if any aggregate function call can't be run in SQL (in that case, they all have to be
-        #   run in Python)
-        candidates = list(exprs.Expr.list_subexprs(
-            analyzer.select_list,
-            filter=lambda e: (
+        # - all grouping exprs
+        # - all stratify exprs
+        candidates = list(
+            exprs.Expr.list_subexprs(
+                analyzer.select_list,
+                filter=lambda e: (
                     sql_elements.contains(e)
                     and not e._contains(cls=exprs.FunctionCall, filter=lambda e: bool(e.is_agg_fn_call))
-            ),
-            traverse_matches=False))
+                ),
+                traverse_matches=False,
+            )
+        )
         if analyzer.filter is not None:
-            candidates.extend(exprs.Expr.subexprs(
-                analyzer.filter, filter=lambda e: sql_elements.contains(e), traverse_matches=False))
-        if is_python_agg and analyzer.group_by_clause is not None:
-            candidates.extend(exprs.Expr.list_subexprs(
-                analyzer.group_by_clause, filter=lambda e: sql_elements.contains(e), traverse_matches=False))
+            candidates.extend(
+                exprs.Expr.subexprs(analyzer.filter, filter=sql_elements.contains, traverse_matches=False)
+            )
+        candidates.extend(
+            exprs.Expr.list_subexprs(analyzer.grouping_exprs, filter=sql_elements.contains, traverse_matches=False)
+        )
+        candidates.extend(
+            exprs.Expr.list_subexprs(analyzer.stratify_exprs, filter=sql_elements.contains, traverse_matches=False)
+        )
         # not isinstance(...): we don't want to materialize Literals via a Select
         sql_exprs = exprs.ExprSet(e for e in candidates if not isinstance(e, exprs.Literal))
 
@@ -701,7 +1049,8 @@ class Planner:
         join_exprs = exprs.ExprSet(
             join_clause.join_predicate
             for join_clause in analyzer.from_clause.join_clauses
-            if join_clause.join_predicate is not None)
+            if join_clause.join_predicate is not None
+        )
         scan_target_exprs = sql_exprs | join_exprs
         tbl_scan_plans: list[exec.SqlScanNode] = []
         plan: exec.ExecNode
@@ -711,16 +1060,28 @@ class Planner:
                 exprs.Expr.list_subexprs(
                     scan_target_exprs,
                     filter=lambda e: e.is_bound_by([tbl]) and not isinstance(e, exprs.Literal),
-                    traverse_matches=False))
+                    traverse_matches=False,
+                )
+            )
+
             plan = exec.SqlScanNode(
-                tbl, row_builder, select_list=tbl_scan_exprs,
-                set_pk=with_pk, exact_version_only=exact_version_only)
+                tbl,
+                row_builder,
+                select_list=tbl_scan_exprs,
+                columns=[c for c in columns if c.get_tbl().id == tbl.tbl_id],
+                set_pk=with_pk,
+                cell_md_col_refs=cls._cell_md_col_refs(tbl_scan_exprs),
+                exact_version_only=exact_version_only,
+            )
             tbl_scan_plans.append(plan)
 
         if len(analyzer.from_clause.join_clauses) > 0:
             plan = exec.SqlJoinNode(
-                row_builder, inputs=tbl_scan_plans, join_clauses=analyzer.from_clause.join_clauses,
-                select_list=sql_exprs)
+                row_builder,
+                inputs=tbl_scan_plans,
+                join_clauses=analyzer.from_clause.join_clauses,
+                select_list=sql_exprs,
+            )
         else:
             plan = tbl_scan_plans[0]
 
@@ -732,7 +1093,17 @@ class Planner:
             # we need to order the input for window functions
             plan.set_order_by(analyzer.get_window_fn_ob_clause())
 
-        plan = cls._insert_prefetch_node(tbl.tbl_version.id, row_builder, plan)
+        if analyzer.sample_clause is not None:
+            plan = exec.SqlSampleNode(
+                row_builder,
+                input=plan,
+                select_list=tbl_scan_exprs,
+                sample_clause=analyzer.sample_clause,
+                stratify_exprs=analyzer.stratify_exprs,
+            )
+
+        plan = cls._add_prefetch_node(tbl.tbl_version.id, row_builder.unique_exprs, plan)
+        plan = cls._add_cell_reconstruction_node(analyzer.all_exprs, plan)
 
         if analyzer.group_by_clause is not None:
             # we're doing grouping aggregation; the input of the AggregateNode are the grouping exprs plus the
@@ -750,36 +1121,57 @@ class Planner:
             ctx.batch_size = 16
 
             # do aggregation in SQL if all agg exprs can be translated
-            if (sql_elements.contains_all(analyzer.select_list)
-                    and sql_elements.contains_all(analyzer.grouping_exprs)
-                    and isinstance(plan, exec.SqlNode)
-                    and plan.to_cte() is not None):
+            if (
+                sql_elements.contains_all(analyzer.select_list)
+                and sql_elements.contains_all(analyzer.grouping_exprs)
+                and isinstance(plan, exec.SqlNode)
+                and plan.to_cte() is not None
+            ):
                 plan = exec.SqlAggregationNode(
-                    row_builder, input=plan, select_list=analyzer.select_list, group_by_items=analyzer.group_by_clause)
+                    row_builder, input=plan, select_list=analyzer.select_list, group_by_items=analyzer.group_by_clause
+                )
             else:
+                input_sql_node = plan.get_node(exec.SqlNode)
+                assert combined_ordering is not None
+                input_sql_node.set_order_by(combined_ordering)
                 plan = exec.AggregationNode(
-                    tbl.tbl_version, row_builder, analyzer.group_by_clause,
-                    analyzer.agg_fn_calls + analyzer.window_fn_calls, agg_input, input=plan)
+                    tbl.tbl_version,
+                    row_builder,
+                    analyzer.group_by_clause,
+                    analyzer.agg_fn_calls + analyzer.window_fn_calls,
+                    agg_input,
+                    input=plan,
+                )
                 typecheck_dummy = analyzer.grouping_exprs + analyzer.agg_fn_calls + analyzer.window_fn_calls
                 agg_output = exprs.ExprSet(typecheck_dummy)
                 if not agg_output.issuperset(exprs.ExprSet(eval_ctx.target_exprs)):
                     # we need an ExprEvalNode to evaluate the remaining output exprs
                     plan = exec.ExprEvalNode(row_builder, eval_ctx.target_exprs, agg_output, input=plan)
+                plan = cls._add_save_node(plan)
         else:
             if not exprs.ExprSet(sql_exprs).issuperset(exprs.ExprSet(eval_ctx.target_exprs)):
                 # we need an ExprEvalNode to evaluate the remaining output exprs
                 plan = exec.ExprEvalNode(row_builder, eval_ctx.target_exprs, sql_exprs, input=plan)
             # we're returning everything to the user, so we might as well do it in a single batch
+            # TODO: return smaller batches in order to increase inter-ExecNode parallelism
             ctx.batch_size = 0
 
+        sql_node = plan.get_node(exec.SqlNode)
         if len(analyzer.order_by_clause) > 0:
             # we have the last SqlNode we created produce the ordering
-            sql_node = plan.get_node(exec.SqlNode)
             assert sql_node is not None
             sql_node.set_order_by(analyzer.order_by_clause)
 
+        # if we don't need an ordered result, tell the ExprEvalNode not to maintain input order (which allows us to
+        # return batches earlier)
+        if sql_node is not None and len(sql_node.order_by_clause) == 0:
+            expr_eval_node = plan.get_node(exec.ExprEvalNode)
+            if expr_eval_node is not None:
+                expr_eval_node.set_input_order(False)
+
         if limit is not None:
-            plan.set_limit(limit)
+            assert isinstance(limit, exprs.Literal)
+            plan.set_limit(limit.val)
 
         plan.set_ctx(ctx)
         return plan
@@ -789,25 +1181,24 @@ class Planner:
         return Analyzer(FromClause(tbls=[tbl]), [], where_clause=where_clause)
 
     @classmethod
-    def create_add_column_plan(
-            cls, tbl: catalog.TableVersionPath, col: catalog.Column
-    ) -> tuple[exec.ExecNode, Optional[int]]:
+    def create_add_column_plan(cls, tbl: catalog.TableVersionPath, col: catalog.Column) -> exec.ExecNode:
         """Creates a plan for InsertableTable.add_column()
         Returns:
             plan: the plan to execute
             value_expr slot idx for the plan output (for computed cols)
         """
         assert isinstance(tbl, catalog.TableVersionPath)
-        row_builder = exprs.RowBuilder(output_exprs=[], columns=[col], input_exprs=[])
+        row_builder = exprs.RowBuilder(output_exprs=[], columns=[col], input_exprs=[], tbl=tbl.tbl_version.get())
         analyzer = Analyzer(FromClause(tbls=[tbl]), row_builder.default_eval_ctx.target_exprs)
         plan = cls._create_query_plan(
-            row_builder=row_builder, analyzer=analyzer, eval_ctx=row_builder.default_eval_ctx, with_pk=True)
+            row_builder=row_builder, analyzer=analyzer, eval_ctx=row_builder.default_eval_ctx, with_pk=True
+        )
+
         plan.ctx.batch_size = 16
         plan.ctx.show_pbar = True
         plan.ctx.ignore_errors = True
+        computed_exprs = row_builder.output_exprs - row_builder.input_exprs
+        plan.ctx.num_computed_exprs = len(computed_exprs)  # we are adding a computed column, so we need to evaluate it
+        plan = cls._add_save_node(plan)
 
-        # we want to flush images
-        if col.is_computed and col.is_stored and col.col_type.is_image_type():
-            plan.set_stored_img_cols(row_builder.output_slot_idxs())
-        value_expr_slot_idx = row_builder.output_slot_idxs()[0].slot_idx if col.is_computed else None
-        return plan, value_expr_slot_idx
+        return plan