pixeltable 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

pixeltable/exprs/function_call.py

@@ -28,7 +28,7 @@ class FunctionCall(Expr):
         if group_by_clause is None:
             group_by_clause = []
         signature = fn.signature
-        super().__init__(signature.get_return_type(bound_args))
+        super().__init__(fn.call_return_type(bound_args))
         self.fn = fn
         self.is_method_call = is_method_call
         self.check_args(signature, bound_args)
@@ -46,9 +46,9 @@ class FunctionCall(Expr):
 
         # Tuple[int, Any]:
         # - for Exprs: (index into components, None)
-        # - otherwise: (-1, val)
-        self.args: List[Tuple[int, Any]] = []
-        self.kwargs: Dict[str, Tuple[int, Any]] = {}
+        # - otherwise: (None, val)
+        self.args: List[Tuple[Optional[int], Optional[Any]]] = []
+        self.kwargs: Dict[str, Tuple[Optional[int], Optional[Any]]] = {}
 
         # we record the types of non-variable parameters for runtime type checks
         self.arg_types: List[ts.ColumnType] = []
@@ -62,7 +62,7 @@ class FunctionCall(Expr):
                 self.args.append((len(self.components), None))
                 self.components.append(arg.copy())
             else:
-                self.args.append((-1, arg))
+                self.args.append((None, arg))
             if param.kind != inspect.Parameter.VAR_POSITIONAL and param.kind != inspect.Parameter.VAR_KEYWORD:
                 self.arg_types.append(signature.parameters[param.name].col_type)
 
@@ -74,7 +74,7 @@ class FunctionCall(Expr):
                 self.kwargs[param_name] = (len(self.components), None)
                 self.components.append(arg.copy())
             else:
-                self.kwargs[param_name] = (-1, arg)
+                self.kwargs[param_name] = (None, arg)
             if fn.py_signature.parameters[param_name].kind != inspect.Parameter.VAR_KEYWORD:
                 self.kwarg_types[param_name] = signature.parameters[param_name].col_type
 
@@ -215,12 +215,12 @@ class FunctionCall(Expr):
 
     def _print_args(self, start_idx: int = 0, inline: bool = True) -> str:
         arg_strs = [
-            str(arg) if idx == -1 else str(self.components[idx]) for idx, arg in self.args[start_idx:]
+            str(arg) if idx is None else str(self.components[idx]) for idx, arg in self.args[start_idx:]
         ]
         def print_arg(arg: Any) -> str:
             return f"'{arg}'" if isinstance(arg, str) else str(arg)
         arg_strs.extend([
-            f'{param_name}={print_arg(arg) if idx == -1 else str(self.components[idx])}'
+            f'{param_name}={print_arg(arg) if idx is None else str(self.components[idx])}'
             for param_name, (idx, arg) in self.kwargs.items()
         ])
         if len(self.order_by) > 0:
@@ -287,7 +287,7 @@ class FunctionCall(Expr):
         """Return args and kwargs, constructed for data_row"""
         kwargs: Dict[str, Any] = {}
         for param_name, (component_idx, arg) in self.kwargs.items():
-            val = arg if component_idx == -1 else data_row[self.components[component_idx].slot_idx]
+            val = arg if component_idx is None else data_row[self.components[component_idx].slot_idx]
             param = self.fn.signature.parameters[param_name]
             if param.kind == inspect.Parameter.VAR_KEYWORD:
                 # expand **kwargs parameter
@@ -298,7 +298,7 @@ class FunctionCall(Expr):
 
         args: List[Any] = []
         for param_idx, (component_idx, arg) in enumerate(self.args):
-            val = arg if component_idx == -1 else data_row[self.components[component_idx].slot_idx]
+            val = arg if component_idx is None else data_row[self.components[component_idx].slot_idx]
             param = self.fn.signature.parameters_by_pos[param_idx]
             if param.kind == inspect.Parameter.VAR_POSITIONAL:
                 # expand *args parameter
@@ -333,7 +333,8 @@ class FunctionCall(Expr):
             # 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):
+        elif 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:
             if self.has_group_by():
