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

pixeltable/exprs/function_call.py

@@ -15,6 +15,7 @@ import pixeltable.type_system as ts
 from .data_row import DataRow
 from .expr import Expr
 from .inline_expr import InlineDict, InlineList
+from .literal import Literal
 from .row_builder import RowBuilder
 from .rowid_ref import RowidRef
 from .sql_element_cache import SqlElementCache
@@ -34,6 +35,7 @@ class FunctionCall(Expr):
 
     arg_types: list[ts.ColumnType]
     kwarg_types: dict[str, ts.ColumnType]
+    return_type: ts.ColumnType
     group_by_start_idx: int
     group_by_stop_idx: int
     fn_expr_idx: int
@@ -43,17 +45,25 @@ class FunctionCall(Expr):
     current_partition_vals: Optional[list[Any]]
 
     def __init__(
-            self, fn: func.Function, bound_args: dict[str, Any], order_by_clause: Optional[list[Any]] = None,
-            group_by_clause: Optional[list[Any]] = None, is_method_call: bool = False):
+        self,
+        fn: func.Function,
+        bound_args: dict[str, Any],
+        return_type: ts.ColumnType,
+        order_by_clause: Optional[list[Any]] = None,
+        group_by_clause: Optional[list[Any]] = None,
+        is_method_call: bool = False
+    ):
         if order_by_clause is None:
             order_by_clause = []
         if group_by_clause is None:
             group_by_clause = []
-        signature = fn.signature
-        return_type = fn.call_return_type(bound_args)
+
+        assert not fn.is_polymorphic
+
         self.fn = fn
         self.is_method_call = is_method_call
-        self.normalize_args(fn.name, signature, bound_args)
+
+        signature = fn.signature
 
         # If `return_type` is non-nullable, but the function call has a nullable input to any of its non-nullable
         # parameters, then we need to make it nullable. This is because Pixeltable defaults a function output to
@@ -67,6 +77,8 @@ class FunctionCall(Expr):
                 return_type = return_type.copy(nullable=True)
                 break
 
+        self.return_type = return_type
+
         super().__init__(return_type)
 
         self.agg_init_args = {}
@@ -74,9 +86,9 @@ class FunctionCall(Expr):
             # we separate out the init args for the aggregator
             assert isinstance(fn, func.AggregateFunction)
             self.agg_init_args = {
-                arg_name: arg for arg_name, arg in bound_args.items() if arg_name in fn.init_param_names
+                arg_name: arg for arg_name, arg in bound_args.items() if arg_name in fn.init_param_names[0]
             }
-            bound_args = {arg_name: arg for arg_name, arg in bound_args.items() if arg_name not in fn.init_param_names}
+            bound_args = {arg_name: arg for arg_name, arg in bound_args.items() if arg_name not in fn.init_param_names[0]}
 
         # construct components, args, kwargs
         self.args = []
@@ -85,8 +97,10 @@ class FunctionCall(Expr):
         # we record the types of non-variable parameters for runtime type checks
         self.arg_types = []
         self.kwarg_types = {}
+
         # the prefix of parameters that are bound can be passed by position
-        for py_param in fn.signature.py_signature.parameters.values():
+        processed_args: set[str] = set()
+        for py_param in signature.py_signature.parameters.values():
             if py_param.name not in bound_args or py_param.kind == inspect.Parameter.KEYWORD_ONLY:
                 break
             arg = bound_args[py_param.name]
@@ -97,18 +111,19 @@ class FunctionCall(Expr):
                 self.args.append((None, arg))
             if py_param.kind != inspect.Parameter.VAR_POSITIONAL and py_param.kind != inspect.Parameter.VAR_KEYWORD:
                 self.arg_types.append(signature.parameters[py_param.name].col_type)
+            processed_args.add(py_param.name)
 
         # the remaining args are passed as keywords
