pixeltable 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

pixeltable/exprs/column_ref.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 from typing import Any, Optional, Sequence
 from uuid import UUID
 
@@ -125,11 +126,46 @@ class ColumnRef(Expr):
 
         return super().__getattr__(name)
 
+    @classmethod
+    def find_embedding_index(
+        cls, col: catalog.Column, idx_name: Optional[str], method_name: str
+    ) -> dict[str, catalog.TableVersion.IndexInfo]:
+        """Return IndexInfo for a column, with an optional given name"""
+        # determine index to use
+        idx_info_dict = col.get_idx_info()
+        from pixeltable import index
+
+        embedding_idx_info = {
+            info: value for info, value in idx_info_dict.items() if isinstance(value.idx, index.EmbeddingIndex)
+        }
+        if len(embedding_idx_info) == 0:
+            raise excs.Error(f'No indices found for {method_name!r} on column {col.name!r}')
+        if idx_name is not None and idx_name not in embedding_idx_info:
+            raise excs.Error(f'Index {idx_name!r} not found for {method_name!r} on column {col.name!r}')
+        if len(embedding_idx_info) > 1:
+            if idx_name is None:
+                raise excs.Error(
+                    f'Column {col.name!r} has multiple indices; use the index name to disambiguate: '
+                    f'`{method_name}(..., idx=<index_name>)`'
+                )
+            idx_info = {idx_name: embedding_idx_info[idx_name]}
+        else:
+            idx_info = embedding_idx_info
+        return idx_info
+
     def similarity(self, item: Any, *, idx: Optional[str] = None) -> Expr:
         from .similarity_expr import SimilarityExpr
 
         return SimilarityExpr(self, item, idx_name=idx)
 
+    def embedding(self, *, idx: Optional[str] = None) -> ColumnRef:
+        idx_info = ColumnRef.find_embedding_index(self.col, idx, 'embedding')
+        assert len(idx_info) == 1
+        col = copy.copy(next(iter(idx_info.values())).val_col)
+        col.name = f'{self.col.name}_embedding_{idx if idx is not None else ""}'
+        col.create_sa_cols()
+        return ColumnRef(col)
+
     def default_column_name(self) -> Optional[str]:
         return str(self)
 

pixeltable/exprs/expr.py

@@ -17,7 +17,7 @@ from typing_extensions import Self, _AnnotatedAlias
 from pixeltable import catalog, exceptions as excs, func, type_system as ts
 
 from .data_row import DataRow
-from .globals import ArithmeticOperator, ComparisonOperator, LiteralPythonTypes, LogicalOperator
+from .globals import ArithmeticOperator, ComparisonOperator, LiteralPythonTypes, LogicalOperator, StringOperator
 
 if TYPE_CHECKING:
     from pixeltable import exprs
@@ -605,10 +605,6 @@ class Expr(abc.ABC):
                 # Return the `MethodRef` object itself; it requires arguments to become a `FunctionCall`
                 return method_ref
 
-    def __rshift__(self, other: object) -> 'exprs.Expr':
-        # Implemented here for type-checking purposes
-        raise excs.Error('The `>>` operator can only be applied to Json expressions')
-
     def __bool__(self) -> bool:
         raise TypeError(
             f'Pixeltable expressions cannot be used in conjunction with Python boolean operators (and/or/not)\n{self!r}'
@@ -658,13 +654,17 @@ class Expr(abc.ABC):
     def __neg__(self) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, -1)
 
-    def __add__(self, other: object) -> 'exprs.ArithmeticExpr':
+    def __add__(self, other: object) -> Union[exprs.ArithmeticExpr, exprs.StringOp]:
+        if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
+            return self._make_string_expr(StringOperator.CONCAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.ADD, other)
 
     def __sub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __mul__(self, other: object) -> 'exprs.ArithmeticExpr':
