pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/index/base.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import abc
-from typing import Any
 
 import sqlalchemy as sql
 
-from pixeltable import catalog, exprs
+import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
+import pixeltable.type_system as ts
 
 
 class IndexBase(abc.ABC):
@@ -18,39 +19,34 @@ class IndexBase(abc.ABC):
     """
 
     @abc.abstractmethod
-    def __init__(self, c: catalog.Column, **kwargs: Any):
-        pass
-
-    @abc.abstractmethod
-    def index_value_expr(self) -> exprs.Expr:
-        """Return expression that computes the value that goes into the index"""
-        pass
+    def create_value_expr(self, c: catalog.Column) -> exprs.Expr:
+        """
+        Validates that the index can be created on column c and returns an expression that computes the index value.
+        """
 
     @abc.abstractmethod
     def records_value_errors(self) -> bool:
         """True if index_value_expr() can raise errors"""
-        pass
 
     @abc.abstractmethod
-    def index_sa_type(self) -> sql.types.TypeEngine:
+    def get_index_sa_type(self, value_col_type: ts.ColumnType) -> sql.types.TypeEngine:
         """Return the sqlalchemy type of the index value column"""
-        pass
 
     @abc.abstractmethod
-    def create_index(self, index_name: str, index_value_col: catalog.Column) -> None:
-        """Create the index on the index value column"""
-        pass
+    def sa_create_stmt(self, store_index_name: str, sa_value_col: sql.Column) -> sql.Compiled:
+        """Return a sqlalchemy statement for creating the index"""
+
+    @abc.abstractmethod
+    def drop_index(self, index_name: str, index_value_col: catalog.Column) -> None:
+        """Drop the index on the index value column"""
 
     @classmethod
     @abc.abstractmethod
-    def display_name(cls) -> str:
-        pass
+    def display_name(cls) -> str: ...
 
     @abc.abstractmethod
-    def as_dict(self) -> dict:
-        pass
+    def as_dict(self) -> dict: ...
 
     @classmethod
     @abc.abstractmethod
-    def from_dict(cls, c: catalog.Column, d: dict) -> IndexBase:
-        pass
+    def from_dict(cls, d: dict) -> IndexBase: ...

pixeltable/index/btree.py

@@ -1,18 +1,18 @@
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import sqlalchemy as sql
 
 # TODO: why does this import result in a circular import, but the one im embedding_index.py doesn't?
 # import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
-from pixeltable import catalog, exprs
-from pixeltable.env import Env
+import pixeltable.exprs as exprs
+import pixeltable.type_system as ts
 from pixeltable.func.udf import udf
 
 from .base import IndexBase
 
 if TYPE_CHECKING:
-    import pixeltable.exprs
+    import pixeltable.catalog as catalog
 
 
 class BtreeIndex(IndexBase):
@@ -22,42 +22,48 @@ class BtreeIndex(IndexBase):
 
     MAX_STRING_LEN = 256
 
-    value_expr: 'pixeltable.exprs.Expr'
-
     @staticmethod
     @udf
-    def str_filter(s: Optional[str]) -> Optional[str]:
+    def str_filter(s: str | None) -> str | None:
         if s is None:
             return None
         return s[: BtreeIndex.MAX_STRING_LEN]
 
-    def __init__(self, c: 'catalog.Column'):
+    def __init__(self) -> None:
+        pass
+
+    def create_value_expr(self, c: 'catalog.Column') -> 'exprs.Expr':
         if not c.col_type.is_scalar_type() and not c.col_type.is_media_type():
             raise excs.Error(f'Index on column {c.name}: B-tree index requires scalar or media type, got {c.col_type}')
+        value_expr: exprs.Expr
         if c.col_type.is_media_type():
             # an index on a media column is an index on the file url
             # no validation for media columns: we're only interested in the string value
-            self.value_expr = exprs.ColumnRef(c, perform_validation=False)
+            value_expr = exprs.ColumnRef(c, perform_validation=False)
         else:
-            self.value_expr = (
+            value_expr = (
                 BtreeIndex.str_filter(exprs.ColumnRef(c)) if c.col_type.is_string_type() else exprs.ColumnRef(c)
             )
-
-    def index_value_expr(self) -> 'exprs.Expr':
-        return self.value_expr
+        return value_expr
 
     def records_value_errors(self) -> bool:
         return False
 
-    def index_sa_type(self) -> sql.types.TypeEngine:
+    def get_index_sa_type(self, val_col_type: ts.ColumnType) -> sql.types.TypeEngine:
         """Return the sqlalchemy type of the index value column"""
-        return self.value_expr.col_type.to_sa_type()
+        return val_col_type.to_sa_type()
+
+    def sa_create_stmt(self, store_index_name: str, sa_value_col: sql.Column) -> sql.Compiled:
+        """Return a sqlalchemy statement for creating the index"""
+        from sqlalchemy.dialects import postgresql
+
+        sa_idx = sql.Index(store_index_name, sa_value_col, postgresql_using='btree')
+        return sql.schema.CreateIndex(sa_idx, if_not_exists=True).compile(dialect=postgresql.dialect())
 
-    def create_index(self, index_name: str, index_value_col: 'catalog.Column') -> None:
-        """Create the index on the index value column"""
-        idx = sql.Index(index_name, index_value_col.sa_col, postgresql_using='btree')
-        conn = Env.get().conn
-        idx.create(bind=conn)
+    def drop_index(self, index_name: str, index_value_col: 'catalog.Column') -> None:
+        """Drop the index on the index value column"""
+        # TODO: implement
+        raise NotImplementedError()
 
     @classmethod
     def display_name(cls) -> str:
@@ -67,5 +73,5 @@ class BtreeIndex(IndexBase):
         return {}
 
     @classmethod
-    def from_dict(cls, c: 'catalog.Column', d: dict) -> 'BtreeIndex':
-        return cls(c)
+    def from_dict(cls, d: dict) -> 'BtreeIndex':
+        return cls()

pixeltable/index/embedding_index.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 import enum
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 import numpy as np
 import pgvector.sqlalchemy  # type: ignore[import-untyped]
-import PIL.Image
 import sqlalchemy as sql
 
+import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
+import pixeltable.exprs as exprs
+import pixeltable.func as func
 import pixeltable.type_system as ts
-from pixeltable import catalog, exprs, func
 from pixeltable.env import Env
 
 from .base import IndexBase
@@ -39,124 +40,105 @@ class EmbeddingIndex(IndexBase):
     }
 
     metric: Metric
-    value_expr: exprs.FunctionCall
-    string_embed: Optional[func.Function]
-    image_embed: Optional[func.Function]
-    string_embed_signature_idx: int
-    image_embed_signature_idx: int
-    index_col_type: pgvector.sqlalchemy.Vector
+    embeddings: dict[ts.ColumnType.Type, func.Function]
 
     def __init__(
         self,
-        c: catalog.Column,
         metric: str,
-        embed: Optional[func.Function] = None,
-        string_embed: Optional[func.Function] = None,
-        image_embed: Optional[func.Function] = None,
+        embed: func.Function | None = None,
+        string_embed: func.Function | None = None,
+        image_embed: func.Function | None = None,
+        audio_embed: func.Function | None = None,
+        video_embed: func.Function | None = None,
     ):
         if embed is None and string_embed is None and image_embed is None:
             raise excs.Error('At least one of `embed`, `string_embed`, or `image_embed` must be specified')
         metric_names = [m.name.lower() for m in self.Metric]
         if metric.lower() not in metric_names:
             raise excs.Error(f'Invalid metric {metric}, must be one of {metric_names}')
-        if not c.col_type.is_string_type() and not c.col_type.is_image_type():
-            raise excs.Error('Embedding index requires string or image column')
-
-        self.string_embed = None
-        self.image_embed = None
-
-        # Resolve the specific embedding functions corresponding to the user-provided `string_embed`, `image_embed`,
-        # and/or `embed`. For string embeddings, `string_embed` will be used if specified; otherwise, `embed` will
-        # be used as a fallback, if it has a matching signature. Likewise for image embeddings.
-
-        if string_embed is not None:
-            # `string_embed` is specified; it MUST be valid.
-            self.string_embed = self._resolve_embedding_fn(string_embed, ts.ColumnType.Type.STRING)
-            if self.string_embed is None:
-                raise excs.Error(
-                    f'The function `{string_embed.name}` is not a valid string embedding: '
-                    'it must take a single string parameter'
-                )
-        elif embed is not None:
-            # `embed` is specified; see if it has a string signature.
-            self.string_embed = self._resolve_embedding_fn(embed, ts.ColumnType.Type.STRING)
-
-        if image_embed is not None:
-            # `image_embed` is specified; it MUST be valid.
-            self.image_embed = self._resolve_embedding_fn(image_embed, ts.ColumnType.Type.IMAGE)
-            if self.image_embed is None:
-                raise excs.Error(
-                    f'The function `{image_embed.name}` is not a valid image embedding: '
-                    'it must take a single image parameter'
-                )
-        elif embed is not None:
-            # `embed` is specified; see if it has an image signature.
-            self.image_embed = self._resolve_embedding_fn(embed, ts.ColumnType.Type.IMAGE)
-
-        if self.string_embed is None and self.image_embed is None:
-            # No string OR image signature was found. This can only happen if `embed` was specified and
-            # contains no matching signatures.
+
+        self.embeddings = {}
+
+        # Resolve the specific embedding functions corresponding to the user-provided embedding functions.
+        # For string embeddings, for example, `string_embed` will be used if specified; otherwise, `embed` will
+        # be used as a fallback, if it has a matching signature.
+
+        for embed_type, embed_fn in (
+            (ts.ColumnType.Type.STRING, string_embed),
+            (ts.ColumnType.Type.IMAGE, image_embed),
+            (ts.ColumnType.Type.AUDIO, audio_embed),
+            (ts.ColumnType.Type.VIDEO, video_embed),
+        ):
+            if embed_fn is not None:
+                # Embedding function for the requisite type is specified directly; it MUST be valid.
+                resolved_fn = self._resolve_embedding_fn(embed_fn, embed_type)
+                if resolved_fn is None:
+                    raise excs.Error(
+                        f'The function `{embed_fn.name}` is not a valid {embed_type.name.lower()} '
+                        f'embedding: it must take a single {embed_type.name.lower()} parameter'
+                    )
+                self.embeddings[embed_type] = resolved_fn
+            elif embed is not None:
+                # General `embed` is specified; see if it has a matching signature.
+                resolved_fn = self._resolve_embedding_fn(embed, embed_type)
+                if resolved_fn is not None:
+                    self.embeddings[embed_type] = resolved_fn
+
+        if len(self.embeddings) == 0:
+            # `embed` was specified and contains no matching signatures.
             assert embed is not None
             raise excs.Error(
-                f'The function `{embed.name}` is not a valid embedding: it must take a single string or image parameter'
+                f'The function `{embed.name}` is not a valid embedding: '
+                'it must take a single string, image, audio, or video parameter'
             )
 
         # Now validate the return types of the embedding functions.
-
-        if self.string_embed is not None:
-            self._validate_embedding_fn(self.string_embed)
-
-        if self.image_embed is not None:
-            self._validate_embedding_fn(self.image_embed)
-
-        if c.col_type.is_string_type() and self.string_embed is None:
-            raise excs.Error(f"Text embedding function is required for column {c.name} (parameter 'string_embed')")
-        if c.col_type.is_image_type() and self.image_embed is None:
-            raise excs.Error(f"Image embedding function is required for column {c.name} (parameter 'image_embed')")
+        for _, embed_fn in self.embeddings.items():
+            self._validate_embedding_fn(embed_fn)
 
         self.metric = self.Metric[metric.upper()]
-        self.value_expr = (
-            self.string_embed(exprs.ColumnRef(c))
-            if c.col_type.is_string_type()
-            else self.image_embed(exprs.ColumnRef(c))
-        )
-        assert isinstance(self.value_expr.col_type, ts.ArrayType)
-        vector_size = self.value_expr.col_type.shape[0]
-        assert vector_size is not None
-        self.index_col_type = pgvector.sqlalchemy.Vector(vector_size)
 
-    def index_value_expr(self) -> exprs.Expr:
-        """Return expression that computes the value that goes into the index"""
-        return self.value_expr
+    def create_value_expr(self, c: catalog.Column) -> exprs.Expr:
+        if c.col_type._type not in (
+            ts.ColumnType.Type.STRING,
+            ts.ColumnType.Type.IMAGE,
+            ts.ColumnType.Type.AUDIO,
+            ts.ColumnType.Type.VIDEO,
+        ):
+            raise excs.Error(f'Type `{c.col_type}` of column {c.name!r} is not a valid type for an embedding index.')
+        if c.col_type._type not in self.embeddings:
+            raise excs.Error(
+                f'The specified embedding function does not support the type `{c.col_type}` of column {c.name!r}.'
+            )
+
+        embed_fn = self.embeddings[c.col_type._type]
+        return embed_fn(exprs.ColumnRef(c))
 
     def records_value_errors(self) -> bool:
         return True
 
-    def index_sa_type(self) -> sql.types.TypeEngine:
-        """Return the sqlalchemy type of the index value column"""
-        return self.index_col_type
-
-    def create_index(self, index_name: str, index_value_col: catalog.Column) -> None:
-        """Create the index on the index value column"""
-        idx = sql.Index(
-            index_name,
-            index_value_col.sa_col,
-            postgresql_using='hnsw',
-            postgresql_with={'m': 16, 'ef_construction': 64},
-            postgresql_ops={index_value_col.sa_col.name: self.PGVECTOR_OPS[self.metric]},
+    def get_index_sa_type(self, val_col_type: ts.ColumnType) -> sql.types.TypeEngine:
+        assert isinstance(val_col_type, ts.ArrayType) and val_col_type.shape is not None
+        vector_size = val_col_type.shape[0]
+        assert vector_size is not None
+        return pgvector.sqlalchemy.Vector(vector_size)
+
+    def sa_create_stmt(self, store_index_name: str, sa_value_col: sql.Column) -> sql.Compiled:
+        """Return a sqlalchemy statement for creating the index"""
+        return Env.get().dbms.create_vector_index_stmt(
+            store_index_name, sa_value_col, metric=self.PGVECTOR_OPS[self.metric]
         )
-        conn = Env.get().conn
-        idx.create(bind=conn)
 
-    def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ColumnElement:
+    def drop_index(self, index_name: str, index_value_col: catalog.Column) -> None:
+        """Drop the index on the index value column"""
+        # TODO: implement
+        raise NotImplementedError()
+
+    def similarity_clause(self, val_column: catalog.Column, item: exprs.Literal) -> sql.ColumnElement:
         """Create a ColumnElement that represents '<val_column> <op> <item>'"""
-        assert isinstance(item, (str, PIL.Image.Image))
-        if isinstance(item, str):
-            assert self.string_embed is not None
-            embedding = self.string_embed.exec([item], {})
-        if isinstance(item, PIL.Image.Image):
-            assert self.image_embed is not None
-            embedding = self.image_embed.exec([item], {})
+        assert item.col_type._type in self.embeddings
+        embedding = self.embeddings[item.col_type._type].exec([item.val], {})
+        assert isinstance(embedding, np.ndarray)
 
         if self.metric == self.Metric.COSINE:
             return val_column.sa_col.cosine_distance(embedding) * -1 + 1
@@ -166,17 +148,11 @@ class EmbeddingIndex(IndexBase):
             assert self.metric == self.Metric.L2
             return val_column.sa_col.l2_distance(embedding)
 
-    def order_by_clause(self, val_column: catalog.Column, item: Any, is_asc: bool) -> sql.ColumnElement:
+    def order_by_clause(self, val_column: catalog.Column, item: exprs.Literal, is_asc: bool) -> sql.ColumnElement:
         """Create a ColumnElement that is used in an ORDER BY clause"""
-        assert isinstance(item, (str, PIL.Image.Image))
-        embedding: Optional[np.ndarray] = None
-        if isinstance(item, str):
-            assert self.string_embed is not None
-            embedding = self.string_embed.exec([item], {})
-        if isinstance(item, PIL.Image.Image):
-            assert self.image_embed is not None
-            embedding = self.image_embed.exec([item], {})
-        assert embedding is not None
+        assert item.col_type._type in self.embeddings
+        embedding = self.embeddings[item.col_type._type].exec([item.val], {})
+        assert isinstance(embedding, np.ndarray)
 
         if self.metric == self.Metric.COSINE:
             result = val_column.sa_col.cosine_distance(embedding)
@@ -194,9 +170,7 @@ class EmbeddingIndex(IndexBase):
         return 'embedding'
 
     @classmethod
-    def _resolve_embedding_fn(
-        cls, embed_fn: func.Function, expected_type: ts.ColumnType.Type
-    ) -> Optional[func.Function]:
+    def _resolve_embedding_fn(cls, embed_fn: func.Function, expected_type: ts.ColumnType.Type) -> func.Function | None:
         """Find an overload resolution for `embed_fn` that matches the given type."""
         assert isinstance(embed_fn, func.Function)
         for resolved_fn in embed_fn._resolved_fns:
@@ -243,14 +217,22 @@ class EmbeddingIndex(IndexBase):
             )
 
     def as_dict(self) -> dict:
-        return {
-            'metric': self.metric.name.lower(),
-            'string_embed': None if self.string_embed is None else self.string_embed.as_dict(),
-            'image_embed': None if self.image_embed is None else self.image_embed.as_dict(),
-        }
+        d: dict[str, Any] = {'metric': self.metric.name.lower()}
+        for embed_type, embed_fn in self.embeddings.items():
+            key = f'{embed_type.name.lower()}_embed'
+            d[key] = embed_fn.as_dict()
+        return d
 
     @classmethod
-    def from_dict(cls, c: catalog.Column, d: dict) -> EmbeddingIndex:
-        string_embed = func.Function.from_dict(d['string_embed']) if d['string_embed'] is not None else None
-        image_embed = func.Function.from_dict(d['image_embed']) if d['image_embed'] is not None else None
-        return cls(c, metric=d['metric'], string_embed=string_embed, image_embed=image_embed)
+    def from_dict(cls, d: dict) -> EmbeddingIndex:
+        string_embed = func.Function.from_dict(d['string_embed']) if d.get('string_embed') is not None else None
+        image_embed = func.Function.from_dict(d['image_embed']) if d.get('image_embed') is not None else None
+        audio_embed = func.Function.from_dict(d['audio_embed']) if d.get('audio_embed') is not None else None
+        video_embed = func.Function.from_dict(d['video_embed']) if d.get('video_embed') is not None else None
+        return cls(
+            metric=d['metric'],
+            string_embed=string_embed,
+            image_embed=image_embed,
+            audio_embed=audio_embed,
+            video_embed=video_embed,
+        )

pixeltable/io/__init__.py

@@ -1,14 +1,16 @@
+"""Functions for importing and exporting Pixeltable data."""
 # ruff: noqa: F401
 
 from .datarows import import_json, import_rows
-from .external_store import ExternalStore, SyncStatus
+from .external_store import ExternalStore
 from .globals import create_label_studio_project, export_images_as_fo_dataset
 from .hf_datasets import import_huggingface_dataset
+from .lancedb import export_lancedb
 from .pandas import import_csv, import_excel, import_pandas
 from .parquet import export_parquet, import_parquet
 
 __default_dir = {symbol for symbol in dir() if not symbol.startswith('_')}
-__removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet', 'datarows'}
+__removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet', 'datarows', 'lancedb'}
 __all__ = sorted(__default_dir - __removed_symbols)
 
 

pixeltable/io/datarows.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Iterable, Optional, Union
+from typing import Any, Iterable
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
@@ -8,7 +8,7 @@ from pixeltable import exceptions as excs
 
 
 def _infer_schema_from_rows(
-    rows: Iterable[dict[str, Any]], schema_overrides: dict[str, Any], primary_key: list[str]
+    rows: Iterable[dict[str, Any]], schema_overrides: dict[str, ts.ColumnType], primary_key: list[str]
 ) -> dict[str, ts.ColumnType]:
     schema: dict[str, ts.ColumnType] = {}
     cols_with_nones: set[str] = set()
@@ -20,6 +20,7 @@ def _infer_schema_from_rows(
                 # in which the column names are encountered in the input data, even if `schema_overrides`
                 # is specified.
                 if col_name not in schema:
+                    assert isinstance(schema_overrides[col_name], ts.ColumnType)
                     schema[col_name] = schema_overrides[col_name]
             elif value is not None:
                 # If `key` is not in `schema_overrides`, then we infer its type from the data.
@@ -33,7 +34,7 @@ def _infer_schema_from_rows(
                 if col_name not in schema:
                     schema[col_name] = col_type
                 else:
-                    supertype = schema[col_name].supertype(col_type)
+                    supertype = schema[col_name].supertype(col_type, for_inference=True)
                     if supertype is None:
                         raise excs.Error(
                             f'Could not infer type of column `{col_name}`; the value in row {n} '
@@ -59,8 +60,8 @@ def import_rows(
     tbl_path: str,
     rows: list[dict[str, Any]],
     *,
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
 ) -> pxt.Table:
@@ -103,8 +104,8 @@ def import_json(
     tbl_path: str,
     filepath_or_url: str,
     *,
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
     **kwargs: Any,