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

pixeltable/exec/sql_node.py

@@ -1,13 +1,14 @@
 import logging
 import warnings
 from decimal import Decimal
-from typing import Iterable, Iterator, NamedTuple, Optional, TYPE_CHECKING, Sequence, AsyncIterator
+from typing import TYPE_CHECKING, AsyncIterator, Iterable, Iterator, NamedTuple, Optional, Sequence
 from uuid import UUID
 
 import sqlalchemy as sql
 
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
@@ -53,10 +54,12 @@ def combine_order_by_clauses(clauses: Iterable[OrderByClause]) -> Optional[Order
 
 
 def print_order_by_clause(clause: OrderByClause) -> str:
-    return ', '.join([
-        f'({item.expr}{", asc=True" if item.asc is True else ""}{", asc=False" if item.asc is False else ""})'
-        for item in clause
-    ])
+    return ', '.join(
+        [
+            f'({item.expr}{", asc=True" if item.asc is True else ""}{", asc=False" if item.asc is False else ""})'
+            for item in clause
+        ]
+    )
 
 
 class SqlNode(ExecNode):
@@ -82,8 +85,12 @@ class SqlNode(ExecNode):
     limit: Optional[int]
 
     def __init__(
-            self, tbl: Optional[catalog.TableVersionPath], row_builder: exprs.RowBuilder,
-            select_list: Iterable[exprs.Expr], sql_elements: exprs.SqlElementCache, set_pk: bool = False
+        self,
+        tbl: Optional[catalog.TableVersionPath],
+        row_builder: exprs.RowBuilder,
+        select_list: Iterable[exprs.Expr],
+        sql_elements: exprs.SqlElementCache,
+        set_pk: bool = False,
     ):
         """
         If row_builder contains references to unstored iter columns, expands the select list to include their
@@ -186,8 +193,11 @@ class SqlNode(ExecNode):
 
     @classmethod
     def create_from_clause(
-            cls, tbl: catalog.TableVersionPath, stmt: sql.Select, refd_tbl_ids: Optional[set[UUID]] = None,
-            exact_version_only: Optional[set[UUID]] = None
+        cls,
+        tbl: catalog.TableVersionPath,
+        stmt: sql.Select,
+        refd_tbl_ids: Optional[set[UUID]] = None,
+        exact_version_only: Optional[set[UUID]] = None,
     ) -> sql.Select:
         """Add From clause to stmt for tables/views referenced by materialized_exprs
         Args:
@@ -220,15 +230,14 @@ class SqlNode(ExecNode):
                 # join tbl to prev_tbl on prev_tbl's rowid cols
                 prev_tbl_rowid_cols = prev_tbl.store_tbl.rowid_columns()
                 tbl_rowid_cols = tbl.store_tbl.rowid_columns()
-                rowid_clauses = \
-                    [c1 == c2 for c1, c2 in zip(prev_tbl_rowid_cols, tbl_rowid_cols[:len(prev_tbl_rowid_cols)])]
+                rowid_clauses = [
+                    c1 == c2 for c1, c2 in zip(prev_tbl_rowid_cols, tbl_rowid_cols[: len(prev_tbl_rowid_cols)])
+                ]
                 stmt = stmt.join(tbl.store_tbl.sa_tbl, sql.and_(*rowid_clauses))
             if tbl.id in exact_version_only:
                 stmt = stmt.where(tbl.store_tbl.v_min_col == tbl.version)
             else:
-                stmt = stmt \
-                    .where(tbl.store_tbl.v_min_col <= tbl.version) \
-                    .where(tbl.store_tbl.v_max_col > tbl.version)
+                stmt = stmt.where(tbl.store_tbl.v_min_col <= tbl.version).where(tbl.store_tbl.v_max_col > tbl.version)
             prev_tbl = tbl
         return stmt
 
@@ -291,7 +300,7 @@ class SqlNode(ExecNode):
 
             # populate output_row
             if self.num_pk_cols > 0:
-                output_row.set_pk(tuple(sql_row[-self.num_pk_cols:]))
+                output_row.set_pk(tuple(sql_row[-self.num_pk_cols :]))
             # copy the output of the SQL query into the output row
             for i, e in enumerate(self.select_list):
                 slot_idx = e.slot_idx
@@ -341,12 +350,16 @@ class SqlScanNode(SqlNode):
 
     Supports filtering and ordering.
     """
+
     exact_version_only: list[catalog.TableVersion]
 
     def __init__(
-        self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
+        self,
+        tbl: catalog.TableVersionPath,
+        row_builder: exprs.RowBuilder,
         select_list: Iterable[exprs.Expr],
-        set_pk: bool = False, exact_version_only: Optional[list[catalog.TableVersion]] = None
+        set_pk: bool = False,
+        exact_version_only: Optional[list[catalog.TableVersion]] = None,
     ):
         """
         Args:
@@ -367,7 +380,8 @@ class SqlScanNode(SqlNode):
         where_clause_tbl_ids = self.where_clause.tbl_ids() if self.where_clause is not None else set()
         refd_tbl_ids = exprs.Expr.all_tbl_ids(self.select_list) | where_clause_tbl_ids | self._ordering_tbl_ids()
         stmt = self.create_from_clause(
-            self.tbl, stmt, refd_tbl_ids, exact_version_only={t.id for t in self.exact_version_only})
+            self.tbl, stmt, refd_tbl_ids, exact_version_only={t.id for t in self.exact_version_only}
+        )
         return stmt
 
 