+    def __mul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+        if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
+            return self._make_string_expr(StringOperator.REPEAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, other)
 
     def __truediv__(self, other: object) -> 'exprs.ArithmeticExpr':
@@ -676,13 +676,17 @@ class Expr(abc.ABC):
     def __floordiv__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.FLOORDIV, other)
 
-    def __radd__(self, other: object) -> 'exprs.ArithmeticExpr':
+    def __radd__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+        if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
+            return self._rmake_string_expr(StringOperator.CONCAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.ADD, other)
 
     def __rsub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._rmake_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __rmul__(self, other: object) -> 'exprs.ArithmeticExpr':
+    def __rmul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+        if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
+            return self._rmake_string_expr(StringOperator.REPEAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.MUL, other)
 
     def __rtruediv__(self, other: object) -> 'exprs.ArithmeticExpr':
@@ -694,6 +698,32 @@ class Expr(abc.ABC):
     def __rfloordiv__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._rmake_arithmetic_expr(ArithmeticOperator.FLOORDIV, other)
 
+    def _make_string_expr(self, op: StringOperator, other: object) -> 'exprs.StringOp':
+        """
+        Make left-handed version of string expression.
+        """
+        from .literal import Literal
+        from .string_op import StringOp
+
+        if isinstance(other, Expr):
+            return StringOp(op, self, other)
+        if isinstance(other, typing.get_args(LiteralPythonTypes)):
+            return StringOp(op, self, Literal(other))
+        raise TypeError(f'Other must be Expr or literal: {type(other)}')
+
+    def _rmake_string_expr(self, op: StringOperator, other: object) -> 'exprs.StringOp':
+        """
+        Right-handed version of _make_string_expr. other must be a literal; if it were an Expr,
+        the operation would have already been evaluated in its left-handed form.
+        """
+        from .literal import Literal
+        from .string_op import StringOp
+
+        assert not isinstance(other, Expr)  # Else the left-handed form would have evaluated first
+        if isinstance(other, typing.get_args(LiteralPythonTypes)):
+            return StringOp(op, Literal(other), self)
+        raise TypeError(f'Other must be Expr or literal: {type(other)}')
+
     def _make_arithmetic_expr(self, op: ArithmeticOperator, other: object) -> 'exprs.ArithmeticExpr':
         """
         other: Union[Expr, LiteralPythonTypes]

pixeltable/exprs/globals.py

@@ -87,3 +87,15 @@ class ArithmeticOperator(enum.Enum):
         if self == self.FLOORDIV:
             return '//'
         raise AssertionError()
+
+
+class StringOperator(enum.Enum):
+    CONCAT = 0
+    REPEAT = 1
+
+    def __str__(self) -> str:
+        if self == self.CONCAT:
+            return '+'
+        if self == self.REPEAT:
+            return '*'
+        raise AssertionError()

pixeltable/exprs/json_mapper.py

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

pixeltable/exprs/json_path.py

@@ -110,12 +110,6 @@ class JsonPath(Expr):
             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)
-        if rhs_expr is None:
-            raise excs.Error(f'>> requires an expression on the right-hand side, found {type(other)}')
-        return JsonMapper(self, rhs_expr)
-
     def default_column_name(self) -> Optional[str]:
         anchor_name = self._anchor.default_column_name() if self._anchor is not None else ''
         ret_name = f'{anchor_name}.{self._json_path()}'

pixeltable/exprs/similarity_expr.py

@@ -23,26 +23,12 @@ class SimilarityExpr(Expr):
 
         self.components = [col_ref, item_expr]
 
-        # determine index to use
-        idx_info = col_ref.col.get_idx_info()
         from pixeltable import index
 
-        embedding_idx_info = {
-            info.name: info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)
-        }
-        if len(embedding_idx_info) == 0:
-            raise excs.Error(f'No index found for column {col_ref.col!r}')
-        if idx_name is not None and idx_name not in embedding_idx_info:
-            raise excs.Error(f'Index {idx_name!r} not found for column {col_ref.col.name!r}')
-        if len(embedding_idx_info) > 1:
-            if idx_name is None:
-                raise excs.Error(
-                    f'Column {col_ref.col.name!r} has multiple indices; use the index name to disambiguate: '
-                    f'`{col_ref.col.name}.similarity(..., idx=<name>)`'
-                )
-            self.idx_info = embedding_idx_info[idx_name]
-        else:
-            self.idx_info = next(iter(embedding_idx_info.values()))
+        # determine index to use
+        idx_dict = ColumnRef.find_embedding_index(col_ref.col, idx_name, 'similarity')
+        assert len(idx_dict) == 1
+        self.idx_info = next(iter(idx_dict.values()))
         idx = self.idx_info.idx
         assert isinstance(idx, index.EmbeddingIndex)
 
@@ -86,8 +72,7 @@ class SimilarityExpr(Expr):
         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
-        raise AssertionError()
+        raise excs.Error('similarity(): cannot be used in a computed column')
 
     def _as_dict(self) -> dict:
         return {'idx_name': self.idx_info.name, **super()._as_dict()}

pixeltable/exprs/string_op.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Any, Optional, Union
+
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+from .data_row import DataRow
+from .expr import Expr
+from .globals import StringOperator
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
+
+
+class StringOp(Expr):
+    """
+    Allows operations on strings
+    """
+
+    operator: StringOperator
+
+    def __init__(self, operator: StringOperator, op1: Expr, op2: Expr):
+        super().__init__(ts.StringType(nullable=op1.col_type.nullable))
+        self.operator = operator
+        self.components = [op1, op2]
+        assert op1.col_type.is_string_type()
+        if operator in {StringOperator.CONCAT, StringOperator.REPEAT}:
+            if operator == StringOperator.CONCAT and not op2.col_type.is_string_type():
+                raise excs.Error(
+                    f'{self}: {operator} on strings requires string type, but {op2} has type {op2.col_type}'
+                )
+            if operator == StringOperator.REPEAT and not op2.col_type.is_int_type():
+                raise excs.Error(f'{self}: {operator} on strings requires int type, but {op2} has type {op2.col_type}')
+        else:
+            raise excs.Error(
+                f'{self}: invalid operation {operator} on strings; '
+                f'only operators {StringOperator.CONCAT} and {StringOperator.REPEAT} are supported'
+            )
+        self.id = self._create_id()
+
+    @property
+    def _op1(self) -> Expr:
+        return self.components[0]
+
+    @property
+    def _op2(self) -> Expr:
+        return self.components[1]
+
+    def __repr__(self) -> str:
+        # add parentheses around operands that are StringOpExpr to express precedence
+        op1_str = f'({self._op1})' if isinstance(self._op1, StringOp) else str(self._op1)
+        op2_str = f'({self._op2})' if isinstance(self._op2, StringOp) else str(self._op2)
+        return f'{op1_str} {self.operator} {op2_str}'
+
+    def _equals(self, other: StringOp) -> bool:
+        return self.operator == other.operator
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
+
+    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+        left = sql_elements.get(self._op1)
+        right = sql_elements.get(self._op2)
+        if left is None or right is None:
+            return None
+        if self.operator == StringOperator.CONCAT:
+            return left.concat(right)
+        if self.operator == StringOperator.REPEAT:
+            return sql.func.repeat(sql.cast(left, sql.String), sql.cast(right, sql.Integer))
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        op1_val = data_row[self._op1.slot_idx]
+        op2_val = data_row[self._op2.slot_idx]
+        data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
+
+    def eval_nullable(self, op1_val: Union[str, None], op2_val: Union[int, str, None]) -> Union[str, None]:
+        """
+        Return the result of evaluating the expression on two nullable int/float operands,
+        None is interpreted as SQL NULL
+        """
+        if op1_val is None or op2_val is None:
+            return None
+        return self.eval_non_null(op1_val, op2_val)
+
+    def eval_non_null(self, op1_val: str, op2_val: Union[int, str]) -> str:
+        """
+        Return the result of evaluating the expression on two int/float operands
+        """
+        assert self.operator in {StringOperator.CONCAT, StringOperator.REPEAT}
+        if self.operator == StringOperator.CONCAT:
+            assert isinstance(op2_val, str)
+            return op1_val + op2_val
+        else:
+            assert isinstance(op2_val, int)
+            return op1_val * op2_val
+
+    def _as_dict(self) -> dict:
+        return {'operator': self.operator.value, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> StringOp:
+        assert 'operator' in d
+        assert len(components) == 2
+        return cls(StringOperator(d['operator']), components[0], components[1])

pixeltable/ext/functions/yolox.py

@@ -1,21 +1,15 @@
 import logging
-from pathlib import Path
-from typing import TYPE_CHECKING, Iterable, Iterator
-from urllib.request import urlretrieve
+from typing import TYPE_CHECKING
 
-import numpy as np
 import PIL.Image
 
 import pixeltable as pxt
-from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.functions.util import normalize_image_mode
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import torch
-    from yolox.exp import Exp  # type: ignore[import-untyped]
-    from yolox.models import YOLOX  # type: ignore[import-untyped]
+    from yolox.models import Yolox, YoloxProcessor  # type: ignore[import-untyped]
 
 _logger = logging.getLogger('pixeltable')
 
@@ -30,7 +24,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
 
     __Requirements__:
 
-    - `pip install git+https://github.com/Megvii-BaseDetection/YOLOX`
+    - `pip install pixeltable-yolox`
 
     Args:
         model_id: one of: `yolox_nano`, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
@@ -46,31 +40,14 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
         >>> tbl.add_computed_column(detections=yolox(tbl.image, model_id='yolox_m', threshold=0.8))
     """
     import torch
