pixeltable 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

pixeltable/exprs/__init__.py

@@ -9,6 +9,7 @@ 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
@@ -17,10 +18,9 @@ from .json_path import RELATIVE_PATH_ROOT, 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 .type_cast import TypeCast
 from .variable import Variable
-from .globals import ComparisonOperator, LogicalOperator, ArithmeticOperator

pixeltable/exprs/arithmetic_expr.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import sqlalchemy as sql
 
 import pixeltable.exceptions as excs
+import pixeltable.exprs as exprs
 import pixeltable.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 +20,7 @@ class ArithmeticExpr(Expr):
     """
     Allows arithmetic exprs on json paths
     """
+
     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,6 +38,14 @@ 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)
@@ -47,14 +58,6 @@ class ArithmeticExpr(Expr):
     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]
-
-    @property
-    def _op2(self) -> Expr:
-        return self.components[1]
-
     def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
         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)
@@ -73,7 +76,7 @@ class ArithmeticExpr(Expr):
             # 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)
+            return sql.sql.expression.cast(left / nullif, self.col_type.to_sa_type())
         if self.operator == ArithmeticOperator.MOD:
             if self.col_type.is_int_type():
                 nullif = sql.sql.func.nullif(right, 0)
@@ -89,9 +92,9 @@ class ArithmeticExpr(Expr):
             # 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)
+                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), sql.Float)
+                return sql.sql.expression.cast(sql.func.floor(left / nullif), self.col_type.to_sa_type())
         assert False
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
@@ -99,30 +102,58 @@ class ArithmeticExpr(Expr):
         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__}')
-
-        # if either operand is None, always return None
+                f'{self.operator} requires numeric types, but {self._op2} has type {type(op2_val).__name__}'
+            )
+
+        data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
+
+    def eval_nullable(
+        self, op1_val: Union[int, float, None], op2_val: Union[int, float, None]
+    ) -> Union[int, 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: Union[int, float], op2_val: Union[int, float]) -> Union[int, 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) -> Optional[Literal]:
+        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

@@ -15,6 +15,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], ...]):
         assert arr.col_type.is_array_type()
         # determine result type
@@ -68,4 +69,3 @@ class ArraySlice(Expr):
             else:
                 index.append(el)
         return cls(components[0], tuple(index))
-

pixeltable/exprs/column_property_ref.py

@@ -7,6 +7,7 @@ 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,6 +20,7 @@ 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
@@ -103,4 +105,3 @@ class ColumnPropertyRef(Expr):
         assert 'prop' in d
         assert isinstance(components[0], ColumnRef)
         return cls(components[0], cls.Property(d['prop']))
-

pixeltable/exprs/column_ref.py

@@ -51,8 +51,9 @@ class ColumnRef(Expr):
         super().__init__(col.col_type)
         assert col.tbl is not None
         self.col = col
-        self.is_unstored_iter_col = \
+        self.is_unstored_iter_col = (
             col.tbl.is_component_view() and col.tbl.is_iterator_column(col) and not col.is_stored
+        )
         self.iter_arg_ctx = None
         # number of rowid columns in the base table
         self.base_rowid_len = col.tbl.base.num_rowid_columns() if self.is_unstored_iter_col else 0
@@ -83,10 +84,11 @@ class ColumnRef(Expr):
         assert len(self.iter_arg_ctx.target_slot_idxs) == 1  # a single inline dict
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return (
-            super()._id_attrs()
-            + [('tbl_id', self.col.tbl.id), ('col_id', self.col.id), ('perform_validation', self.perform_validation)]
-        )
+        return super()._id_attrs() + [
+            ('tbl_id', self.col.tbl.id),
+            ('col_id', self.col.id),
+            ('perform_validation', self.perform_validation),
+        ]
 
     # override
     def _retarget(self, tbl_versions: dict[UUID, catalog.TableVersion]) -> ColumnRef:
@@ -99,14 +101,18 @@ class ColumnRef(Expr):
         from .column_property_ref import ColumnPropertyRef
 
         # resolve column properties
-        if name == ColumnPropertyRef.Property.ERRORTYPE.name.lower() \
-                or name == ColumnPropertyRef.Property.ERRORMSG.name.lower():
+        if (
+            name == ColumnPropertyRef.Property.ERRORTYPE.name.lower()
+            or name == ColumnPropertyRef.Property.ERRORMSG.name.lower()
+        ):
             property_is_present = self.col.is_stored and (self.col.is_computed or self.col_type.is_media_type())
             if not property_is_present:
                 raise excs.Error(f'{name} only valid for a stored computed or media column: {self}')
             return ColumnPropertyRef(self, ColumnPropertyRef.Property[name.upper()])
