pixeltable 0.2.24__py3-none-any.whl → 0.3.0__py3-none-any.whl

pixeltable/func/aggregate_function.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import abc
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, overload
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
@@ -35,26 +35,73 @@ class AggregateFunction(Function):
     GROUP_BY_PARAM = 'group_by'
     RESERVED_PARAMS = {ORDER_BY_PARAM, GROUP_BY_PARAM}
 
+    agg_classes: list[type[Aggregator]]  # classes for each signature, in signature order
+    init_param_names: list[list[str]]  # names of the __init__ parameters for each signature
+
     def __init__(
-            self, aggregator_class: type[Aggregator], self_path: str,
-            init_types: list[ts.ColumnType], update_types: list[ts.ColumnType], value_type: ts.ColumnType,
-            requires_order_by: bool, allows_std_agg: bool, allows_window: bool):
-        self.agg_cls = aggregator_class
+        self,
+        agg_class: type[Aggregator],
+        type_substitutions: Optional[Sequence[dict]],
+        self_path: str,
+        requires_order_by: bool,
+        allows_std_agg: bool,
+        allows_window: bool
+    ) -> None:
+        if type_substitutions is None:
+            type_substitutions = [None]  # single signature with no substitutions
+            self.agg_classes = [agg_class]
+        else:
+            self.agg_classes = [agg_class] * len(type_substitutions)
+        self.init_param_names = []
         self.requires_order_by = requires_order_by
         self.allows_std_agg = allows_std_agg
         self.allows_window = allows_window
-        self.__doc__ = aggregator_class.__doc__
+        self.__doc__ = agg_class.__doc__
+
+        signatures: list[Signature] = []
+
+        # If no type_substitutions were provided, construct a single signature for the class.
+        # Otherwise, construct one signature for each type substitution instance.
+        for subst in type_substitutions:
+            signature, init_param_names = self.__cls_to_signature(agg_class, subst)
+            signatures.append(signature)
+            self.init_param_names.append(init_param_names)
+
+        super().__init__(signatures, self_path=self_path)
+
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        self.agg_classes = [self.agg_classes[signature_idx]]
+        self.init_param_names = [self.init_param_names[signature_idx]]
+
+    def __cls_to_signature(
+        self, cls: type[Aggregator], type_substitutions: Optional[dict] = None
+    ) -> tuple[Signature, list[str]]:
+        """Inspects the Aggregator class to infer the corresponding function signature. Returns the
+        inferred signature along with the list of init_param_names (for downstream error handling).
+        """
+        # infer type parameters; set return_type=InvalidType() because it has no meaning here
+        init_sig = Signature.create(py_fn=cls.__init__, return_type=ts.InvalidType(), is_cls_method=True, type_substitutions=type_substitutions)
+        update_sig = Signature.create(py_fn=cls.update, return_type=ts.InvalidType(), is_cls_method=True, type_substitutions=type_substitutions)
+        value_sig = Signature.create(py_fn=cls.value, is_cls_method=True, type_substitutions=type_substitutions)
+
+        init_types = [p.col_type for p in init_sig.parameters.values()]
+        update_types = [p.col_type for p in update_sig.parameters.values()]
+        value_type = value_sig.return_type
+        assert value_type is not None
+
+        if len(update_types) == 0:
+            raise excs.Error('update() must have at least one parameter')
 
         # our signature is the signature of 'update', but without self,
         # plus the parameters of 'init' as keyword-only parameters
-        py_update_params = list(inspect.signature(self.agg_cls.update).parameters.values())[1:]  # leave out self
+        py_update_params = list(inspect.signature(cls.update).parameters.values())[1:]  # leave out self
         assert len(py_update_params) == len(update_types)
         update_params = [
             Parameter(p.name, col_type=update_types[i], kind=p.kind, default=p.default)
             for i, p in enumerate(py_update_params)
         ]
         # starting at 1: leave out self
