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

pixeltable/globals.py

@@ -0,0 +1,425 @@
+import dataclasses
+import logging
+from typing import Any, Optional, Union, Type
+
+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.catalog import Catalog
+from pixeltable.env import Env
+from pixeltable.exprs import Predicate
+from pixeltable.iterators import ComponentIterator
+from pixeltable.metadata import schema
+
+_logger = logging.getLogger('pixeltable')
+
+
+def init() -> None:
+    """Initializes the Pixeltable environment."""
+    _ = Catalog.get()
+
+
+def create_table(
+    path_str: str,
+    schema: dict[str, Any],
+    *,
+    primary_key: Optional[Union[str, list[str]]] = None,
+    num_retained_versions: int = 10,
+    comment: str = '',
+) -> catalog.InsertableTable:
+    """Create a new `InsertableTable`.
+
+    Args:
+        path_str: Path to the table.
+        schema: dictionary mapping column names to column types, value expressions, or to column specifications.
+        num_retained_versions: Number of versions of the table to retain.
+
+    Returns:
+        The newly created table.
+
+    Raises:
+        Error: if the path already exists or is invalid.
+
+    Examples:
+        Create a table with an int and a string column:
+
+        >>> table = cl.create_table('my_table', schema={'col1': IntType(), 'col2': StringType()})
+    """
+    path = catalog.Path(path_str)
+    Catalog.get().paths.check_is_valid(path, expected=None)
+    dir = Catalog.get().paths[path.parent]
+
+    if len(schema) == 0:
+        raise excs.Error(f'Table schema is empty: `{path_str}`')
+
+    if primary_key is None:
+        primary_key = []
+    elif isinstance(primary_key, str):
+        primary_key = [primary_key]
+    else:
+        if not isinstance(primary_key, list) or not all(isinstance(pk, str) for pk in primary_key):
+            raise excs.Error('primary_key must be a single column name or a list of column names')
+
+    tbl = catalog.InsertableTable.create(
+        dir._id,
+        path.name,
+        schema,
+        primary_key=primary_key,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+    )
+    Catalog.get().paths[path] = tbl
+    _logger.info(f'Created table `{path_str}`.')
+    return tbl
+
+
+def create_view(
+    path_str: str,
+    base: catalog.Table,
+    *,
+    schema: Optional[dict[str, Any]] = None,
+    filter: Optional[Predicate] = None,
+    is_snapshot: bool = False,
+    iterator: Optional[tuple[type[ComponentIterator], dict[str, Any]]] = None,
+    num_retained_versions: int = 10,
+    comment: str = '',
+    ignore_errors: bool = False,
+) -> catalog.View:
+    """Create a new `View`.
+
+    Args:
+        path_str: Path to the view.
+        base: Table (ie, table or view or snapshot) 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.
+        iterator_class: Class of the iterator to use for the view.
+        iterator_args: Arguments to pass to the iterator class.
+        num_retained_versions: Number of versions of the view to retain.
+        ignore_errors: if True, fail silently if the path already exists or is invalid.
+
+    Returns:
+        The newly created view.
+
+    Raises:
+        Error: if the path already exists or is invalid.
+
+    Examples:
+        Create a view with an additional int and a string column and a filter:
+
+        >>> view = cl.create_view(
+            'my_view', base, schema={'col3': IntType(), 'col4': StringType()}, filter=base.col1 > 10)
+
+        Create a table snapshot:
+
+        >>> snapshot_view = cl.create_view('my_snapshot_view', base, is_snapshot=True)
+
+        Create an immutable view with additional computed columns and a filter:
+
+        >>> 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)
+    path = catalog.Path(path_str)
+    try:
+        Catalog.get().paths.check_is_valid(path, expected=None)
+    except Exception as e:
+        if ignore_errors:
+            return
+        else:
+            raise e
+    dir = Catalog.get().paths[path.parent]
+
+    if schema is None:
+        schema = {}
+    if iterator is None:
+        iterator_class, iterator_args = None, None
+    else:
+        iterator_class, iterator_args = iterator
+    view = catalog.View.create(
+        dir._id,
+        path.name,
+        base=base,
+        schema=schema,
+        predicate=filter,
+        is_snapshot=is_snapshot,
+        iterator_cls=iterator_class,
+        iterator_args=iterator_args,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+    )
+    Catalog.get().paths[path] = view
+    _logger.info(f'Created view `{path_str}`.')
+    return view
+
+
+def get_table(path: str) -> catalog.Table:
+    """Get a handle to a table (including views and snapshots).
+
+    Args:
+        path: Path to the table.
+
+    Returns:
+        A `InsertableTable` or `View` object.
+
+    Raises:
+        Error: If the path does not exist or does not designate a table.
+
+    Examples:
+        Get handle for a table in the top-level directory:
+
+        >>> table = cl.get_table('my_table')
+
+        For a table in a subdirectory:
+
+        >>> table = cl.get_table('subdir.my_table')
+
+        For a snapshot in the top-level directory:
+
+        >>> table = cl.get_table('my_snapshot')
+    """
+    p = catalog.Path(path)
+    Catalog.get().paths.check_is_valid(p, expected=catalog.Table)
+    obj = Catalog.get().paths[p]
+    return obj
+
+
+def move(path: str, new_path: str) -> None:
+    """Move a schema object to a new directory and/or rename a schema object.
+
+    Args:
+        path: absolute path to the existing schema object.
+        new_path: absolute new path for the schema object.
+
+    Raises:
+        Error: If path does not exist or new_path already exists.
+
+    Examples:
+        Move a table to a different directory:
+
+        >>>> cl.move('dir1.my_table', 'dir2.my_table')
+
+        Rename a table:
+
+        >>>> cl.move('dir1.my_table', 'dir1.new_name')
+    """
+    p = catalog.Path(path)
+    Catalog.get().paths.check_is_valid(p, expected=catalog.SchemaObject)
+    new_p = catalog.Path(new_path)
+    Catalog.get().paths.check_is_valid(new_p, expected=None)
+    obj = Catalog.get().paths[p]
+    Catalog.get().paths.move(p, new_p)
+    new_dir = Catalog.get().paths[new_p.parent]
+    obj.move(new_p.name, new_dir._id)
+
+
+def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> None:
+    """Drop a table.
+
+    Args:
+        path: Path to the table.
+        force: Whether to drop the table even if it has unsaved changes.
+        ignore_errors: Whether to ignore errors if the table does not exist.
+
+    Raises:
+        Error: If the path does not exist or does not designate a table and ignore_errors is False.
+
+    Examples:
+        >>> cl.drop_table('my_table')
+    """
+    path_obj = catalog.Path(path)
+    try:
+        Catalog.get().paths.check_is_valid(path_obj, expected=catalog.Table)
+    except Exception as e:
+        if ignore_errors:
+            _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._drop()
+    del Catalog.get().paths[path_obj]
+    _logger.info(f'Dropped table `{path}`.')
+
+
+def list_tables(dir_path: str = '', recursive: bool = True) -> list[str]:
+    """List the tables in a directory.
+
+    Args:
+        dir_path: Path to the directory. Defaults to the root directory.
+        recursive: Whether to list tables in subdirectories as well.
+
+    Returns:
+        A list of table paths.
+
+    Raises:
+        Error: If the path does not exist or does not designate a directory.
+
+    Examples:
+        List tables in top-level directory:
+
+        >>> cl.list_tables()
+        ['my_table', ...]
+
+        List tables in 'dir1':
+
+        >>> cl.list_tables('dir1')
+        [...]
+    """
+    assert dir_path is not None
+    path = catalog.Path(dir_path, empty_is_valid=True)
+    Catalog.get().paths.check_is_valid(path, expected=catalog.Dir)
+    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:
+    """Create a directory.
+
+    Args:
+        path_str: Path to the directory.
+        ignore_errors: if True, silently returns on error
+
+    Raises:
+        Error: If the path already exists or the parent is not a directory.
+
+    Examples:
+        >>> cl.create_dir('my_dir')
+
+        Create a subdirectory:
+
+        >>> cl.create_dir('my_dir.sub_dir')
+    """
+    try:
+        path = catalog.Path(path_str)
+        Catalog.get().paths.check_is_valid(path, expected=None)
+        parent = Catalog.get().paths[path.parent]
+        assert parent is not None
+        with orm.Session(Env.get().engine, future=True) as session:
+            dir_md = schema.DirMd(name=path.name)
+            dir_record = schema.Dir(parent_id=parent._id, md=dataclasses.asdict(dir_md))
+            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)
+            session.commit()
+            _logger.info(f'Created directory `{path_str}`.')
+            print(f'Created directory `{path_str}`.')
+    except excs.Error as e:
+        if ignore_errors:
+            return
+        else:
+            raise e
+
+
+def rm_dir(path_str: str) -> None:
+    """Remove a directory.
+
+    Args:
+        path_str: Path to the directory.
+
+    Raises:
+        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')
+
+        Remove a subdirectory:
+
+        >>> cl.rm_dir('my_dir.sub_dir')
+    """
+    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)
+
+    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}')
+
+
+def list_dirs(path_str: str = '', recursive: bool = True) -> list[str]:
+    """List the directories in a directory.
+
+    Args:
+        path_str: Path to the directory.
+        recursive: Whether to list subdirectories recursively.
+
+    Returns:
+        List of directory paths.
+
+    Raises:
+        Error: If the path does not exist or does not designate a directory.
+
+    Examples:
+        >>> cl.list_dirs('my_dir', recursive=True)
+        ['my_dir', 'my_dir.sub_dir1']
+    """
+    path = catalog.Path(path_str, empty_is_valid=True)
+    Catalog.get().paths.check_is_valid(path, expected=catalog.Dir)
+    return [str(p) for p in Catalog.get().paths.get_children(path, child_type=catalog.Dir, recursive=recursive)]
+
+
+def list_functions() -> pd.DataFrame:
+    """Returns information about all registered functions.
+
+    Returns:
+        Pandas DataFrame with columns 'Path', 'Name', 'Parameters', 'Return Type', 'Is Agg', 'Library'
+    """
+    functions = func.FunctionRegistry.get().list_functions()
+    paths = ['.'.join(f.self_path.split('.')[:-1]) for f in functions]
+    names = [f.name for f in functions]
+    params = [
+        ', '.join([param_name + ': ' + str(param_type) for param_name, param_type in f.signature.parameters.items()])
+        for f in functions
+    ]
+    pd_df = pd.DataFrame(
+        {
+            'Path': paths,
+            'Function Name': names,
+            'Parameters': params,
+            'Return Type': [str(f.signature.get_return_type()) for f in functions],
+        }
+    )
+    pd_df = pd_df.style.set_properties(**{'text-align': 'left'}).set_table_styles(
+        [dict(selector='th', props=[('text-align', 'center')])]
+    )  # center-align headings
+    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)

pixeltable/index/base.py

@@ -13,8 +13,10 @@ class IndexBase(abc.ABC):
     Internal interface used by the catalog and runtime system to interact with indices:
     - types and expressions needed to create and populate the index value column
     - creating/dropping the index
-    - TODO: translating queries into sqlalchemy predicates
+    This doesn't cover querying the index, which is dependent on the index semantics and handled by
+    the specific subclass.
     """