-    from yolox.utils import postprocess  # type: ignore[import-untyped]
-
-    model, exp = _lookup_model(model_id, 'cpu')
-    image_tensors = list(_images_to_tensors(images, exp))
-    batch_tensor = torch.stack(image_tensors)
 
+    model = _lookup_model(model_id, 'cpu')
+    processor = _lookup_processor(model_id)
+    normalized_images = [normalize_image_mode(image) for image in images]
     with torch.no_grad():
-        output_tensor = model(batch_tensor)
-
-    outputs = postprocess(output_tensor, 80, threshold, exp.nmsthre, class_agnostic=False)
-
-    results: list[dict] = []
-    for image in images:
-        ratio = min(exp.test_size[0] / image.height, exp.test_size[1] / image.width)
-        if outputs[0] is None:
-            results.append({'bboxes': [], 'scores': [], 'labels': []})
-        else:
-            results.append(
-                {
-                    'bboxes': [(output[:4] / ratio).tolist() for output in outputs[0]],
-                    'scores': [output[4].item() * output[5].item() for output in outputs[0]],
-                    'labels': [int(output[6]) for output in outputs[0]],
-                }
-            )
-    return results
+        tensor = processor(normalized_images)
+        output = model(tensor)
+        return processor.postprocess(normalized_images, output, threshold=threshold)
 
 
 @pxt.udf