-        py_init_params = list(inspect.signature(self.agg_cls.__init__).parameters.values())[1:]
+        py_init_params = list(inspect.signature(cls.__init__).parameters.values())[1:]
         assert len(py_init_params) == len(init_types)
         init_params = [
             Parameter(p.name, col_type=init_types[i], kind=inspect.Parameter.KEYWORD_ONLY, default=p.default)
@@ -67,23 +114,39 @@ class AggregateFunction(Function):
                 f'{", ".join(duplicate_params)}'
             )
         params = update_params + init_params  # init_params are keyword-only and come last
+        init_param_names = [p.name for p in init_params]
 
-        signature = Signature(value_type, params)
-        super().__init__(signature, self_path=self_path)
-        self.init_param_names = [p.name for p in init_params]
+        return Signature(value_type, params), init_param_names
 
-        # make sure the signature doesn't contain reserved parameter names;
-        # do this after super().__init__(), otherwise self.name is invalid
-        for param in signature.parameters:
-            if param.lower() in self.RESERVED_PARAMS:
-                raise excs.Error(f'{self.name}(): parameter name {param} is reserved')
+    @property
+    def agg_class(self) -> type[Aggregator]:
+        assert not self.is_polymorphic
+        return self.agg_classes[0]
 
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
+    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
         raise NotImplementedError
 
+    def overload(self, cls: type[Aggregator]) -> AggregateFunction:
+        if not isinstance(cls, type) or not issubclass(cls, Aggregator):
+            raise excs.Error(f'Invalid argument to @overload decorator: {cls}')
+        if self._has_resolved_fns:
+            raise excs.Error('New `overload` not allowed after the UDF has already been called')
+        if self._conditional_return_type is not None:
+            raise excs.Error('New `overload` not allowed after a conditional return type has been specified')
+        sig, init_param_names = self.__cls_to_signature(cls)
+        self.signatures.append(sig)
+        self.agg_classes.append(cls)
+        self.init_param_names.append(init_param_names)
+        return self
+
+    def _docstring(self) -> Optional[str]:
+        return inspect.getdoc(self.agg_classes[0])
+
     def help_str(self) -> str:
         res = super().help_str()
-        res += '\n\n' + inspect.getdoc(self.agg_cls.update)
+        # We need to reference agg_classes[0] rather than agg_class here, because we want this to work even if the
+        # aggregator is polymorphic (in which case we use the docstring of the originally decorated UDA).
+        res += '\n\n' + inspect.getdoc(self.agg_classes[0].update)
         return res
 
     def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.FunctionCall':
@@ -121,18 +184,24 @@ class AggregateFunction(Function):
                     f'{self.display_name}(): group_by invalid with an aggregate function that does not allow windows')
             group_by_clause = kwargs.pop(self.GROUP_BY_PARAM)
 
-        bound_args = self.signature.py_signature.bind(*args, **kwargs)
-        self.validate_call(bound_args.arguments)
+        resolved_fn, bound_args = self._bind_to_matching_signature(args, kwargs)
+        return_type = resolved_fn.call_return_type(args, kwargs)
         return exprs.FunctionCall(
-            self, bound_args.arguments,
+            resolved_fn,
+            bound_args,
+            return_type,
             order_by_clause=[order_by_clause] if order_by_clause is not None else [],
-            group_by_clause=[group_by_clause] if group_by_clause is not None else [])
+            group_by_clause=[group_by_clause] if group_by_clause is not None else []
+        )
 
     def validate_call(self, bound_args: dict[str, Any]) -> None:
         # check that init parameters are not Exprs
         # TODO: do this in the planner (check that init parameters are either constants or only refer to grouping exprs)
