pixeltable 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

pixeltable/exprs/expr.py

@@ -518,6 +518,38 @@ class Expr(abc.ABC):
             return ArithmeticExpr(op, self, Literal(other))  # type: ignore[arg-type]
         raise TypeError(f'Other must be Expr or literal: {type(other)}')
 
+    def __and__(self, other: object) -> Expr:
+        if not isinstance(other, Expr):
+            raise TypeError(f'Other needs to be an expression: {type(other)}')
+        if not other.col_type.is_bool_type():
+            raise TypeError(f'Other needs to be an expression that returns a boolean: {other.col_type}')
+        from .compound_predicate import CompoundPredicate
+        return CompoundPredicate(LogicalOperator.AND, [self, other])
+
+    def __or__(self, other: object) -> Expr:
+        if not isinstance(other, Expr):
+            raise TypeError(f'Other needs to be an expression: {type(other)}')
+        if not other.col_type.is_bool_type():
+            raise TypeError(f'Other needs to be an expression that returns a boolean: {other.col_type}')
+        from .compound_predicate import CompoundPredicate
+        return CompoundPredicate(LogicalOperator.OR, [self, other])
+
+    def __invert__(self) -> Expr:
+        from .compound_predicate import CompoundPredicate
+        return CompoundPredicate(LogicalOperator.NOT, [self])
+
+    def split_conjuncts(
+            self, condition: Callable[[Expr], bool]) -> tuple[list[Expr], Optional[Expr]]:
+        """
+        Returns clauses of a conjunction that meet condition in the first element.
+        The second element contains remaining clauses, rolled into a conjunction.
+        """
+        assert self.col_type.is_bool_type()  # only valid for predicates
+        if condition(self):
+            return [self], None
+        else:
+            return [], self
+
     def _make_applicator_function(self, fn: Callable, col_type: Optional[ts.ColumnType]) -> 'pixeltable.func.Function':
         """
         Creates a unary pixeltable `Function` that encapsulates a python `Callable`. The result type of

pixeltable/exprs/in_predicate.py

@@ -5,20 +5,20 @@ from typing import Optional, List, Any, Dict, Tuple, Iterable
 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 .predicate import Predicate
 from .row_builder import RowBuilder
 
 
-class InPredicate(Predicate):
+class InPredicate(Expr):
     """Predicate corresponding to the SQL IN operator."""
 
     def __init__(self, lhs: Expr, value_set_literal: Optional[Iterable] = None, value_set_expr: Optional[Expr] = None):
         assert (value_set_literal is None) != (value_set_expr is None)
         if not lhs.col_type.is_scalar_type():
             raise excs.Error(f'isin(): only supported for scalar types, not {lhs.col_type}')
-        super().__init__()
+        super().__init__(ts.BoolType())
 
         self.value_list: Optional[list] = None  # only contains values of the correct type
         if value_set_expr is not None:

pixeltable/exprs/is_null.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
+
 from typing import Optional, List, Dict
 
 import sqlalchemy as sql
 
-from .predicate import Predicate
-from .expr import Expr
+import pixeltable.type_system as ts
 from .data_row import DataRow
+from .expr import Expr
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
 
 
-class IsNull(Predicate):
+class IsNull(Expr):
     def __init__(self, e: Expr):
-        super().__init__()
+        super().__init__(ts.BoolType())
         self.components = [e]
         self.id = self._create_id()
 

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,4 @@
-from typing import Optional, List
+from typing import Optional, List, Any
 
 import sqlalchemy as sql
 import PIL.Image
@@ -14,33 +14,44 @@ from .row_builder import RowBuilder
 
 class SimilarityExpr(Expr):
 
-    def __init__(self, col_ref: ColumnRef, item: Expr):
+    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: Optional[str] = None):
         super().__init__(ts.FloatType())
-        self.components = [col_ref, item]
+        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()
-        assert item.col_type.is_string_type() or item.col_type.is_image_type()
 
         # determine index to use
         idx_info = col_ref.col.get_idx_info()
         import pixeltable.index as index
-        embedding_idx_info = [info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)]
+        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}')
+            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:
-            raise excs.Error(
-                f'Column {col_ref.col.name} has multiple indices; use the index name to disambiguate, '
-                f'e.g., `{col_ref.col.name}.<index-name>.similarity(...)`')
-        self.idx_info = embedding_idx_info[0]
+            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
 
-        if item.col_type.is_string_type() and idx.txt_embed is None:
+        if item_expr.col_type.is_string_type() and idx.string_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
-                f'text_embed parameter and does not support text queries')
-        if item.col_type.is_image_type() and idx.img_embed is None:
+                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:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
-                f'img_embed parameter and does not support image queries')
+                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")
 
     def __str__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'

pixeltable/func/aggregate_function.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
 import abc
-import importlib
 import inspect
-from typing import Optional, Any, Type, List, Dict, Callable
-import itertools
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+
 from .function import Function
-from .signature import Signature, Parameter
 from .globals import validate_symbol_path
+from .signature import Parameter, Signature
+
+if TYPE_CHECKING:
+    import pixeltable
 
 
 class Aggregator(abc.ABC):
@@ -40,6 +42,7 @@ class AggregateFunction(Function):
         self.requires_order_by = requires_order_by
         self.allows_std_agg = allows_std_agg
         self.allows_window = allows_window
+        self.__doc__ = aggregator_class.__doc__
 
         # our signature is the signature of 'update', but without self,
         # plus the parameters of 'init' as keyword-only parameters
@@ -135,6 +138,9 @@ class AggregateFunction(Function):
                     f'expression'
                 )
 
+    def __repr__(self) -> str:
+        return f'<Pixeltable Aggregator {self.name}>'
+
 
 def uda(
         *,

pixeltable/func/callable_function.py

@@ -25,6 +25,7 @@ class CallableFunction(Function):
         self.py_fn = py_fn
         self.self_name = self_name
         self.batch_size = batch_size
+        self.__doc__ = py_fn.__doc__
         super().__init__(signature, self_path=self_path)
 
     @property
@@ -113,3 +114,6 @@ class CallableFunction(Function):
                         f'{self.display_name}(): '
                         f'parameter {param.name} must be a constant value, not a Pixeltable expression'
                     )
+
+    def __repr__(self) -> str:
+        return f'<Pixeltable UDF {self.name}>'

pixeltable/func/function_registry.py

@@ -66,6 +66,8 @@ class FunctionRegistry:
     #             self.module_fns[fn_path] = obj
 
     def register_function(self, fqn: str, fn: Function) -> None:
+        if fqn in self.module_fns:
+            raise excs.Error(f'A UDF with that name already exists: {fqn}')
         self.module_fns[fqn] = fn
 
     def list_functions(self) -> List[Function]:

pixeltable/functions/globals.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 import pixeltable.func as func
 import pixeltable.type_system as ts
@@ -14,6 +14,7 @@ def cast(expr: exprs.Expr, target_type: ts.ColumnType) -> exprs.Expr:
 
 @func.uda(update_types=[ts.IntType()], value_type=ts.IntType(), allows_window=True, requires_order_by=False)
 class sum(func.Aggregator):
+    """Sums the selected integers or floats."""
     def __init__(self):
         self.sum: Union[int, float] = 0
 
@@ -38,6 +39,40 @@ class count(func.Aggregator):
         return self.count
 
 
+@func.uda(update_types=[ts.FloatType()], value_type=ts.FloatType(nullable=True), allows_window=True, requires_order_by=False)
+class max(func.Aggregator):
+    def __init__(self):
+        self.val = None
+
+    def update(self, val: Optional[float]) -> None:
+        if val is not None:
+            if self.val is None:
+                self.val = val
+            else:
+                import builtins
+                self.val = builtins.max(self.val, val)
+
+    def value(self) -> Optional[float]:
+        return self.val
+
+
+@func.uda(update_types=[ts.FloatType()], value_type=ts.FloatType(nullable=True), allows_window=True, requires_order_by=False)
+class min(func.Aggregator):
+    def __init__(self):
+        self.val = None
+
+    def update(self, val: Optional[float]) -> None:
+        if val is not None:
+            if self.val is None:
+                self.val = val
+            else:
+                import builtins
+                self.val = builtins.min(self.val, val)
+
+    def value(self) -> Optional[float]:
+        return self.val
+
+
 @func.uda(update_types=[ts.IntType()], value_type=ts.FloatType(), allows_window=False, requires_order_by=False)
 class mean(func.Aggregator):
     def __init__(self):

pixeltable/functions/huggingface.py

@@ -1,3 +1,12 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various models from the Hugging Face `transformers` package.
+
+These UDFs will cause Pixeltable to invoke the relevant models locally. In order to use them, you must
+first `pip install transformers` (or in some cases, `sentence-transformers`, as noted in the specific
+UDFs).
+"""
+
 from typing import Callable, TypeVar, Optional, Any
 
 import PIL.Image
