pixeltable 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

pixeltable/plan.py

@@ -1,8 +1,9 @@
-from typing import Tuple, Optional, List, Set, Any, Dict
+from typing import Any, Iterable, List, Optional, Sequence, Set, Tuple
 from uuid import UUID
 
 import sqlalchemy as sql
 
+import pixeltable as pxt
 import pixeltable.exec as exec
 import pixeltable.func as func
 from pixeltable import catalog
@@ -39,7 +40,7 @@ class Analyzer:
     """Class to perform semantic analysis of a query and to store the analysis state"""
 
     def __init__(
-            self, tbl: catalog.TableVersionPath, select_list: List[exprs.Expr],
+            self, tbl: catalog.TableVersionPath, 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):
         if group_by_clause is None:
@@ -68,7 +69,7 @@ class Analyzer:
         # all exprs that are evaluated in Python; not executable
         self.all_exprs = self.select_list.copy()
         self.all_exprs.extend(self.group_by_clause)
-        self.all_exprs.extend([e for e, _ in self.order_by_clause])
+        self.all_exprs.extend(e for e, _ in self.order_by_clause)
         if self.filter is not None:
             self.all_exprs.append(self.filter)
         self.sql_exprs = list(exprs.Expr.list_subexprs(
@@ -84,7 +85,7 @@ class Analyzer:
 
     def _analyze_agg(self) -> None:
         """Check semantic correctness of aggregation and fill in agg-specific fields of Analyzer"""
-        self.agg_fn_calls = [e for e in self.all_exprs if _is_agg_fn_call(e)]
+        self.agg_fn_calls = [e for e in self.all_exprs if isinstance(e, exprs.FunctionCall) and _is_agg_fn_call(e)]
         if len(self.agg_fn_calls) == 0:
             # nothing to do
             return
@@ -98,7 +99,7 @@ class Analyzer:
 
         # check that filter doesn't contain aggregates
         if self.filter is not None:
-            agg_fn_calls = [e for e in self.filter.subexprs(filter=lambda e: _is_agg_fn_call(e))]
+            agg_fn_calls = [e for e in self.filter.subexprs(expr_class=exprs.FunctionCall, filter=lambda e: _is_agg_fn_call(e))]
             if len(agg_fn_calls) > 0:
                 raise excs.Error(f'Filter cannot contain aggregate functions: {self.filter}')
 
@@ -111,7 +112,7 @@ class Analyzer:
                 raise excs.Error(f'Grouping expression contains aggregate function: {e}')
 
         # check that agg fn calls don't have contradicting ordering requirements
-        order_by: List[exprs.Exprs] = []
+        order_by: list[exprs.Expr] = []
         order_by_origin: Optional[exprs.Expr] = None  # the expr that determines the ordering
         for agg_fn_call in self.agg_fn_calls:
             fn_call_order_by = agg_fn_call.get_agg_order_by()
@@ -185,7 +186,7 @@ class Planner:
     def create_count_stmt(
             cls, tbl: catalog.TableVersionPath, where_clause: Optional[exprs.Expr] = None
     ) -> sql.Select:
-        stmt = sql.select(sql.func.count('*'))
+        stmt = sql.select(sql.func.count())
         refd_tbl_ids: Set[UUID] = set()
         if where_clause is not None:
             analyzer = cls.analyze(tbl, where_clause)
@@ -200,7 +201,7 @@ class Planner:
 
     @classmethod
     def create_insert_plan(
-            cls, tbl: catalog.TableVersion, rows: List[Dict[str, Any]], ignore_errors: bool
+        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()
@@ -214,12 +215,12 @@ class Planner:
         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()]
         input_col_info = [info for info in stored_col_info if not info.col.is_computed]
-        plan = exec.InMemoryDataNode(tbl, rows, row_builder, tbl.next_rowid)
+        plan: exec.ExecNode = exec.InMemoryDataNode(tbl, rows, row_builder, tbl.next_rowid)
 
         media_input_cols = [info for info in input_col_info if info.col.col_type.is_media_type()]
         if len(media_input_cols) > 0:
             # prefetch external files for all input column refs for validation
-            plan = exec.CachePrefetchNode(tbl.id, media_input_cols, plan)
+            plan = exec.CachePrefetchNode(tbl.id, media_input_cols, input=plan)
             plan = exec.MediaValidationNode(row_builder, media_input_cols, input=plan)
 
         computed_exprs = [e for e in row_builder.default_eval_ctx.target_exprs if not isinstance(e, exprs.ColumnRef)]
@@ -234,6 +235,34 @@ class Planner:
                 ignore_errors=ignore_errors))
         return plan
 
