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

pixeltable/__init__.py

@@ -21,7 +21,7 @@ from .type_system import (
 )
 from .utils.help import help
 
-from . import functions, io, iterators
+from . import ext, functions, io, iterators
 from .__version__ import __version__, __version_tuple__
 
 # This is the safest / most maintainable way to do this: start with the default and "blacklist" stuff that

pixeltable/__version__.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced during build.
-__version__ = "0.2.12"
-__version_tuple__ = (0, 2, 12)
+__version__ = "0.2.14"
+__version_tuple__ = (0, 2, 14)

pixeltable/catalog/column.py

@@ -152,6 +152,11 @@ class Column:
             return self._records_errors
         return self.is_stored and (self.is_computed or self.col_type.is_media_type())
 
+    @property
+    def qualified_name(self) -> str:
+        assert self.tbl is not None
+        return f'{self.tbl.name}.{self.name}'
+
     def source(self) -> None:
         """
         If this is a computed col and the top-level expr is a function call, print the source, if possible.

pixeltable/catalog/globals.py

@@ -19,6 +19,14 @@ class UpdateStatus:
     updated_cols: List[str] = dataclasses.field(default_factory=list)
     cols_with_excs: List[str] = dataclasses.field(default_factory=list)
 
+    def __iadd__(self, other: 'UpdateStatus') -> 'UpdateStatus':
+        self.num_rows += other.num_rows
+        self.num_computed_values += other.num_computed_values
+        self.num_excs += other.num_excs
+        self.updated_cols = list(dict.fromkeys(self.updated_cols + other.updated_cols))
+        self.cols_with_excs = list(dict.fromkeys(self.cols_with_excs + other.cols_with_excs))
+        return self
+
 def is_valid_identifier(name: str) -> bool:
     return name.isidentifier() and not name.startswith('_')
 

pixeltable/catalog/insertable_table.py

@@ -129,11 +129,11 @@ class InsertableTable(Table):
                     msg = str(e)
                     raise excs.Error(f'Error in column {col.name}: {msg[0].lower() + msg[1:]}\nRow: {row}')
 
-    def delete(self, where: Optional['pixeltable.exprs.Predicate'] = None) -> UpdateStatus:
+    def delete(self, where: Optional['pixeltable.exprs.Expr'] = None) -> UpdateStatus:
         """Delete rows in this table.
 
         Args:
-            where: a Predicate to filter rows to delete.
+            where: a predicate to filter rows to delete.
 
         Examples:
             Delete all rows in a table:

pixeltable/catalog/table.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import logging
 from pathlib import Path
-from typing import Union, Any, Optional, Callable, Set, Tuple, Iterable, overload, Type
+from typing import Union, Any, Optional, Callable, Set, Tuple, Iterable, overload, Type, Literal
 from uuid import UUID
 import abc
 
@@ -113,7 +113,7 @@ class Table(SchemaObject):
         from pixeltable.dataframe import DataFrame
         return DataFrame(self._tbl_version_path).select(*items, **named_items)
 
-    def where(self, pred: 'exprs.Predicate') -> 'pixeltable.dataframe.DataFrame':
+    def where(self, pred: 'exprs.Expr') -> 'pixeltable.dataframe.DataFrame':
         """Return a DataFrame for this table.
         """
         # local import: avoid circular imports