-        if name == ColumnPropertyRef.Property.FILEURL.name.lower() \
-                or name == ColumnPropertyRef.Property.LOCALPATH.name.lower():
+        if (
+            name == ColumnPropertyRef.Property.FILEURL.name.lower()
+            or name == ColumnPropertyRef.Property.LOCALPATH.name.lower()
+        ):
             if not self.col.col_type.is_media_type():
                 raise excs.Error(f'{name} only valid for image/video/audio/document columns: {self}')
             if self.col.is_computed and not self.col.is_stored:
@@ -115,12 +121,14 @@ class ColumnRef(Expr):
 
         if self.col_type.is_json_type():
             from .json_path import JsonPath
+
             return JsonPath(self, [name])
 
         return super().__getattr__(name)
 
     def similarity(self, item: Any, *, idx: Optional[str] = None) -> Expr:
         from .similarity_expr import SimilarityExpr
+
         return SimilarityExpr(self, item, idx_name=idx)
 
     def default_column_name(self) -> Optional[str]:
@@ -206,11 +214,11 @@ class ColumnRef(Expr):
             return
 
         # if this is a new base row, we need to instantiate a new iterator
-        if self.base_rowid != data_row.pk[:self.base_rowid_len]:
+        if self.base_rowid != data_row.pk[: self.base_rowid_len]:
             row_builder.eval(data_row, self.iter_arg_ctx)
             iterator_args = data_row[self.iter_arg_ctx.target_slot_idxs[0]]
             self.iterator = self.col.tbl.iterator_cls(**iterator_args)
-            self.base_rowid = data_row.pk[:self.base_rowid_len]
+            self.base_rowid = data_row.pk[: self.base_rowid_len]
         self.iterator.set_pos(data_row.pk[self.pos_idx])
         res = next(self.iterator)
         data_row[self.slot_idx] = res[self.col.name]
@@ -224,7 +232,7 @@ class ColumnRef(Expr):
             'tbl_id': str(tbl.id),
             'tbl_version': version,
             'col_id': self.col.id,
-            'perform_validation': self.perform_validation
+            'perform_validation': self.perform_validation,
         }
 
     @classmethod
@@ -240,6 +248,3 @@ class ColumnRef(Expr):
         col = cls.get_column(d)
         perform_validation = d['perform_validation']
         return cls(col, perform_validation=perform_validation)
-
-    def is_constant(self) -> bool:
-        return False

pixeltable/exprs/comparison.py

@@ -39,8 +39,12 @@ class Comparison(Expr):
             self.components = [op1, op2]
 
         import pixeltable.index as index
-        if self.is_search_arg_comparison and self._op2.col_type.is_string_type() \
-                and len(self._op2.val) >= index.BtreeIndex.MAX_STRING_LEN:
+
+        if (
+            self.is_search_arg_comparison
+            and self._op2.col_type.is_string_type()
+            and len(self._op2.val) >= index.BtreeIndex.MAX_STRING_LEN
+        ):
             # we can't use an index for this after all
             raise excs.Error(
                 f'String literal too long for comparison against indexed column {self._op1.col.name!r} '

pixeltable/exprs/compound_predicate.py

@@ -60,8 +60,7 @@ class CompoundPredicate(Expr):
     def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('operator', self.operator.value)]
 
-    def split_conjuncts(
-            self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Optional[Expr]]:
+    def split_conjuncts(self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Optional[Expr]]:
         if self.operator == LogicalOperator.OR or self.operator == LogicalOperator.NOT:
             return super().split_conjuncts(condition)
         matches = [op for op in self.components if condition(op)]
@@ -97,4 +96,3 @@ class CompoundPredicate(Expr):
     def _from_dict(cls, d: dict, components: list[Expr]) -> CompoundPredicate:
         assert 'operator' in d
         return cls(LogicalOperator(d['operator']), components)
-

pixeltable/exprs/data_row.py

@@ -6,10 +6,10 @@ import urllib.parse
 import urllib.request
 from typing import Any, Optional
 
-import PIL
-import PIL.Image
 import numpy as np
 import pgvector.sqlalchemy  # type: ignore[import-untyped]
+import PIL
+import PIL.Image
 import sqlalchemy as sql
 
 from pixeltable import env