pixeltable 0.4.15__py3-none-any.whl → 0.4.17__py3-none-any.whl

pixeltable/exec/object_store_save_node.py

@@ -7,7 +7,6 @@ from collections import defaultdict, deque
 from concurrent import futures
 from pathlib import Path
 from typing import AsyncIterator, Iterator, NamedTuple, Optional
-from uuid import UUID
 
 from pixeltable import exprs
 from pixeltable.utils.object_stores import ObjectOps, ObjectPath, StorageTarget
@@ -81,9 +80,7 @@ class ObjectStoreSaveNode(ExecNode):
         num_missing: int  # number of references to media files in this row
         delete_destinations: list[Path]  # paths to delete after all copies are complete
 
-    def __init__(
-        self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True
-    ):
+    def __init__(self, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True):
         # input_/output_exprs=[]: we don't have anything to evaluate
         super().__init__(input.row_builder, [], [], input)
         self.retain_input_order = retain_input_order

pixeltable/exec/row_update_node.py

@@ -14,10 +14,18 @@ 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.
+    Populates the slots of the columns present in the update list.
+    Assumptions:
+    - all update dicts contain the same keys
+    - the input node populates DataRow.cell_vals for all primary key columns
     """
 
+    updates: dict[tuple, dict[catalog.Column, Any]]
+    is_rowid_key: bool  # if True, key_vals_batch contains rowids rather than primary key values
+    col_slot_idxs: dict[catalog.Column, int]
+    pk_columns: list[catalog.Column]
+    matched_key_vals: set[tuple]
+
     def __init__(
         self,
         tbl: catalog.TableVersionPath,
@@ -37,16 +45,16 @@ class RowUpdateNode(ExecNode):
             for col_ref in row_builder.unique_exprs
             if isinstance(col_ref, exprs.ColumnRef)
         }
+        # all update target columns should have assigned slot idxs
+        assert all(col in all_col_slot_idxs for col in col_vals_batch[0])
         self.col_slot_idxs = {col: all_col_slot_idxs[col] for col in col_vals_batch[0]}
-        self.key_slot_idxs = {col: all_col_slot_idxs[col] for col in tbl.tbl_version.get().primary_key_columns()}
-        self.matched_key_vals: set[tuple] = set()
+        self.pk_columns = tbl.tbl_version.get().primary_key_columns()
+        self.matched_key_vals = set()
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         async for batch in 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())
-                )
+                key_vals = row.rowid if self.is_rowid_key else tuple(row.cell_vals[col.id] for col in self.pk_columns)
                 if key_vals not in self.updates:
                     continue
                 self.matched_key_vals.add(key_vals)
@@ -59,11 +67,10 @@ class RowUpdateNode(ExecNode):
     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 = {col.name: val for col, val in zip(self.pk_columns, key_vals)}
             row.update({col.name: val for col, val in col_vals.items()})
             result.append(row)
         return result

pixeltable/exec/sql_node.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import warnings
 from decimal import Decimal
@@ -65,7 +66,7 @@ def print_order_by_clause(clause: OrderByClause) -> str:
 
 class SqlNode(ExecNode):
     """
-    Materializes data from the store via an SQL statement.
+    Materializes data from the store via a SQL statement.
     This only provides the select list. The subclasses are responsible for the From clause and any additional clauses.
     The pk columns are not included in the select list.
     If set_pk is True, they are added to the end of the result set when creating the SQL statement
@@ -82,6 +83,8 @@ class SqlNode(ExecNode):
 
     tbl: Optional[catalog.TableVersionPath]
     select_list: exprs.ExprSet
+    columns: list[catalog.Column]  # for which columns to populate DataRow.cell_vals/cell_md
+    cell_md_refs: list[exprs.ColumnPropertyRef]  # of ColumnRefs which also need DataRow.slot_cellmd for evaluation
     set_pk: bool
     num_pk_cols: int
     py_filter: Optional[exprs.Expr]  # a predicate that can only be run in Python
@@ -89,6 +92,12 @@ class SqlNode(ExecNode):
     cte: Optional[sql.CTE]
     sql_elements: exprs.SqlElementCache
 
