pixeltable 0.4.6__py3-none-any.whl → 0.4.7__py3-none-any.whl

pixeltable/__init__.py

@@ -1,11 +1,12 @@
 # ruff: noqa: F401
 
 from .__version__ import __version__, __version_tuple__
-from .catalog import Column, InsertableTable, Table, UpdateStatus, View
+from .catalog import Column, ColumnMetadata, IndexMetadata, InsertableTable, Table, TableMetadata, UpdateStatus, View
 from .dataframe import DataFrame
 from .exceptions import Error, ExprEvalError, PixeltableWarning
-from .func import Aggregator, Function, expr_udf, mcp_udfs, query, retrieval_udf, uda, udf
+from .func import Aggregator, Function, Tool, ToolChoice, Tools, expr_udf, mcp_udfs, query, retrieval_udf, uda, udf
 from .globals import (
+    DirContents,
     array,
     configure_logging,
     create_dir,
@@ -15,6 +16,7 @@ from .globals import (
     create_view,
     drop_dir,
     drop_table,
+    get_dir_contents,
     get_table,
     init,
     list_dirs,

pixeltable/catalog/__init__.py

@@ -8,7 +8,7 @@ from .insertable_table import InsertableTable
 from .named_function import NamedFunction
 from .path import Path
 from .schema_object import SchemaObject
-from .table import Table
+from .table import ColumnMetadata, IndexMetadata, Table, TableMetadata
 from .table_version import TableVersion
 from .table_version_handle import ColumnHandle, TableVersionHandle
 from .table_version_path import TableVersionPath

pixeltable/catalog/catalog.py

@@ -103,7 +103,7 @@ def retry_loop(
                 except PendingTableOpsError as e:
                     Env.get().console_logger.debug(f'retry_loop(): finalizing pending ops for {e.tbl_id}')
                     Catalog.get()._finalize_pending_ops(e.tbl_id)
-                except sql.exc.DBAPIError as e:
+                except (sql.exc.DBAPIError, sql.exc.OperationalError) as e:
                     # TODO: what other exceptions should we be looking for?
                     if isinstance(e.orig, (psycopg.errors.SerializationFailure, psycopg.errors.LockNotAvailable)):
                         if num_retries < _MAX_RETRIES or _MAX_RETRIES == -1:
@@ -356,7 +356,7 @@ class Catalog:
                             # raise to abort the transaction
                             raise
 
-                        except sql.exc.DBAPIError as e:
+                        except (sql.exc.DBAPIError, sql.exc.OperationalError) as e:
                             has_exc = True
                             if isinstance(
                                 e.orig, (psycopg.errors.SerializationFailure, psycopg.errors.LockNotAvailable)
@@ -380,7 +380,7 @@ class Catalog:
                     # we got this exception after getting the initial table locks and therefore need to abort
                     raise
 
-            except sql.exc.DBAPIError as e:
+            except (sql.exc.DBAPIError, sql.exc.OperationalError) as e:
                 has_exc = True
                 # we got some db error during the actual operation (not just while trying to get locks on the metadata
                 # records): we convert these into Errors, if asked to do so, and abort

pixeltable/catalog/column.py

@@ -10,6 +10,7 @@ import sqlalchemy as sql
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable import exprs
+from pixeltable.metadata import schema
 
 from .globals import MediaValidation, is_valid_identifier
 
@@ -126,6 +127,54 @@ class Column:
         # computed cols also have storage columns for the exception string and type
         self.sa_cellmd_col = None
 
+    def to_md(self, pos: Optional[int] = None) -> tuple[schema.ColumnMd, Optional[schema.SchemaColumn]]:
+        """Returns the Column and optional SchemaColumn metadata for this Column."""
+        assert self.is_pk is not None
+        col_md = schema.ColumnMd(
+            id=self.id,
+            col_type=self.col_type.as_dict(),
+            is_pk=self.is_pk,
+            schema_version_add=self.schema_version_add,
+            schema_version_drop=self.schema_version_drop,
+            value_expr=self.value_expr.as_dict() if self.value_expr is not None else None,
+            stored=self.stored,
+        )
+        if pos is None:
+            return col_md, None
+        assert self.name is not None, 'Column name must be set for user-facing columns'
+        sch_md = schema.SchemaColumn(
+            name=self.name,
+            pos=pos,
+            media_validation=self._media_validation.name.lower() if self._media_validation is not None else None,
+        )
+        return col_md, sch_md
+
+    @classmethod
+    def from_md(
+        cls, col_md: schema.ColumnMd, tbl: TableVersion, schema_col_md: Optional[schema.SchemaColumn]
+    ) -> Column:
+        """Create a Column from a ColumnMd."""
+        assert col_md.id is not None
+        col_name = schema_col_md.name if schema_col_md is not None else None
+        media_val = (
+            MediaValidation[schema_col_md.media_validation.upper()]
+            if schema_col_md is not None and schema_col_md.media_validation is not None
+            else None
+        )
+        col = cls(
+            col_id=col_md.id,
+            name=col_name,
+            col_type=ts.ColumnType.from_dict(col_md.col_type),
+            is_pk=col_md.is_pk,
+            stored=col_md.stored,
+            media_validation=media_val,
+            schema_version_add=col_md.schema_version_add,
+            schema_version_drop=col_md.schema_version_drop,
+            value_expr_dict=col_md.value_expr,
+            tbl=tbl,
+        )
+        return col
+
     def init_value_expr(self) -> None:
         from pixeltable import exprs
 

pixeltable/catalog/insertable_table.py

@@ -105,13 +105,6 @@ class InsertableTable(Table):
         Env.get().console_logger.info(f'Created table {name!r}.')
         return tbl
 
-    def _get_metadata(self) -> dict[str, Any]:
-        md = super()._get_metadata()
-        md['base'] = None
-        md['is_view'] = False
-        md['is_snapshot'] = False
-        return md
-
     @overload
     def insert(
         self,

pixeltable/catalog/schema_object.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
 from uuid import UUID
 
 if TYPE_CHECKING:
@@ -41,19 +41,6 @@ class SchemaObject:
             path = Catalog.get().get_dir_path(self._dir_id)
             return str(path.append(self._name))
 
-    def get_metadata(self) -> dict[str, Any]:
-        """Returns metadata associated with this schema object."""
-        from pixeltable.catalog import retry_loop
-
-        @retry_loop(for_write=False)
-        def op() -> dict[str, Any]:
-            return self._get_metadata()
-
-        return op()
-
-    def _get_metadata(self) -> dict[str, Any]:
-        return {'name': self._name, 'path': self._path()}
-
     @abstractmethod
     def _display_name(self) -> str:
         """

pixeltable/catalog/table.py

@@ -6,7 +6,7 @@ import json
 import logging
 from keyword import iskeyword as is_python_keyword
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Literal, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Literal, Optional, TypedDict, overload
 
 from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
 import datetime
@@ -80,49 +80,70 @@ class Table(SchemaObject):
         conn.execute(stmt, {'new_dir_id': new_dir_id, 'new_name': json.dumps(new_name), 'id': self._id})
 
     # this is duplicated from SchemaObject so that our API docs show the docstring for Table
-    def get_metadata(self) -> dict[str, Any]:
+    def get_metadata(self) -> 'TableMetadata':
         """
         Retrieves metadata associated with this table.
 
         Returns:
-            A dictionary containing the metadata, in the following format:
-
-                ```python
-                {
-                    'name': 'my_table',
-                    'path': 'my_dir.my_subdir.my_table',
-                    'base': None,  # If this is a view or snapshot, will contain the name of its base table
-                    'schema': {
-                        'col1': StringType(),
-                        'col2': IntType(),
-                    },
-                    'is_replica': False,
-                    'version': 22,
-                    'version_created': datetime.datetime(...),
-                    'schema_version': 1,
-                    'comment': '',
-                    'num_retained_versions': 10,
-                    'is_view': False,
-                    'is_snapshot': False,
-                    'media_validation': 'on_write',
-                }
-                ```
+            A [TableMetadata][pixeltable.TableMetadata] instance containing this table's metadata.
         """
-        return super().get_metadata()
-
-    def _get_metadata(self) -> dict[str, Any]:
-        md = super()._get_metadata()
-        md['schema'] = self._get_schema()
-        md['is_replica'] = self._tbl_version_path.is_replica()
-        md['version'] = self._get_version()
-        md['version_created'] = datetime.datetime.fromtimestamp(
-            self._tbl_version_path.tbl_version.get().created_at, tz=datetime.timezone.utc
+        from pixeltable.catalog import retry_loop
+
+        @retry_loop(for_write=False)
+        def op() -> 'TableMetadata':
+            return self._get_metadata()
+
+        return op()
+
+    def _get_metadata(self) -> 'TableMetadata':
+        columns = self._tbl_version_path.columns()
+        column_info: dict[str, ColumnMetadata] = {}
+        for col in columns:
+            column_info[col.name] = ColumnMetadata(
+                name=col.name,
+                type_=col.col_type._to_str(as_schema=True),
+                version_added=col.schema_version_add,
+                is_stored=col.is_stored,
+                is_primary_key=col.is_pk,
+                media_validation=col.media_validation.name.lower() if col.media_validation is not None else None,  # type: ignore[typeddict-item]
+                computed_with=col.value_expr.display_str(inline=False) if col.value_expr is not None else None,
+            )
+        # Pure snapshots have no indices
+        indices = self._tbl_version.get().idxs_by_name.values() if self._tbl_version is not None else {}
+        index_info: dict[str, IndexMetadata] = {}
+        for info in indices:
+            if isinstance(info.idx, index.EmbeddingIndex):
+                embeddings: list[str] = []
+                if info.idx.string_embed is not None:
+                    embeddings.append(str(info.idx.string_embed))
+                if info.idx.image_embed is not None:
+                    embeddings.append(str(info.idx.image_embed))
+                index_info[info.name] = IndexMetadata(
+                    name=info.name,
+                    columns=[info.col.name],
+                    index_type='embedding',
+                    parameters=EmbeddingIndexParams(
+                        metric=info.idx.metric.name.lower(),  # type: ignore[typeddict-item]
+                        embeddings=embeddings,
+                    ),
+                )
+        return TableMetadata(
+            name=self._name,
+            path=self._path(),
+            columns=column_info,
+            indices=index_info,
+            is_replica=self._tbl_version_path.is_replica(),
+            is_view=False,
+            is_snapshot=False,
+            version=self._get_version(),
+            version_created=datetime.datetime.fromtimestamp(
+                self._tbl_version_path.tbl_version.get().created_at, tz=datetime.timezone.utc
+            ),
+            schema_version=self._tbl_version_path.schema_version(),
+            comment=self._get_comment(),
+            media_validation=self._get_media_validation().name.lower(),  # type: ignore[typeddict-item]
+            base=None,
         )
-        md['schema_version'] = self._tbl_version_path.schema_version()
-        md['comment'] = self._get_comment()
-        md['num_retained_versions'] = self._get_num_retained_versions()
-        md['media_validation'] = self._get_media_validation().name.lower()
-        return md
 
     def _get_version(self) -> int:
         """Return the version of this table. Used by tests to ascertain version changes."""
@@ -455,7 +476,7 @@ class Table(SchemaObject):
 
     def add_columns(
         self,
-        schema: dict[str, Union[ts.ColumnType, builtins.type, _GenericAlias]],
+        schema: dict[str, ts.ColumnType | builtins.type | _GenericAlias],
         if_exists: Literal['error', 'ignore', 'replace', 'replace_force'] = 'error',
     ) -> UpdateStatus:
         """
@@ -529,7 +550,7 @@ class Table(SchemaObject):
         self,
         *,
         if_exists: Literal['error', 'ignore', 'replace', 'replace_force'] = 'error',
-        **kwargs: Union[ts.ColumnType, builtins.type, _GenericAlias, exprs.Expr],
+        **kwargs: ts.ColumnType | builtins.type | _GenericAlias | exprs.Expr,
     ) -> UpdateStatus:
         """
         Adds an ordinary (non-computed) column to the table.
@@ -774,7 +795,7 @@ class Table(SchemaObject):
             cls._verify_column(col)
             column_names.add(col.name)
 
-    def drop_column(self, column: Union[str, ColumnRef], if_not_exists: Literal['error', 'ignore'] = 'error') -> None:
+    def drop_column(self, column: str | ColumnRef, if_not_exists: Literal['error', 'ignore'] = 'error') -> None:
         """Drop a column from the table.
 
         Args:
@@ -896,7 +917,7 @@ class Table(SchemaObject):
 
     def add_embedding_index(
         self,
-        column: Union[str, ColumnRef],
+        column: str | ColumnRef,
         *,
         idx_name: Optional[str] = None,
         embedding: Optional[pxt.Function] = None,
@@ -1023,7 +1044,7 @@ class Table(SchemaObject):
     def drop_embedding_index(
         self,
         *,
-        column: Union[str, ColumnRef, None] = None,
+        column: str | ColumnRef | None = None,
         idx_name: Optional[str] = None,
         if_not_exists: Literal['error', 'ignore'] = 'error',
     ) -> None:
@@ -1083,7 +1104,7 @@ class Table(SchemaObject):
 
             self._drop_index(col=col, idx_name=idx_name, _idx_class=index.EmbeddingIndex, if_not_exists=if_not_exists)
 
-    def _resolve_column_parameter(self, column: Union[str, ColumnRef]) -> Column:
+    def _resolve_column_parameter(self, column: str | ColumnRef) -> Column:
         """Resolve a column parameter to a Column object"""
         col: Column = None
         if isinstance(column, str):
@@ -1102,7 +1123,7 @@ class Table(SchemaObject):
     def drop_index(
         self,
         *,
-        column: Union[str, ColumnRef, None] = None,
+        column: str | ColumnRef | None = None,
         idx_name: Optional[str] = None,
         if_not_exists: Literal['error', 'ignore'] = 'error',
     ) -> None:
@@ -1421,7 +1442,7 @@ class Table(SchemaObject):
             return result
 
     def recompute_columns(
-        self, *columns: Union[str, ColumnRef], errors_only: bool = False, cascade: bool = True
+        self, *columns: str | ColumnRef, errors_only: bool = False, cascade: bool = True
     ) -> UpdateStatus:
         """Recompute the values in one or more computed columns of this table.
 
@@ -1533,11 +1554,7 @@ class Table(SchemaObject):
             env.Env.get().console_logger.info(f'Linked external store `{store.name}` to table `{self._name}`.')
 
     def unlink_external_stores(
-        self,
-        stores: Optional[str | list[str]] = None,
-        *,
-        delete_external_data: bool = False,
-        ignore_errors: bool = False,
+        self, stores: str | list[str] | None = None, *, delete_external_data: bool = False, ignore_errors: bool = False
     ) -> None:
         """
         Unlinks this table's external stores.
@@ -1579,7 +1596,7 @@ class Table(SchemaObject):
                 env.Env.get().console_logger.info(f'Unlinked external store from table `{self._name}`: {store_str}')
 
     def sync(
-        self, stores: Optional[str | list[str]] = None, *, export_data: bool = True, import_data: bool = True
+        self, stores: str | list[str] | None = None, *, export_data: bool = True, import_data: bool = True
     ) -> UpdateStatus:
         """
         Synchronizes this table with its linked external stores.
@@ -1657,7 +1674,7 @@ class Table(SchemaObject):
         from pixeltable.catalog import Catalog
 
         if n is None:
-            n = 1000_000_000
+            n = 1_000_000_000
         if not isinstance(n, int) or n < 1:
             raise excs.Error(f'Invalid value for n: {n}')
 
@@ -1709,3 +1726,72 @@ class Table(SchemaObject):
             raise excs.Error(f'{self._display_str()}: Cannot {op_descr} a snapshot.')
         if self._tbl_version_path.is_replica():
             raise excs.Error(f'{self._display_str()}: Cannot {op_descr} a {self._display_name()}.')
+
+
+class ColumnMetadata(TypedDict):
+    """Metadata for a column of a Pixeltable table."""
+
+    name: str
+    """The name of the column."""
+    type_: str
+    """The type specifier of the column."""
+    version_added: int
+    """The table version when this column was added."""
+    is_stored: bool
+    """`True` if this is a stored column; `False` if it is dynamically computed."""
+    is_primary_key: bool
+    """`True` if this column is part of the table's primary key."""
+    media_validation: Optional[Literal['on_read', 'on_write']]
+    """The media validation policy for this column."""
+    computed_with: Optional[str]
+    """Expression used to compute this column; `None` if this is not a computed column."""
+
+
+class IndexMetadata(TypedDict):
+    """Metadata for a column of a Pixeltable table."""
+
+    name: str
+    """The name of the index."""
+    columns: list[str]
+    """The table columns that are indexed."""
+    index_type: Literal['embedding']
+    """The type of index (currently only `'embedding'` is supported, but others will be added in the future)."""
+    parameters: EmbeddingIndexParams
+
+
+class EmbeddingIndexParams(TypedDict):
+    metric: Literal['cosine', 'ip', 'l2']
+    """Index metric."""
+    embeddings: list[str]
+    """List of embeddings defined for this index."""
+
+
+class TableMetadata(TypedDict):
+    """Metadata for a Pixeltable table."""
+
+    name: str
+    """The name of the table (ex: `'my_table'`)."""
+    path: str
+    """The full path of the table (ex: `'my_dir.my_subdir.my_table'`)."""
+    columns: dict[str, ColumnMetadata]
+    """Column metadata for all of the visible columns of the table."""
+    indices: dict[str, IndexMetadata]
+    """Index metadata for all of the indices of the table."""
+    is_replica: bool
+    """`True` if this table is a replica of another (shared) table."""
+    is_view: bool
+    """`True` if this table is a view."""
+    is_snapshot: bool
+    """`True` if this table is a snapshot."""
+    version: int
+    """The current version of the table."""
+    version_created: datetime.datetime
+    """The timestamp when this table version was created."""
+    schema_version: int
+    """The current schema version of the table."""
+    comment: Optional[str]
+    """User-provided table comment, if one exists."""
+    media_validation: Literal['on_read', 'on_write']
+    """The media validation policy for this table."""
+    base: Optional[str]
+    """If this table is a view or snapshot, the full path of its base table; otherwise `None`."""