@@ -716,13 +716,13 @@ class Table(SchemaObject):
         raise NotImplementedError
 
     def update(
-            self, value_spec: dict[str, Any], where: Optional['pixeltable.exprs.Predicate'] = None, cascade: bool = True
+            self, value_spec: dict[str, Any], where: Optional['pixeltable.exprs.Expr'] = None, cascade: bool = True
     ) -> UpdateStatus:
         """Update rows in this table.
 
         Args:
             value_spec: a dictionary mapping column names to literal values or Pixeltable expressions.
-            where: a Predicate to filter rows to update.
+            where: a predicate to filter rows to update.
             cascade: if True, also update all computed columns that transitively depend on the updated columns.
 
         Examples:
@@ -745,18 +745,34 @@ class Table(SchemaObject):
         self._check_is_dropped()
         return self._tbl_version.update(value_spec, where, cascade)
 
-    def batch_update(self, rows: Iterable[dict[str, Any]], cascade: bool = True) -> UpdateStatus:
+    def batch_update(
+            self, rows: Iterable[dict[str, Any]], cascade: bool = True,
+            if_not_exists: Literal['error', 'ignore', 'insert'] = 'error'
+    ) -> UpdateStatus:
         """Update rows in this table.
 
         Args:
             rows: an Iterable of dictionaries containing values for the updated columns plus values for the primary key
                   columns.
             cascade: if True, also update all computed columns that transitively depend on the updated columns.
+            if_not_exists: Specifies the behavior if a row to update does not exist:
+
+                - `'error'`: Raise an error.
+                - `'ignore'`: Skip the row silently.
+                - `'insert'`: Insert the row.
 
         Examples:
-            Update the 'name' and 'age' columns for the rows with ids 1 and 2 (assuming 'id' is the primary key):
+            Update the `name` and `age` columns for the rows with ids 1 and 2 (assuming `id` is the primary key).
+            If either row does not exist, this raises an error:
 
             >>> tbl.update([{'id': 1, 'name': 'Alice', 'age': 30}, {'id': 2, 'name': 'Bob', 'age': 40}])
+
+            Update the `name` and `age` columns for the row with `id` 1 (assuming `id` is the primary key) and insert
+            the row with new `id` 3 (assuming this key does not exist):
+
+            >>> tbl.update(
+                [{'id': 1, 'name': 'Alice', 'age': 30}, {'id': 3, 'name': 'Bob', 'age': 40}],
+                if_not_exists='insert')
         """
         if self._tbl_version_path.is_snapshot():
             raise excs.Error('Cannot update a snapshot')
@@ -784,13 +800,15 @@ class Table(SchemaObject):
                     missing_cols = pk_col_names - set(col.name for col in col_vals.keys())
                     raise excs.Error(f'Primary key columns ({", ".join(missing_cols)}) missing in {row_spec}')
             row_updates.append(col_vals)
-        return self._tbl_version.batch_update(row_updates, rowids, cascade)
+        return self._tbl_version.batch_update(
+            row_updates, rowids, error_if_not_exists=if_not_exists == 'error',
+            insert_if_not_exists=if_not_exists == 'insert', cascade=cascade)
 
-    def delete(self, where: Optional['pixeltable.exprs.Predicate'] = None) -> UpdateStatus:
+    def delete(self, where: Optional['pixeltable.exprs.Expr'] = None) -> UpdateStatus:
         """Delete rows in this table.
 
         Args:
-            where: a Predicate to filter rows to delete.
+            where: a predicate to filter rows to delete.
 
         Examples:
             Delete all rows in a table:

pixeltable/catalog/table_version.py

@@ -678,12 +678,12 @@ class TableVersion:
         return result
 
     def update(
-        self, value_spec: dict[str, Any], where: Optional['exprs.Predicate'] = None, cascade: bool = True
+        self, value_spec: dict[str, Any], where: Optional['exprs.Expr'] = None, cascade: bool = True
     ) -> UpdateStatus:
         """Update rows in this TableVersionPath.
         Args:
             value_spec: a list of (column, value) pairs specifying the columns to update and their new values.
-            where: a Predicate to filter rows to update.
+            where: a predicate to filter rows to update.
             cascade: if True, also update all computed columns that transitively depend on the updated columns,
                 including within views.
         """
@@ -694,18 +694,26 @@ class TableVersion:
 
         update_spec = self._validate_update_spec(value_spec, allow_pk=False, allow_exprs=True)
         if where is not None:
-            if not isinstance(where, exprs.Predicate):
-                raise excs.Error(f"'where' argument must be a Predicate, got {type(where)}")
+            if not isinstance(where, exprs.Expr):
+                raise excs.Error(f"'where' argument must be a predicate, got {type(where)}")
             analysis_info = Planner.analyze(self.path, where)
             # for now we require that the updated rows can be identified via SQL, rather than via a Python filter
             if analysis_info.filter is not None:
                 raise excs.Error(f'Filter {analysis_info.filter} not expressible in SQL')
 
         with Env.get().engine.begin() as conn:
-            return self._update(conn, update_spec, where, cascade)
+            plan, updated_cols, recomputed_cols = (
+                Planner.create_update_plan(self.path, update_spec, [], where, cascade)
+            )
+            result = self.propagate_update(
+                plan, where.sql_expr() if where is not None else None, recomputed_cols,
+                base_versions=[], conn=conn, timestamp=time.time(), cascade=cascade, show_progress=True)
+            result.updated_cols = updated_cols
+            return result
 
     def batch_update(
-            self, batch: list[dict[Column, 'exprs.Expr']], rowids: list[tuple[int, ...]], cascade: bool = True
+            self, batch: list[dict[Column, 'exprs.Expr']], rowids: list[tuple[int, ...]], insert_if_not_exists: bool,
+            error_if_not_exists: bool, cascade: bool = True,
     ) -> UpdateStatus:
         """Update rows in batch.
         Args:
@@ -714,62 +722,26 @@ class TableVersion:
         """
         # if we do lookups of rowids, we must have one for each row in the batch
         assert len(rowids) == 0 or len(rowids) == len(batch)
-        result_status = UpdateStatus()
         cols_with_excs: set[str] = set()
-        updated_cols: set[str] = set()
-        pk_cols = self.primary_key_columns()
-        use_rowids = len(rowids) > 0
 
         with Env.get().engine.begin() as conn:
-            for i, row in enumerate(batch):
-                where_clause: Optional[exprs.Expr] = None
-                if use_rowids:
-                    # construct Where clause to match rowid
-                    num_rowid_cols = len(self.store_tbl.rowid_columns())
-                    for col_idx in range(num_rowid_cols):
-                        assert len(rowids[i]) == num_rowid_cols, f'len({rowids[i]}) != {num_rowid_cols}'
-                        clause = exprs.RowidRef(self, col_idx) == rowids[i][col_idx]
-                        if where_clause is None:
-                            where_clause = clause
-                        else:
-                            where_clause = where_clause & clause
-                else:
-                    # construct Where clause for primary key columns
-                    for col in pk_cols:
-                        assert col in row
-                        clause = exprs.ColumnRef(col) == row[col]
-                        if where_clause is None:
-                            where_clause = clause
-                        else:
-                            where_clause = where_clause & clause
-
-                update_targets = {col: row[col] for col in row if col not in pk_cols}
-                status = self._update(conn, update_targets, where_clause, cascade, show_progress=False)
-                result_status.num_rows += status.num_rows
-                result_status.num_excs += status.num_excs
-                result_status.num_computed_values += status.num_computed_values
-                cols_with_excs.update(status.cols_with_excs)
-                updated_cols.update(status.updated_cols)
-
-            result_status.cols_with_excs = list(cols_with_excs)
-            result_status.updated_cols = list(updated_cols)
-            return result_status
-
-    def _update(
-            self, conn: sql.engine.Connection, update_targets: dict[Column, 'pixeltable.exprs.Expr'],
-            where_clause: Optional['pixeltable.exprs.Predicate'] = None, cascade: bool = True,
-            show_progress: bool = True
-    ) -> UpdateStatus:
-        from pixeltable.plan import Planner
+            from pixeltable.plan import Planner
 
