pixeltable 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

pixeltable/exprs/expr_dict.py

@@ -0,0 +1,55 @@
+from typing import Generic, TypeVar, Optional, Iterator, Iterable
+
+T = TypeVar('T')
+
+from .expr import Expr
+
+class ExprDict(Generic[T]):
+    """
+    A dictionary that maps Expr instances to values of type T.
+
+    We cannot use dict[Expr, T] because Expr.__eq__() serves a different purpose than the default __eq__.
+    """
+
+    _data: dict[int, tuple[Expr, T]]
+
+    def __init__(self, iterable: Optional[Iterable[tuple[Expr, T]]] = None):
+        self._data = {}
+
+        if iterable is not None:
+            for key, value in iterable:
+                self[key] = value
+
+    def __setitem__(self, key: Expr, value: T) -> None:
+        self._data[key.id] = (key, value)
+
+    def __getitem__(self, key: Expr) -> T:
+        return self._data[key.id][1]
+
+    def __delitem__(self, key: Expr) -> None:
+        del self._data[key.id]
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __iter__(self) -> Iterator[Expr]:
+        return (expr for expr, _ in self._data.values())
+
+    def __contains__(self, key: Expr) -> bool:
+        return key.id in self._data
+
+    def get(self, key: Expr, default: Optional[T] = None) -> Optional[T]:
+        item = self._data.get(key.id)
+        return item[1] if item is not None else default
+
+    def clear(self) -> None:
+        self._data.clear()
+
+    def keys(self) -> Iterator[Expr]:
+        return self.__iter__()
+
+    def values(self) -> Iterator[T]:
+        return (value for _, value in self._data.values())
+
+    def items(self) -> Iterator[tuple[Expr, T]]:
+        return ((expr, value) for expr, value in self._data.values())

pixeltable/exprs/expr_set.py

@@ -1,25 +1,26 @@
 from __future__ import annotations
 
-from typing import Optional, Iterable, Iterator
+from typing import Optional, Iterable, Iterator, TypeVar, Generic
 
 from .expr import Expr
 
+T = TypeVar('T', bound='Expr')
 
-class ExprSet:
+class ExprSet(Generic[T]):
     """
     A set that also supports indexed lookup (by slot_idx and Expr.id). Exprs are uniquely identified by Expr.id.
     """
-    exprs: dict[int, Expr]  # key: Expr.id
-    exprs_by_idx: dict[int, Expr]  # key: slot_idx
+    exprs: dict[int, T]  # key: Expr.id
+    exprs_by_idx: dict[int, T]  # key: slot_idx
 
-    def __init__(self, elements: Optional[Iterable[Expr]] = None):
+    def __init__(self, elements: Optional[Iterable[T]] = None):
         self.exprs = {}
         self.exprs_by_idx = {}
         if elements is not None:
             for e in elements:
                 self.add(e)
 
-    def add(self, expr: Expr) -> None:
+    def add(self, expr: T) -> None:
         if expr.id in self.exprs:
             return
         self.exprs[expr.id] = expr
@@ -27,24 +28,22 @@ class ExprSet:
             return
         self.exprs_by_idx[expr.slot_idx] = expr
 
-    def update(self, *others: Iterable[Expr]) -> None:
+    def update(self, *others: Iterable[T]) -> None:
         for other in others:
             for e in other:
                 self.add(e)
 
-    def __contains__(self, item: Expr) -> bool:
+    def __contains__(self, item: T) -> bool:
         return item.id in self.exprs
 
     def __len__(self) -> int:
         return len(self.exprs)
 
-    def __iter__(self) -> Iterator[Expr]:
+    def __iter__(self) -> Iterator[T]:
         return iter(self.exprs.values())
 
-    def __getitem__(self, index: object) -> Optional[Expr]:
+    def __getitem__(self, index: object) -> Optional[T]:
         """Indexed lookup by slot_idx or Expr.id."""
-        if not isinstance(index, int) and not isinstance(index, Expr):
-            pass
         assert isinstance(index, int) or isinstance(index, Expr)
         if isinstance(index, int):
             # return expr with matching slot_idx
@@ -52,11 +51,18 @@ class ExprSet:
         else:
             return self.exprs.get(index.id)
 
-    def issuperset(self, other: ExprSet) -> bool:
+    def issuperset(self, other: ExprSet[T]) -> bool:
         return self.exprs.keys() >= other.exprs.keys()
 