+    # execution state
+    cellmd_item_idxs: exprs.ExprDict[int]  # cellmd expr -> idx in sql select list
+    column_item_idxs: dict[catalog.Column, int]  # column -> idx in sql select list
+    column_cellmd_item_idxs: dict[catalog.Column, int]  # column -> idx in sql select list
+    result_cursor: sql.engine.CursorResult | None
+
     # where_clause/-_element: allow subclass to set one or the other (but not both)
     where_clause: Optional[exprs.Expr]
     where_clause_element: Optional[sql.ColumnElement]
@@ -101,12 +110,22 @@ class SqlNode(ExecNode):
         tbl: Optional[catalog.TableVersionPath],
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
         sql_elements: exprs.SqlElementCache,
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
         set_pk: bool = False,
     ):
         # create Select stmt
         self.sql_elements = sql_elements
         self.tbl = tbl
+        self.columns = columns
+        if cell_md_col_refs is not None:
+            assert all(ref.col.stores_cellmd for ref in cell_md_col_refs)
+            self.cell_md_refs = [
+                exprs.ColumnPropertyRef(ref, exprs.ColumnPropertyRef.Property.CELLMD) for ref in cell_md_col_refs
+            ]
+        else:
+            self.cell_md_refs = []
         self.select_list = exprs.ExprSet(select_list)
         # unstored iter columns: we also need to retrieve whatever is needed to materialize the iter args
         for iter_arg in row_builder.unstored_iter_args.values():
@@ -129,6 +148,9 @@ class SqlNode(ExecNode):
             assert self.num_pk_cols > 1
 
         # additional state
+        self.cellmd_item_idxs = exprs.ExprDict()
+        self.column_item_idxs = {}
+        self.column_cellmd_item_idxs = {}
         self.result_cursor = None
         # the filter is provided by the subclass
         self.py_filter = None
@@ -144,10 +166,9 @@ class SqlNode(ExecNode):
             if tv is not None:
                 assert tv.is_validated
 
-    def _create_pk_cols(self) -> list[sql.Column]:
-        """Create a list of pk columns"""
-        # we need to retrieve the pk columns
+    def _pk_col_items(self) -> list[sql.Column]:
         if self.set_pk:
+            # we need to retrieve the pk columns
             assert self.tbl is not None
             assert self.tbl.tbl_version.get().is_validated
             return self.tbl.tbl_version.get().store_tbl.pk_columns()
@@ -157,7 +178,19 @@ class SqlNode(ExecNode):
         """Create Select from local state"""
 
         assert self.sql_elements.contains_all(self.select_list)
-        sql_select_list = [self.sql_elements.get(e) for e in self.select_list] + self._create_pk_cols()
+        sql_select_list_exprs = exprs.ExprSet(self.select_list)
+        self.cellmd_item_idxs = exprs.ExprDict((ref, sql_select_list_exprs.add(ref)) for ref in self.cell_md_refs)
+        column_refs = [exprs.ColumnRef(col) for col in self.columns]
+        self.column_item_idxs = {col_ref.col: sql_select_list_exprs.add(col_ref) for col_ref in column_refs}
+        column_cellmd_refs = [
+            exprs.ColumnPropertyRef(col_ref, exprs.ColumnPropertyRef.Property.CELLMD)
+            for col_ref in column_refs
+            if col_ref.col.stores_cellmd
+        ]
+        self.column_cellmd_item_idxs = {
+            cellmd_ref.col_ref.col: sql_select_list_exprs.add(cellmd_ref) for cellmd_ref in column_cellmd_refs
+        }
+        sql_select_list = [self.sql_elements.get(e) for e in sql_select_list_exprs] + self._pk_col_items()
         stmt = sql.select(*sql_select_list)
 
         where_clause_element = (
@@ -198,9 +231,7 @@ class SqlNode(ExecNode):
             if not keep_pk:
                 self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
             self.cte = self._create_stmt().cte()
-        pk_count = self.num_pk_cols if self.set_pk else 0
-        assert len(self.select_list) + pk_count == len(self.cte.c)
-        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))  # skip pk cols
+        return self.cte, exprs.ExprDict(zip(list(self.select_list) + self.cell_md_refs, self.cte.c))  # skip pk cols
 
     @classmethod
     def retarget_rowid_refs(cls, target: catalog.TableVersionPath, expr_seq: Iterable[exprs.Expr]) -> None:
@@ -318,24 +349,53 @@ class SqlNode(ExecNode):
         output_batch = DataRowBatch(self.row_builder)
         output_row: Optional[exprs.DataRow] = None
         num_rows_returned = 0
