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

pixeltable/exprs/arithmetic_expr.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 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
@@ -58,29 +58,36 @@ class ArithmeticExpr(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('operator', self.operator.value)]
 
-    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
         assert self.col_type.is_int_type() or self.col_type.is_float_type() or self.col_type.is_json_type()
         left = sql_elements.get(self._op1)
         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:
@@ -113,7 +118,7 @@ class ArithmeticExpr(Expr):
 
         data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
 
-    def eval_nullable(self, op1_val: Optional[float], op2_val: Optional[float]) -> Optional[float]:
+    def eval_nullable(self, op1_val: float | None, op2_val: float | None) -> float | None:
         """
         Return the result of evaluating the expression on two nullable int/float operands,
         None is interpreted as SQL NULL
@@ -139,7 +144,7 @@ class ArithmeticExpr(Expr):
         elif self.operator == ArithmeticOperator.FLOORDIV:
             return op1_val // op2_val
 
-    def as_literal(self) -> Optional[Literal]:
+    def as_literal(self) -> Literal | None:
         op1_lit = self._op1.as_literal()
         if op1_lit is None:
             return None

pixeltable/exprs/array_slice.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -16,7 +16,7 @@ class ArraySlice(Expr):
     Slice operation on an array, eg, t.array_col[:, 1:2].
     """
 
-    def __init__(self, arr: Expr, index: tuple[Union[int, slice], ...]):
+    def __init__(self, arr: Expr, index: tuple[int | slice, ...]):
         assert arr.col_type.is_array_type()
         # determine result type
         super().__init__(arr.col_type)
@@ -43,7 +43,7 @@ class ArraySlice(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return [*super()._id_attrs(), ('index', self.index)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/column_property_ref.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import enum
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -26,6 +26,7 @@ class ColumnPropertyRef(Expr):
         ERRORMSG = 1
         FILEURL = 2
         LOCALPATH = 3
+        CELLMD = 4  # JSON metadata for the cell, e.g. errortype, errormsg for media columns
 
     def __init__(self, col_ref: ColumnRef, prop: Property):
         super().__init__(ts.StringType(nullable=True))
@@ -33,7 +34,7 @@ class ColumnPropertyRef(Expr):
         self.prop = prop
         self.id = self._create_id()
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         return str(self).replace('.', '_')
 
     def _equals(self, other: ColumnPropertyRef) -> bool:
@@ -43,57 +44,69 @@ 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_error_prop(self) -> bool:
-        return self.prop in (self.Property.ERRORTYPE, self.Property.ERRORMSG)
+    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.is_stored:
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
+        if not self.col_ref.col_handle.get().is_stored:
             return None
+        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 (
-            self._col_ref.col.col_type.is_media_type()
-            and self._col_ref.col.media_validation == catalog.MediaValidation.ON_READ
-            and self.is_error_prop()
+            col.col_type.is_media_type()
+            and col.media_validation == catalog.MediaValidation.ON_READ
+            and self.is_cellmd_prop()
         ):
             return None
 
         if self.prop == self.Property.ERRORTYPE:
-            assert self._col_ref.col.sa_errortype_col is not None
-            return self._col_ref.col.sa_errortype_col
+            return col.sa_cellmd_col.op('->>')('errortype')
         if self.prop == self.Property.ERRORMSG:
-            assert self._col_ref.col.sa_errormsg_col is not None
-            return self._col_ref.col.sa_errormsg_col
+            return col.sa_cellmd_col.op('->>')('errormsg')
+        if self.prop == self.Property.CELLMD:
+            assert col.sa_cellmd_col is not None
+            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
+    def create_cellmd_exc(cls, exc: Exception) -> dict[str, str]:
+        """Create a cellmd value from an exception."""
+        return {'errortype': type(exc).__name__, 'errormsg': str(exc)}
+
     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_error_prop():
-            exc = data_row.get_exc(self._col_ref.slot_idx)
+        elif self.is_cellmd_prop():
+            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:
                 data_row[self.slot_idx] = type(exc).__name__
-            else:
+            elif self.prop == self.Property.ERRORMSG:
                 data_row[self.slot_idx] = str(exc)
+            elif self.prop == self.Property.CELLMD:
+                data_row[self.slot_idx] = self.create_cellmd_exc(exc)
+            else:
+                raise AssertionError(f'Unknown property {self.prop}')
+            return
         else:
             raise AssertionError()