-        kw_param_names = set(bound_args.keys()) - set(list(fn.signature.py_signature.parameters.keys())[:len(self.args)])
-        for param_name in kw_param_names:
-            arg = bound_args[param_name]
-            if isinstance(arg, Expr):
-                self.kwargs[param_name] = (len(self.components), None)
-                self.components.append(arg.copy())
-            else:
-                self.kwargs[param_name] = (None, arg)
-            if fn.signature.py_signature.parameters[param_name].kind != inspect.Parameter.VAR_KEYWORD:
-                self.kwarg_types[param_name] = signature.parameters[param_name].col_type
+        for param_name in bound_args.keys():
+            if param_name not in processed_args:
+                arg = bound_args[param_name]
+                if isinstance(arg, Expr):
+                    self.kwargs[param_name] = (len(self.components), None)
+                    self.components.append(arg.copy())
+                else:
+                    self.kwargs[param_name] = (None, arg)
+                if signature.py_signature.parameters[param_name].kind != inspect.Parameter.VAR_KEYWORD:
+                    self.kwarg_types[param_name] = signature.parameters[param_name].col_type
 
         # window function state:
         # self.components[self.group_by_start_idx:self.group_by_stop_idx] contains group_by exprs
@@ -126,7 +141,7 @@ class FunctionCall(Expr):
 
         if isinstance(self.fn, func.ExprTemplateFunction):
             # we instantiate the template to create an Expr that can be evaluated and record that as a component
-            fn_expr = self.fn.instantiate(**bound_args)
+            fn_expr = self.fn.instantiate([], bound_args)
             self.components.append(fn_expr)
             self.fn_expr_idx = len(self.components) - 1
         else:
@@ -184,11 +199,6 @@ class FunctionCall(Expr):
                     pass
 
             if not isinstance(arg, Expr):
-                # make sure that non-Expr args are json-serializable and are literals of the correct type
-                try:
-                    _ = json.dumps(arg)
-                except TypeError:
-                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg} (of type {type(arg)})')
                 if arg is not None:
                     try:
                         param_type = param.col_type
@@ -255,7 +265,7 @@ class FunctionCall(Expr):
             ('order_by_start_idx', self.order_by_start_idx)
         ]
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         return self.display_str()
 
     def display_str(self, inline: bool = True) -> str:
@@ -357,7 +367,7 @@ class FunctionCall(Expr):
         """
         assert self.is_agg_fn_call
         assert isinstance(self.fn, func.AggregateFunction)
-        self.aggregator = self.fn.agg_cls(**self.agg_init_args)
+        self.aggregator = self.fn.agg_class(**self.agg_init_args)
 
     def update(self, data_row: DataRow) -> None:
         """
@@ -429,27 +439,32 @@ class FunctionCall(Expr):
             data_row[self.slot_idx] = self.fn.py_fn(*args, **kwargs)
         elif self.is_window_fn_call:
             assert isinstance(self.fn, func.AggregateFunction)
+            agg_cls = self.fn.agg_class
             if self.has_group_by():
                 if self.current_partition_vals is None:
                     self.current_partition_vals = [None] * len(self.group_by)
                 partition_vals = [data_row[e.slot_idx] for e in self.group_by]
                 if partition_vals != self.current_partition_vals:
                     # new partition
-                    self.aggregator = self.fn.agg_cls(**self.agg_init_args)
+                    self.aggregator = agg_cls(**self.agg_init_args)
                     self.current_partition_vals = partition_vals
             elif self.aggregator is None:
-                self.aggregator = self.fn.agg_cls(**self.agg_init_args)
+                self.aggregator = agg_cls(**self.agg_init_args)
             self.aggregator.update(*args)
             data_row[self.slot_idx] = self.aggregator.value()
         else:
-            data_row[self.slot_idx] = self.fn.exec(*args, **kwargs)
+            data_row[self.slot_idx] = self.fn.exec(args, kwargs)
 
     def _as_dict(self) -> dict:
         result = {
-            'fn': self.fn.as_dict(), 'args': self.args, 'kwargs': self.kwargs,
-            'group_by_start_idx': self.group_by_start_idx, 'group_by_stop_idx': self.group_by_stop_idx,
+            'fn': self.fn.as_dict(),
+            'args': self.args,
+            'kwargs': self.kwargs,
+            'return_type': self.return_type.as_dict(),
+            'group_by_start_idx': self.group_by_start_idx,
+            'group_by_stop_idx': self.group_by_stop_idx,
             'order_by_start_idx': self.order_by_start_idx,
-            **super()._as_dict()
+            **super()._as_dict(),
         }
         return result
 
@@ -458,15 +473,108 @@ class FunctionCall(Expr):
         assert 'fn' in d
         assert 'args' in d
         assert 'kwargs' in d
-        # reassemble bound args
+
         fn = func.Function.from_dict(d['fn'])