-    def __ge__(self, other: ExprSet) -> bool:
+    def __ge__(self, other: ExprSet[T]) -> bool:
         return self.issuperset(other)
 
-    def __le__(self, other: ExprSet) -> bool:
+    def __le__(self, other: ExprSet[T]) -> bool:
         return other.issuperset(self)
+
+    def difference(self, *others: Iterable[T]) -> ExprSet[T]:
+        id_diff = set(self.exprs.keys()).difference(e.id for other_set in others for e in other_set)
+        return ExprSet(self.exprs[id] for id in id_diff)
+
+    def __sub__(self, other: ExprSet[T]) -> ExprSet[T]:
+        return self.difference(other)

pixeltable/exprs/function_call.py

@@ -294,6 +294,10 @@ class FunctionCall(Expr):
     def get_window_sort_exprs(self) -> tuple[list[Expr], list[Expr]]:
         return self.group_by, self.order_by
 
+    def get_window_ordering(self) -> list[tuple[Expr, bool]]:
+        # ordering is implicitly ascending
+        return [(e, None) for e in self.group_by] + [(e, True) for e in self.order_by]
+
     @property
     def is_agg_fn_call(self) -> bool:
         return isinstance(self.fn, func.AggregateFunction)
@@ -303,6 +307,10 @@ class FunctionCall(Expr):
         return self.order_by
 
     def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+        # we currently can't translate aggregate functions with grouping and/or ordering to SQL
+        if self.has_group_by() or len(self.order_by) > 0:
+            return None
+
         # try to construct args and kwargs to call self.fn._to_sql()
         kwargs: dict[str, sql.ColumnElement] = {}
         for param_name, (component_idx, arg) in self.kwargs.items():
@@ -374,6 +382,18 @@ class FunctionCall(Expr):
         return args, kwargs
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        if isinstance(self.fn, func.ExprTemplateFunction):
+            # we need to evaluate the template
+            # TODO: can we get rid of this extra copy?
+            fn_expr = self.components[self.fn_expr_idx]
+            data_row[self.slot_idx] = data_row[fn_expr.slot_idx]
+            return
+        elif self.is_agg_fn_call and not self.is_window_fn_call:
+            if self.aggregator is None:
+                pass
+            data_row[self.slot_idx] = self.aggregator.value()
+            return
+
         args, kwargs = self._make_args(data_row)
         signature = self.fn.signature
         if signature.parameters is not None:
@@ -389,12 +409,7 @@ class FunctionCall(Expr):
                     data_row[self.slot_idx] = None
                     return
 
-        if isinstance(self.fn, func.ExprTemplateFunction):
-            # we need to evaluate the template
-            # TODO: can we get rid of this extra copy?
-            fn_expr = self.components[self.fn_expr_idx]
-            data_row[self.slot_idx] = data_row[fn_expr.slot_idx]
-        elif isinstance(self.fn, func.CallableFunction) and not self.fn.is_batched:
+        if isinstance(self.fn, func.CallableFunction) and not self.fn.is_batched:
             # optimization: avoid additional level of indirection we'd get from calling Function.exec()
             data_row[self.slot_idx] = self.fn.py_fn(*args, **kwargs)
         elif self.is_window_fn_call:
@@ -410,8 +425,6 @@ class FunctionCall(Expr):
                 self.aggregator = self.fn.agg_cls(**self.agg_init_args)
             self.aggregator.update(*args)
             data_row[self.slot_idx] = self.aggregator.value()
-        elif self.is_agg_fn_call:
-            data_row[self.slot_idx] = self.aggregator.value()
         else:
             data_row[self.slot_idx] = self.fn.exec(*args, **kwargs)
 

pixeltable/exprs/json_path.py

@@ -105,12 +105,9 @@ class JsonPath(Expr):
         return JsonPath(self._anchor, self.path_elements + [name])
 
     def __getitem__(self, index: object) -> 'JsonPath':
-        if isinstance(index, str):
-            if index != '*':
-                raise excs.Error(f'Invalid json list index: {index}')
-        elif not isinstance(index, (int, slice)):
-                raise excs.Error(f'Invalid json list index: {index}')
-        return JsonPath(self._anchor, self.path_elements + [index])
+        if isinstance(index, (int, slice, str)):
+            return JsonPath(self._anchor, self.path_elements + [index])
+        raise excs.Error(f'Invalid json list index: {index}')
 
     def __rshift__(self, other: object) -> 'JsonMapper':
         rhs_expr = Expr.from_object(other)