@@ -348,9 +349,10 @@ 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()
-        else:
-            assert self.is_agg_fn_call
+        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)
 
     def _as_dict(self) -> Dict:
         result = {
@@ -369,9 +371,9 @@ class FunctionCall(Expr):
         # 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 == -1 else components[idx] for i, (idx, arg) in enumerate(d['args'])}
+        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 == -1 else components[idx] for param_name, (idx, val) in d['kwargs'].items()})
+            {param_name: val if idx is None else components[idx] for param_name, (idx, val) in d['kwargs'].items()})
         group_by_exprs = components[d['group_by_start_idx']:d['group_by_stop_idx']]
         order_by_exprs = components[d['order_by_start_idx']:]
         fn_call = cls(

pixeltable/exprs/image_member_access.py

@@ -1,19 +1,17 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple, Union
+
+from typing import Optional, List, Any, Dict, Tuple
 
 import PIL
 import sqlalchemy as sql
 
+import pixeltable.exceptions as excs
+import pixeltable.func as func
+import pixeltable.type_system as ts
+from .data_row import DataRow
 from .expr import Expr
-from .column_ref import ColumnRef
 from .function_call import FunctionCall
-from .image_similarity_predicate import ImageSimilarityPredicate
-from .data_row import DataRow
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
-import pixeltable.func as func
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
 
 
 # TODO: this doesn't dig up all attrs for actual jpeg images
@@ -43,9 +41,7 @@ class ImageMemberAccess(Expr):
     attr_info = _create_pil_attr_info()
 
     def __init__(self, member_name: str, caller: Expr):
-        if member_name == 'nearest':
-            super().__init__(ts.InvalidType())  # requires FunctionCall to return value
-        elif member_name in self.attr_info:
+        if member_name in self.attr_info:
             super().__init__(self.attr_info[member_name])
         else:
             candidates = func.FunctionRegistry.get().get_type_methods(member_name, ts.ColumnType.Type.IMAGE)
@@ -78,22 +74,8 @@ class ImageMemberAccess(Expr):
         assert len(components) == 1
         return cls(d['member_name'], components[0])
 
-    def __call__(self, *args, **kwargs) -> Union[FunctionCall, ImageSimilarityPredicate]:
-        caller = self._caller
-        call_signature = f'({",".join([type(arg).__name__ for arg in args])})'
-        if self.member_name == 'nearest':
-            # - caller must be ColumnRef
-            # - signature is (Union[PIL.Image.Image, str])
-            if not isinstance(caller, ColumnRef):
-                raise excs.Error(f'nearest(): caller must be an image column')
-            if len(args) != 1 or (not isinstance(args[0], PIL.Image.Image) and not isinstance(args[0], str)):
-                raise excs.Error(f'nearest(): requires a PIL.Image.Image or str, got {call_signature} instead')
-            return ImageSimilarityPredicate(
-                caller,
-                img=args[0] if isinstance(args[0], PIL.Image.Image) else None,
-                text=args[0] if isinstance(args[0], str) else None)
-
-        result = self.img_method(*[caller, *args], **kwargs)
+    def __call__(self, *args, **kwargs) -> FunctionCall:
+        result = self.img_method(*[self._caller, *args], **kwargs)
         result.is_method_call = True
         return result
 
@@ -112,4 +94,3 @@ class ImageMemberAccess(Expr):
             data_row[self.slot_idx] = getattr(caller_val, self.member_name)
         except AttributeError:
             data_row[self.slot_idx] = None
-

pixeltable/exprs/in_predicate.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import Optional, List, Any, Dict, Tuple, Iterable
+
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+from .data_row import DataRow
+from .expr import Expr
+from .predicate import Predicate
+from .row_builder import RowBuilder
+
+
+class InPredicate(Predicate):
+    """Predicate corresponding to the SQL IN operator."""
+
+    def __init__(self, lhs: Expr, value_set_literal: Optional[Iterable] = None, value_set_expr: Optional[Expr] = None):
+        assert (value_set_literal is None) != (value_set_expr is None)
+        if not lhs.col_type.is_scalar_type():
+            raise excs.Error(f'isin(): only supported for scalar types, not {lhs.col_type}')
+        super().__init__()
+
+        self.value_list: Optional[list] = None  # only contains values of the correct type
+        if value_set_expr is not None:
+            if not value_set_expr.col_type.is_json_type():
+                raise excs.Error(
+                    f'isin(): argument must have a JSON type, but {value_set_expr} has type {value_set_expr.col_type}')
+            self.components = [lhs.copy(), value_set_expr.copy()]
+        else:
+            assert value_set_literal is not None
+            self.components = [lhs.copy()]
+            self.value_list = self._normalize_value_set(value_set_literal)
+
+        self.id = self._create_id()
+
+    @property
+    def _lhs(self) -> Expr:
+        return self.components[0]
+
+    @property
+    def _value_set_expr(self) -> Expr:
+        assert len(self.components) == 2
+        return self.components[1]
+
+    def _normalize_value_set(self, value_set: Any, filter_type_mismatches: bool = True) -> Iterable:
+        if not isinstance(value_set, Iterable):
+            raise excs.Error(f'isin(): argument must be an Iterable (eg, list, dict, ...), not {value_set!r}')
+        value_list = list(value_set)
+        if not filter_type_mismatches:
+            return value_list
+
+        # ignore elements of the wrong type
+        result = []
+        for val in value_list:
+            try:
+                self._lhs.col_type.validate_literal(val)
+                result.append(val)
+            except TypeError:
+                pass
+        return result
+
+    def __str__(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]})'
+
+    def _equals(self, other: InPredicate) -> bool:
+        return self.value_list == other.value_list
+
+    def _id_attrs(self) -> List[Tuple[str, Any]]:
+        return super()._id_attrs() + [('value_list', self.value_list)]
+
+    def sql_expr(self) -> Optional[sql.ClauseElement]:
+        lhs_sql_exprs = self.components[0].sql_expr()
+        if lhs_sql_exprs is None or self.value_list is None:
+            return None
+        return lhs_sql_exprs.in_(self.value_list)
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        lhs_val = data_row[self._lhs.slot_idx]
+        if self.value_list is not None:
+            data_row[self.slot_idx] = lhs_val in self.value_list
+        else:
+            value_set = data_row[self._value_set_expr.slot_idx]
+            value_list = self._normalize_value_set(value_set, filter_type_mismatches=False)
+            data_row[self.slot_idx] = lhs_val in value_list
+
+    def _as_dict(self) -> Dict:
+        return {'value_list': self.value_list, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
+        assert 'value_list' in d
+        assert len(components) <= 2
+        return cls(components[0], d['value_list'], components[1] if len(components) == 2 else None)
+

pixeltable/exprs/inline_array.py

@@ -1,16 +1,16 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple
+
 import copy
+from typing import Optional, List, Any, Dict, Tuple
 
-import sqlalchemy as sql
 import numpy as np
+import sqlalchemy as sql
 
-from .expr import Expr
+import pixeltable.type_system as ts
 from .data_row import DataRow
+from .expr import Expr
 from .inline_dict import InlineDict
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
-import pixeltable.type_system as ts
 
 
 class InlineArray(Expr):
@@ -27,8 +27,8 @@ class InlineArray(Expr):
 
         # elements contains
         # - for Expr elements: (index into components, None)
-        # - for non-Expr elements: (-1, value)
-        self.elements: List[Tuple[int, Any]] = []
+        # - for non-Expr elements: (None, value)
+        self.elements: List[Tuple[Optional[int], Any]] = []
         for el in elements:
             el = copy.deepcopy(el)
             if isinstance(el, list):
@@ -41,11 +41,11 @@ class InlineArray(Expr):
                 self.elements.append((len(self.components), None))
                 self.components.append(el)
             else:
-                self.elements.append((-1, el))
+                self.elements.append((None, el))
 
         inferred_element_type = ts.InvalidType()
         for idx, val in self.elements:
-            if idx >= 0:
+            if idx is not None:
                 inferred_element_type = ts.ColumnType.supertype(inferred_element_type, self.components[idx].col_type)
             else:
                 inferred_element_type = ts.ColumnType.supertype(inferred_element_type, ts.ColumnType.infer_literal_type(val))
@@ -83,7 +83,7 @@ class InlineArray(Expr):
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         result = [None] * len(self.elements)
         for i, (child_idx, val) in enumerate(self.elements):
-            if child_idx >= 0:
+            if child_idx is not None:
                 result[i] = data_row[self.components[child_idx].slot_idx]
             else:
                 result[i] = copy.deepcopy(val)
@@ -100,7 +100,9 @@ class InlineArray(Expr):
         assert 'elements' in d
         arg: List[Any] = []
         for idx, val in d['elements']:
-            if idx >= 0:
+            # TODO Normalize idx -1 to None via schema migrations.
+            # Long-term we should not be allowing idx == -1.
+            if idx is not None and idx >= 0:  # Older schemas might have -1 instead of None
                 arg.append(components[idx])
             else:
                 arg.append(val)

pixeltable/exprs/inline_dict.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple
+
 import copy
+from typing import Optional, List, Any, Dict, Tuple
 
 import sqlalchemy as sql
 
-from .expr import Expr
-from .data_row import DataRow
-from .row_builder import RowBuilder
 import pixeltable.exceptions as excs
-import pixeltable.catalog as catalog
 import pixeltable.type_system as ts
+from .data_row import DataRow
+from .expr import Expr
+from .row_builder import RowBuilder
 
 
 class InlineDict(Expr):
@@ -21,8 +21,8 @@ class InlineDict(Expr):
         super().__init__(ts.JsonType())  # we need to call this in order to populate self.components
         # dict_items contains
         # - for Expr fields: (key, index into components, None)
-        # - for non-Expr fields: (key, -1, value)
-        self.dict_items: List[Tuple[str, int, Any]] = []
+        # - for non-Expr fields: (key, None, value)
+        self.dict_items: List[Tuple[str, Optional[int], Any]] = []
         for key, val in d.items():
             if not isinstance(key, str):
                 raise excs.Error(f'Dictionary requires string keys, {key} has type {type(key)}')
@@ -35,11 +35,11 @@ class InlineDict(Expr):
                 self.dict_items.append((key, len(self.components), None))
                 self.components.append(val)
             else:
-                self.dict_items.append((key, -1, val))
+                self.dict_items.append((key, None, val))
 
         self.type_spec: Optional[Dict[str, ts.ColumnType]] = {}
         for key, idx, _ in self.dict_items:
-            if idx == -1:
+            if idx is None:
                 # TODO: implement type inference for values
                 self.type_spec = None
                 break
@@ -56,7 +56,7 @@ class InlineDict(Expr):
                 return f"'{val}'"
             return str(val)
         for key, idx, val in self.dict_items:
-            if idx != -1:
+            if idx is not None:
                 item_strs.append(f"'{key}': {str(self.components[i])}")
                 i += 1
             else:
@@ -71,7 +71,7 @@ class InlineDict(Expr):
 
     def to_dict(self) -> Dict[str, Any]:
         """Return the original dict used to construct this"""
-        return {key: val if idx == -1 else self.components[idx] for key, idx, val in self.dict_items}
+        return {key: val if idx is None else self.components[idx] for key, idx, val in self.dict_items}
 
     def sql_expr(self) -> Optional[sql.ClauseElement]:
         return None
@@ -80,7 +80,7 @@ class InlineDict(Expr):
         result = {}
         for key, idx, val in self.dict_items:
             assert isinstance(key, str)
-            if idx >= 0:
+            if idx is not None:
                 result[key] = data_row[self.components[idx].slot_idx]
             else:
                 result[key] = copy.deepcopy(val)
@@ -94,7 +94,9 @@ class InlineDict(Expr):
         assert 'dict_items' in d
         arg: Dict[str, Any] = {}
         for key, idx, val in d['dict_items']:
-            if idx >= 0:
+            # TODO Normalize idx -1 to None via schema migrations.
+            # Long-term we should not be allowing idx == -1.
+            if idx is not None and idx >= 0:  # Older schemas might have -1 instead of None
                 arg[key] = components[idx]
             else:
                 arg[key] = val

pixeltable/exprs/literal.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
+
+import datetime
 from typing import Optional, List, Any, Dict, Tuple
 
 import sqlalchemy as sql
 
-from .expr import Expr
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from .data_row import DataRow
+from .expr import Expr
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
-import pixeltable.type_system as ts
+
 
 class Literal(Expr):
     def __init__(self, val: Any, col_type: Optional[ts.ColumnType] = None):
@@ -46,9 +49,18 @@ class Literal(Expr):
         data_row[self.slot_idx] = self.val
 
     def _as_dict(self) -> Dict:
-        return {'val': self.val, **super()._as_dict()}
+        # For some types, we need to explictly record their type, because JSON does not know
+        # how to interpret them unambiguously
+        if self.col_type.is_timestamp_type():
+            return {'val': self.val.isoformat(), 'val_t': self.col_type._type.name, **super()._as_dict()}
+        else:
+            return {'val': self.val, **super()._as_dict()}
 
     @classmethod
     def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
         assert 'val' in d
+        if 'val_t' in d:
+            val_t = d['val_t']
+            assert val_t == ts.ColumnType.Type.TIMESTAMP.name
+            return cls(datetime.datetime.fromisoformat(d['val']))
         return cls(d['val'])

pixeltable/exprs/row_builder.py

@@ -54,14 +54,14 @@ class RowBuilder:
         target_exprs: List[Expr]  # exprs corresponding to target_slot_idxs
 
     def __init__(
-            self, output_exprs: List[Expr], columns: List[catalog.Column],
-            indices: List[Tuple[catalog.Column, func.Function]], input_exprs: List[Expr]
+            self, output_exprs: List[Expr], columns: List[catalog.Column], input_exprs: List[Expr]
     ):
         """
         Args:
             output_exprs: list of Exprs to be evaluated
             columns: list of columns to be materialized
-            indices: list of embeddings to be materialized (Tuple[indexed column, embedding function])
+            input_exprs: list of Exprs that are excluded from evaluation (because they're already materialized)
+        TODO: enforce that output_exprs doesn't overlap with input_exprs?
         """
         self.unique_exprs = ExprSet()  # dependencies precede their dependents
         self.next_slot_idx = 0
@@ -73,7 +73,6 @@ class RowBuilder:
         # output exprs: all exprs the caller wants to materialize
         # - explicitly requested output_exprs
         # - values for computed columns
-        # - embedding values for indices
         resolve_cols = set(columns)
         self.output_exprs = [
             self._record_unique_expr(e.copy().resolve_computed_cols(resolve_cols=resolve_cols), recursive=True)
@@ -97,21 +96,6 @@ class RowBuilder:
                 ref = self._record_unique_expr(ref, recursive=False)
                 self.add_table_column(col, ref.slot_idx)
 
-        # record indices; indexed by slot_idx
-        self.index_columns: List[catalog.Column] = []
-        for col, embedding_fn in indices:
-            # we assume that the parameter of the embedding function is a ref to an image column
-            assert col.col_type.is_image_type()
-            # construct expr to compute embedding; explicitly resize images to the required size
-            target_img_type = next(iter(embedding_fn.signature.parameters.values())).col_type
-            expr = embedding_fn(ColumnRef(col).resize(target_img_type.size))
-            expr = self._record_unique_expr(expr, recursive=True)
-            self.output_exprs.append(expr)
-            if len(self.index_columns) <= expr.slot_idx:
-                # pad to slot_idx
-                self.index_columns.extend([None] * (expr.slot_idx - len(self.index_columns) + 1))
-            self.index_columns[expr.slot_idx] = col
-
         # default eval ctx: all output exprs
         self.default_eval_ctx = self.create_eval_ctx(self.output_exprs, exclude=unique_input_exprs)
 
@@ -170,13 +154,6 @@ class RowBuilder:
         """Return ColumnSlotIdx for output columns"""
         return self.table_columns
 
-    def index_slot_idxs(self) -> List[ColumnSlotIdx]:
-        """Return ColumnSlotIdx for index columns"""
-        return [
-            ColumnSlotIdx(self.output_columns[i], i) for i in range(len(self.index_columns))
-            if self.output_columns[i] is not None
-        ]
-
     @property
     def num_materialized(self) -> int:
         return self.next_slot_idx
@@ -204,12 +181,16 @@ class RowBuilder:
             for i, c in enumerate(expr.components):
                 # make sure we only refer to components that have themselves been recorded
                 expr.components[i] = self._record_unique_expr(c, True)
-        assert expr.slot_idx < 0
+        assert expr.slot_idx is None
         expr.slot_idx = self._next_slot_idx()
         self.unique_exprs.append(expr)
         return expr
 
     def _record_output_expr_id(self, e: Expr, output_expr_id: int) -> None:
+        assert e.slot_idx is not None
+        assert output_expr_id is not None
+        if e.slot_idx in self.input_expr_slot_idxs:
+            return
         self.output_expr_ids[e.slot_idx].add(output_expr_id)
         for d in e.dependencies():
             self._record_output_expr_id(d, output_expr_id)
@@ -334,22 +315,15 @@ class RowBuilder:
                 exc = data_row.get_exc(slot_idx)
                 num_excs += 1
                 exc_col_ids.add(col.id)
-                table_row[col.storage_name()] = None
-                table_row[col.errortype_storage_name()] = type(exc).__name__
-                table_row[col.errormsg_storage_name()] = str(exc)
+                table_row[col.store_name()] = None
+                table_row[col.errortype_store_name()] = type(exc).__name__
+                table_row[col.errormsg_store_name()] = str(exc)
             else:
-                val = data_row.get_stored_val(slot_idx)
-                table_row[col.storage_name()] = val
+                val = data_row.get_stored_val(slot_idx, col.sa_col.type)
+                table_row[col.store_name()] = val
                 # we unfortunately need to set these, even if there are no errors
-                table_row[col.errortype_storage_name()] = None
-                table_row[col.errormsg_storage_name()] = None
-
-        for slot_idx, col in enumerate(self.index_columns):
-            if col is None:
-                continue
-            # don't use get_stored_val() here, we need to pass in the ndarray
-            val = data_row[slot_idx]
-            table_row[col.index_storage_name()] = val
+                table_row[col.errortype_store_name()] = None
+                table_row[col.errormsg_store_name()] = None
 
         return table_row, num_excs
 

pixeltable/exprs/similarity_expr.py

@@ -0,0 +1,65 @@
+from typing import Optional, List
+
+import sqlalchemy as sql
+import PIL.Image
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .column_ref import ColumnRef
+from .data_row import DataRow
+from .expr import Expr
+from .literal import Literal
+from .row_builder import RowBuilder
+
+
+class SimilarityExpr(Expr):
+
+    def __init__(self, col_ref: ColumnRef, item: Expr):
+        super().__init__(ts.FloatType())
+        self.components = [col_ref, item]
+        self.id = self._create_id()
+        assert isinstance(item, Literal)
+        assert item.col_type.is_string_type() or item.col_type.is_image_type()
+
+        # determine index to use
+        idx_info = col_ref.col.get_idx_info()
+        if len(idx_info) == 0:
+            raise excs.Error(f'No index found for column {col_ref.col}')
+        if len(idx_info) > 1:
+            raise excs.Error(
+                f'Column {col_ref.col.name} has multiple indices; use the index name to disambiguate, '
+                f'e.g., `{col_ref.col.name}.<index-name>.similarity(...)`')
+        self.idx_info = next(iter(idx_info.values()))
+        idx = self.idx_info.idx
+
+        if item.col_type.is_string_type() and idx.txt_embed is None:
+            raise excs.Error(
+                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
+                f'text_embed parameter and does not support text queries')
+        if item.col_type.is_image_type() and idx.img_embed is None:
+            raise excs.Error(
+                f'Embedding index {self.idx_info.name} on column {self.idx_info.col.name} was created without the '
+                f'img_embed parameter and does not support image queries')
+
+    def __str__(self) -> str:
+        return f'{self.components[0]}.similarity({self.components[1]})'
+
+    def sql_expr(self) -> Optional[sql.ClauseElement]:
+        assert isinstance(self.components[1], Literal)
+        item = self.components[1].val
+        return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
+
+    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ClauseElement]:
+        assert isinstance(self.components[1], Literal)
+        item = self.components[1].val
+        return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        # this should never get called
+        assert False
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: List[Expr]) -> Expr:
+        assert len(components) == 2
+        assert isinstance(components[0], ColumnRef)
+        return cls(components[0], components[1])

pixeltable/ext/__init__.py

@@ -0,0 +1,5 @@
+"""
+Extended integrations for Pixeltable. This package contains experimental or demonstration features that
+are not intended for production use. Long-term support cannot be guaranteed, usually because the features
+have dependencies whose future support is unclear.
+"""