pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__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,44 +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"""
-        pass
 
     @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,43 @@ 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
 
-    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)
+        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 drop_index(self, index_name: str, index_value_col: 'catalog.Column') -> None:
         """Drop the index on the index value column"""
@@ -72,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,18 @@
 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,28 +41,23 @@ class EmbeddingIndex(IndexBase):
     }
 
     metric: Metric
-    value_expr: exprs.FunctionCall
-    string_embed: Optional[func.Function]
-    image_embed: Optional[func.Function]
+    string_embed: func.Function | None
+    image_embed: func.Function | None
     string_embed_signature_idx: int
     image_embed_signature_idx: int
-    index_col_type: pgvector.sqlalchemy.Vector
 
     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,
     ):
         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
@@ -102,51 +99,43 @@ class EmbeddingIndex(IndexBase):
             )
 
         # 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)
 
+        self.metric = self.Metric[metric.upper()]
+
+    def create_value_expr(self, c: catalog.Column) -> exprs.Expr:
+        if not c.col_type.is_string_type() and not c.col_type.is_image_type():
+            raise excs.Error(
+                f'Embedding index requires string or image column, column {c.name!r} has type {c.col_type}'
+            )
         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')")
 
-        self.metric = self.Metric[metric.upper()]
-        self.value_expr = (
+        return (
             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 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 drop_index(self, index_name: str, index_value_col: catalog.Column) -> None:
         """Drop the index on the index value column"""
@@ -156,6 +145,7 @@ class EmbeddingIndex(IndexBase):
     def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ColumnElement:
         """Create a ColumnElement that represents '<val_column> <op> <item>'"""
         assert isinstance(item, (str, PIL.Image.Image))
+        embedding: np.ndarray
         if isinstance(item, str):
             assert self.string_embed is not None
             embedding = self.string_embed.exec([item], {})
@@ -174,7 +164,7 @@ class EmbeddingIndex(IndexBase):
     def order_by_clause(self, val_column: catalog.Column, item: Any, 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
+        embedding: np.ndarray | None = None
         if isinstance(item, str):
             assert self.string_embed is not None
             embedding = self.string_embed.exec([item], {})
@@ -199,9 +189,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:
@@ -255,7 +243,7 @@ class EmbeddingIndex(IndexBase):
         }
 
     @classmethod
-    def from_dict(cls, c: catalog.Column, d: dict) -> EmbeddingIndex:
+    def from_dict(cls, 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)
+        return cls(metric=d['metric'], string_embed=string_embed, image_embed=image_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.
@@ -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,