pixeltable/exprs/rowid_ref.py

@@ -68,8 +68,8 @@ class RowidRef(Expr):
         """
         if self.tbl_id == tbl.tbl_version.id:
             return
-        tbl_version_ids = [tbl_version.id for tbl_version in tbl.get_tbl_versions()]
-        assert self.tbl_id in tbl_version_ids
+        base_ids = [tbl_version.id for tbl_version in tbl.get_tbl_versions()]
+        assert self.tbl_id in base_ids  # our current TableVersion is a base of the new TableVersion
         self.tbl = tbl.tbl_version
         self.tbl_id = self.tbl.id
 

pixeltable/exprs/sql_element_cache.py

@@ -3,6 +3,7 @@ from typing import Iterable, Union, Optional
 import sqlalchemy as sql
 
 from .expr import Expr
+from .expr_dict import ExprDict
 
 
 class SqlElementCache:
@@ -10,8 +11,11 @@ class SqlElementCache:
 
     cache: dict[int, Optional[sql.ColumnElement]]  # key: Expr.id
 
-    def __init__(self):
+    def __init__(self, elements: Optional[ExprDict[sql.ColumnElement]] = None):
         self.cache = {}
+        if elements is not None:
+            for e, el in elements.items():
+                self.cache[e.id] = el
 
     def get(self, e: Expr) -> Optional[sql.ColumnElement]:
         """Returns the sql.ColumnElement for the given Expr, or None if Expr.to_sql() returns None."""

pixeltable/ext/functions/whisperx.py

@@ -8,9 +8,14 @@ if TYPE_CHECKING:
 import pixeltable as pxt
 
 
-@pxt.udf(param_types=[pxt.AudioType(), pxt.StringType(), pxt.StringType(), pxt.StringType(), pxt.IntType()])
+@pxt.udf
 def transcribe(
-    audio: str, *, model: str, compute_type: Optional[str] = None, language: Optional[str] = None, chunk_size: int = 30
+    audio: pxt.Audio,
+    *,
+    model: str,
+    compute_type: Optional[str] = None,
+    language: Optional[str] = None,
+    chunk_size: int = 30
 ) -> dict:
     """
     Transcribe an audio file using WhisperX.

pixeltable/func/callable_function.py

@@ -4,7 +4,7 @@ import inspect
 from typing import Any, Callable, Optional
 from uuid import UUID
 
-import cloudpickle
+import cloudpickle  # type: ignore[import-untyped]
 
 from .function import Function
 from .signature import Signature
@@ -108,7 +108,7 @@ class CallableFunction(Function):
     @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)
+        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)

pixeltable/func/function_registry.py

@@ -4,7 +4,7 @@ import dataclasses
 import importlib
 import logging
 import sys
-from typing import Optional, Dict, List
+from typing import Optional
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -14,7 +14,6 @@ import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable.metadata import schema
 from .function import Function
-from .globals import get_caller_module_path
 
 _logger = logging.getLogger('pixeltable')
 
@@ -32,15 +31,15 @@ class FunctionRegistry:
         return cls._instance
 
     def __init__(self):
-        self.stored_fns_by_id: Dict[UUID, Function] = {}
-        self.module_fns: Dict[str, Function] = {}  # fqn -> Function
+        self.stored_fns_by_id: dict[UUID, Function] = {}
+        self.module_fns: dict[str, Function] = {}  # fqn -> Function
         self.type_methods: dict[ts.ColumnType.Type, dict[str, Function]] = {}
 
     def clear_cache(self) -> None:
         """
         Useful during testing
         """
-        self.stored_fns_by_id: Dict[UUID, Function] = {}
+        self.stored_fns_by_id = {}
 
     # def register_std_modules(self) -> None:
     #     """Register all submodules of pixeltable.functions"""
@@ -76,7 +75,7 @@ class FunctionRegistry:
                 raise excs.Error(f'Duplicate method name for type {base_type}: {fn.name}')
             self.type_methods[base_type][fn.name] = fn
 