-        import pixeltable.exprs as exprs
-        for param_name in self.init_param_names:
+        from pixeltable import exprs
+
+        assert not self.is_polymorphic
+
+        for param_name in self.init_param_names[0]:
             if param_name in bound_args and isinstance(bound_args[param_name], exprs.Expr):
                 raise excs.Error(
                     f'{self.display_name}(): init() parameter {param_name} needs to be a constant, not a Pixeltable '
@@ -143,13 +212,23 @@ class AggregateFunction(Function):
         return f'<Pixeltable Aggregator {self.name}>'
 
 
+# Decorator invoked without parentheses: @pxt.uda
+@overload
+def uda(decorated_fn: Callable) -> AggregateFunction: ...
+
+
+# Decorator schema invoked with parentheses: @pxt.uda(**kwargs)
+@overload
 def uda(
-        *,
-        value_type: ts.ColumnType,
-        update_types: list[ts.ColumnType],
-        init_types: Optional[list[ts.ColumnType]] = None,
-        requires_order_by: bool = False, allows_std_agg: bool = True, allows_window: bool = False,
-) -> Callable[[type[Aggregator]], AggregateFunction]:
+    *,
+    requires_order_by: bool = False,
+    allows_std_agg: bool = True,
+    allows_window: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None
+) -> Callable[[type[Aggregator]], AggregateFunction]: ...
+
+
+def uda(*args, **kwargs):
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:
@@ -161,46 +240,50 @@ def uda(
     to the module where the class is defined.
 
     Parameters:
-    - init_types: list of types for the __init__() parameters; must match the number of parameters
-    - update_types: list of types for the update() parameters; must match the number of parameters
-    - value_type: return type of the aggregator
     - requires_order_by: if True, the first parameter to the function is the order-by expression
     - allows_std_agg: if True, the function can be used as a standard aggregate function w/o a window
     - allows_window: if True, the function can be used with a window
     """
-    if init_types is None:
-        init_types = []
-
-    def decorator(cls: type[Aggregator]) -> AggregateFunction:
-        # validate type parameters
-        num_init_params = len(inspect.signature(cls.__init__).parameters) - 1
-        if num_init_params > 0:
-            if len(init_types) != num_init_params:
-                raise excs.Error(
-                    f'init_types must be a list of {num_init_params} types, one for each parameter of __init__()')
-        num_update_params = len(inspect.signature(cls.update).parameters) - 1
-        if num_update_params == 0:
-            raise excs.Error('update() must have at least one parameter')
-        if len(update_types) != num_update_params:
-            raise excs.Error(
-                f'update_types must be a list of {num_update_params} types, one for each parameter of update()')
-        assert value_type is not None
-
-        # the AggregateFunction instance resides in the same module as cls
-        class_path = f'{cls.__module__}.{cls.__qualname__}'
-        # nonlocal name
-        # name = name or cls.__name__
-        # instance_path_elements = class_path.split('.')[:-1] + [name]
-        # instance_path = '.'.join(instance_path_elements)
-
-        # create the corresponding AggregateFunction instance
-        instance = AggregateFunction(
-            cls, class_path, init_types, update_types, value_type, requires_order_by, allows_std_agg, allows_window)
-        # do the path validation at the very end, in order to be able to write tests for the other failure cases
-        validate_symbol_path(class_path)
-        #module = importlib.import_module(cls.__module__)
-        #setattr(module, name, instance)
-
-        return instance
+    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+
+        # Decorator invoked without parentheses: @pxt.uda
+        # Simply call make_aggregator with defaults.
+        return make_aggregator(cls=args[0])
+
+    else:
+
+        # Decorator schema invoked with parentheses: @pxt.uda(**kwargs)
+        # Create a decorator for the specified schema.
+        requires_order_by = kwargs.pop('requires_order_by', False)
+        allows_std_agg = kwargs.pop('allows_std_agg', True)
+        allows_window = kwargs.pop('allows_window', False)
+        type_substitutions = kwargs.pop('type_substitutions', None)
+        if len(kwargs) > 0:
+            raise excs.Error(f'Invalid @uda decorator kwargs: {", ".join(kwargs.keys())}')
+        if len(args) > 0:
+            raise excs.Error('Unexpected @uda decorator arguments.')
+
+        def decorator(cls: type[Aggregator]) -> AggregateFunction:
+            return make_aggregator(
+                cls,
+                requires_order_by=requires_order_by,
+                allows_std_agg=allows_std_agg,
+                allows_window=allows_window,
+                type_substitutions=type_substitutions
+            )
 
-    return decorator
+        return decorator
+
+
+def make_aggregator(
+    cls: type[Aggregator],
+    requires_order_by: bool = False,
+    allows_std_agg: bool = True,
+    allows_window: bool = False,
+    type_substitutions: Optional[Sequence[dict]] = None
+) -> AggregateFunction:
+    class_path = f'{cls.__module__}.{cls.__qualname__}'
+    instance = AggregateFunction(cls, type_substitutions, class_path, requires_order_by, allows_std_agg, allows_window)
+    # do the path validation at the very end, in order to be able to write tests for the other failure cases
+    validate_symbol_path(class_path)
+    return instance

pixeltable/func/callable_function.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Sequence
 from uuid import UUID
 
 import cloudpickle  # type: ignore[import-untyped]
 
+import pixeltable.exceptions as excs
+
 from .function import Function
 from .signature import Signature
 
@@ -20,26 +22,42 @@ class CallableFunction(Function):
 
     def __init__(
         self,
-        signature: Signature,
-        py_fn: Callable,
+        signatures: list[Signature],
+        py_fns: list[Callable],
         self_path: Optional[str] = None,
         self_name: Optional[str] = None,
         batch_size: Optional[int] = None,
         is_method: bool = False,
         is_property: bool = False
     ):
-        assert py_fn is not None
-        self.py_fn = py_fn
+        assert len(signatures) > 0
+        assert len(signatures) == len(py_fns)
+        if self_path is None and len(signatures) > 1:
+            raise excs.Error('Multiple signatures are only allowed for module UDFs (not locally defined UDFs)')
+        self.py_fns = py_fns
         self.self_name = self_name
         self.batch_size = batch_size
-        self.__doc__ = py_fn.__doc__
-        super().__init__(signature, self_path=self_path, is_method=is_method, is_property=is_property)
+        self.__doc__ = self.py_fns[0].__doc__
+        super().__init__(signatures, self_path=self_path, is_method=is_method, is_property=is_property)
+
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        assert len(self.py_fns) > signature_idx
+        self.py_fns = [self.py_fns[signature_idx]]
 
     @property
     def is_batched(self) -> bool:
         return self.batch_size is not None
 
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
+    def _docstring(self) -> Optional[str]:
+        return inspect.getdoc(self.py_fns[0])
+
+    @property
+    def py_fn(self) -> Callable:
+        assert not self.is_polymorphic
+        return self.py_fns[0]
+
+    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
+        assert not self.is_polymorphic
         if self.is_batched:
             # Pack the batched parameters into singleton lists
             constant_param_names = [p.name for p in self.signature.constant_parameters]
@@ -52,13 +70,14 @@ class CallableFunction(Function):
         else:
             return self.py_fn(*args, **kwargs)
 
-    def exec_batch(self, *args: Any, **kwargs: Any) -> list:
+    def exec_batch(self, args: list[Any], kwargs: dict[str, Any]) -> list:
         """Execute the function with the given arguments and return the result.
         The arguments are expected to be batched: if the corresponding parameter has type T,
         then the argument should have type T if it's a constant parameter, or list[T] if it's
         a batched parameter.
         """
         assert self.is_batched
+        assert not self.is_polymorphic
         # Unpack the constant parameters
         constant_param_names = [p.name for p in self.signature.constant_parameters]
         constant_kwargs = {k: v[0] for k, v in kwargs.items() if k in constant_param_names}
@@ -77,10 +96,19 @@ class CallableFunction(Function):
     def name(self) -> str:
         return self.self_name
 
-    def help_str(self) -> str:
-        res = super().help_str()
-        res += '\n\n' + inspect.getdoc(self.py_fn)
-        return res
+    def overload(self, fn: Callable) -> CallableFunction:
+        if self.self_path is None:
+            raise excs.Error('`overload` can only be used with module UDFs (not locally defined UDFs)')
+        if self.is_method or self.is_property:
+            raise excs.Error('`overload` cannot be used with `is_method` or `is_property`')
+        if self._has_resolved_fns:
+            raise excs.Error('New `overload` not allowed after the UDF has already been called')
+        if self._conditional_return_type is not None:
+            raise excs.Error('New `overload` not allowed after a conditional return type has been specified')
+        sig = Signature.create(fn)
+        self.signatures.append(sig)
+        self.py_fns.append(fn)
+        return self
 
     def _as_dict(self) -> dict:
         if self.self_path is None:
@@ -99,6 +127,7 @@ class CallableFunction(Function):
         return super()._from_dict(d)
 
     def to_store(self) -> tuple[dict, bytes]:
+        assert not self.is_polymorphic  # multi-signature UDFs not allowed for stored fns
         md = {
             'signature': self.signature.as_dict(),
             'batch_size': self.batch_size,
@@ -111,12 +140,15 @@ class CallableFunction(Function):
         assert callable(py_fn)
         sig = Signature.from_dict(md['signature'])
         batch_size = md['batch_size']
-        return CallableFunction(sig, py_fn, self_name=name, batch_size=batch_size)
+        return CallableFunction([sig], [py_fn], self_name=name, batch_size=batch_size)
 
     def validate_call(self, bound_args: dict[str, Any]) -> None:
-        import pixeltable.exprs as exprs
+        from pixeltable import exprs
+
+        assert not self.is_polymorphic
         if self.is_batched:
-            for param in self.signature.constant_parameters:
+            signature = self.signatures[0]
+            for param in signature.constant_parameters:
                 if param.name in bound_args and isinstance(bound_args[param.name], exprs.Expr):
                     raise ValueError(
                         f'{self.display_name}(): '

pixeltable/func/expr_template_function.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 
 import pixeltable
 import pixeltable.exceptions as excs
@@ -8,15 +8,22 @@ from .function import Function
 from .signature import Signature
 
 
-class ExprTemplateFunction(Function):
-    """A parameterized expression from which an executable Expr is created with a function call."""
+class ExprTemplate:
+    """
+    Encapsulates a single signature of an `ExprTemplateFunction` and its associated parameterized expression,
+    along with various precomputed metadata. (This is analogous to a `Callable`-`Signature` pair in a
+    `CallableFunction`.)
+    """
+    expr: 'pixeltable.exprs.Expr'
+    signature: Signature
+    param_exprs: list['pixeltable.exprs.Variable']
+
+    def __init__(self, expr: 'pixeltable.exprs.Expr', signature: Signature):
+        from pixeltable import exprs
 
-    def __init__(
-            self, expr: 'pixeltable.exprs.Expr', signature: Signature, self_path: Optional[str] = None,
-            name: Optional[str] = None):
-        import pixeltable.exprs as exprs
         self.expr = expr
-        self.self_name = name
+        self.signature = signature
+
         self.param_exprs = list(set(expr.subexprs(expr_class=exprs.Variable)))
         # make sure there are no duplicate names
         assert len(self.param_exprs) == len(set(p.name for p in self.param_exprs))
@@ -24,7 +31,7 @@ class ExprTemplateFunction(Function):
 
         # verify default values
         self.defaults: dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
-        for param in signature.parameters.values():
+        for param in self.signature.parameters.values():
             if param.default is inspect.Parameter.empty:
                 continue
             param_expr = self.param_exprs_by_name[param.name]
@@ -35,18 +42,39 @@ class ExprTemplateFunction(Function):
                 msg = str(e)
                 raise excs.Error(f"Default value for parameter '{param.name}': {msg[0].lower() + msg[1:]}")
 
-        super().__init__(signature, self_path=self_path)
 
-    def instantiate(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+class ExprTemplateFunction(Function):
+    """A parameterized expression from which an executable Expr is created with a function call."""
+    templates: list[ExprTemplate]
+    self_name: str
+
+    def __init__(self, templates: list[ExprTemplate], self_path: Optional[str] = None, name: Optional[str] = None):
+        self.templates = templates
+        self.self_name = name
+
+        super().__init__([t.signature for t in templates], self_path=self_path)
+
+    def _update_as_overload_resolution(self, signature_idx: int) -> None:
+        self.templates = [self.templates[signature_idx]]
+
+    @property
+    def template(self) -> ExprTemplate:
+        assert not self.is_polymorphic
+        return self.templates[0]
+
+    def instantiate(self, args: Sequence[Any], kwargs: dict[str, Any]) -> 'pixeltable.exprs.Expr':
+        from pixeltable import exprs
+
+        assert not self.is_polymorphic
+        template = self.template
         bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         # apply defaults, otherwise we might have Parameters left over
         bound_args.update(
-            {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
-        result = self.expr.copy()
-        import pixeltable.exprs as exprs
+            {param_name: default for param_name, default in template.defaults.items() if param_name not in bound_args})
+        result = template.expr.copy()
         arg_exprs: dict[exprs.Expr, exprs.Expr] = {}
         for param_name, arg in bound_args.items():
-            param_expr = self.param_exprs_by_name[param_name]
+            param_expr = template.param_exprs_by_name[param_name]
             if not isinstance(arg, exprs.Expr):
                 # TODO: use the available param_expr.col_type
                 arg_expr = exprs.Expr.from_object(arg)
@@ -56,15 +84,22 @@ class ExprTemplateFunction(Function):
                 arg_expr = arg
             arg_exprs[param_expr] = arg_expr
         result = result.substitute(arg_exprs)
-        import pixeltable.exprs as exprs
         assert not result._contains(exprs.Variable)
         return result
 
-    def exec(self, *args: Any, **kwargs: Any) -> Any:
-        expr = self.instantiate(*args, **kwargs)
-        import pixeltable.exprs as exprs
+    def _docstring(self) -> Optional[str]:
+        from pixeltable import exprs
+
+        if isinstance(self.templates[0].expr, exprs.FunctionCall):
+            return self.templates[0].expr.fn._docstring()
+        return None
+
+    def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
+        from pixeltable import exec, exprs
+
+        assert not self.is_polymorphic
+        expr = self.instantiate(args, kwargs)
         row_builder = exprs.RowBuilder(output_exprs=[expr], columns=[], input_exprs=[])
-        import pixeltable.exec as exec
         row_batch = exec.DataRowBatch(tbl=None, row_builder=row_builder, len=1)
         row = row_batch[0]
         row_builder.eval(row, ctx=row_builder.default_eval_ctx)
@@ -78,11 +113,16 @@ class ExprTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
+    def __str__(self) -> str:
+        return str(self.templates[0].expr)
+
     def _as_dict(self) -> dict:
         if self.self_path is not None:
             return super()._as_dict()
+        assert not self.is_polymorphic
+        assert len(self.templates) == 1
         return {
-            'expr': self.expr.as_dict(),
+            'expr': self.template.expr.as_dict(),
             'signature': self.signature.as_dict(),
             'name': self.name,
         }
@@ -90,7 +130,8 @@ class ExprTemplateFunction(Function):
     @classmethod
     def _from_dict(cls, d: dict) -> Function:
         if 'expr' not in d:
-            return super()._from_dict(d)
+             return super()._from_dict(d)
         assert 'signature' in d and 'name' in d
         import pixeltable.exprs as exprs
-        return cls(exprs.Expr.from_dict(d['expr']), Signature.from_dict(d['signature']), name=d['name'])
+        template = ExprTemplate(exprs.Expr.from_dict(d['expr']), Signature.from_dict(d['signature']))
+        return cls([template], name=d['name'])