@@ -377,8 +391,12 @@ class SqlLookupNode(SqlNode):
     """
 
     def __init__(
-        self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
-        select_list: Iterable[exprs.Expr], sa_key_cols: list[sql.Column], key_vals: list[tuple],
+        self,
+        tbl: catalog.TableVersionPath,
+        row_builder: exprs.RowBuilder,
+        select_list: Iterable[exprs.Expr],
+        sa_key_cols: list[sql.Column],
+        key_vals: list[tuple],
     ):
         """
         Args:
@@ -406,11 +424,13 @@ class SqlAggregationNode(SqlNode):
     group_by_items: Optional[list[exprs.Expr]]
 
     def __init__(
-        self, row_builder: exprs.RowBuilder,
+        self,
+        row_builder: exprs.RowBuilder,
         input: SqlNode,
         select_list: Iterable[exprs.Expr],
         group_by_items: Optional[list[exprs.Expr]] = None,
-        limit: Optional[int] = None, exact_version_only: Optional[list[catalog.TableVersion]] = None
+        limit: Optional[int] = None,
+        exact_version_only: Optional[list[catalog.TableVersion]] = None,
     ):
         """
         Args:
@@ -436,12 +456,16 @@ class SqlJoinNode(SqlNode):
     """
     Materializes data from the store via a Select ... From ... that contains joins
     """
+
     input_ctes: list[sql.CTE]
     join_clauses: list['pixeltable.plan.JoinClause']
 
     def __init__(
-        self, row_builder: exprs.RowBuilder,
-        inputs: Sequence[SqlNode], join_clauses: list['pixeltable.plan.JoinClause'], select_list: Iterable[exprs.Expr]
+        self,
+        row_builder: exprs.RowBuilder,
+        inputs: Sequence[SqlNode],
+        join_clauses: list['pixeltable.plan.JoinClause'],
+        select_list: Iterable[exprs.Expr],
     ):
         assert len(inputs) > 1
         assert len(inputs) == len(join_clauses) + 1
@@ -456,16 +480,21 @@ class SqlJoinNode(SqlNode):
 
     def _create_stmt(self) -> sql.Select:
         from pixeltable import plan
+
         stmt = super()._create_stmt()
         stmt = stmt.select_from(self.input_ctes[0])
         for i in range(len(self.join_clauses)):
             join_clause = self.join_clauses[i]
             on_clause = (
-                self.sql_elements.get(join_clause.join_predicate) if join_clause.join_type != plan.JoinType.CROSS
+                self.sql_elements.get(join_clause.join_predicate)
+                if join_clause.join_type != plan.JoinType.CROSS
                 else sql.sql.expression.literal(True)
             )
             is_outer = join_clause.join_type == plan.JoinType.LEFT or join_clause.join_type == plan.JoinType.FULL_OUTER
             stmt = stmt.join(
-                self.input_ctes[i + 1], onclause=on_clause, isouter=is_outer,
-                full=join_clause == plan.JoinType.FULL_OUTER)
+                self.input_ctes[i + 1],
+                onclause=on_clause,
+                isouter=is_outer,
+                full=join_clause == plan.JoinType.FULL_OUTER,
+            )
         return stmt

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