-        param_names = list(fn.signature.parameters.keys())
-        bound_args = {param_names[i]: arg if idx is None else components[idx] for i, (idx, arg) in enumerate(d['args'])}
-        bound_args.update(
-            {param_name: val if idx is None else components[idx] for param_name, (idx, val) in d['kwargs'].items()})
+        assert not fn.is_polymorphic
+        return_type = ts.ColumnType.from_dict(d['return_type']) if 'return_type' in d else None
         group_by_exprs = components[d['group_by_start_idx']:d['group_by_stop_idx']]
         order_by_exprs = components[d['order_by_start_idx']:]
+
+        args = [
+            expr if idx is None else components[idx]
+            for idx, expr in d['args']
+        ]
+        kwargs = {
+            param_name: (expr if idx is None else components[idx])
+            for param_name, (idx, expr) in d['kwargs'].items()
+        }
+
+        # `Function.from_dict()` does signature matching, so it is safe to assume that `args` and `kwargs` are
+        # consistent with its signature.
+
+        # Reassemble bound_args. Note that args and kwargs represent "already bound arguments": they are not bindable
+        # in the Python sense, because variable args (such as *args and **kwargs) have already been condensed.
+        param_names = list(fn.signature.parameters.keys())
+        bound_args = {param_names[i]: arg for i, arg in enumerate(args)}
+        bound_args.update(kwargs.items())
+
+        # TODO: In order to properly invoke call_return_type, we need to ensure that any InlineLists or InlineDicts
+        # in bound_args are unpacked into Python lists/dicts. There is an open task to ensure this is true in general;
+        # for now, as a hack, we do the unpacking here for the specific case of an InlineList of Literals (the only
+        # case where this is necessary to support existing conditional_return_type implementations). Once the general
+        # pattern is implemented, we can remove this hack.
+        unpacked_bound_args = {
+            param_name: cls.__unpack_bound_arg(arg) for param_name, arg in bound_args.items()
+        }
+
+        # Evaluate the call_return_type as defined in the current codebase.
+        call_return_type = fn.call_return_type([], unpacked_bound_args)
+
+        if return_type is None:
+            # Schema versions prior to 25 did not store the return_type in metadata, and there is no obvious way to
+            # infer it during DB migration, so we might encounter a stored return_type of None. In that case, we use
+            # the call_return_type that we just inferred (which matches the deserialization behavior prior to
+            # version 25).
+            return_type = call_return_type
+        else:
+            # There is a return_type stored in metadata (schema version >= 25).
+            # Check that the stored return_type of the UDF call matches the column type of the FunctionCall, and
+            # fail-fast if it doesn't (otherwise we risk getting downstream database errors).
+            # TODO: Handle this more gracefully (instead of failing the DB load, allow the DB load to succeed, but
+            #       mark this FunctionCall as unusable). It's the same issue as dealing with a renamed UDF or Function
+            #       signature mismatch.
+            if not return_type.is_supertype_of(call_return_type, ignore_nullable=True):
+                raise excs.Error(
+                    f'The return type stored in the database for a UDF call to `{fn.self_path}` no longer matches the '
+                    f'return type of the UDF as currently defined in the code.\nThis probably means that the code for '
+                    f'`{fn.self_path}` has changed in a backward-incompatible way.\n'
+                    f'Return type in database: `{return_type}`\n'
+                    f'Return type as currently defined: `{call_return_type}`'
+                )
+
         fn_call = cls(
-            func.Function.from_dict(d['fn']), bound_args, group_by_clause=group_by_exprs,
-            order_by_clause=order_by_exprs)
+            fn,
+            bound_args,
+            return_type,
+            group_by_clause=group_by_exprs,
+            order_by_clause=order_by_exprs
+        )
         return fn_call
