pixeltable 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

pixeltable/func/signature.py

@@ -111,6 +111,38 @@ class Signature:
         parameters = [Parameter.from_dict(param_dict) for param_dict in d['parameters']]
         return cls(ts.ColumnType.from_dict(d['return_type']), parameters, d['is_batched'])
 
+    def is_consistent_with(self, other: Signature) -> bool:
+        """
+        Returns True if this signature is consistent with the other signature.
+        S is consistent with T if we could safely replace S by T in any call where S is used. Specifically:
+        (i) S.return_type is a supertype of T.return_type
+        (ii) For each parameter p in S, there is a parameter q in T such that:
+            - p and q have the same name and kind
+            - q.col_type is a supertype of p.col_type
+        (iii) For each *required* parameter q in T, there is a parameter p in S with the same name (in which
+            case the kinds and types must also match, by condition (ii)).
+        """
+        # Check (i)
+        if not self.get_return_type().is_supertype_of(other.get_return_type(), ignore_nullable=True):
+            return False
+
+        # Check (ii)
+        for param_name, param in self.parameters.items():
+            if param_name not in other.parameters:
+                return False
+            other_param = other.parameters[param_name]
+            if (param.kind != other_param.kind or
+                (param.col_type is None) != (other_param.col_type is None) or  # this can happen if they are varargs
+                param.col_type is not None and not other_param.col_type.is_supertype_of(param.col_type, ignore_nullable=True)):
+                return False
+
+        # Check (iii)
+        for other_param in other.required_parameters:
+            if other_param.name not in self.parameters:
+                return False
+
+        return True
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Signature):
             return False
@@ -156,8 +188,12 @@ class Signature:
 
     @classmethod
     def create_parameters(
-            cls, py_fn: Optional[Callable] = None, py_params: Optional[list[inspect.Parameter]] = None,
-            param_types: Optional[list[ts.ColumnType]] = None
+        cls,
+        py_fn: Optional[Callable] = None,
+        py_params: Optional[list[inspect.Parameter]] = None,
+        param_types: Optional[list[ts.ColumnType]] = None,
+        type_substitutions: Optional[dict] = None,
+        is_cls_method: bool = False
     ) -> list[Parameter]:
         assert (py_fn is None) != (py_params is None)
         if py_fn is not None:
@@ -165,7 +201,12 @@ class Signature:
             py_params = list(sig.parameters.values())
         parameters: list[Parameter] = []
 
+        if type_substitutions is None:
+            type_substitutions = {}
+
         for idx, param in enumerate(py_params):
+            if is_cls_method and idx == 0:
+                continue  # skip 'self' or 'cls' parameter
             if param.name in cls.SPECIAL_PARAM_NAMES:
                 raise excs.Error(f"'{param.name}' is a reserved parameter name")
             if param.kind == inspect.Parameter.VAR_POSITIONAL or param.kind == inspect.Parameter.VAR_KEYWORD:
@@ -179,7 +220,12 @@ class Signature:
                 param_type = param_types[idx]
                 is_batched = False
             else:
-                param_type, is_batched = cls._infer_type(param.annotation)
+                py_type: Optional[type]
+                if param.annotation in type_substitutions:
+                    py_type = type_substitutions[param.annotation]
+                else:
+                    py_type = param.annotation
+                param_type, is_batched = cls._infer_type(py_type)
                 if param_type is None:
                     raise excs.Error(f'Cannot infer pixeltable type for parameter {param.name}')
 
@@ -190,18 +236,29 @@ class Signature:
 
     @classmethod
     def create(
-        cls, py_fn: Callable,
+        cls,
+        py_fn: Callable,
         param_types: Optional[list[ts.ColumnType]] = None,
-        return_type: Optional[ts.ColumnType] = None
+        return_type: Optional[ts.ColumnType] = None,
+        type_substitutions: Optional[dict] = None,
+        is_cls_method: bool = False
     ) -> Signature:
         """Create a signature for the given Callable.
         Infer the parameter and return types, if none are specified.
         Raises an exception if the types cannot be inferred.
         """
-        parameters = cls.create_parameters(py_fn=py_fn, param_types=param_types)
+        if type_substitutions is None:
+            type_substitutions = {}
+
+        parameters = cls.create_parameters(py_fn=py_fn, param_types=param_types, is_cls_method=is_cls_method, type_substitutions=type_substitutions)
         sig = inspect.signature(py_fn)
         if return_type is None:
-            return_type, return_is_batched = cls._infer_type(sig.return_annotation)
+            py_type: Optional[type]
+            if sig.return_annotation in type_substitutions:
+                py_type = type_substitutions[sig.return_annotation]
+            else:
+                py_type = sig.return_annotation
+            return_type, return_is_batched = cls._infer_type(py_type)
             if return_type is None:
                 raise excs.Error('Cannot infer pixeltable return type')
         else:

pixeltable/func/udf.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Optional, overload
+from typing import Any, Callable, Optional, Sequence, overload
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
 from .callable_function import CallableFunction
-from .expr_template_function import ExprTemplateFunction
+from .expr_template_function import ExprTemplateFunction, ExprTemplate
 from .function import Function
 from .function_registry import FunctionRegistry
 from .globals import validate_symbol_path
@@ -21,13 +21,14 @@ def udf(decorated_fn: Callable) -> Function: ...
 # Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
 @overload
 def udf(
-        *,
-        batch_size: Optional[int] = None,
-        substitute_fn: Optional[Callable] = None,
-        is_method: bool = False,
-        is_property: bool = False,
-        _force_stored: bool = False
-) -> Callable[[Callable], Function]: ...
+    *,
+    batch_size: Optional[int] = None,
+    substitute_fn: Optional[Callable] = None,
+    is_method: bool = False,
+    is_property: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None,
+    _force_stored: bool = False
+) -> Callable[[Callable], CallableFunction]: ...
 
 
 def udf(*args, **kwargs):
@@ -52,6 +53,7 @@ def udf(*args, **kwargs):
         substitute_fn = kwargs.pop('substitute_fn', None)
         is_method = kwargs.pop('is_method', None)
         is_property = kwargs.pop('is_property', None)
+        type_substitutions = kwargs.pop('type_substitutions', None)
         force_stored = kwargs.pop('_force_stored', False)
         if len(kwargs) > 0:
             raise excs.Error(f'Invalid @udf decorator kwargs: {", ".join(kwargs.keys())}')
@@ -65,6 +67,7 @@ def udf(*args, **kwargs):
                 substitute_fn=substitute_fn,
                 is_method=is_method,
                 is_property=is_property,
+                type_substitutions=type_substitutions,
                 force_stored=force_stored
             )
 