-    def list_functions(self) -> List[Function]:
+    def list_functions(self) -> list[Function]:
         # retrieve Function.Metadata data for all existing stored functions from store directly
         # (self.stored_fns_by_id isn't guaranteed to contain all functions)
         # TODO: have the client do this, once the client takes over the Db functionality
@@ -85,7 +84,7 @@ class FunctionRegistry:
         #         schema.Db.name, schema.Dir.path, sql_func.length(schema.Function.init_obj))\
         #     .where(schema.Function.db_id == schema.Db.id)\
         #     .where(schema.Function.dir_id == schema.Dir.id)
-        # stored_fn_md: List[Function.Metadata] = []
+        # stored_fn_md: list[Function.Metadata] = []
         # with Env.get().engine.begin() as conn:
         #     rows = conn.execute(stmt)
         #     for name, md_dict, db_name, dir_path, init_obj_len in rows:

pixeltable/func/query_template_function.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
+
 import inspect
-from typing import Dict, Optional, Any, Callable
+from typing import Any, Callable, Optional
 
 import sqlalchemy as sql
 
-import pixeltable
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+import pixeltable as pxt
+from pixeltable import exprs
+
 from .function import Function
-from .signature import Signature, Parameter
+from .signature import Signature
 
 
 class QueryTemplateFunction(Function):