+
     @abc.abstractmethod
     def __init__(self, c: catalog.Column, **kwargs: Any):
         pass

pixeltable/index/embedding_index.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Any
+import enum
 
+import PIL.Image
+import numpy as np
 import pgvector.sqlalchemy
+import PIL.Image
 import sqlalchemy as sql
 
 import pixeltable.catalog as catalog
@@ -14,15 +18,31 @@ from .base import IndexBase
 
 class EmbeddingIndex(IndexBase):
     """
-    Internal interface used by the catalog and runtime system to interact with (embedding) indices:
-    - types and expressions needed to create and populate the index value column
-    - creating/dropping the index
-    - translating 'matches' queries into sqlalchemy predicates
+    Interface to the pgvector access method in Postgres.
+    - pgvector converts the cosine metric to '1 - metric' and the inner product metric to '-metric', in order to
+      satisfy the Postgres requirement that an index scan requires an ORDER BY ... ASC clause
+    - 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
     """
 
+    class Metric(enum.Enum):
+        COSINE = 1
+        IP = 2
+        L2 = 3
+
+    PGVECTOR_OPS = {
+        Metric.COSINE: 'vector_cosine_ops',
+        Metric.IP: 'vector_ip_ops',
+        Metric.L2: 'vector_l2_ops'
+    }
+
     def __init__(
-            self, c: catalog.Column, text_embed: Optional[func.Function] = None,
+            self, c: catalog.Column, metric: str, text_embed: Optional[func.Function] = None,
             img_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:
@@ -36,6 +56,7 @@ class EmbeddingIndex(IndexBase):
             # verify signature
             self._validate_embedding_fn(img_embed, 'img_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))
         assert self.value_expr.col_type.is_array_type()
@@ -59,10 +80,51 @@ class EmbeddingIndex(IndexBase):
             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: 'vector_cosine_ops'}
+            postgresql_ops={index_value_col.sa_col.name: self.PGVECTOR_OPS[self.metric]}
         )
         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>'"""