@@ -79,9 +82,10 @@ def make_function(
     substitute_fn: Optional[Callable] = None,
     is_method: bool = False,
     is_property: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None,
     function_name: Optional[str] = None,
     force_stored: bool = False
-) -> Function:
+) -> CallableFunction:
     """
     Constructs a `CallableFunction` from the specified parameters.
     If `substitute_fn` is specified, then `decorated_fn`
@@ -104,25 +108,43 @@ def make_function(
     # Display name to use for error messages
     errmsg_name = function_name if function_path is None else function_path
 
-    sig = Signature.create(decorated_fn, param_types, return_type)
-
-    # batched functions must have a batched return type
-    # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
-    if batch_size is not None and not sig.is_batched:
-        raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
-    if batch_size is not None and len(sig.batched_parameters) == 0:
-        raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
-    if batch_size is None and len(sig.batched_parameters) > 0:
-        raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
-
-    if is_method and is_property:
-        raise excs.Error(f'Cannot specify both `is_method` and `is_property` (in function `{function_name}`)')
-    if is_property and len(sig.parameters) != 1:
-        raise excs.Error(
-            f"`is_property=True` expects a UDF with exactly 1 parameter, but `{function_name}` has {len(sig.parameters)}"
-        )
-    if (is_method or is_property) and function_path is None:
-        raise excs.Error('Stored functions cannot be declared using `is_method` or `is_property`')
+    signatures: list[Signature]
+    if type_substitutions is None:
+        sig = Signature.create(decorated_fn, param_types, return_type)
+
+        # batched functions must have a batched return type
+        # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
+        if batch_size is not None and not sig.is_batched:
+            raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
+        if batch_size is not None and len(sig.batched_parameters) == 0:
+            raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
+        if batch_size is None and len(sig.batched_parameters) > 0:
+            raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
+
+        if is_method and is_property:
+            raise excs.Error(f'Cannot specify both `is_method` and `is_property` (in function `{function_name}`)')
+        if is_property and len(sig.parameters) != 1:
+            raise excs.Error(
+                f"`is_property=True` expects a UDF with exactly 1 parameter, but `{function_name}` has {len(sig.parameters)}"
+            )
+        if (is_method or is_property) and function_path is None:
+            raise excs.Error('Stored functions cannot be declared using `is_method` or `is_property`')
+
+        signatures = [sig]
+    else:
+        if function_path is None:
+            raise excs.Error(
+                f'{errmsg_name}(): type substitutions can only be used with module UDFs (not locally defined UDFs)'
+            )
+        if batch_size is not None:
+            raise excs.Error(f'{errmsg_name}(): type substitutions cannot be used with batched functions')
+        if is_method is not None or is_property is not None:
+            # TODO: Support this for `is_method`?
+            raise excs.Error(f'{errmsg_name}(): type substitutions cannot be used with `is_method` or `is_property`')
+        signatures = [
+            Signature.create(decorated_fn, param_types, return_type, type_substitutions=subst)
+            for subst in type_substitutions
+        ]
 
     if substitute_fn is None:
         py_fn = decorated_fn
@@ -132,8 +154,8 @@ def make_function(
         py_fn = substitute_fn
 
     result = CallableFunction(
-        signature=sig,
-        py_fn=py_fn,
+        signatures=signatures,
+        py_fns=[py_fn] * len(signatures),  # All signatures share the same Python function
         self_path=function_path,
         self_name=function_name,
         batch_size=batch_size,
@@ -171,12 +193,12 @@ def expr_udf(*args: Any, **kwargs: Any) -> Any:
         import pixeltable.exprs as exprs
         var_exprs = [exprs.Variable(param.name, param.col_type) for param in sig.parameters.values()]
         # call the function with the parameter expressions to construct an Expr with parameters
-        template = py_fn(*var_exprs)
-        assert isinstance(template, exprs.Expr)
-        sig.return_type = template.col_type
+        expr = py_fn(*var_exprs)
+        assert isinstance(expr, exprs.Expr)
+        sig.return_type = expr.col_type
         if function_path is not None:
             validate_symbol_path(function_path)
-        return ExprTemplateFunction(template, sig, self_path=function_path, name=py_fn.__name__)
+        return ExprTemplateFunction([ExprTemplate(expr, sig)], self_path=function_path, name=py_fn.__name__)
 
     if len(args) == 1:
         assert len(kwargs) == 0 and callable(args[0])

pixeltable/functions/globals.py

@@ -1,6 +1,7 @@
 import builtins
 from typing import _GenericAlias  # type: ignore[attr-defined]
 from typing import Optional, Union
+import typing
 
 import sqlalchemy as sql
 
@@ -16,23 +17,24 @@ def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias
     return expr
 
 
-@func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=False),
-    allows_window=True, requires_order_by=False)
-class sum(func.Aggregator):
+T = typing.TypeVar('T')
+
+
+@func.uda(allows_window=True, type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
     def __init__(self):
-        self.sum: Optional[int] = None
+        self.sum: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
 
-    def value(self) -> Union[int, float]:
+    def value(self) -> T:
         return self.sum
 
 
@@ -43,12 +45,22 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.sum(val)
 
 
-@func.uda(update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(), allows_window=True, requires_order_by=False)
-class count(func.Aggregator):
+@func.uda(
+    allows_window=True,
+    # Allow counting non-null values of any type
+    # TODO: I couldn't include "Array" because we don't have a way to represent a generic array (of arbitrary dimension).
+    # TODO: should we have an "Any" type that can be used here?
+    type_substitutions=tuple(
+        {T: Optional[t]}  # type: ignore[misc]
+        for t in (ts.String, ts.Int, ts.Float, ts.Bool, ts.Timestamp,
+                  ts.Json, ts.Image, ts.Video, ts.Audio, ts.Document)
+    ),
+)
+class count(func.Aggregator, typing.Generic[T]):
     def __init__(self):
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is not None:
             self.count += 1
 
@@ -62,74 +74,82 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 
 @func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class min(func.Aggregator):
+    allows_window=True,
+    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp))  # type: ignore[misc]
+)
+class min(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.val: Optional[int] = None
+        self.val: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.min(self.val, val)
+            self.val = builtins.min(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @min.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+    if val.type.python_type == bool:
+        # TODO: min/max aggregation of booleans is not supported in Postgres (but it is in Python).
+        # Right now we simply force the computation to be done in Python; we might consider implementing an alternate
+        # way of doing it in SQL. (min/max of booleans is simply logical and/or, respectively.)
+        return None
     return sql.sql.func.min(val)
 
 
 @func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.IntType(nullable=True), allows_window=True,
-    requires_order_by=False)
-class max(func.Aggregator):
+    allows_window=True,
+    type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp))  # type: ignore[misc]
+)
+class max(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.val: Optional[int] = None
+        self.val: T = None
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.val is None:
             self.val = val
         else:
-            self.val = builtins.max(self.val, val)
+            self.val = builtins.max(self.val, val)  # type: ignore[call-overload]
 
-    def value(self) -> Optional[int]:
+    def value(self) -> T:
         return self.val
 
 
 @max.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
+    if val.type.python_type == bool:
+        # TODO: see comment in @min.to_sql.
+        return None
     return sql.sql.func.max(val)
 
 
-@func.uda(
-    update_types=[ts.IntType(nullable=True)], value_type=ts.FloatType(nullable=True), allows_window=False,
-    requires_order_by=False)
-class mean(func.Aggregator):
+@func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
+class mean(func.Aggregator, typing.Generic[T]):
     def __init__(self):
-        self.sum: Optional[int] = None
+        self.sum: T = None
         self.count = 0
 
-    def update(self, val: Optional[int]) -> None:
+    def update(self, val: T) -> None:
         if val is None:
             return
         if self.sum is None:
             self.sum = val
         else:
-            self.sum += val
+            self.sum += val  # type: ignore[operator]
         self.count += 1
 
-    def value(self) -> Optional[float]:
+    def value(self) -> Optional[float]:  # Always a float
         if self.count == 0:
             return None
-        return self.sum / self.count
+        return self.sum / self.count  # type: ignore[operator]
 
 
 @mean.to_sql

pixeltable/functions/json.py

@@ -16,20 +16,15 @@ import pixeltable as pxt
 from pixeltable.utils.code import local_public_names
 
 
-@pxt.uda(
-    update_types=[pxt.JsonType(nullable=True)],
-    value_type=pxt.JsonType(),
-    requires_order_by=False,
-    allows_window=False,
-)
+@pxt.uda
 class make_list(pxt.Aggregator):
     """
     Collects arguments into a list.
     """
-    def __init__(self):
+    def __init__(self) -> None:
         self.output: list[Any] = []
 
-    def update(self, obj: Any) -> None:
+    def update(self, obj: pxt.Json) -> None:
         if obj is None:
             return
         self.output.append(obj)

pixeltable/functions/ollama.py

@@ -34,7 +34,7 @@ def generate(
     template: str = '',
     context: Optional[list[int]] = None,
     raw: bool = False,
-    format: str = '',
+    format: Optional[str] = None,
     options: Optional[dict] = None,
 ) -> dict:
     """