+
+    @classmethod
+    def __find_matching_signature(cls, fn: func.Function, args: list[Any], kwargs: dict[str, Any]) -> Optional[int]:
+        for idx, sig in enumerate(fn.signatures):
+            if cls.__signature_matches(sig, args, kwargs):
+                return idx
+        return None
+
+    @classmethod
+    def __signature_matches(cls, sig: func.Signature, args: list[Any], kwargs: dict[str, Any]) -> bool:
+        unbound_parameters = set(sig.parameters.keys())
+        for i, arg in enumerate(args):
+            if i >= len(sig.parameters_by_pos):
+                return False
+            param = sig.parameters_by_pos[i]
+            arg_type = arg.col_type if isinstance(arg, Expr) else ts.ColumnType.infer_literal_type(arg)
+            if param.col_type is not None and not param.col_type.is_supertype_of(arg_type, ignore_nullable=True):
+                return False
+            unbound_parameters.remove(param.name)
+        for param_name, arg in kwargs.items():
+            if param_name not in unbound_parameters:
+                return False
+            param = sig.parameters[param_name]
+            arg_type = arg.col_type if isinstance(arg, Expr) else ts.ColumnType.infer_literal_type(arg)
+            if param.col_type is not None and not param.col_type.is_supertype_of(arg_type, ignore_nullable=True):
+                return False
+            unbound_parameters.remove(param_name)
+        for param_name in unbound_parameters:
+            param = sig.parameters[param_name]
+            if not param.has_default:
+                return False
+        return True
+
+    @classmethod
+    def __unpack_bound_arg(cls, arg: Any) -> Any:
+        if isinstance(arg, InlineList) and all(isinstance(el, Literal) for el in arg.components):
+            return [el.val for el in arg.components]
+        return arg

pixeltable/exprs/in_predicate.py

@@ -61,7 +61,7 @@ class InPredicate(Expr):
                 pass
         return result
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         if self.value_list is not None:
             return f'{self.components[0]}.isin({self.value_list})'
         return f'{self.components[0]}.isin({self.components[1]})'

pixeltable/exprs/inline_expr.py

@@ -56,7 +56,7 @@ class InlineArray(Expr):
         self.components.extend(exprs)
         self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         elem_strs = [str(expr) for expr in self.components]
         return f'[{", ".join(elem_strs)}]'
 
@@ -101,11 +101,17 @@ class InlineList(Expr):
             else:
                 exprs.append(Literal(el))
 
-        super().__init__(ts.JsonType())
+        json_schema = {
+            'type': 'array',
+            'prefixItems': [expr.col_type.to_json_schema() for expr in exprs],
+            'items': False  # No additional items (fixed length)
+        }
+
+        super().__init__(ts.JsonType(json_schema))
         self.components.extend(exprs)
         self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         elem_strs = [str(expr) for expr in self.components]
         return f'[{", ".join(elem_strs)}]'
 
@@ -149,11 +155,25 @@ class InlineDict(Expr):
             else:
                 exprs.append(Literal(val))
 
-        super().__init__(ts.JsonType())
+        json_schema: Optional[dict[str, Any]]
+        try:
+            json_schema = {
+                'type': 'object',
+                'properties': {
+                    key: expr.col_type.to_json_schema()
+                    for key, expr in zip(self.keys, exprs)
+                },
+            }
+        except excs.Error:
+            # InlineDicts are used to store iterator arguments, which are not required to be valid JSON types,
+            # so we can't always construct a valid schema.
+            json_schema = None
+
+        super().__init__(ts.JsonType(json_schema))
         self.components.extend(exprs)
         self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         item_strs = list(f"'{key}': {str(expr)}" for key, expr in zip(self.keys, self.components))
         return '{' + ', '.join(item_strs) + '}'
 

pixeltable/exprs/is_null.py

@@ -18,7 +18,7 @@ class IsNull(Expr):
         self.components = [e]
         self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         return f'{str(self.components[0])} == None'
 
     def _equals(self, other: IsNull) -> bool:

pixeltable/exprs/json_mapper.py

@@ -69,7 +69,7 @@ class JsonMapper(Expr):
             return False
         return self._src_expr.equals(other._src_expr) and self._target_expr.equals(other._target_expr)
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         return f'{str(self._src_expr)} >> {str(self._target_expr)}'
 
     @property

pixeltable/exprs/json_path.py

@@ -42,7 +42,7 @@ class JsonPath(Expr):
         # this is not a problem, because _create_id() shouldn't be called after init()
         self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         # else "R": the anchor is RELATIVE_PATH_ROOT
         return (f'{str(self._anchor) if self._anchor is not None else "R"}'
             f'{"." if isinstance(self.path_elements[0], str) else ""}{self._json_path()}')

pixeltable/exprs/method_ref.py

@@ -60,5 +60,5 @@ class MethodRef(Expr):
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         assert False, 'MethodRef cannot be evaluated directly'
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         return f'{self.base_expr}.{self.method_name}'

pixeltable/exprs/row_builder.py

