pixeltable 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/func/aggregate_function.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import abc
+import importlib
+import inspect
+from typing import Optional, Any, Type, List, Dict
+import itertools
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .function import Function
+from .signature import Signature, Parameter
+
+
+class Aggregator(abc.ABC):
+    def update(self, *args: Any, **kwargs: Any) -> None:
+        pass
+    def value(self) -> Any:
+        pass
+
+
+class AggregateFunction(Function):
+    """Function interface for an aggregation operation.
+
+    requires_order_by: if True, the first parameter to an aggregate function defines the order in which the function
+    sees rows in update()
+    allows_std_agg: if True, the aggregate function can be used as a standard aggregate function w/o a window
+    allows_window: if True, the aggregate function can be used with a window
+    """
+    ORDER_BY_PARAM = 'order_by'
+    GROUP_BY_PARAM = 'group_by'
+    RESERVED_PARAMS = {ORDER_BY_PARAM, GROUP_BY_PARAM}
+
+    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.requires_order_by = requires_order_by
+        self.allows_std_agg = allows_std_agg
+        self.allows_window = allows_window
+
+        # our signature is the signature of 'update', but without self,
+        # plus the parameters of 'init' as keyword-only parameters
+        update_params = list(inspect.signature(self.agg_cls.update).parameters.values())[1:]  # leave out self
+        assert len(update_params) == len(update_types)
+        init_params = [
+            inspect.Parameter(p.name, inspect.Parameter.KEYWORD_ONLY, default=p.default)
+            # starting at 1: leave out self
+            for p in itertools.islice(inspect.signature(self.agg_cls.__init__).parameters.values(), 1, None)
+        ]
+        assert len(init_params) == len(init_types)
+        duplicate_params = set(p.name for p in init_params) & set(p.name for p in update_params)
+        if len(duplicate_params) > 0:
+            raise excs.Error(
+                f'__init__() and update() cannot have parameters with the same name: '
+                f'{", ".join(duplicate_params)}'
+            )
+        py_params = update_params + init_params  # init_params are keyword-only and come last
+        py_signature = inspect.Signature(py_params)
+
+        params = [Parameter(p.name, update_types[i], p.kind, is_batched=False) for i, p in enumerate(update_params)]
+        params.extend([Parameter(p.name, init_types[i], p.kind, is_batched=False) for i, p in enumerate(init_params)])
+        signature = Signature(value_type, params)
+        super().__init__(signature, py_signature=py_signature, self_path=self_path)
+        self.init_param_names = [p.name for p in init_params]
+
+        # 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')
+
+    def help_str(self) -> str:
+        res = super().help_str()
+        res += '\n\n' + inspect.getdoc(self.agg_cls.update)
+        return res
+
+    def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+        from pixeltable import exprs
+
+        # perform semantic analysis of special parameters 'order_by' and 'group_by'
+        order_by_clause: Optional[Any] = None
+        if self.ORDER_BY_PARAM in kwargs:
+            if self.requires_order_by:
+                raise excs.Error(
+                    f'{self.display_name}(): order_by invalid, this function requires the first argument to be the '
+                    f'ordering expression'
+                )
+            if not self.allows_window:
+                raise excs.Error(
+                    f'{self.display_name}(): order_by invalid with an aggregate function that does not allow windows')
+            order_by_clause = kwargs.pop(self.ORDER_BY_PARAM)
+        elif self.requires_order_by:
+            # the first argument is the order-by expr
+            if len(args) == 0:
+                raise excs.Error(f'{self.display_name}(): requires an ordering expression as its first argument')
+            order_by_clause = args[0]
+            if not isinstance(order_by_clause, exprs.Expr):
+                raise excs.Error(
+                    f'{self.display_name}(): the first argument needs to be a Pixeltable expression, but instead is a '
+                    f'{type(order_by_clause)}'
+                )
+            # don't pass the first parameter on, the Function doesn't get to see it
+            args = args[1:]
+
+        group_by_clause: Optional[Any] = None
+        if self.GROUP_BY_PARAM in kwargs:
+            if not self.allows_window:
+                raise excs.Error(
+                    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.py_signature.bind(*args, **kwargs)
+        self.validate_call(bound_args.arguments)
+        return exprs.FunctionCall(
+            self, bound_args.arguments,
+            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 [])
+
+    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:
+            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 '
+                    f'expression'
+                )
+
+
+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,
+        name: Optional[str] = None
+) -> Type[Aggregator]:
+    """Decorator for user-defined aggregate functions.
+
+    The decorated class must inherit from Aggregator and implement the following methods:
+    - __init__(self, ...) to initialize the aggregator
+    - update(self, ...) to update the aggregator with a new value
+    - value(self) to return the final result
+
+    The decorator creates an AggregateFunction instance from the class and adds it
+    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
+    - name: name of the AggregateFunction instance; if None, the class name is used
+    """
+    if name is not None and not name.isidentifier():
+        raise excs.Error(f'Invalid name: {name}')
+    if init_types is None:
+        init_types = []
+
+    def decorator(cls: Type[Aggregator]) -> Type[Aggregator]:
+        # 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
+        module_path = cls.__module__
+        nonlocal name
+        name = name or cls.__name__
+        instance_path = f'{module_path}.{name}'
+
+        # create the corresponding AggregateFunction instance
+        instance = AggregateFunction(
+            cls, instance_path, init_types, update_types, value_type, requires_order_by, allows_std_agg, allows_window)
+        module = importlib.import_module(module_path)
+        setattr(module, name, instance)
+
+        return cls
+
+    return decorator

