pixeltable 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

pixeltable/exprs/similarity_expr.py

@@ -1,4 +1,4 @@
-from typing import Optional, List
+from typing import Optional, List, Any
 
 import sqlalchemy as sql
 import PIL.Image
@@ -14,33 +14,44 @@ from .row_builder import RowBuilder
 
 class SimilarityExpr(Expr):
 
-    def __init__(self, col_ref: ColumnRef, item: Expr):
+    def __init__(self, col_ref: ColumnRef, item: Any, idx_name: Optional[str] = None):
         super().__init__(ts.FloatType())
-        self.components = [col_ref, item]
+        item_expr = Expr.from_object(item)
+        if item_expr is None or not(item_expr.col_type.is_string_type() or item_expr.col_type.is_image_type()):
+            raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not a {type(item)}')
+        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()
-        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()
         import pixeltable.index as index
-        embedding_idx_info = [info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)]
+        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}')
+            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:
-            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 = embedding_idx_info[0]
+            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()))
         idx = self.idx_info.idx
 
-        if item.col_type.is_string_type() and idx.txt_embed is None:
+        if item_expr.col_type.is_string_type() and idx.string_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:
+                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")
+        if item_expr.col_type.is_image_type() and idx.image_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')
+                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")
 
     def __str__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'

pixeltable/functions/openai.py