@@ -107,47 +84,27 @@ def yolo_to_coco(detections: dict) -> list:
     return result
 
 
-def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: 'Exp') -> Iterator['torch.Tensor']:
-    import torch
-    from yolox.data import ValTransform  # type: ignore[import-untyped]
-
-    val_transform = ValTransform(legacy=False)
-    for image in images:
-        normalized_image = normalize_image_mode(image)
-        image_transform, _ = val_transform(np.array(normalized_image), None, exp.test_size)
-        yield torch.from_numpy(image_transform)
-
-
-def _lookup_model(model_id: str, device: str) -> tuple['YOLOX', 'Exp']:
-    import torch
-    from yolox.exp import get_exp
+def _lookup_model(model_id: str, device: str) -> 'Yolox':
+    from yolox.models import Yolox
 
     key = (model_id, device)
-    if key in _model_cache:
-        return _model_cache[key]
+    if key not in _model_cache:
+        _model_cache[key] = Yolox.from_pretrained(model_id, device=device)
 
-    weights_url = f'https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/{model_id}.pth'
-    weights_file = Path(f'{env.Env.get().tmp_dir}/{model_id}.pth')
-    if not weights_file.exists():
-        _logger.info(f'Downloading weights for YOLOX model {model_id}: from {weights_url} -> {weights_file}')
-        urlretrieve(weights_url, weights_file)
+    return _model_cache[key]
 
