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

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,6 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import sqlalchemy as sql
 
@@ -12,77 +14,90 @@ from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
+if TYPE_CHECKING:
+    from pixeltable.catalog.table_version import TableVersion
+
 
 class SimilarityExpr(Expr):
+    """
+    A similarity expression against an embedding index.
+    """
+
+    idx_id: int
+    idx_name: str
+
+    def __init__(self, col_ref: ColumnRef, item: Expr, idx_name: str | None = None):
+        from pixeltable.index import EmbeddingIndex
 
-    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: Optional[str] = None):
         super().__init__(ts.FloatType())
-        item_expr = Expr.from_object(item)
-        if item_expr is None or not(item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()):
-            raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not a {type(item)}')
-        assert item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()
 
-        self.components = [col_ref, item_expr]
-        self.id = self._create_id()
+        self.components = [col_ref, item]
 
         # determine index to use
-        idx_info = col_ref.col.get_idx_info()
-        from pixeltable import index
-        embedding_idx_info = {
-            info.name: info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)
-        }
-        if len(embedding_idx_info) == 0:
-            raise excs.Error(f'No index found for column {col_ref.col!r}')
-        if idx_name is not None and idx_name not in embedding_idx_info:
-            raise excs.Error(f'Index {idx_name!r} not found for column {col_ref.col.name!r}')
-        if len(embedding_idx_info) > 1:
-            if idx_name is None:
-                raise excs.Error(
-                    f'Column {col_ref.col.name!r} has multiple indices; use the index name to disambiguate: '
-                    f'`{col_ref.col.name}.similarity(..., idx=<name>)`')
-            self.idx_info = embedding_idx_info[idx_name]
-        else:
-            self.idx_info = next(iter(embedding_idx_info.values()))
-        idx = self.idx_info.idx
-        assert isinstance(idx, index.EmbeddingIndex)
-
-        if item_expr.col_type.is_string_type() and idx.string_embed is None:
-            raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} was created without the '
-                f"'string_embed' parameter and does not support string queries")
-        if item_expr.col_type.is_image_type() and idx.image_embed is None:
+        idx_info = col_ref.tbl.get().get_idx(col_ref.col, idx_name, EmbeddingIndex)
+        self.idx_id = idx_info.id
+        self.idx_name = idx_info.name
+        idx = idx_info.idx
+        assert isinstance(idx, EmbeddingIndex)
+
+        if item.col_type._type not in idx.embeddings:
+            type_str = item.col_type._type.name.lower()
+            article = 'an' if type_str[0] in 'aeiou' else 'a'
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} was created without the '
-                f"'image_embed' parameter and does not support image queries")
+                f'Embedding index {idx_info.name!r} on column {idx_info.col.name!r} does not have {article} '
+                f'{type_str} embedding and does not support {type_str} queries'
+            )
+        self.id = self._create_id()
 
     def __repr__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'
 
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return [*super()._id_attrs(), ('idx_id', self.idx_id)]
+
     def default_column_name(self) -> str:
         return 'similarity'
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
+
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
-             raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not an expression')
-        item = self.components[1].val
-        from pixeltable import index
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
+            raise excs.Error('similarity(): requires a value, not an expression')
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.similarity_clause(idx_info.val_col, self.components[1])
+
+    def as_order_by_clause(self, is_asc: bool) -> sql.ColumnElement | None:
+        from pixeltable.index import EmbeddingIndex
 
-    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ColumnElement]:
+        # check for a literal here, instead of the c'tor: needed for ExprTemplateFunctions
         if not isinstance(self.components[1], Literal):