+    @classmethod
+    def create_df_insert_plan(
+        cls,
+        tbl: catalog.TableVersion,
+        df: 'pxt.DataFrame',
+        ignore_errors: bool
+    ) -> exec.ExecNode:
+        assert not tbl.is_view()
+        plan = df._create_query_plan()  # ExecNode constructed by the DataFrame
+
+        # Modify the plan RowBuilder to register the output columns
+        for col_name, expr in zip(df.schema.keys(), df._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)
+
+        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)
+
+        plan.set_ctx(
+            exec.ExecContext(
+                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,
@@ -258,7 +287,7 @@ class Planner:
         target = tbl.tbl_version  # the one we need to update
         updated_cols = list(update_targets.keys())
         if len(recompute_targets) > 0:
-            recomputed_cols = recompute_targets.copy()
+            recomputed_cols = set(recompute_targets)
         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
@@ -270,13 +299,13 @@ class Planner:
         copied_cols = [
             col for col in target.cols if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
         ]
-        select_list = [exprs.ColumnRef(col) for col in copied_cols]
+        select_list: list[exprs.Expr] = [exprs.ColumnRef(col) for col in copied_cols]
         select_list.extend(update_targets.values())
 
         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 = {exprs.ColumnRef(col): e for col, e in update_targets.items()}
+        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)
 
@@ -284,16 +313,17 @@ class Planner:
         plan = cls.create_query_plan(tbl, select_list, where_clause=where_clause, with_pk=True, ignore_errors=True)
         all_base_cols = copied_cols + updated_cols + list(recomputed_base_cols)  # same order as select_list
         # update row builder with column information
-        [plan.row_builder.add_table_column(col, select_list[i].slot_idx) for i, col in enumerate(all_base_cols)]
+        for i, col in enumerate(all_base_cols):
+            plan.row_builder.add_table_column(col, select_list[i].slot_idx)
         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
 
     @classmethod
     def create_batch_update_plan(
-            cls, tbl: catalog.TableVersionPath,
-            batch: list[dict[catalog.Column, exprs.Expr]], rowids: list[tuple[int, ...]],
-            cascade: bool
-    ) -> Tuple[exec.ExecNode, exec.RowUpdateNode, sql.ClauseElement, List[catalog.Column], List[catalog.Column]]:
+        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:
         - root node of the plan to produce the updated rows
@@ -327,7 +357,7 @@ class Planner:
             col for col in target.cols if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
         ]
         select_list = [exprs.ColumnRef(col) for col in copied_cols]
-        select_list.extend([exprs.ColumnRef(col) for col in updated_cols])
+        select_list.extend(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]
@@ -341,10 +371,11 @@ class Planner:
         analyzer = Analyzer(tbl, select_list)
         row_builder = exprs.RowBuilder(analyzer.all_exprs, [], analyzer.sql_exprs)
         analyzer.finalize(row_builder)
-        plan = exec.SqlLookupNode(tbl, row_builder, analyzer.sql_exprs, sa_key_cols, key_vals)
-        delete_where_clause = plan.where_clause
+        sql_lookup_node = exec.SqlLookupNode(tbl, row_builder, analyzer.sql_exprs, sa_key_cols, key_vals)
+        delete_where_clause = sql_lookup_node.where_clause
         col_vals = [{col: row[col].val for col in updated_cols} for row in batch]