-        plan, updated_cols, recomputed_cols = (
-            Planner.create_update_plan(self.path, update_targets, [], where_clause, cascade)
-        )
-        result = self.propagate_update(
-            plan, where_clause.sql_expr() if where_clause is not None else None, recomputed_cols,
-            base_versions=[], conn=conn, timestamp=time.time(), cascade=cascade, show_progress=show_progress)
-        result.updated_cols = updated_cols
-        return result
+            plan, row_update_node, delete_where_clause, updated_cols, recomputed_cols = \
+                Planner.create_batch_update_plan(self.path, batch, rowids, cascade=cascade)
+            result = self.propagate_update(
+                plan, delete_where_clause, recomputed_cols, base_versions=[], conn=conn, timestamp=time.time(),
+                cascade=cascade)
+            result.updated_cols = [c.qualified_name for c in updated_cols]
+
+            unmatched_rows = row_update_node.unmatched_rows()
+            if len(unmatched_rows) > 0:
+                if error_if_not_exists:
+                    raise excs.Error(f'batch_update(): {len(unmatched_rows)} row(s) not found')
+                if insert_if_not_exists:
+                    insert_status = self.insert(unmatched_rows, print_stats=False, fail_on_exception=False)
+                    result += insert_status
+            return result
 
     def _validate_update_spec(
             self, value_spec: dict[str, Any], allow_pk: bool, allow_exprs: bool
@@ -779,7 +751,10 @@ class TableVersion:
             if not isinstance(col_name, str):
                 raise excs.Error(f'Update specification: dict key must be column name, got {col_name!r}')
             if col_name == _ROWID_COLUMN_NAME:
-                # ignore pseudo-column _rowid
+                # a valid rowid is a list of ints, one per rowid column
+                assert len(val) == len(self.store_tbl.rowid_columns())
+                for el in val:
+                    assert isinstance(el, int)
                 continue
             col = self.path.get_column(col_name, include_bases=False)
             if col is None:
@@ -789,8 +764,6 @@ class TableVersion:
                 raise excs.Error(f'Column {col_name} is computed and cannot be updated')
             if col.is_pk and not allow_pk:
                 raise excs.Error(f'Column {col_name} is a primary key column and cannot be updated')
-            if col.col_type.is_media_type():
-                raise excs.Error(f'Column {col_name} has type image/video/audio/document and cannot be updated')
 
             # make sure that the value is compatible with the column type
             try:
@@ -848,17 +821,17 @@ class TableVersion:
         result.cols_with_excs = list(dict.fromkeys(result.cols_with_excs).keys())  # remove duplicates
         return result
 
-    def delete(self, where: Optional['exprs.Predicate'] = None) -> UpdateStatus:
+    def delete(self, where: Optional['exprs.Expr'] = None) -> UpdateStatus:
         """Delete rows in this table.
         Args:
-            where: a Predicate to filter rows to delete.
+            where: a predicate to filter rows to delete.
         """
         assert self.is_insertable()
-        from pixeltable.exprs import Predicate
+        from pixeltable.exprs import Expr
         from pixeltable.plan import Planner
         if where is not None:
-            if not isinstance(where, Predicate):
-                raise excs.Error(f"'where' argument must be a Predicate, got {type(where)}")
+            if not isinstance(where, Expr):
+                raise excs.Error(f"'where' argument must be a predicate, got {type(where)}")
             analysis_info = Planner.analyze(self.path, where)
             # for now we require that the updated rows can be identified via SQL, rather than via a Python filter
             if analysis_info.filter is not None:
@@ -872,11 +845,11 @@ class TableVersion:
         return status
 
     def propagate_delete(
-            self, where: Optional['exprs.Predicate'], base_versions: List[Optional[int]],
+            self, where: Optional['exprs.Expr'], base_versions: List[Optional[int]],
             conn: sql.engine.Connection, timestamp: float) -> int:
         """Delete rows in this table and propagate to views.
         Args:
-            where: a Predicate to filter rows to delete.
+            where: a predicate to filter rows to delete.
         Returns:
             number of deleted rows
         """

pixeltable/catalog/view.py

@@ -51,7 +51,7 @@ class View(Table):
     @classmethod
     def create(
             cls, dir_id: UUID, name: str, base: TableVersionPath, schema: Dict[str, Any],
-            predicate: 'pxt.exprs.Predicate', is_snapshot: bool, num_retained_versions: int, comment: str,
+            predicate: 'pxt.exprs.Expr', is_snapshot: bool, num_retained_versions: int, comment: str,
             iterator_cls: Optional[Type[ComponentIterator]], iterator_args: Optional[Dict]
     ) -> View:
         columns = cls._create_columns(schema)
@@ -92,7 +92,7 @@ class View(Table):
                 ]
                 sig = func.Signature(InvalidType(), params)
                 from pixeltable.exprs import FunctionCall