@@ -13,15 +22,39 @@ from pixeltable.utils.code import local_public_names
 
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
 def sentence_transformer(
-    sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
+    sentence: Batch[str], *, model_id: str, normalize_embeddings: bool = False
 ) -> Batch[np.ndarray]:
-    """Runs the specified sentence transformer model."""
+    """
+    Runs the specified pretrained sentence-transformers model. `model_id` should be a pretrained model, as described
+    in the [Sentence Transformers Pretrained Models](https://sbert.net/docs/sentence_transformer/pretrained_models.html)
+    documentation.
+
+    __Requirements:__
+
+    - `pip install sentence-transformers`
+
+    Args:
+        sentence: The sentence to embed.
+        model_id: The pretrained model to use for the encoding.
+        normalize_embeddings: If `True`, normalizes embeddings to length 1; see the
+            [Sentence Transformers API Docs](https://sbert.net/docs/package_reference/sentence_transformer/SentenceTransformer.html)
+            for more details
+
+    Returns:
+        An array containing the output of the embedding model.
+
+    Examples:
+        Add a computed column that applies the model `all-mpnet-base-2` to an existing Pixeltable column `tbl.sentence`
+        of the table `tbl`:
+
+        >>> tbl['result'] = sentence_transformer(tbl.sentence, model_id='all-mpnet-base-v2')
+    """
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import SentenceTransformer
 
     model = _lookup_model(model_id, SentenceTransformer)
 