+        is_using_cockroachdb = Env.get().is_using_cockroachdb
+        tzinfo = Env.get().default_time_zone
 
         for sql_row in result_cursor:
             output_row = output_batch.add_row(output_row)
 
             # populate output_row
+
             if self.num_pk_cols > 0:
                 output_row.set_pk(tuple(sql_row[-self.num_pk_cols :]))
+
+            # column copies
+            for col, item_idx in self.column_item_idxs.items():
+                output_row.cell_vals[col.id] = sql_row[item_idx]
+            for col, item_idx in self.column_cellmd_item_idxs.items():
+                cell_md_dict = sql_row[item_idx]
+                output_row.cell_md[col.id] = exprs.CellMd(**cell_md_dict) if cell_md_dict is not None else None
+
+            # populate DataRow.slot_cellmd, where requested
+            for cellmd_ref, item_idx in self.cellmd_item_idxs.items():
+                cell_md_dict = sql_row[item_idx]
+                output_row.slot_md[cellmd_ref.col_ref.slot_idx] = (
+                    exprs.CellMd.from_dict(cell_md_dict) if cell_md_dict is not None else None
+                )
+
             # copy the output of the SQL query into the output row
             for i, e in enumerate(self.select_list):
                 slot_idx = e.slot_idx
-                # certain numerical operations can produce Decimals (eg, SUM(<int column>)); we need to convert them
                 if isinstance(sql_row[i], Decimal):
+                    # certain numerical operations can produce Decimals (eg, SUM(<int column>)); we need to convert them
                     if e.col_type.is_int_type():
                         output_row[slot_idx] = int(sql_row[i])
                     elif e.col_type.is_float_type():
                         output_row[slot_idx] = float(sql_row[i])
                     else:
                         raise RuntimeError(f'Unexpected Decimal value for {e}')
+                elif is_using_cockroachdb and isinstance(sql_row[i], datetime.datetime):
+                    # Ensure that the datetime is timezone-aware and in the session time zone
+                    # cockroachDB returns timestamps in the session time zone, with numeric offset,
+                    # convert to the session time zone with the requested tzinfo for DST handling
+                    if e.col_type.is_timestamp_type():
+                        if isinstance(sql_row[i].tzinfo, datetime.timezone):
+                            output_row[slot_idx] = sql_row[i].astimezone(tz=tzinfo)
+                        else:
+                            output_row[slot_idx] = sql_row[i]
+                    else:
+                        raise RuntimeError(f'Unexpected datetime value for {e}')
                 else:
                     output_row[slot_idx] = sql_row[i]
 
@@ -387,11 +447,21 @@ class SqlScanNode(SqlNode):
         tbl: catalog.TableVersionPath,
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
         set_pk: bool = False,
         exact_version_only: Optional[list[catalog.TableVersionHandle]] = None,
     ):
         sql_elements = exprs.SqlElementCache()
-        super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=set_pk)
+        super().__init__(
+            tbl,
+            row_builder,
+            select_list,
+            columns=columns,
+            sql_elements=sql_elements,
+            set_pk=set_pk,
+            cell_md_col_refs=cell_md_col_refs,
+        )
         # create Select stmt
         if exact_version_only is None:
             exact_version_only = []
@@ -423,11 +493,21 @@ class SqlLookupNode(SqlNode):
         tbl: catalog.TableVersionPath,
         row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
+        columns: list[catalog.Column],
         sa_key_cols: list[sql.Column],
         key_vals: list[tuple],
+        cell_md_col_refs: list[exprs.ColumnRef] | None = None,
     ):
         sql_elements = exprs.SqlElementCache()
-        super().__init__(tbl, row_builder, select_list, sql_elements, set_pk=True)
+        super().__init__(
+            tbl,
+            row_builder,
+            select_list,
+            columns=columns,
+            sql_elements=sql_elements,
+            set_pk=True,
+            cell_md_col_refs=cell_md_col_refs,
+        )
         # Where clause: (key-col-1, key-col-2, ...) IN ((val-1, val-2, ...), ...)
         self.where_clause_element = sql.tuple_(*sa_key_cols).in_(key_vals)
 
@@ -460,9 +540,10 @@ class SqlAggregationNode(SqlNode):
         limit: Optional[int] = None,
         exact_version_only: Optional[list[catalog.TableVersion]] = None,
     ):