-                FunctionCall.check_args(sig, bound_args)
+                FunctionCall.normalize_args(sig, bound_args)
             except TypeError as e:
                 raise Error(f'Cannot instantiate iterator with given arguments: {e}')
 
@@ -213,5 +213,5 @@ class View(Table):
     ) -> UpdateStatus:
         raise excs.Error(f'{self.display_name()} {self._name!r}: cannot insert into view')
 
-    def delete(self, where: Optional['pixeltable.exprs.Predicate'] = None) -> UpdateStatus:
+    def delete(self, where: Optional['pixeltable.exprs.Expr'] = None) -> UpdateStatus:
         raise excs.Error(f'{self.display_name()} {self._name!r}: cannot delete from view')

pixeltable/dataframe.py

@@ -153,7 +153,7 @@ class DataFrame:
         self,
         tbl: catalog.TableVersionPath,
         select_list: Optional[List[Tuple[exprs.Expr, Optional[str]]]] = None,
-        where_clause: Optional[exprs.Predicate] = None,
+        where_clause: Optional[exprs.Expr] = None,
         group_by_clause: Optional[List[exprs.Expr]] = None,
         grouping_tbl: Optional[catalog.TableVersion] = None,
         order_by_clause: Optional[List[Tuple[exprs.Expr, bool]]] = None,  # List[(expr, asc)]
@@ -530,7 +530,11 @@ class DataFrame:
             limit=self.limit_val,
         )
 