-            raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not an expression')
-        item = self.components[1].val
-        from pixeltable import index
-        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
-        return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
+            raise excs.Error('similarity(): requires a value, not an expression')
+        idx_info = self._resolve_idx()
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info.idx.order_by_clause(idx_info.val_col, self.components[1], is_asc)
+
+    def _resolve_idx(self) -> 'TableVersion.IndexInfo':
+        from pixeltable.index import EmbeddingIndex
+
+        # resolve idx_id
+        col_ref = self.components[0]
+        if self.idx_id not in col_ref.tbl.get().idxs:
+            raise excs.Error(f'Index {self.idx_name!r} not found')
+        idx_info = col_ref.tbl.get().idxs[self.idx_id]
+        assert isinstance(idx_info.idx, EmbeddingIndex)
+        return idx_info
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        # this should never get called
-        assert False
+        raise excs.Error('similarity(): cannot be used in a computed column')
+
+    def _as_dict(self) -> dict:
+        return {'idx_name': self.idx_name, **super()._as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> 'SimilarityExpr':
+        idx_name = d.get('idx_name')
         assert len(components) == 2
         assert isinstance(components[0], ColumnRef)
-        return cls(components[0], components[1])
+        return cls(components[0], components[1], idx_name=idx_name)

pixeltable/exprs/sql_element_cache.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Union, Optional, cast
+from typing import Iterable
 
 import sqlalchemy as sql
 
@@ -9,19 +9,19 @@ from .expr_dict import ExprDict
 class SqlElementCache:
     """Cache of sql.ColumnElements for exprs"""
 
-    cache: dict[int, Optional[sql.ColumnElement]]  # key: Expr.id
+    cache: dict[int, sql.ColumnElement | None]  # key: Expr.id
 
-    def __init__(self, elements: Optional[ExprDict[sql.ColumnElement]] = None):
+    def __init__(self, elements: ExprDict[sql.ColumnElement] | None = None):
         self.cache = {}
         if elements is not None:
             for e, el in elements.items():
                 self.cache[e.id] = el
 
-    def extend(self, elements: ExprDict[sql.ColumnElement]):
+    def extend(self, elements: ExprDict[sql.ColumnElement]) -> None:
         for e, el in elements.items():
             self.cache[e.id] = el
 
-    def get(self, e: Expr) -> Optional[sql.ColumnElement]:
+    def get(self, e: Expr) -> sql.ColumnElement | None:
         """Returns the sql.ColumnElement for the given Expr, or None if Expr.to_sql() returns None."""
         try:
             return self.cache[e.id]

pixeltable/exprs/string_op.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Any
+
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+from .data_row import DataRow
+from .expr import Expr
+from .globals import StringOperator
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
+
+
+class StringOp(Expr):
+    """
+    Allows operations on strings
+    """
+
+    operator: StringOperator
+
+    def __init__(self, operator: StringOperator, op1: Expr, op2: Expr):
+        super().__init__(ts.StringType(nullable=op1.col_type.nullable))
+        self.operator = operator
+        self.components = [op1, op2]
+        assert op1.col_type.is_string_type()
+        if operator in (StringOperator.CONCAT, StringOperator.REPEAT):
+            if operator == StringOperator.CONCAT and not op2.col_type.is_string_type():
+                raise excs.Error(
+                    f'{self}: {operator} on strings requires string type, but {op2} has type {op2.col_type}'
+                )
+            if operator == StringOperator.REPEAT and not op2.col_type.is_int_type():
+                raise excs.Error(f'{self}: {operator} on strings requires int type, but {op2} has type {op2.col_type}')
+        else:
+            raise excs.Error(
+                f'{self}: invalid operation {operator} on strings; '
+                f'only operators {StringOperator.CONCAT} and {StringOperator.REPEAT} are supported'
+            )
+        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 StringOpExpr to express precedence
+        op1_str = f'({self._op1})' if isinstance(self._op1, StringOp) else str(self._op1)
+        op2_str = f'({self._op2})' if isinstance(self._op2, StringOp) else str(self._op2)
+        return f'{op1_str} {self.operator} {op2_str}'
+
+    def _equals(self, other: StringOp) -> bool:
+        return self.operator == other.operator
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
+
+    def sql_expr(self, sql_elements: SqlElementCache) -> sql.ColumnElement | None:
+        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 == StringOperator.CONCAT:
+            return left.concat(right)
+        if self.operator == StringOperator.REPEAT:
+            return sql.func.repeat(left.cast(sql.String), right.cast(sql.Integer))
+        return None
+
+    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]
+        data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
+
+    def eval_nullable(self, op1_val: str | None, op2_val: int | str | None) -> str | 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:
+            return None
+        return self.eval_non_null(op1_val, op2_val)
+
+    def eval_non_null(self, op1_val: str, op2_val: int | str) -> str:
+        """
+        Return the result of evaluating the expression on two int/float operands
+        """
+        assert self.operator in (StringOperator.CONCAT, StringOperator.REPEAT)
+        if self.operator == StringOperator.CONCAT:
+            assert isinstance(op2_val, str)
+            return op1_val + op2_val
+        else:
+            assert isinstance(op2_val, int)
+            return op1_val * op2_val
+
+    def _as_dict(self) -> dict:
+        return {'operator': self.operator.value, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> StringOp:
+        assert 'operator' in d
+        assert len(components) == 2
+        return cls(StringOperator(d['operator']), components[0], components[1])

pixeltable/exprs/type_cast.py

@@ -1,10 +1,9 @@
-from typing import Any, Optional
-
 import sqlalchemy as sql
 
-import pixeltable.type_system as ts
+from pixeltable import type_system as ts
 
 from .expr import DataRow, Expr
+from .literal import Literal
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
@@ -14,23 +13,21 @@ class TypeCast(Expr):
     An `Expr` that represents a type conversion from an underlying `Expr` to
     a specified `ColumnType`.
     """
+
     def __init__(self, underlying: Expr, new_type: ts.ColumnType):
         super().__init__(new_type)
         self.components: list[Expr] = [underlying]
-        self.id: Optional[int] = self._create_id()
-
-    @property
-    def _underlying(self):
-        return self.components[0]
+        self.id: int | None = self._create_id()
 
     def _equals(self, other: 'TypeCast') -> bool:
         # `TypeCast` has no properties beyond those captured by `Expr`.
         return True
 
-    def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('new_type', self.col_type)]
+    @property
+    def _op1(self) -> Expr:
+        return self.components[0]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         """
         sql_expr() is unimplemented for now, in order to sidestep potentially thorny
         questions about consistency of doing type conversions in both Python and Postgres.
@@ -38,9 +35,24 @@ class TypeCast(Expr):
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        original_val = data_row[self._underlying.slot_idx]
+        original_val = data_row[self._op1.slot_idx]
         data_row[self.slot_idx] = self.col_type.create_literal(original_val)
 
+    def as_literal(self) -> Literal | None:
+        op1_lit = self._op1.as_literal()
+        if op1_lit is None:
+            return None
+        if not (
+            self.col_type.is_numeric_type() and (op1_lit.col_type.is_numeric_type() or op1_lit.col_type.is_bool_type())
+        ):
+            return None
+
+        op1_val = op1_lit.val
+        if self.col_type.is_int_type():
+            return Literal(int(op1_val), self.col_type)
+        elif self.col_type.is_float_type():
+            return Literal(float(op1_val), self.col_type)
+        return None
 
     def _as_dict(self) -> dict:
         return {'new_type': self.col_type.as_dict(), **super()._as_dict()}
@@ -52,4 +64,4 @@ class TypeCast(Expr):
         return cls(components[0], ts.ColumnType.from_dict(d['new_type']))
 
     def __repr__(self) -> str:
-        return f'{self._underlying}.astype({self.col_type._to_str(as_schema=True)})'
+        return f'{self._op1}.astype({self.col_type._to_str(as_schema=True)})'

pixeltable/exprs/variable.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any, NoReturn
 
-import pixeltable.type_system as ts
+from pixeltable import type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
@@ -22,7 +22,7 @@ class Variable(Expr):
         self.id = self._create_id()
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('name', self.name)]
+        return [*super()._id_attrs(), ('name', self.name)]
 
     def default_column_name(self) -> NoReturn:
         raise NotImplementedError()

pixeltable/func/__init__.py

@@ -1,8 +1,12 @@
-from .aggregate_function import Aggregator, AggregateFunction, uda
+# ruff: noqa: F401
+
+from .aggregate_function import AggregateFunction, Aggregator, uda
 from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
-from .function import Function
+from .function import Function, InvalidFunction
 from .function_registry import FunctionRegistry
-from .query_template_function import QueryTemplateFunction
-from .signature import Signature, Parameter, Batch
-from .udf import udf, make_function, expr_udf
+from .mcp import mcp_udfs
+from .query_template_function import QueryTemplateFunction, query, retrieval_udf
+from .signature import Batch, Parameter, Signature
+from .tools import Tool, ToolChoice, Tools
+from .udf import expr_udf, make_function, udf