pixeltable/func/batched_function.py

@@ -0,0 +1,53 @@
+import inspect
+from typing import List, Dict, Any, Optional, Callable
+import abc
+
+from .function import Function
+from .signature import Signature
+
+
+class BatchedFunction(Function):
+    """Base class for functions that can run on batches"""
+
+    @abc.abstractmethod
+    def get_batch_size(self, *args: Any, **kwargs: Any) -> Optional[int]:
+        """Return the batch size for the given arguments, or None if the batch size is unknown.
+        args/kwargs might be empty
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def invoke(self, arg_batches: List[List[Any]], kwarg_batches: Dict[str, List[Any]]) -> List[Any]:
+        """Invoke the function for the given batch and return a batch of results"""
+        raise NotImplementedError
+
+
+class ExplicitBatchedFunction(BatchedFunction):
+    """
+    A `BatchedFunction` that is defined by a signature and an explicit python
+    `Callable`.
+    """
+    def __init__(self, signature: Signature, batch_size: Optional[int], invoker_fn: Callable, self_path: str):
+        super().__init__(signature=signature, py_signature=inspect.signature(invoker_fn), self_path=self_path)
+        self.batch_size = batch_size
+        self.invoker_fn = invoker_fn
+
+    def get_batch_size(self, *args: Any, **kwargs: Any) -> Optional[int]:
+        return self.batch_size
+
+    def invoke(self, arg_batches: List[List[Any]], kwarg_batches: Dict[str, List[Any]]) -> List[Any]:
+        """Invoke the function for the given batch and return a batch of results"""
+        constant_param_names = [p.name for p in self.signature.constant_parameters]
+        kwargs = {k: v[0] for k, v in kwarg_batches.items() if k in constant_param_names}
+        kwarg_batches = {k: v for k, v in kwarg_batches.items() if k not in constant_param_names}
+        return self.invoker_fn(*arg_batches, **kwargs, **kwarg_batches)
+
+    def validate_call(self, bound_args: Dict[str, Any]) -> None:
+        """Verify constant parameters"""
+        import pixeltable.exprs as exprs
+        for param in self.signature.constant_parameters:
+            if param.name in bound_args and isinstance(bound_args[param.name], exprs.Expr):
+                raise ValueError(
+                    f'{self.display_name}(): '
+                    f'parameter {param.name} must be a constant value, not a Pixeltable expression'
+                )

pixeltable/func/callable_function.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import inspect
+import sys
+from typing import Optional, Dict, Callable, List, Tuple
+from uuid import UUID
+import cloudpickle
+
+import pixeltable.type_system as ts
+import pixeltable.exceptions as excs
+from .function import Function
+from .function_registry import FunctionRegistry
+from .globals import get_caller_module_path
+from .signature import Signature
+
+
+class CallableFunction(Function):
+    """Pixeltable Function backed by a Python Callable.
+
+    CallableFunctions come in two flavors:
+    - references to lambdas and functions defined in notebooks, which are pickled and serialized to the store
+    - functions that are defined in modules are serialized via the default mechanism
+    """
+
+    def __init__(
+            self, signature: Signature, py_fn: Callable, self_path: Optional[str] = None,
+            self_name: Optional[str] = None):
+        assert py_fn is not None
+        self.py_fn = py_fn
+        self.self_name = self_name
+        py_signature = inspect.signature(self.py_fn)
+        super().__init__(signature, py_signature, self_path=self_path)
+
+    @property
+    def display_name(self) -> str:
+        return self.self_name
+
+    @property
+    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 _as_dict(self) -> Dict:
+        if self.self_path is None:
+            # this is not a module function
+            from .function_registry import FunctionRegistry
+            id = FunctionRegistry.get().create_stored_function(self)
+            return {'id': id.hex}
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, d: Dict) -> Function:
+        if 'id' in d:
+            from .function_registry import FunctionRegistry
+            return FunctionRegistry.get().get_stored_function(UUID(hex=d['id']))
+        return super()._from_dict(d)
+
+    def to_store(self) -> Tuple[Dict, bytes]:
+        return (self.signature.as_dict(), cloudpickle.dumps(self.py_fn))
+
+    @classmethod
+    def from_store(cls, name: Optional[str], md: Dict, binary_obj: bytes) -> Function:
+        py_fn = cloudpickle.loads(binary_obj)
+        assert isinstance(py_fn, Callable)
+        return CallableFunction(Signature.from_dict(md), py_fn, self_name=name)

pixeltable/func/expr_template_function.py

@@ -0,0 +1,82 @@
+import inspect
+from typing import Dict, Optional, Callable, List
+
+import pixeltable
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .function import Function
+from .signature import Signature, Parameter
+
+
+class ExprTemplateFunction(Function):
+    """A parameterized expression from which an executable Expr is created with a function call."""
+
+    def __init__(
+            self, expr: 'pixeltable.exprs.Expr', py_signature: inspect.Signature, self_path: Optional[str] = None,
+            name: Optional[str] = None):
+        import pixeltable.exprs as exprs
+        self.expr = expr
+        self.self_name = name
+        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))
+        self.param_exprs_by_name = {p.name: p for p in self.param_exprs}
+
+        # verify default values
+        self.defaults: Dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
+        for py_param in py_signature.parameters.values():
+            if py_param.default is inspect.Parameter.empty:
+                continue
+            param_expr = self.param_exprs_by_name[py_param.name]
+            try:
+                literal_default = exprs.Literal(py_param.default, col_type=param_expr.col_type)
+                self.defaults[py_param.name] = literal_default
+            except TypeError as e:
+                msg = str(e)
+                raise excs.Error(f"Default value for parameter '{py_param.name}': {msg[0].lower() + msg[1:]}")
+        # construct signature
+        assert len(self.param_exprs) == len(py_signature.parameters)
+        fn_params = [
+            Parameter(p.name, self.param_exprs_by_name[p.name].col_type, p.kind)
+            for p in py_signature.parameters.values()
+        ]
+        signature = Signature(return_type=expr.col_type, parameters=fn_params)
+
+        super().__init__(signature, py_signature=py_signature, self_path=self_path)
+
+    def instantiate(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+        bound_args = self.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()
+        for param_name, arg in bound_args.items():
+            param_expr = self.param_exprs_by_name[param_name]
+            result = result.substitute(param_expr, arg)
+        import pixeltable.exprs as exprs
+        assert not result.contains(exprs.Variable)
+        return result
+
+    @property
+    def display_name(self) -> str:
+        return self.self_name
+
+    @property
+    def name(self) -> str:
+        return self.self_name
+
+    def _as_dict(self) -> Dict:
+        if self.self_path is not None:
+            return super()._as_dict()
+        return {
+            'name': self.name,
+            'expr': self.expr.as_dict(),
+            **super()._as_dict()
+        }
+
+    @classmethod
+    def _from_dict(cls, d: Dict) -> Function:
+        if 'expr' not in d:
+            return super()._from_dict(d)
+        import pixeltable.exprs as exprs
+        return cls(exprs.Expr.from_dict(d['expr']), name=d['name'])

pixeltable/func/function.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import abc
+import importlib
+import inspect
+import pixeltable
+from typing import Optional, Dict, Any, Tuple
+
+from .globals import resolve_symbol
+from .signature import Signature
+
+
+class Function(abc.ABC):
+    """Base class for Pixeltable's function interface.
+
+    A function in Pixeltable is an object that has a signature and implements __call__().
+    This base class provides a default serialization mechanism for Function instances provided by Python modules,
+    via the member self_path.
+    """
+
+    def __init__(self, signature: Signature, py_signature: inspect.Signature, self_path: Optional[str] = None):
+        self.signature = signature
+        self.py_signature = py_signature
+        self.self_path = self_path  # fully-qualified path to self
+
+    @property
+    def name(self) -> str:
+        assert self.self_path is not None
+        return self.self_path.split('.')[-1]
+
+    @property
+    def display_name(self) -> str:
+        if self.self_path is None:
+            return '<anonymous>'
+        ptf_prefix = 'pixeltable.functions.'
+        if self.self_path.startswith(ptf_prefix):
+            return self.self_path[len(ptf_prefix):]
+        return self.self_path
+
+    def help_str(self) -> str:
+        return self.display_name + str(self.signature)
+
+    def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+        from pixeltable import exprs
+        bound_args = self.py_signature.bind(*args, **kwargs)
+        self.validate_call(bound_args.arguments)
+        return exprs.FunctionCall(self, bound_args.arguments)
+
+    def validate_call(self, bound_args: Dict[str, Any]) -> None:
+        """Override this to do custom validation of the arguments"""
+        pass
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        return self.self_path == other.self_path
+
+    def source(self) -> None:
+        """Print source code"""
+        print('source not available')
+
+    def as_dict(self) -> Dict:
+        """
+        Return a serialized reference to the instance that can be passed to json.dumps() and converted back
+        to an instance with from_dict().
+        Subclasses can override _as_dict().
+        """
+        classpath = f'{self.__class__.__module__}.{self.__class__.__qualname__}'
+        return {'_classpath': classpath, **self._as_dict()}
+
+    def _as_dict(self) -> Dict:
+        """Default serialization: store the path to self (which includes the module path)"""
+        assert self.self_path is not None
+        return {'path': self.self_path}
+
+    @classmethod
+    def from_dict(cls, d: Dict) -> Function:
+        """
+        Turn dict that was produced by calling as_dict() into an instance of the correct Function subclass.
+        """
+        assert '_classpath' in d
+        module_path, class_name = d['_classpath'].rsplit('.', 1)
+        class_module = importlib.import_module(module_path)
+        func_class = getattr(class_module, class_name)
+        return func_class._from_dict(d)
+
+    @classmethod
+    def _from_dict(cls, d: Dict) -> Function:
+        """Default deserialization: load the symbol indicated by the stored symbol_path"""
+        assert 'path' in d and d['path'] is not None
+        instance = resolve_symbol(d['path'])
+        assert isinstance(instance, Function)
+        return instance
+
+    def to_store(self) -> Tuple[Dict, bytes]:
+        """
+        Serialize the function to a format that can be stored in the Pixeltable store
+        Returns:
+            - a dict that can be passed to json.dumps()
+            - additional binary data
+        Only Function subclasses that can be stored need to override this.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def from_store(cls, name: Optional[str], md: Dict, binary_obj: bytes) -> Function:
+        """
+        Create a Function instance from the serialized representation returned by to_store()
+        """
+        raise NotImplementedError()