+        assert len(input.cell_md_refs) == 0  # there's no aggregation over json or arrays in SQL
         self.input_cte, input_col_map = input.to_cte()
         sql_elements = exprs.SqlElementCache(input_col_map)
-        super().__init__(None, row_builder, select_list, sql_elements)
+        super().__init__(None, row_builder, select_list, columns=[], sql_elements=sql_elements)
         self.group_by_items = group_by_items
 
     def _create_stmt(self) -> sql.Select:
@@ -498,7 +579,10 @@ class SqlJoinNode(SqlNode):
             input_cte, input_col_map = input_node.to_cte()
             self.input_ctes.append(input_cte)
             sql_elements.extend(input_col_map)
-        super().__init__(None, row_builder, select_list, sql_elements)
+        cell_md_col_refs = [cell_md_ref.col_ref for input in inputs for cell_md_ref in input.cell_md_refs]
+        super().__init__(
+            None, row_builder, select_list, columns=[], sql_elements=sql_elements, cell_md_col_refs=cell_md_col_refs
+        )
 
     def _create_stmt(self) -> sql.Select:
         from pixeltable import plan
@@ -552,7 +636,16 @@ class SqlSampleNode(SqlNode):
         assert self.pk_count > 1
         sql_elements = exprs.SqlElementCache(input_col_map)
         assert sql_elements.contains_all(stratify_exprs)
-        super().__init__(input.tbl, row_builder, select_list, sql_elements, set_pk=True)
+        cell_md_col_refs = [cell_md_ref.col_ref for cell_md_ref in input.cell_md_refs]
+        super().__init__(
+            input.tbl,
+            row_builder,
+            select_list,
+            columns=[],
+            sql_elements=sql_elements,
+            cell_md_col_refs=cell_md_col_refs,
+            set_pk=True,
+        )
         self.stratify_exprs = stratify_exprs
         self.sample_clause = sample_clause
         assert isinstance(self.sample_clause.seed, int)

pixeltable/exprs/__init__.py

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

pixeltable/exprs/arithmetic_expr.py

@@ -4,7 +4,7 @@ from typing import Any, Optional
 
 import sqlalchemy as sql
 
-from pixeltable import exceptions as excs, type_system as ts
+from pixeltable import env, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
@@ -64,23 +64,30 @@ class ArithmeticExpr(Expr):
         right = sql_elements.get(self._op2)
         if left is None or right is None:
             return None
-        if self.operator == ArithmeticOperator.ADD:
-            return left + right
-        if self.operator == ArithmeticOperator.SUB:
-            return left - right
-        if self.operator == ArithmeticOperator.MUL:
-            return left * right
+        if self.operator in (ArithmeticOperator.ADD, ArithmeticOperator.SUB, ArithmeticOperator.MUL):
+            if env.Env.get().is_using_cockroachdb and self._op1.col_type != self._op2.col_type:
+                if self._op1.col_type != self.col_type:
+                    left = sql.cast(left, self.col_type.to_sa_type())
+                if self._op2.col_type != self.col_type:
+                    right = sql.cast(right, self.col_type.to_sa_type())
+            if self.operator == ArithmeticOperator.ADD:
+                return left + right
+            if self.operator == ArithmeticOperator.SUB:
+                return left - right
+            if self.operator == ArithmeticOperator.MUL:
+                return left * right
         if self.operator == ArithmeticOperator.DIV:
             assert self.col_type.is_float_type()
-            # Avoid DivisionByZero: if right is 0, make this a NULL
+            # Avoid division by zero errors by converting any zero divisor to NULL.
             # TODO: Should we cast the NULLs to NaNs when they are retrieved back into Python?
-            nullif = sql.sql.func.nullif(right, 0)
-            # We have to cast to a `float`, or else we'll get a `Decimal`
-            return sql.sql.expression.cast(left / nullif, self.col_type.to_sa_type())
+            # These casts cause the computation to take place in float units, rather than DECIMAL.
+            nullif = sql.cast(sql.func.nullif(right, 0), self.col_type.to_sa_type())
+            return sql.cast(left, self.col_type.to_sa_type()) / nullif
         if self.operator == ArithmeticOperator.MOD:
             if self.col_type.is_int_type():