+        assert isinstance(item, (str, PIL.Image.Image))
+        if isinstance(item, str):
+            assert self.txt_embed is not None
+            embedding = self.txt_embed.exec(item)
+        if isinstance(item, PIL.Image.Image):
+            assert self.img_embed is not None
+            embedding = self.img_embed.exec(item)
+
+        if self.metric == self.Metric.COSINE:
+            return val_column.sa_col.cosine_distance(embedding) * -1 + 1
+        elif self.metric == self.Metric.IP:
+            return val_column.sa_col.max_inner_product(embedding) * -1
+        else:
+            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.ClauseElement:
+        """Create a ClauseElement 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.txt_embed is not None
+            embedding = self.txt_embed.exec(item)
+        if isinstance(item, PIL.Image.Image):
+            assert self.img_embed is not None
+            embedding = self.img_embed.exec(item)
+        assert embedding is not None
+
+        if self.metric == self.Metric.COSINE:
+            result = val_column.sa_col.cosine_distance(embedding)
+            result = result.desc() if is_asc else result
+        elif self.metric == self.Metric.IP:
+            result = val_column.sa_col.max_inner_product(embedding)
+            result = result.desc() if is_asc else result
+        else:
+            assert self.metric == self.Metric.L2
+            result = val_column.sa_col.l2_distance(embedding)
+        return result
+
     @classmethod
     def display_name(cls) -> str:
         return 'embedding'
@@ -72,18 +134,29 @@ class EmbeddingIndex(IndexBase):
         """Validate the signature"""
         assert isinstance(embed_fn, func.Function)
         sig = embed_fn.signature
-        if not sig.return_type.is_array_type():
-            raise excs.Error(f'{name} must return an array, but returns {sig.return_type}')
-        else:
-            shape = sig.return_type.shape
-            if len(shape) != 1 or shape[0] == None:
-                raise excs.Error(f'{name} must return a 1D array of a specific length, but returns {sig.return_type}')
         if len(sig.parameters) != 1 or sig.parameters_by_pos[0].col_type.type_enum != expected_type:
             raise excs.Error(
                 f'{name} must take a single {expected_type.name.lower()} parameter, but has signature {sig}')
 
+        # validate return type
+        param_name = sig.parameters_by_pos[0].name
+        if expected_type == ts.ColumnType.Type.STRING:
+            return_type = embed_fn.call_return_type({param_name: 'dummy'})
+        else:
+            assert expected_type == ts.ColumnType.Type.IMAGE
+            img = PIL.Image.new('RGB', (512, 512))
+            return_type = embed_fn.call_return_type({param_name: img})
+        assert return_type is not None
+        if not return_type.is_array_type():
+            raise excs.Error(f'{name} must return an array, but returns {return_type}')
+        else:
+            shape = return_type.shape
+            if len(shape) != 1 or shape[0] == None:
+                raise excs.Error(f'{name} must return a 1D array of a specific length, but returns {return_type}')
+
     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()
         }
@@ -92,4 +165,4 @@ class EmbeddingIndex(IndexBase):
     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, text_embed=txt_embed, img_embed=img_embed)
+        return cls(c, metric=d['metric'], text_embed=txt_embed, img_embed=img_embed)

pixeltable/io/__init__.py

@@ -0,0 +1,3 @@
+from .hf_datasets import import_huggingface_dataset
+from .pandas import import_csv, import_excel, import_pandas
+from .parquet import import_parquet