-    array = model.encode(sentences, normalize_embeddings=normalize_embeddings)
+    array = model.encode(sentence, normalize_embeddings=normalize_embeddings)
     return [array[i] for i in range(array.shape[0])]
 
 
@@ -49,7 +82,32 @@ def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embed
 
 @pxt.udf(batch_size=32)
 def cross_encoder(sentences1: Batch[str], sentences2: Batch[str], *, model_id: str) -> Batch[float]:
-    """Runs the specified cross-encoder model."""
+    """
+    Runs the specified cross-encoder model to compute similarity scores for pairs of sentences.
+    `model_id` should be a pretrained model, as described in the
+    [Cross-Encoder Pretrained Models](https://www.sbert.net/docs/cross_encoder/pretrained_models.html)
+    documentation.
+
+    __Requirements:__
+
+    - `pip install sentence-transformers`
+
+    Parameters:
+        sentences1: The first sentence to be paired.
+        sentences2: The second sentence to be paired.
+        model_id: The identifier of the cross-encoder model to use.
+
+    Returns:
+        The similarity score between the inputs.
+
+    Examples:
+        Add a computed column that applies the model `ms-marco-MiniLM-L-4-v2` to the sentences in
+        columns `tbl.sentence1` and `tbl.sentence2`:
+
+        >>> tbl['result'] = sentence_transformer(
+                tbl.sentence1, tbl.sentence2, model_id='ms-marco-MiniLM-L-4-v2'
+            )
+    """
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import CrossEncoder
 

pixeltable/functions/image.py

@@ -1,3 +1,15 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `ImageType`.
+
+Example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.img_col.convert('L')).collect()
+```
+"""
+
 import base64
 from typing import Optional, Tuple
 
@@ -41,6 +53,11 @@ def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.
 # Image.convert()
 @func.udf
 def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
+    """
+    Convert the image to a different mode.
+
+    Equivalent to [`PIL.Image.Image.convert()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.convert).
+    """
     return self.convert(mode)
 
 

pixeltable/functions/openai.py

@@ -141,7 +141,7 @@ def chat_completions(
 
 
 @pxt.udf
-def vision(prompt: str, image: PIL.Image.Image, *, model: str = 'gpt-4-vision-preview') -> str:
+def vision(prompt: str, image: PIL.Image.Image, *, model: str) -> str:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     bytes_arr = io.BytesIO()
     image.save(bytes_arr, format='png')