@@ -368,7 +368,7 @@ class RowBuilder:
                 if not ignore_errors:
                     input_vals = [data_row[d.slot_idx] for d in expr.dependencies()]
                     raise excs.ExprEvalError(
-                        expr, f'expression {expr}', data_row.get_exc(expr.slot_idx), exc_tb, input_vals, 0)
+                        expr, f'expression {expr}', data_row.get_exc(expr.slot_idx), exc_tb, input_vals, 0) from exc
 
     def create_table_row(self, data_row: DataRow, exc_col_ids: set[int]) -> tuple[dict[str, Any], int]:
         """Create a table row from the slots that have an output column assigned

pixeltable/exprs/rowid_ref.py

@@ -55,7 +55,7 @@ class RowidRef(Expr):
         return super()._id_attrs() +\
             [('normalized_base_id', self.normalized_base_id), ('idx', self.rowid_component_idx)]
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         # check if this is the pos column of a component view
         tbl = self.tbl if self.tbl is not None else catalog.Catalog.get().tbl_versions[(self.tbl_id, None)]
         if tbl.is_component_view() and self.rowid_component_idx == tbl.store_tbl.pos_col_idx:  # type: ignore[attr-defined]

pixeltable/exprs/similarity_expr.py

@@ -23,7 +23,6 @@ class SimilarityExpr(Expr):
         assert item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()
 
         self.components = [col_ref, item_expr]
-        self.id = self._create_id()
 
         # determine index to use
         idx_info = col_ref.col.get_idx_info()
@@ -48,16 +47,20 @@ class SimilarityExpr(Expr):
 
         if item_expr.col_type.is_string_type() and idx.string_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} was created without the '
-                f"'string_embed' parameter and does not support string queries")
+                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have a '
+                f"string embedding and does not support string queries")
         if item_expr.col_type.is_image_type() and idx.image_embed is None:
             raise excs.Error(
-                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} was created without the '
-                f"'image_embed' parameter and does not support image queries")
+                f'Embedding index {self.idx_info.name!r} on column {self.idx_info.col.name!r} does not have an '
+                f"image embedding and does not support image queries")
+        self.id = self._create_id()
 
-    def __str__(self) -> str:
+    def __repr__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'
 
+    def _id_attrs(self):
+        return super()._id_attrs() + [('idx_name', self.idx_info.name)]
+
     def default_column_name(self) -> str:
         return 'similarity'
 
@@ -81,8 +84,12 @@ class SimilarityExpr(Expr):
         # this should never get called
         assert False
 
+    def _as_dict(self) -> dict:
+        return {'idx_name': self.idx_info.name, **super()._as_dict()}
+
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> 'SimilarityExpr':
+        iname = d['idx_name'] if 'idx_name' in d else None
         assert len(components) == 2
         assert isinstance(components[0], ColumnRef)
-        return cls(components[0], components[1])
+        return cls(components[0], components[1], idx_name=iname)

pixeltable/exprs/sql_element_cache.py

@@ -17,6 +17,10 @@ class SqlElementCache:
             for e, el in elements.items():
                 self.cache[e.id] = el
 
+    def extend(self, elements: ExprDict[sql.ColumnElement]):
+        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."""
         try:

pixeltable/exprs/type_cast.py

@@ -51,5 +51,5 @@ class TypeCast(Expr):
         assert len(components) == 1
         return cls(components[0], ts.ColumnType.from_dict(d['new_type']))
 
-    def __str__(self) -> str:
-        return f'{self._underlying}.astype({self.col_type})'
+    def __repr__(self) -> str:
+        return f'{self._underlying}.astype({self.col_type._to_str(as_schema=True)})'

pixeltable/exprs/variable.py

@@ -33,6 +33,9 @@ class Variable(Expr):
     def __str__(self) -> str:
         return self.name
 
+    def __repr__(self) -> str:
+        return f"Variable('{self.name}')"
+
     def sql_expr(self, _: SqlElementCache) -> NoReturn:
         raise NotImplementedError()
 

pixeltable/func/__init__.py

@@ -1,8 +1,9 @@
-from .aggregate_function import Aggregator, AggregateFunction, uda
+from .aggregate_function import AggregateFunction, Aggregator, uda
 from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function
 from .function_registry import FunctionRegistry
-from .query_template_function import QueryTemplateFunction
-from .signature import Signature, Parameter, Batch
-from .udf import udf, make_function, expr_udf
+from .query_template_function import QueryTemplateFunction, query
+from .signature import Batch, Parameter, Signature
+from .tools import Tool, ToolChoice, Tools
+from .udf import expr_udf, make_function, udf