-    def where(self, pred: exprs.Predicate) -> DataFrame:
+    def where(self, pred: exprs.Expr) -> DataFrame:
+        if not isinstance(pred, exprs.Expr):
+            raise excs.Error(f'Where() requires a Pixeltable expression, but instead got {type(pred)}')
+        if not pred.col_type.is_bool_type():
+            raise excs.Error(f'Where(): expression needs to return bool, but instead returns {pred.col_type}')
         return DataFrame(
             self.tbl,
             select_list=self.select_list,
@@ -628,12 +632,9 @@ class DataFrame:
     def __getitem__(self, index: object) -> DataFrame:
         """
         Allowed:
-        - [<Predicate>]: filter operation
         - [List[Expr]]/[Tuple[Expr]]: setting the select list
         - [Expr]: setting a single-col select list
         """
-        if isinstance(index, exprs.Predicate):
-            return self.where(index)
         if isinstance(index, tuple):
             index = list(index)
         if isinstance(index, exprs.Expr):
@@ -668,7 +669,7 @@ class DataFrame:
         tbl = catalog.TableVersionPath.from_dict(d['tbl'])
         select_list = [(exprs.Expr.from_dict(e), name) for e, name in d['select_list']] \
             if d['select_list'] is not None else None
-        where_clause = exprs.Predicate.from_dict(d['where_clause']) \
+        where_clause = exprs.Expr.from_dict(d['where_clause']) \
             if d['where_clause'] is not None else None
         group_by_clause = [exprs.Expr.from_dict(e) for e in d['group_by_clause']] \
             if d['group_by_clause'] is not None else None

pixeltable/exec/__init__.py

@@ -5,6 +5,7 @@ from .exec_context import ExecContext
 from .exec_node import ExecNode
 from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
-from .sql_scan_node import SqlScanNode
+from .sql_node import SqlScanNode, SqlLookupNode
+from .row_update_node import RowUpdateNode
 from .media_validation_node import MediaValidationNode
 from .data_row_batch import DataRowBatch

pixeltable/exec/expr_eval_node.py

@@ -50,7 +50,14 @@ class ExprEvalNode(ExecNode):
 
     def _open(self) -> None:
         warnings.simplefilter("ignore", category=TqdmWarning)
-        if self.ctx.show_pbar:
+        # This is a temporary hack. When B-tree indices on string columns were implemented (via computed columns
+        # that invoke the `BtreeIndex.str_filter` udf), it resulted in frivolous progress bars appearing on every
+        # insertion. This special-cases the `str_filter` call to suppress the corresponding progress bar.
+        # TODO(aaron-siegel) Remove this hack once we clean up progress bars more generally.
+        is_str_filter_node = all(
+            isinstance(expr, exprs.FunctionCall) and expr.fn.name == 'str_filter' for expr in self.output_exprs
+        )
+        if self.ctx.show_pbar and not is_str_filter_node:
             self.pbar = tqdm(
                 total=len(self.target_exprs) * self.ctx.num_rows,
                 desc='Computing cells',

pixeltable/exec/row_update_node.py

@@ -0,0 +1,61 @@
+import logging
+from typing import Any
+
+import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
+from pixeltable.utils.media_store import MediaStore
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+
+_logger = logging.getLogger('pixeltable')
+
+class RowUpdateNode(ExecNode):
+    """
+    Update individual rows in the input batches, identified by key columns.
+
+    The updates for a row are provided as a dict of column names to new values.
+    The node assumes that all update dicts contain the same keys, and it populates the slots of the columns present in
+    the update list.
+    """
+    def __init__(
+            self, tbl: catalog.TableVersionPath, key_vals_batch: list[tuple], is_rowid_key: bool,
+            col_vals_batch: list[dict[catalog.Column, Any]], row_builder: exprs.RowBuilder, input: ExecNode,
+    ):
+        super().__init__(row_builder, [], [], input)
+        self.updates = {key_vals: col_vals for key_vals, col_vals in zip(key_vals_batch, col_vals_batch)}
+        self.is_rowid_key = is_rowid_key
+        # determine slot idxs of all columns we need to read or write
+        # retrieve ColumnRefs from the RowBuilder (has slot_idx set)
+        all_col_slot_idxs = {
+            col_ref.col: col_ref.slot_idx
+            for col_ref in row_builder.unique_exprs if isinstance(col_ref, exprs.ColumnRef)
+        }
+        self.col_slot_idxs = {col: all_col_slot_idxs[col] for col in col_vals_batch[0].keys()}
+        self.key_slot_idxs = {col: all_col_slot_idxs[col] for col in tbl.tbl_version.primary_key_columns()}
+        self.matched_key_vals: set[tuple] = set()
+
+    def __next__(self) -> DataRowBatch:
+        batch = next(self.input)
+        for row in batch:
+            key_vals = row.rowid if self.is_rowid_key else \
+                tuple(row[slot_idx] for slot_idx in self.key_slot_idxs.values())
+            if key_vals not in self.updates:
+                continue
+            self.matched_key_vals.add(key_vals)
+            col_vals = self.updates[key_vals]
+            for col, val in col_vals.items():
+                slot_idx = self.col_slot_idxs[col]
+                row[slot_idx] = val
+        return batch
+
+    def unmatched_rows(self) -> list[dict[str, Any]]:
+        """Return rows that didn't get used in the updates as a list of dicts compatible with TableVersion.insert()."""
+        result: list[dict[str, Any]] = []
+        key_cols = self.key_slot_idxs.keys()
+        for key_vals, col_vals in self.updates.items():
+            if key_vals in self.matched_key_vals:
+                continue
+            row = {col.name: val for col, val in zip(key_cols, key_vals)}
+            row.update({col.name: val for col, val in col_vals.items()})
+            result.append(row)
+        return result