@@ -44,7 +44,7 @@ def generate(
         prompt: The prompt to generate a response for.
         model: The model name.
         suffix: The text after the model response.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        format: The format of the response; must be one of `'json'` or `None`.
         system: System message.
         template: Prompt template to use.
         context: The context parameter returned from a previous call to `generate()`.
@@ -77,7 +77,7 @@ def chat(
     *,
     model: str,
     tools: Optional[list[dict]] = None,
-    format: str = '',
+    format: Optional[str] = None,
     options: Optional[dict] = None,
 ) -> dict:
     """
@@ -87,7 +87,7 @@ def chat(
         messages: The messages of the chat.
         model: The model name.
         tools: Tools for the model to use.
-        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        format: The format of the response; must be one of `'json'` or `None`.
         options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
             For details, see the
             [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)

pixeltable/functions/timestamp.py

@@ -232,7 +232,7 @@ def _(
         sql.cast(day, sql.Integer),
         sql.cast(hour, sql.Integer),
         sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Double))
+        sql.cast(second + microsecond / 1000000.0, sql.Float))
 
 # @pxt.udf
 # def date(self: datetime) -> datetime:

pixeltable/functions/video.py

@@ -47,13 +47,7 @@ _format_defaults = {  # format -> (codec, ext)
 #         output_container.mux(packet)
 
 
-@pxt.uda(
-    init_types=[pxt.IntType()],
-    update_types=[pxt.ImageType()],
-    value_type=pxt.VideoType(),
-    requires_order_by=True,
-    allows_window=False,
-)
+@pxt.uda(requires_order_by=True)
 class make_video(pxt.Aggregator):
     """
     Aggregator that creates a video from a sequence of images.
@@ -80,7 +74,7 @@ class make_video(pxt.Aggregator):
         for packet in self.stream.encode(av_frame):
             self.container.mux(packet)
 
-    def value(self) -> str:
+    def value(self) -> pxt.Video:
         for packet in self.stream.encode():
             self.container.mux(packet)
         self.container.close()

pixeltable/functions/vision.py

@@ -220,7 +220,7 @@ def eval_detections(
     return result
 
 
-@pxt.uda(update_types=[pxt.JsonType()], value_type=pxt.JsonType(), allows_std_agg=True, allows_window=False)
+@pxt.uda
 class mean_ap(pxt.Aggregator):
     """
     Calculates the mean average precision (mAP) over