-    exp = get_exp(exp_name=model_id)
-    model = exp.get_model().to(device)
 
-    model.eval()
-    model.head.training = False
-    model.training = False
+def _lookup_processor(model_id: str) -> 'YoloxProcessor':
+    from yolox.models import YoloxProcessor
 
-    # Load in the weights from training
-    weights = torch.load(weights_file, map_location=torch.device(device))
-    model.load_state_dict(weights['model'])
+    if model_id not in _processor_cache:
+        _processor_cache[model_id] = YoloxProcessor(model_id)
 
-    _model_cache[key] = (model, exp)
-    return model, exp
+    return _processor_cache[model_id]
 
 
-_model_cache: dict[tuple[str, str], tuple['YOLOX', 'Exp']] = {}
+_model_cache: dict[tuple[str, str], 'Yolox'] = {}
+_processor_cache: dict[str, 'YoloxProcessor'] = {}
 
 
 __all__ = local_public_names(__name__)

pixeltable/func/tools.py

@@ -51,10 +51,10 @@ class Tool(pydantic.BaseModel):
     # The output of `tool_calls` must be a dict in standardized tool invocation format:
     # {tool_name: [{'args': {name1: value1, name2: value2, ...}}, ...], ...}
     def invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.Expr':
-        from pixeltable import exprs
+        import pixeltable.functions as pxtf
 
         func_name = self.name or self.fn.name
-        return exprs.JsonMapper(tool_calls[func_name]['*'], self.__invoke_kwargs(exprs.RELATIVE_PATH_ROOT.args))
+        return pxtf.map(tool_calls[func_name]['*'], lambda x: self.__invoke_kwargs(x.args))
 
     def __invoke_kwargs(self, kwargs: 'exprs.Expr') -> 'exprs.FunctionCall':
         kwargs = {param.name: self.__extract_tool_arg(param, kwargs) for param in self.parameters.values()}

pixeltable/functions/__init__.py

@@ -24,7 +24,7 @@ from . import (
     vision,
     whisper,
 )
-from .globals import count, max, mean, min, sum
+from .globals import count, map, max, mean, min, sum
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 

pixeltable/functions/globals.py

@@ -1,15 +1,14 @@
 import builtins
 import typing
-
-from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
-from typing import Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import sqlalchemy as sql
 
-import pixeltable.type_system as ts
-from pixeltable import exprs, func
+from pixeltable import exceptions as excs, exprs, func, type_system as ts
 from pixeltable.utils.code import local_public_names
 
+from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
+
 
 # TODO: remove and replace calls with astype()
 def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias]) -> exprs.Expr:
@@ -168,6 +167,18 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.avg(val)
 
 
+def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
+    target_expr: exprs.Expr
+    try:
+        target_expr = exprs.Expr.from_object(fn(exprs.json_path.RELATIVE_PATH_ROOT))
+    except Exception as e:
+        raise excs.Error(
+            'Failed to evaluate map function. '
+            '(The `fn` argument to `map()` must produce a valid Pixeltable expression.)'
+        ) from e
+    return exprs.JsonMapper(expr, target_expr)
+
+
 __all__ = local_public_names(__name__)