-        plan = row_update_node = exec.RowUpdateNode(tbl, key_vals, len(rowids) > 0, col_vals, row_builder, plan)
+        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, analyzer.sql_exprs):
             # we need an ExprEvalNode to evaluate the remaining output exprs
             plan = exec.ExprEvalNode(row_builder, analyzer.select_list, analyzer.sql_exprs, input=plan)
@@ -388,7 +419,7 @@ class Planner:
         # retrieve all stored cols and all target exprs
         recomputed_cols = set(recompute_targets.copy())
         copied_cols = [col for col in target.cols if col.is_stored and not col in recomputed_cols]
-        select_list = [exprs.ColumnRef(col) for col in copied_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]
@@ -398,10 +429,8 @@ class Planner:
         plan = cls.create_query_plan(
             view, select_list, where_clause=target.predicate, with_pk=True, ignore_errors=True,
             exact_version_only=view.get_bases())
-        [
+        for i, col in enumerate(copied_cols + list(recomputed_cols)):  # same order as select_list
             plan.row_builder.add_table_column(col, select_list[i].slot_idx)
-            for i, col in enumerate(copied_cols + list(recomputed_cols))  # same order as select_list
-        ]
         # 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()]
@@ -532,7 +561,7 @@ class Planner:
                 return False
             tbl = e.col.tbl
             return tbl.is_component_view() and tbl.is_iterator_column(e.col) and not e.col.is_stored
-        unstored_iter_col_refs = list(exprs.Expr.list_subexprs(analyzer.all_exprs, filter=refs_unstored_iter_col))
+        unstored_iter_col_refs = list(exprs.Expr.list_subexprs(analyzer.all_exprs, expr_class=exprs.ColumnRef, filter=refs_unstored_iter_col))
         if len(unstored_iter_col_refs) > 0 and len(order_by_items) == 0:
             # we don't already have a user-requested ordering and we access unstored iterator columns:
             # order by the primary key of the component view, which minimizes the number of iterator instantiations
@@ -554,9 +583,9 @@ class Planner:
         return order_by_items
 
     @classmethod
-    def _is_contained_in(cls, l1: List[exprs.Expr], l2: List[exprs.Expr]) -> bool:
+    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])
+        s1, s2 = set(e.id for e in l1), set(e.id for e in l2)
         return s1 <= s2
 
     @classmethod

pixeltable/store.py

@@ -7,7 +7,7 @@ import sys
 import urllib.parse
 import urllib.request
 import warnings
-from typing import Optional, Dict, Any, List, Tuple, Set
+from typing import Optional, Dict, Any, List, Tuple, Set, Union
 
 import sqlalchemy as sql
 from tqdm import tqdm, TqdmWarning
@@ -15,10 +15,8 @@ from tqdm import tqdm, TqdmWarning
 import pixeltable.catalog as catalog
 import pixeltable.env as env
 from pixeltable import exprs
-import pixeltable.exceptions as excs
 from pixeltable.exec import ExecNode
 from pixeltable.metadata import schema
-from pixeltable.type_system import StringType
 from pixeltable.utils.media_store import MediaStore
 from pixeltable.utils.sql import log_stmt, log_explain
 