@@ -16,24 +17,23 @@ class QueryTemplateFunction(Function):
 
     @classmethod
     def create(
-            cls, template_callable: Callable, param_types: Optional[list[ts.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: Optional[list[pxt.ColumnType]], path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
         py_params = list(py_sig.parameters.values())
         params = Signature.create_parameters(py_params=py_params, param_types=param_types)
         # invoke template_callable with parameter expressions to construct a DataFrame with parameters
-        import pixeltable.exprs as exprs
         var_exprs = [exprs.Variable(param.name, param.col_type) for param in params]
         template_df = template_callable(*var_exprs)
         from pixeltable import DataFrame
         assert isinstance(template_df, DataFrame)
         # we take params and return json
-        sig = Signature(return_type=ts.JsonType(), parameters=params)
+        sig = Signature(return_type=pxt.JsonType(), parameters=params)
         return QueryTemplateFunction(template_df, sig, path=path, name=name)
 
     def __init__(
-            self, template_df: Optional['pixeltable.DataFrame'], sig: Optional[Signature], path: Optional[str] = None,
+            self, template_df: Optional['pxt.DataFrame'], sig: Optional[Signature], path: Optional[str] = None,
             name: Optional[str] = None,
     ):
         super().__init__(sig, self_path=path)
@@ -46,7 +46,6 @@ class QueryTemplateFunction(Function):
         self.conn: Optional[sql.engine.Connection] = None
 
         # convert defaults to Literals
-        import pixeltable.exprs as exprs
         self.defaults: dict[str, exprs.Literal] = {}  # key: param name, value: default value converted to a Literal
         param_types = self.template_df.parameters()
         for param in [p for p in self.signature.parameters.values() if p.has_default()]:
@@ -75,10 +74,10 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> Function:
+    def _from_dict(cls, d: dict) -> Function:
         from pixeltable.dataframe import DataFrame
         return cls(DataFrame.from_dict(d['df']), Signature.from_dict(d['signature']), name=d['name'])

pixeltable/func/signature.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-import json
 import dataclasses
 import enum
 import inspect
+import json
 import logging
 import typing
-from typing import Optional, Callable, Dict, List, Any, Union, Tuple
+from typing import Any, Callable, Optional, Union
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
@@ -18,7 +18,7 @@ _logger = logging.getLogger('pixeltable')
 class Parameter:
     name: str
     col_type: Optional[ts.ColumnType]  # None for variable parameters
-    kind: enum.Enum  # inspect.Parameter.kind; inspect._ParameterKind is private
+    kind: inspect._ParameterKind
     # for some reason, this needs to precede is_batched in the dataclass definition,
     # otherwise Python complains that an argument with a default is followed by an argument without a default
     default: Any = inspect.Parameter.empty  # default value for the parameter
@@ -82,7 +82,7 @@ class Signature:
     """
     SPECIAL_PARAM_NAMES = ['group_by', 'order_by']
 
-    def __init__(self, return_type: ts.ColumnType, parameters: List[Parameter], is_batched: bool = False):
+    def __init__(self, return_type: ts.ColumnType, parameters: list[Parameter], is_batched: bool = False):
         assert isinstance(return_type, ts.ColumnType)
         self.return_type = return_type
         self.is_batched = is_batched
@@ -97,7 +97,7 @@ class Signature:
         assert isinstance(self.return_type, ts.ColumnType)
         return self.return_type
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         result = {
             'return_type': self.get_return_type().as_dict(),
             'parameters': [p.as_dict() for p in self.parameters.values()],
@@ -106,11 +106,13 @@ class Signature:
         return result
 
     @classmethod
-    def from_dict(cls, d: Dict[str, Any]) -> Signature:
+    def from_dict(cls, d: dict[str, Any]) -> 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 __eq__(self, other: Signature) -> bool:
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Signature):
+            return False
         if self.get_return_type() != other.get_return_type():
             return False
         if len(self.parameters) != len(other.parameters):
@@ -122,7 +124,7 @@ class Signature:
         return True
 
     def __str__(self) -> str:
-        param_strs: List[str] = []
+        param_strs: list[str] = []
         for p in self.parameters.values():
             if p.kind == inspect.Parameter.VAR_POSITIONAL:
                 param_strs.append(f'*{p.name}')
@@ -133,7 +135,7 @@ class Signature:
         return f'({", ".join(param_strs)}) -> {str(self.get_return_type())}'
 
     @classmethod
-    def _infer_type(cls, annotation: Optional[type]) -> Tuple[Optional[ts.ColumnType], Optional[bool]]:
+    def _infer_type(cls, annotation: Optional[type]) -> tuple[Optional[ts.ColumnType], Optional[bool]]:
         """Returns: (column type, is_batched) or (None, ...) if the type cannot be inferred"""
         if annotation is None:
             return (None, None)
@@ -154,13 +156,13 @@ 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
-    ) -> List[Parameter]:
+            param_types: Optional[list[ts.ColumnType]] = None
+    ) -> list[Parameter]:
         assert (py_fn is None) != (py_params is None)
         if py_fn is not None:
             sig = inspect.signature(py_fn)
             py_params = list(sig.parameters.values())
-        parameters: List[Parameter] = []
+        parameters: list[Parameter] = []
 
         for idx, param in enumerate(py_params):
             if param.name in cls.SPECIAL_PARAM_NAMES:
@@ -187,9 +189,9 @@ class Signature:
 
     @classmethod
     def create(
-            cls, py_fn: Callable,
-            param_types: Optional[List[ts.ColumnType]] = None,
-            return_type: Optional[Union[ts.ColumnType, Callable]] = None
+        cls, py_fn: Callable,
+        param_types: Optional[list[ts.ColumnType]] = None,
+        return_type: Optional[ts.ColumnType] = None
     ) -> Signature:
         """Create a signature for the given Callable.
         Infer the parameter and return types, if none are specified.

pixeltable/func/udf.py

@@ -38,10 +38,6 @@ def udf(*args, **kwargs):
         >>> @pxt.udf
         ... def my_function(x: int) -> int:
         ...    return x + 1
-
-        >>> @pxt.udf(param_types=[pxt.IntType()], return_type=pxt.IntType())
-        ... def my_function(x):
-        ...    return x + 1
     """
     if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
 

pixeltable/functions/__init__.py

@@ -1,7 +1,7 @@
 from pixeltable.utils.code import local_public_names
 
 from . import (anthropic, audio, fireworks, huggingface, image, json, mistralai, openai, string, timestamp, together,
-               video, vision)
+               video, vision, whisper)
 from .globals import *
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)

pixeltable/functions/audio.py

@@ -11,18 +11,16 @@ t.select(pxtf.audio.get_metadata()).collect()
 ```
 """
 
-import pixeltable.func as func
-import pixeltable.type_system as ts
+import pixeltable as pxt
 from pixeltable.utils.code import local_public_names
 
 
-@func.udf(return_type=ts.JsonType(nullable=False), param_types=[ts.AudioType(nullable=False)], is_method=True)
-def get_metadata(audio: str) -> dict:
+@pxt.udf(is_method=True)
+def get_metadata(audio: pxt.Audio) -> dict:
     """
     Gets various metadata associated with an audio file and returns it as a dictionary.
     """
-    import pixeltable.functions as pxtf
-    return pxtf.video._get_metadata(audio)
+    return pxt.functions.video._get_metadata(audio)
 
 
 __all__ = local_public_names(__name__)