@@ -141,7 +141,7 @@ def chat_completions(
 
 
 @pxt.udf
-def vision(prompt: str, image: PIL.Image.Image, *, model: str = 'gpt-4-vision-preview') -> str:
+def vision(prompt: str, image: PIL.Image.Image, *, model: str) -> str:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     bytes_arr = io.BytesIO()
     image.save(bytes_arr, format='png')

pixeltable/globals.py

@@ -1,13 +1,13 @@
 import dataclasses
 import logging
-from typing import Any, Optional, Union, Type
+from typing import Any, Optional, Union
 
 import pandas as pd
 import sqlalchemy as sql
 from sqlalchemy.util.preloaded import orm
 
 import pixeltable.exceptions as excs
-from pixeltable import catalog, func
+from pixeltable import catalog, func, DataFrame
 from pixeltable.catalog import Catalog
 from pixeltable.env import Env
 from pixeltable.exprs import Predicate
@@ -78,7 +78,7 @@ def create_table(
 
 def create_view(
     path_str: str,
-    base: catalog.Table,
+    base: Union[catalog.Table, DataFrame],
     *,
     schema: Optional[dict[str, Any]] = None,
     filter: Optional[Predicate] = None,
@@ -92,7 +92,7 @@ def create_view(
 
     Args:
         path_str: Path to the view.
-        base: Table (ie, table or view or snapshot) to base the view on.
+        base: Table (i.e., table or view or snapshot) or DataFrame to base the view on.
         schema: dictionary mapping column names to column types, value expressions, or to column specifications.
         filter: Predicate to filter rows of the base table.
         is_snapshot: Whether the view is a snapshot.
@@ -122,7 +122,19 @@ def create_view(
         >>> snapshot_view = cl.create_view(
             'my_snapshot', base, schema={'col3': base.col2 + 1}, filter=base.col1 > 10, is_snapshot=True)
     """
-    assert isinstance(base, catalog.Table)
+    if isinstance(base, catalog.Table):
+        tbl_version_path = base._tbl_version_path
+    elif isinstance(base, DataFrame):
+        base._validate_mutable('create_view')
+        tbl_version_path = base.tbl
+        if base.where_clause is not None and filter is not None:
+            raise excs.Error(
+                'Cannot specify a `filter` directly if one is already declared in a `DataFrame.where` clause'
+            )
+        filter = base.where_clause
+    else:
+        raise excs.Error('`base` must be an instance of `Table` or `DataFrame`')
+    assert isinstance(base, catalog.Table) or isinstance(base, DataFrame)
     path = catalog.Path(path_str)
     try:
         Catalog.get().paths.check_is_valid(path, expected=None)
@@ -139,10 +151,11 @@ def create_view(
         iterator_class, iterator_args = None, None
     else:
         iterator_class, iterator_args = iterator
+
     view = catalog.View.create(
         dir._id,
         path.name,
-        base=base,
+        base=tbl_version_path,
         schema=schema,
         predicate=filter,
         is_snapshot=is_snapshot,
@@ -221,7 +234,7 @@ def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> N
 
     Args:
         path: Path to the table.
-        force: Whether to drop the table even if it has unsaved changes.
+        force: If `True`, will also drop all views or sub-views of this table.
         ignore_errors: Whether to ignore errors if the table does not exist.
 
     Raises:
@@ -230,21 +243,27 @@ def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> N
     Examples:
         >>> cl.drop_table('my_table')
     """
+    cat = Catalog.get()
     path_obj = catalog.Path(path)
     try:
-        Catalog.get().paths.check_is_valid(path_obj, expected=catalog.Table)
+        cat.paths.check_is_valid(path_obj, expected=catalog.Table)
     except Exception as e:
-        if ignore_errors:
+        if ignore_errors or force:
             _logger.info(f'Skipped table `{path}` (does not exist).')
             return
         else:
             raise e
-    tbl = Catalog.get().paths[path_obj]
-    if len(Catalog.get().tbl_dependents[tbl._id]) > 0:
-        dependent_paths = [get_path(dep) for dep in Catalog.get().tbl_dependents[tbl._id]]
-        raise excs.Error(f'Table {path} has dependents: {", ".join(dependent_paths)}')
+    tbl = cat.paths[path_obj]
+    assert isinstance(tbl, catalog.Table)
+    if len(cat.tbl_dependents[tbl._id]) > 0:
+        dependent_paths = [dep.path for dep in cat.tbl_dependents[tbl._id]]
+        if force:
+            for dependent_path in dependent_paths:
+                drop_table(dependent_path, force=True)
+        else:
+            raise excs.Error(f'Table {path} has dependents: {", ".join(dependent_paths)}')
     tbl._drop()
-    del Catalog.get().paths[path_obj]
+    del cat.paths[path_obj]
     _logger.info(f'Dropped table `{path}`.')
 
 
@@ -278,7 +297,7 @@ def list_tables(dir_path: str = '', recursive: bool = True) -> list[str]:
     return [str(p) for p in Catalog.get().paths.get_children(path, child_type=catalog.Table, recursive=recursive)]
 
 
-def create_dir(path_str: str, ignore_errors: bool = False) -> None:
+def create_dir(path_str: str, ignore_errors: bool = False) -> catalog.Dir:
     """Create a directory.
 
     Args:
@@ -306,10 +325,12 @@ def create_dir(path_str: str, ignore_errors: bool = False) -> None:
             session.add(dir_record)
             session.flush()
             assert dir_record.id is not None
-            Catalog.get().paths[path] = catalog.Dir(dir_record.id, parent._id, path.name)
+            dir = catalog.Dir(dir_record.id, parent._id, path.name)
+            Catalog.get().paths[path] = dir
             session.commit()
             _logger.info(f'Created directory `{path_str}`.')
             print(f'Created directory `{path_str}`.')
+            return dir
     except excs.Error as e:
         if ignore_errors:
             return
@@ -317,7 +338,7 @@ def create_dir(path_str: str, ignore_errors: bool = False) -> None:
             raise e
 
 
-def rm_dir(path_str: str) -> None:
+def drop_dir(path_str: str, force: bool = False, ignore_errors: bool = False) -> None:
     """Remove a directory.
 
     Args:
@@ -327,31 +348,49 @@ def rm_dir(path_str: str) -> None:
         Error: If the path does not exist or does not designate a directory or if the directory is not empty.
 
     Examples:
-        >>> cl.rm_dir('my_dir')
+        >>> cl.drop_dir('my_dir')
 
         Remove a subdirectory:
 
-        >>> cl.rm_dir('my_dir.sub_dir')
+        >>> cl.drop_dir('my_dir.sub_dir')
     """
+    cat = Catalog.get()
     path = catalog.Path(path_str)
-    Catalog.get().paths.check_is_valid(path, expected=catalog.Dir)
 
-    # make sure it's empty
-    if len(Catalog.get().paths.get_children(path, child_type=None, recursive=True)) > 0:
-        raise excs.Error(f'Directory {path_str} is not empty')
-    # TODO: figure out how to make force=True work in the presence of snapshots
-    #        # delete tables
-    #        for tbl_path in self.paths.get_children(path, child_type=MutableTable, recursive=True):
-    #            self.drop_table(str(tbl_path), force=True)
-    #        # rm subdirs
-    #        for dir_path in self.paths.get_children(path, child_type=Dir, recursive=False):
-    #            self.rm_dir(str(dir_path), force=True)
+    try:
+        cat.paths.check_is_valid(path, expected=catalog.Dir)
+    except Exception as e:
+        if ignore_errors or force:
+            _logger.info(f'Skipped directory `{path}` (does not exist).')
+            return
+        else:
+            raise e
+
+    children = cat.paths.get_children(path, child_type=None, recursive=True)
+
+    if len(children) > 0 and not force:
+        raise excs.Error(f'Directory `{path_str}` is not empty.')
+
+    for child in children:
+        assert isinstance(child, catalog.Path)
+        # We need to check that the child is still in `cat.paths`, since it is possible it was
+        # already deleted as a dependent of a preceding child in the iteration.
+        try:
+            obj = cat.paths[child]
+        except excs.Error:
+            continue
+        if isinstance(obj, catalog.Dir):
+            drop_dir(str(child), force=True)
+        else:
+            assert isinstance(obj, catalog.Table)
+            assert not obj._is_dropped  # else it should have been removed from `cat.paths` already
+            drop_table(str(child), force=True)
 
     with Env.get().engine.begin() as conn:
         dir = Catalog.get().paths[path]
         conn.execute(sql.delete(schema.Dir.__table__).where(schema.Dir.id == dir._id))
     del Catalog.get().paths[path]
-    _logger.info(f'Removed directory {path_str}')
+    _logger.info(f'Removed directory `{path_str}`.')
 
 
 def list_dirs(path_str: str = '', recursive: bool = True) -> list[str]:
@@ -403,28 +442,6 @@ def list_functions() -> pd.DataFrame:
     return pd_df.hide(axis='index')
 
 
-def get_path(schema_obj: catalog.SchemaObject) -> str:
-    """Returns the path to a SchemaObject.
-
-    Args:
-        schema_obj: SchemaObject to get the path for.
-
-    Returns:
-        Path to the SchemaObject.
-    """
-    path_elements: list[str] = []
-    dir_id = schema_obj._dir_id
-    while dir_id is not None:
-        dir = Catalog.get().paths.get_schema_obj(dir_id)
-        if dir._dir_id is None:
-            # this is the root dir with name '', which we don't want to include in the path
-            break
-        path_elements.insert(0, dir._name)
-        dir_id = dir._dir_id
-    path_elements.append(schema_obj._name)
-    return '.'.join(path_elements)
-
-
 def configure_logging(
     *,
     to_stdout: Optional[bool] = None,

pixeltable/index/embedding_index.py

@@ -24,6 +24,7 @@ class EmbeddingIndex(IndexBase):
     - similarity_clause() converts those metrics back to their original form; it is used in expressions outside
       the Order By clause
     - order_by_clause() is used exclusively in the ORDER BY clause
+    - embedding function parameters are named '<type-name>_embed', where type-name is ColumnType.Type.name
     """
 
     class Metric(enum.Enum):
@@ -38,30 +39,30 @@ class EmbeddingIndex(IndexBase):
     }
 
     def __init__(
-            self, c: catalog.Column, metric: str, text_embed: Optional[func.Function] = None,
-            img_embed: Optional[func.Function] = None):
+            self, c: catalog.Column, metric: str, string_embed: Optional[func.Function] = None,
+            image_embed: Optional[func.Function] = None):
         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(f'Embedding index requires string or image column')
-        if c.col_type.is_string_type() and text_embed is None:
-                raise excs.Error(f'Text embedding function is required for column {c.name} (parameter `txt_embed`)')
-        if c.col_type.is_image_type() and img_embed is None:
-            raise excs.Error(f'Image embedding function is required for column {c.name} (parameter `img_embed`)')
-        if text_embed is not None:
+        if c.col_type.is_string_type() and 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 image_embed is None:
+            raise excs.Error(f"Image embedding function is required for column {c.name} (parameter 'image_embed')")
+        if string_embed is not None:
             # verify signature
-            self._validate_embedding_fn(text_embed, 'txt_embed', ts.ColumnType.Type.STRING)
-        if img_embed is not None:
+            self._validate_embedding_fn(string_embed, 'string_embed', ts.ColumnType.Type.STRING)
+        if image_embed is not None:
             # verify signature
-            self._validate_embedding_fn(img_embed, 'img_embed', ts.ColumnType.Type.IMAGE)
+            self._validate_embedding_fn(image_embed, 'image_embed', ts.ColumnType.Type.IMAGE)
 
         self.metric = self.Metric[metric.upper()]
         from pixeltable.exprs import ColumnRef
-        self.value_expr = text_embed(ColumnRef(c)) if c.col_type.is_string_type() else img_embed(ColumnRef(c))
+        self.value_expr = string_embed(ColumnRef(c)) if c.col_type.is_string_type() else image_embed(ColumnRef(c))
         assert self.value_expr.col_type.is_array_type()
-        self.txt_embed = text_embed
-        self.img_embed = img_embed
+        self.string_embed = string_embed
+        self.image_embed = image_embed
         vector_size = self.value_expr.col_type.shape[0]
         assert vector_size is not None
         self.index_col_type = pgvector.sqlalchemy.Vector(vector_size)
@@ -88,14 +89,14 @@ class EmbeddingIndex(IndexBase):
         idx.create(bind=conn)
 
     def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ClauseElement:
-        """Create a ClauseElement to that represents '<val_column> <op> <item>'"""
+        """Create a ClauseElement that represents '<val_column> <op> <item>'"""
         assert isinstance(item, (str, PIL.Image.Image))
         if isinstance(item, str):
-            assert self.txt_embed is not None
-            embedding = self.txt_embed.exec(item)
+            assert self.string_embed is not None
+            embedding = self.string_embed.exec(item)
         if isinstance(item, PIL.Image.Image):
-            assert self.img_embed is not None
-            embedding = self.img_embed.exec(item)
+            assert self.image_embed is not None
+            embedding = self.image_embed.exec(item)
 
         if self.metric == self.Metric.COSINE:
             return val_column.sa_col.cosine_distance(embedding) * -1 + 1
@@ -110,11 +111,11 @@ class EmbeddingIndex(IndexBase):
         assert isinstance(item, (str, PIL.Image.Image))
         embedding: Optional[np.ndarray] = None
         if isinstance(item, str):
-            assert self.txt_embed is not None
-            embedding = self.txt_embed.exec(item)
+            assert self.string_embed is not None
+            embedding = self.string_embed.exec(item)
         if isinstance(item, PIL.Image.Image):
-            assert self.img_embed is not None
-            embedding = self.img_embed.exec(item)
+            assert self.image_embed is not None
+            embedding = self.image_embed.exec(item)
         assert embedding is not None
 
         if self.metric == self.Metric.COSINE:
@@ -160,12 +161,12 @@ class EmbeddingIndex(IndexBase):
     def as_dict(self) -> dict:
         return {
             'metric': self.metric.name.lower(),
-            'txt_embed': None if self.txt_embed is None else self.txt_embed.as_dict(),
-            'img_embed': None if self.img_embed is None else self.img_embed.as_dict()
+            '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()
         }
 
     @classmethod
     def from_dict(cls, c: catalog.Column, d: dict) -> EmbeddingIndex:
-        txt_embed = func.Function.from_dict(d['txt_embed']) if d['txt_embed'] is not None else None
-        img_embed = func.Function.from_dict(d['img_embed']) if d['img_embed'] is not None else None
-        return cls(c, metric=d['metric'], text_embed=txt_embed, img_embed=img_embed)
+        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)

pixeltable/io/external_store.py

@@ -222,12 +222,12 @@ class Project(ExternalStore, abc.ABC):
             if t_col not in t_cols:
                 if is_user_specified_col_mapping:
                     raise excs.Error(
-                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table.get_name()}` '
+                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table.name}` '
                         'contains no such column.'
                     )
                 else:
                     raise excs.Error(
-                        f'Column `{t_col}` does not exist in Table `{table.get_name()}`. Either add a column `{t_col}`, '
+                        f'Column `{t_col}` does not exist in Table `{table.name}`. Either add a column `{t_col}`, '
                         f'or specify a `col_mapping` to associate a different column with the external field `{ext_col}`.'
                     )
             if ext_col not in export_cols and ext_col not in import_cols:

pixeltable/io/globals.py

@@ -50,7 +50,7 @@ def create_label_studio_project(
             `ls_project_0`, `ls_project_1`, etc.
         title: An optional title for the Label Studio project. This is the title that annotators
             will see inside Label Studio. Unlike `name`, it does not need to be an identifier and
-            does not need to be unique. If not specified, the table name `t.get_name()` will be used.
+            does not need to be unique. If not specified, the table name `t.name` will be used.
         media_import_method: The method to use when transferring media files to Label Studio:
             - `post`: Media will be sent to Label Studio via HTTP post. This should generally only be used for
                 prototyping; due to restrictions in Label Studio, it can only be used with projects that have

pixeltable/io/label_studio.py

@@ -95,7 +95,7 @@ class LabelStudioProject(Project):
         return {ANNOTATIONS_COLUMN: pxt.JsonType(nullable=True)}
 
     def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
-        _logger.info(f'Syncing Label Studio project "{self.project_title}" with table `{t.get_name()}`'
+        _logger.info(f'Syncing Label Studio project "{self.project_title}" with table `{t.name}`'
                      f' (export: {export_data}, import: {import_data}).')
         # Collect all existing tasks into a dict with entries `rowid: task`
         tasks = {tuple(task['meta']['rowid']): task for task in self.__fetch_all_tasks()}
@@ -386,7 +386,7 @@ class LabelStudioProject(Project):
         updates = [{'_rowid': rowid, local_annotations_col.name: ann} for rowid, ann in annotations.items()]
         if len(updates) > 0:
             _logger.info(
-                f'Updating table `{t.get_name()}`, column `{local_annotations_col.name}` with {len(updates)} total annotations.'
+                f'Updating table `{t.name}`, column `{local_annotations_col.name}` with {len(updates)} total annotations.'
             )
             # batch_update currently doesn't propagate from views to base tables. As a workaround, we call
             # batch_update on the actual ancestor table that holds the annotations column.
@@ -554,7 +554,7 @@ class LabelStudioProject(Project):
 
         if title is None:
             # `title` defaults to table name
-            title = t.get_name()
+            title = t.name
 
         # Create a column to hold the annotations, if one does not yet exist
         if col_mapping is None or ANNOTATIONS_COLUMN in col_mapping.values():

pixeltable/metadata/__init__.py

@@ -10,7 +10,7 @@ import sqlalchemy.orm as orm
 from .schema import SystemInfo, SystemInfoMd
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 17
+VERSION = 18
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:

pixeltable/metadata/converters/convert_17.py

@@ -0,0 +1,26 @@
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=17)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(
+        engine,
+        table_md_updater=__update_table_md
+    )
+
+
+def __update_table_md(table_md: dict) -> None:
+    # key changes in IndexMd.init_args: img_embed -> image_embed, txt_embed -> string_embed
+    if len(table_md['index_md']) == 0:
+        return
+    for idx_md in table_md['index_md'].values():
+        if not idx_md['class_fqn'].endswith('.EmbeddingIndex'):
+            continue
+        init_dict = idx_md['init_args']
+        init_dict['image_embed'] = init_dict['img_embed']
+        del init_dict['img_embed']
+        init_dict['string_embed'] = init_dict['txt_embed']
+        del init_dict['txt_embed']

pixeltable/tool/create_test_db_dump.py

@@ -253,7 +253,7 @@ class Dumper:
         add_column('c6_to_string', t.c6.apply(json.dumps))
         add_column('c6_back_to_json', t[f'{col_prefix}_c6_to_string'].apply(json.loads))
 
-        t.add_embedding_index(f'{col_prefix}_function_call', text_embed=embed_udf.clip_text_embed)
+        t.add_embedding_index(f'{col_prefix}_function_call', string_embed=embed_udf.clip_text_embed)
 
         # query()
         @t.query