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

pixeltable/exprs/__init__.py

@@ -1,26 +1,29 @@
+# ruff: noqa: F401
+
 from .arithmetic_expr import ArithmeticExpr
 from .array_slice import ArraySlice
 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, BinaryMd, CellMd, DataRow
 from .expr import Expr
 from .expr_dict import ExprDict
 from .expr_set import ExprSet
 from .function_call import FunctionCall
+from .globals import ArithmeticOperator, ComparisonOperator, LogicalOperator
 from .in_predicate import InPredicate
 from .inline_expr import InlineArray, InlineDict, InlineList
 from .is_null import IsNull
-from .json_mapper import JsonMapper
-from .json_path import RELATIVE_PATH_ROOT, JsonPath
+from .json_mapper import JsonMapper, JsonMapperDispatch
+from .json_path import JsonPath
 from .literal import Literal
 from .method_ref import MethodRef
 from .object_ref import ObjectRef
-from .row_builder import RowBuilder, ColumnSlotIdx, ExecProfile
+from .row_builder import ColumnSlotIdx, ExecProfile, RowBuilder
 from .rowid_ref import RowidRef
 from .similarity_expr import SimilarityExpr
 from .sql_element_cache import SqlElementCache
+from .string_op import StringOp
 from .type_cast import TypeCast
 from .variable import Variable
-from .globals import ComparisonOperator, LogicalOperator, ArithmeticOperator

pixeltable/exprs/arithmetic_expr.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import env, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
 from .globals import ArithmeticOperator
+from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
@@ -18,6 +18,9 @@ class ArithmeticExpr(Expr):
     """
     Allows arithmetic exprs on json paths
     """
+
+    operator: ArithmeticOperator
+
     def __init__(self, operator: ArithmeticOperator, op1: Expr, op2: Expr):
         if op1.col_type.is_json_type() or op2.col_type.is_json_type() or operator == ArithmeticOperator.DIV:
             # we assume it's a float
@@ -35,49 +38,56 @@ class ArithmeticExpr(Expr):
 
         self.id = self._create_id()
 
+    @property
+    def _op1(self) -> Expr:
+        return self.components[0]
+
+    @property
+    def _op2(self) -> Expr:
+        return self.components[1]
+
     def __repr__(self) -> str:
         # add parentheses around operands that are ArithmeticExprs to express precedence
         op1_str = f'({self._op1})' if isinstance(self._op1, ArithmeticExpr) else str(self._op1)
         op2_str = f'({self._op2})' if isinstance(self._op2, ArithmeticExpr) else str(self._op2)
-        return f'{op1_str} {str(self.operator)} {op2_str}'
+        return f'{op1_str} {self.operator} {op2_str}'
 
     def _equals(self, other: ArithmeticExpr) -> bool:
         return self.operator == other.operator
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('operator', self.operator.value)]
-
-    @property
-    def _op1(self) -> Expr:
-        return self.components[0]
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
 
-    @property
-    def _op2(self) -> Expr:
-        return self.components[1]
-
-    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, sql.Float)
+            # 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
@@ -87,42 +97,66 @@ 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), sql.Integer)
-            if self.col_type.is_float_type():
-                return sql.sql.expression.cast(sql.func.floor(left / nullif), sql.Float)
-        assert False
+            # 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:
         op1_val = data_row[self._op1.slot_idx]
         op2_val = data_row[self._op2.slot_idx]
 
         # if one or both columns is JsonTyped, we need a dynamic check that they are numeric