@@ -34,10 +32,15 @@ class StoreBase:
     - v_max: version at which the row was deleted (or MAX_VERSION if it's still live)
     """
 
+    __INSERT_BATCH_SIZE = 1000
+
     def __init__(self, tbl_version: catalog.TableVersion):
         self.tbl_version = tbl_version
         self.sa_md = sql.MetaData()
         self.sa_tbl: Optional[sql.Table] = None
+        # We need to declare a `base` variable here, even though it's only defined for instances of `StoreView`,
+        # since it's referenced by various methods of `StoreBase`
+        self.base = None if tbl_version.base is None else tbl_version.base.store_tbl
         self.create_sa_tbl()
 
     def pk_columns(self) -> List[sql.Column]:
@@ -49,7 +52,6 @@ class StoreBase:
     @abc.abstractmethod
     def _create_rowid_columns(self) -> List[sql.Column]:
         """Create and return rowid columns"""
-        pass
 
     @abc.abstractmethod
     def _create_system_columns(self) -> List[sql.Column]:
@@ -61,7 +63,6 @@ class StoreBase:
         self._pk_columns = [*rowid_cols, self.v_min_col]
         return [*rowid_cols, self.v_min_col, self.v_max_col]
 
-
     def create_sa_tbl(self) -> None:
         """Create self.sa_tbl from self.tbl_version."""
         system_cols = self._create_system_columns()
@@ -96,14 +97,12 @@ class StoreBase:
         self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
 
     @abc.abstractmethod
-    def _rowid_join_predicate(self) -> sql.ClauseElement:
+    def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         """Return predicate for rowid joins to all bases"""
-        pass
 
     @abc.abstractmethod
     def _storage_name(self) -> str:
         """Return the name of the data store table"""
-        pass
 
     def _move_tmp_media_file(self, file_url: Optional[str], col: catalog.Column, v_min: int) -> str:
         """Move tmp media file with given url to Env.media_dir and return new url, or given url if not a tmp_dir file"""
@@ -158,10 +157,12 @@ class StoreBase:
 
     def count(self, conn: Optional[sql.engine.Connection] = None) -> int:
         """Return the number of rows visible in self.tbl_version"""
-        stmt = sql.select(sql.func.count('*'))\
-            .select_from(self.sa_tbl)\
-            .where(self.v_min_col <= self.tbl_version.version)\
+        stmt = (
+            sql.select(sql.func.count('*'))  # type: ignore
+            .select_from(self.sa_tbl)
+            .where(self.v_min_col <= self.tbl_version.version)
             .where(self.v_max_col > self.tbl_version.version)
+        )
         if conn is None:
             with env.Env.get().engine.connect() as conn:
                 result = conn.execute(stmt).scalar_one()
@@ -191,12 +192,12 @@ class StoreBase:
         added_storage_cols = [col.store_name()]
         if col.records_errors:
             # we also need to create the errormsg and errortype storage cols
-            stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL')
-            conn.execute(sql.text(stmt))
-            stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL')
-            conn.execute(sql.text(stmt))
+            stmt = sql.text(f'ALTER TABLE {self._storage_name()} '
+                            f'ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL')
+            conn.execute(stmt)
+            stmt = sql.text(f'ALTER TABLE {self._storage_name()} '
+                            f'ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL')
+            conn.execute(stmt)
             added_storage_cols.extend([col.errormsg_store_name(), col.errortype_store_name()])
         self.create_sa_tbl()
         _logger.info(f'Added columns {added_storage_cols} to storage table {self._storage_name()}')
@@ -223,35 +224,70 @@ class StoreBase:
         """
         num_excs = 0
         num_rows = 0
-        for row_batch in exec_plan:
-            num_rows += len(row_batch)
-            for result_row in row_batch:
-                values_dict: Dict[sql.Column, Any] = {}
-
-                if col.is_computed:
-                    if result_row.has_exc(value_expr_slot_idx):
-                        num_excs += 1
-                        value_exc = result_row.get_exc(value_expr_slot_idx)
-                        # we store a NULL value and record the exception/exc type
-                        error_type = type(value_exc).__name__
-                        error_msg = str(value_exc)
-                        values_dict = {
-                            col.sa_col: None,
-                            col.sa_errortype_col: error_type,
-                            col.sa_errormsg_col: error_msg
-                        }
-                    else:
-                        val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
-                        if col.col_type.is_media_type():
-                            val = self._move_tmp_media_file(val, col, result_row.pk[-1])
-                        values_dict = {col.sa_col: val}
-
-                update_stmt = sql.update(self.sa_tbl).values(values_dict)
-                for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
-                    update_stmt = update_stmt.where(pk_col == pk_val)
-                log_stmt(_logger, update_stmt)
-                conn.execute(update_stmt)
 
+        # create temp table to store output of exec_plan, with the same primary key as the store table
+        tmp_name = f'temp_{self._storage_name()}'
+        tmp_pk_cols = [sql.Column(col.name, col.type, primary_key=True) for col in self.pk_columns()]
+        tmp_cols = tmp_pk_cols.copy()
+        tmp_val_col = sql.Column(col.sa_col.name, col.sa_col.type)
+        tmp_cols.append(tmp_val_col)
+        # add error columns if the store column records errors
+        if col.records_errors:
+            tmp_errortype_col = sql.Column(col.sa_errortype_col.name, col.sa_errortype_col.type)
+            tmp_cols.append(tmp_errortype_col)
+            tmp_errormsg_col = sql.Column(col.sa_errormsg_col.name, col.sa_errormsg_col.type)
+            tmp_cols.append(tmp_errormsg_col)
+        tmp_tbl = sql.Table(tmp_name, self.sa_md, *tmp_cols, prefixes=['TEMPORARY'])
+        tmp_tbl.create(bind=conn)
+
+        try:
+            # insert rows from exec_plan into temp table
+            for row_batch in exec_plan:
+                num_rows += len(row_batch)
+                tbl_rows: list[dict[str, Any]] = []
+                for result_row in row_batch:
+                    tbl_row: dict[str, Any] = {}
+                    for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
+                        tbl_row[pk_col.name] = pk_val
+
+                    if col.is_computed:
+                        if result_row.has_exc(value_expr_slot_idx):
+                            num_excs += 1
+                            value_exc = result_row.get_exc(value_expr_slot_idx)
+                            # we store a NULL value and record the exception/exc type
+                            error_type = type(value_exc).__name__
+                            error_msg = str(value_exc)
+                            tbl_row[col.sa_col.name] = None
+                            tbl_row[col.sa_errortype_col.name] = error_type
+                            tbl_row[col.sa_errormsg_col.name] = error_msg
+                        else:
+                            val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
+                            if col.col_type.is_media_type():
+                                val = self._move_tmp_media_file(val, col, result_row.pk[-1])
+                            tbl_row[col.sa_col.name] = val
+                            if col.records_errors:
+                                tbl_row[col.sa_errortype_col.name] = None
+                                tbl_row[col.sa_errormsg_col.name] = None
+
+                    tbl_rows.append(tbl_row)
+                conn.execute(sql.insert(tmp_tbl), tbl_rows)
+
+            # update store table with values from temp table
+            update_stmt = sql.update(self.sa_tbl)
+            for pk_col, tmp_pk_col in zip(self.pk_columns(), tmp_pk_cols):
+                update_stmt = update_stmt.where(pk_col == tmp_pk_col)
+            update_stmt = update_stmt.values({col.sa_col: tmp_val_col})
+            if col.records_errors:
+                update_stmt = update_stmt.values({
+                    col.sa_errortype_col: tmp_errortype_col,
+                    col.sa_errormsg_col: tmp_errormsg_col
+                })
+            log_explain(_logger, update_stmt, conn)
+            conn.execute(update_stmt)
+
+        finally:
+            tmp_tbl.drop(bind=conn)
+            self.sa_md.remove(tmp_tbl)
         return num_excs
 
     def insert_rows(
@@ -264,7 +300,6 @@ class StoreBase:
         """
         assert v_min is not None
         exec_plan.ctx.set_conn(conn)
-        batch_size = 16  # TODO: is this a good batch size?
         # TODO: total?
         num_excs = 0
         num_rows = 0
@@ -276,10 +311,10 @@ class StoreBase:
             exec_plan.open()
             for row_batch in exec_plan:
                 num_rows += len(row_batch)
-                for batch_start_idx in range(0, len(row_batch), batch_size):
+                for batch_start_idx in range(0, len(row_batch), self.__INSERT_BATCH_SIZE):
                     # compute batch of rows and convert them into table rows
                     table_rows: List[Dict[str, Any]] = []
-                    for row_idx in range(batch_start_idx, min(batch_start_idx + batch_size, len(row_batch))):
+                    for row_idx in range(batch_start_idx, min(batch_start_idx + self.__INSERT_BATCH_SIZE, len(row_batch))):
                         row = row_batch[row_idx]
                         table_row, num_row_exc = \
                             self._create_table_row(row, row_builder, media_cols, cols_with_excs, v_min=v_min)
@@ -295,6 +330,8 @@ class StoreBase:
                                     file=sys.stdout
                                 )
                             progress_bar.update(1)
+
+                    # insert batch of rows
                     self._move_tmp_media_files(table_rows, media_cols, v_min)
                     conn.execute(sql.insert(self.sa_tbl), table_rows)
             if progress_bar is not None:
@@ -303,7 +340,7 @@ class StoreBase:
         finally:
             exec_plan.close()
 
-    def _versions_clause(self, versions: List[Optional[int]], match_on_vmin: bool) -> sql.ClauseElement:
+    def _versions_clause(self, versions: list[Optional[int]], match_on_vmin: bool) -> sql.ColumnElement[bool]:
         """Return filter for base versions"""
         v = versions[0]
         if v is None:
@@ -318,7 +355,7 @@ class StoreBase:
 
     def delete_rows(
             self, current_version: int, base_versions: List[Optional[int]], match_on_vmin: bool,
-            where_clause: Optional[sql.ClauseElement], conn: sql.engine.Connection) -> int:
+            where_clause: Optional[sql.ColumnElement[bool]], conn: sql.engine.Connection) -> int:
         """Mark rows as deleted that are live and were created prior to current_version.
         Also: populate the undo columns
         Args:
@@ -338,17 +375,19 @@ class StoreBase:
         rowid_join_clause = self._rowid_join_predicate()
         base_versions_clause = sql.true() if len(base_versions) == 0 \
             else self.base._versions_clause(base_versions, match_on_vmin)
-        set_clause = {self.v_max_col: current_version}
+        set_clause: dict[sql.Column, Union[int, sql.Column]] = {self.v_max_col: current_version}
         for index_info in self.tbl_version.idxs_by_name.values():
             # copy value column to undo column
             set_clause[index_info.undo_col.sa_col] = index_info.val_col.sa_col
             # set value column to NULL
             set_clause[index_info.val_col.sa_col] = None
-        stmt = sql.update(self.sa_tbl) \
-            .values(set_clause) \
-            .where(where_clause) \
-            .where(rowid_join_clause) \
+        stmt = (
+            sql.update(self.sa_tbl)
+            .values(set_clause)
+            .where(where_clause)
+            .where(rowid_join_clause)
             .where(base_versions_clause)
+        )
         log_explain(_logger, stmt, conn)
         status = conn.execute(stmt)
         return status.rowcount
@@ -366,14 +405,13 @@ class StoreTable(StoreBase):
     def _storage_name(self) -> str:
         return f'tbl_{self.tbl_version.id.hex}'
 
-    def _rowid_join_predicate(self) -> sql.ClauseElement:
+    def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.true()
 
 
 class StoreView(StoreBase):
     def __init__(self, catalog_view: catalog.TableVersion):
         assert catalog_view.is_view()
-        self.base = catalog_view.base.store_tbl
         super().__init__(catalog_view)
 
     def _create_rowid_columns(self) -> List[sql.Column]:
@@ -384,7 +422,7 @@ class StoreView(StoreBase):
     def _storage_name(self) -> str:
         return f'view_{self.tbl_version.id.hex}'
 
-    def _rowid_join_predicate(self) -> sql.ClauseElement:
+    def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.and_(
             self.base._rowid_join_predicate(),
             *[c1 == c2 for c1, c2 in zip(self.rowid_columns(), self.base.rowid_columns())])
@@ -411,7 +449,7 @@ class StoreComponentView(StoreView):
         # we need to fix up the 'pos' column in TableVersion
         self.tbl_version.cols_by_name['pos'].sa_col = self.pos_col
 
-    def _rowid_join_predicate(self) -> sql.ClauseElement:
+    def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.and_(
             self.base._rowid_join_predicate(),
             *[c1 == c2 for c1, c2 in zip(self.rowid_columns()[:-1], self.base.rowid_columns())])

pixeltable/tool/create_test_db_dump.py

@@ -6,7 +6,7 @@ import pathlib
 import subprocess
 from typing import Any
 
-import pgserver
+import pixeltable_pgserver
 import toml
 
 import pixeltable as pxt
@@ -41,22 +41,27 @@ class Dumper:
         md_version = metadata.VERSION
         dump_file = self.output_dir / f'pixeltable-v{md_version:03d}-test.dump.gz'
         _logger.info(f'Creating database dump at: {dump_file}')
-        pg_package_dir = os.path.dirname(pgserver.__file__)
+        pg_package_dir = os.path.dirname(pixeltable_pgserver.__file__)
         pg_dump_binary = f'{pg_package_dir}/pginstall/bin/pg_dump'
         _logger.info(f'Using pg_dump binary at: {pg_dump_binary}')
+        # We need the raw DB URL, without a driver qualifier.  (The driver qualifier is needed by
+        # SQLAlchemy, but command-line Postgres won't know how to interpret it.)
+        db_url = Env.get()._db_server.get_uri(Env.get()._db_name)
         with open(dump_file, 'wb') as dump:
             pg_dump_process = subprocess.Popen(
-                [pg_dump_binary, Env.get().db_url, '-U', 'postgres', '-Fc'],
+                (pg_dump_binary, db_url, '-U', 'postgres', '-Fc'),
                 stdout=subprocess.PIPE
             )
             subprocess.run(
-                ["gzip", "-9"],
+                ('gzip', '-9'),
                 stdin=pg_dump_process.stdout,
                 stdout=dump,
                 check=True
             )
+            if pg_dump_process.poll() != 0:
+                raise RuntimeError(f'pg_dump failed with return code {pg_dump_process.returncode}')
         info_file = self.output_dir / f'pixeltable-v{md_version:03d}-test-info.toml'
-        git_sha = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip()
+        git_sha = subprocess.check_output(('git', 'rev-parse', 'HEAD')).decode('ascii').strip()
         user = os.environ.get('USER', os.environ.get('USERNAME'))
         info_dict = {'pixeltable-dump': {
             'metadata-version': md_version,
@@ -177,8 +182,8 @@ class Dumper:
         assert t.base_table_image_rot.col in project.stored_proxies
 
     def __add_expr_columns(self, t: pxt.Table, col_prefix: str, include_expensive_functions=False) -> None:
-        def add_column(col_name: str, col_expr: Any) -> None:
-            t.add_column(**{f'{col_prefix}_{col_name}': col_expr})
+        def add_column(col_name: str, col_expr: Any, stored: bool = True) -> None:
+            t.add_column(**{f'{col_prefix}_{col_name}': col_expr}, stored=stored)
 
         # arithmetic_expr
         add_column('plus', t.c2 + 6)
@@ -187,9 +192,6 @@ class Dumper:
         add_column('div', t.c3 / 1.7)
         add_column('mod', t.c2 % 11)
 
-        # array_slice
-        add_column('array_slice_1', t.c6[5])
-
         # column_property_ref
         add_column('fileurl', t.c8.fileurl)
         add_column('localpath', t.c8.localpath)
@@ -217,7 +219,7 @@ class Dumper:
 
         # image_member_access
         add_column('image_mode', t.c8.mode)
-        add_column('image_rot', t.c8.rotate(180))
+        add_column('image_rot', t.c8.rotate(180), stored=False)
 
         # in_predicate
         add_column('isin_1', t.c1.isin(['test string 1', 'test string 2', 'test string 3']))
@@ -237,6 +239,10 @@ class Dumper:
         # json_mapper and json_path
         add_column('json_mapper', t.c6[3])
         add_column('json_path', t.c6.f1)
+        add_column('json_path_nested', t.c6.f6.f7)
+        add_column('json_path_star', t.c6.f5['*'])
+        add_column('json_path_idx', t.c6.f5[3])
+        add_column('json_path_slice', t.c6.f5[1:3:2])
 
         # literal
         add_column('str_const', 'str')