-                nullif = sql.sql.func.nullif(right, 0)
-                return left % nullif
+                # Avoid division by zero errors by converting any zero divisor to NULL.
+                nullif1 = sql.cast(sql.func.nullif(right, 0), self.col_type.to_sa_type())
+                return left % nullif1
             if self.col_type.is_float_type():
                 # Postgres does not support modulus for floats
                 return None
@@ -90,11 +97,9 @@ class ArithmeticExpr(Expr):
             # We need the behavior to be consistent, so that expressions will evaluate the same way
             # whether or not their operands can be translated to SQL. These SQL clauses should
             # mimic the behavior of Python's // operator.
-            nullif = sql.sql.func.nullif(right, 0)
-            if self.col_type.is_int_type():
-                return sql.sql.expression.cast(sql.func.floor(left / nullif), self.col_type.to_sa_type())
-            if self.col_type.is_float_type():
-                return sql.sql.expression.cast(sql.func.floor(left / nullif), self.col_type.to_sa_type())
+            # Avoid division by zero errors by converting any zero divisor to NULL.
+            nullif = sql.cast(sql.func.nullif(right, 0), self.col_type.to_sa_type())
+            return sql.func.floor(sql.cast(left, self.col_type.to_sa_type()) / nullif)
         raise AssertionError()
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/column_property_ref.py

@@ -44,21 +44,21 @@ class ColumnPropertyRef(Expr):
         return [*super()._id_attrs(), ('prop', self.prop.value)]
 
     @property
-    def _col_ref(self) -> ColumnRef:
+    def col_ref(self) -> ColumnRef:
         col_ref = self.components[0]
         assert isinstance(col_ref, ColumnRef)
         return col_ref
 
     def __repr__(self) -> str:
-        return f'{self._col_ref}.{self.prop.name.lower()}'
+        return f'{self.col_ref}.{self.prop.name.lower()}'
 
     def is_cellmd_prop(self) -> bool:
         return self.prop in (self.Property.ERRORTYPE, self.Property.ERRORMSG, self.Property.CELLMD)
 
     def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
-        if not self._col_ref.col_handle.get().is_stored:
+        if not self.col_ref.col_handle.get().is_stored:
             return None
-        col = self._col_ref.col_handle.get()
+        col = self.col_ref.col_handle.get()
 
         # the errortype/-msg properties of a read-validated media column need to be extracted from the DataRow
         if (
@@ -77,7 +77,7 @@ class ColumnPropertyRef(Expr):
             return col.sa_cellmd_col
         if self.prop == self.Property.FILEURL:
             # the file url is stored as the column value
-            return sql_elements.get(self._col_ref)
+            return sql_elements.get(self.col_ref)
         return None
 
     @classmethod
@@ -87,15 +87,15 @@ class ColumnPropertyRef(Expr):
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         if self.prop == self.Property.FILEURL:
-            assert data_row.has_val[self._col_ref.slot_idx]
-            data_row[self.slot_idx] = data_row.file_urls[self._col_ref.slot_idx]
+            assert data_row.has_val[self.col_ref.slot_idx]
+            data_row[self.slot_idx] = data_row.file_urls[self.col_ref.slot_idx]
             return
         elif self.prop == self.Property.LOCALPATH:
-            assert data_row.has_val[self._col_ref.slot_idx]
-            data_row[self.slot_idx] = data_row.file_paths[self._col_ref.slot_idx]
+            assert data_row.has_val[self.col_ref.slot_idx]
+            data_row[self.slot_idx] = data_row.file_paths[self.col_ref.slot_idx]
             return
         elif self.is_cellmd_prop():
-            exc = data_row.get_exc(self._col_ref.slot_idx)
+            exc = data_row.get_exc(self.col_ref.slot_idx)
             if exc is None:
                 data_row[self.slot_idx] = None
             elif self.prop == self.Property.ERRORTYPE:

pixeltable/exprs/column_ref.py

@@ -123,8 +123,8 @@ class ColumnRef(Expr):
             name == ColumnPropertyRef.Property.ERRORTYPE.name.lower()
             or name == ColumnPropertyRef.Property.ERRORMSG.name.lower()
         ):
-            property_is_present = self.col.stores_cellmd
-            if not property_is_present:
+            is_valid = (self.col.is_computed or self.col.col_type.is_media_type()) and self.col.is_stored
+            if not is_valid:
                 raise excs.Error(f'{name} only valid for a stored computed or media column: {self}')
             return ColumnPropertyRef(self, ColumnPropertyRef.Property[name.upper()])
         if (