-        if self._op1.col_type.is_json_type() and not isinstance(op1_val, int) and not isinstance(op1_val, float):
+        if self._op1.col_type.is_json_type() and op1_val is not None and not isinstance(op1_val, (int, float)):
             raise excs.Error(
-                f'{self.operator} requires numeric type, but {self._op1} has type {type(op1_val).__name__}')
-        if self._op2.col_type.is_json_type() and not isinstance(op2_val, int) and not isinstance(op2_val, float):
+                f'{self.operator} requires numeric types, but {self._op1} has type {type(op1_val).__name__}'
+            )
+        if self._op2.col_type.is_json_type() and op2_val is not None and not isinstance(op2_val, (int, float)):
             raise excs.Error(
-                f'{self.operator} requires numeric type, but {self._op2} has type {type(op2_val).__name__}')
+                f'{self.operator} requires numeric types, but {self._op2} has type {type(op2_val).__name__}'
+            )
 
-        # if either operand is None, always return None
+        data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
+
+    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
+        """
         if op1_val is None or op2_val is None:
-            data_row[self.slot_idx] = None
-            return
+            return None
+        return self.eval_non_null(op1_val, op2_val)
 
+    def eval_non_null(self, op1_val: float, op2_val: float) -> float:
+        """
+        Return the result of evaluating the expression on two int/float operands
+        """
         if self.operator == ArithmeticOperator.ADD:
-            data_row[self.slot_idx] = op1_val + op2_val
+            return op1_val + op2_val
         elif self.operator == ArithmeticOperator.SUB:
-            data_row[self.slot_idx] = op1_val - op2_val
+            return op1_val - op2_val
         elif self.operator == ArithmeticOperator.MUL:
-            data_row[self.slot_idx] = op1_val * op2_val
+            return op1_val * op2_val
         elif self.operator == ArithmeticOperator.DIV:
-            data_row[self.slot_idx] = op1_val / op2_val
+            return op1_val / op2_val
         elif self.operator == ArithmeticOperator.MOD:
-            data_row[self.slot_idx] = op1_val % op2_val
+            return op1_val % op2_val
         elif self.operator == ArithmeticOperator.FLOORDIV:
-            data_row[self.slot_idx] = op1_val // op2_val
+            return op1_val // op2_val
+
+    def as_literal(self) -> Literal | None:
+        op1_lit = self._op1.as_literal()
+        if op1_lit is None:
+            return None
+        op2_lit = self._op2.as_literal()
+        if op2_lit is None:
+            return None
+        op1_val = op1_lit.val
+        assert op1_lit.col_type.is_numeric_type() or op1_val is None
+        op2_val = op2_lit.val
+        assert op2_lit.col_type.is_numeric_type() or op2_val is None
+
+        return Literal(self.eval_nullable(op1_val, op2_val), self.col_type)  # type: ignore[arg-type]
 
     def _as_dict(self) -> dict:
         return {'operator': self.operator.value, **super()._as_dict()}

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
 
@@ -15,7 +15,8 @@ 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)
@@ -40,9 +41,9 @@ class ArraySlice(Expr):
         return self.index == other.index
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('index', self.index)]
+        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:
@@ -68,4 +69,3 @@ class ArraySlice(Expr):
             else:
                 index.append(el)
         return cls(components[0], tuple(index))
-

pixeltable/exprs/column_property_ref.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import enum
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
 import pixeltable.type_system as ts
 from pixeltable import catalog
+
 from .column_ref import ColumnRef
 from .data_row import DataRow
 from .expr import Expr
@@ -19,11 +20,13 @@ class ColumnPropertyRef(Expr):
 
     The properties themselves are type-specific and may or may not need to reference the underlying column data.
     """
+
     class Property(enum.Enum):
         ERRORTYPE = 0
         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))
@@ -31,69 +34,81 @@ 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:
         return self.prop == other.prop
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('prop', self.prop.value)]
+        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 == self.Property.ERRORTYPE or self.prop == 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:
-            assert False
+            raise AssertionError()
 
     def _as_dict(self) -> dict:
         return {'prop': self.prop.value, **super()._as_dict()}
@@ -103,4 +118,3 @@ class ColumnPropertyRef(Expr):
         assert 'prop' in d
         assert isinstance(components[0], ColumnRef)
         return cls(components[0], cls.Property(d['prop']))
-