pixeltable 0.3.11__py3-none-any.whl → 0.3.12__py3-none-any.whl

pixeltable/catalog/table.py

@@ -95,6 +95,7 @@ class Table(SchemaObject):
                         'col1': StringType(),
                         'col2': IntType(),
                     },
+                    'is_replica': False,
                     'version': 22,
                     'schema_version': 1,
                     'comment': '',
@@ -110,6 +111,7 @@ class Table(SchemaObject):
             md = super().get_metadata()
             md['base'] = self._base._path() if self._base is not None else None
             md['schema'] = self._schema
+            md['is_replica'] = self._tbl_version.get().is_replica
             md['version'] = self._version
             md['schema_version'] = self._tbl_version.get().schema_version
             md['comment'] = self._comment
@@ -139,14 +141,14 @@ class Table(SchemaObject):
         if self._is_dropped:
             raise excs.Error(f'{self._display_name()} {self._name} has been dropped')
 
-    def __getattr__(self, name: str) -> 'pxt.exprs.ColumnRef':
+    def __getattr__(self, name: str) -> 'exprs.ColumnRef':
         """Return a ColumnRef for the given name."""
         col = self._tbl_version_path.get_column(name)
         if col is None:
             raise AttributeError(f'Column {name!r} unknown')
         return ColumnRef(col)
 
-    def __getitem__(self, name: str) -> 'pxt.exprs.ColumnRef':
+    def __getitem__(self, name: str) -> 'exprs.ColumnRef':
         """Return a ColumnRef for the given name."""
         return getattr(self, name)
 
@@ -689,7 +691,7 @@ class Table(SchemaObject):
         for name, spec in schema.items():
             col_type: Optional[ts.ColumnType] = None
             value_expr: Optional[exprs.Expr] = None
-            primary_key: Optional[bool] = None
+            primary_key: bool = False
             media_validation: Optional[catalog.MediaValidation] = None
             stored = True
 
@@ -711,7 +713,7 @@ class Table(SchemaObject):
                     value_expr = value_expr.copy()
                     value_expr.bind_rel_paths()
                 stored = spec.get('stored', True)
-                primary_key = spec.get('primary_key')
+                primary_key = spec.get('primary_key', False)
                 media_validation_str = spec.get('media_validation')
                 media_validation = (
                     catalog.MediaValidation[media_validation_str.upper()] if media_validation_str is not None else None
@@ -1282,7 +1284,7 @@ class Table(SchemaObject):
         raise NotImplementedError
 
     def update(
-        self, value_spec: dict[str, Any], where: Optional['pxt.exprs.Expr'] = None, cascade: bool = True
+        self, value_spec: dict[str, Any], where: Optional['exprs.Expr'] = None, cascade: bool = True
     ) -> UpdateStatus:
         """Update rows in this table.
 
@@ -1383,7 +1385,7 @@ class Table(SchemaObject):
             FileCache.get().emit_eviction_warnings()
             return status
 
-    def delete(self, where: Optional['pxt.exprs.Expr'] = None) -> UpdateStatus:
+    def delete(self, where: Optional['exprs.Expr'] = None) -> UpdateStatus:
         """Delete rows in this table.
 
         Args:

pixeltable/catalog/table_version.py

@@ -5,7 +5,7 @@ import importlib
 import logging
 import time
 import uuid
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional, Tuple
 from uuid import UUID
 
 import jsonschema.exceptions
@@ -18,6 +18,7 @@ from pixeltable import exprs, index
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
 from pixeltable.metadata import schema
+from pixeltable.utils.exception_handler import run_cleanup_on_exception
 from pixeltable.utils.filecache import FileCache
 from pixeltable.utils.media_store import MediaStore
 
@@ -55,6 +56,7 @@ class TableVersion:
     name: str
     user: Optional[str]
     effective_version: Optional[int]
+    is_replica: bool
     version: int
     comment: str
     media_validation: MediaValidation
@@ -111,6 +113,7 @@ class TableVersion:
         self.user = tbl_md.user
         self.effective_version = effective_version
         self.version = tbl_md.current_version if effective_version is None else effective_version
+        self.is_replica = tbl_md.is_replica
         self.comment = schema_version_md.comment
         self.num_retained_versions = schema_version_md.num_retained_versions
         self.schema_version = schema_version_md.schema_version
@@ -232,6 +235,7 @@ class TableVersion:
             tbl_id=str(tbl_id),
             name=name,
             user=user,
+            is_replica=False,
             current_version=0,
             current_schema_version=0,
             next_col_id=len(cols),
@@ -310,24 +314,16 @@ class TableVersion:
         session.add(schema_version_record)
         return tbl_record.id, tbl_version
 
-    @classmethod
-    def delete_md(cls, tbl_id: UUID) -> None:
-        conn = Env.get().conn
-        conn.execute(sql.delete(schema.TableSchemaVersion.__table__).where(schema.TableSchemaVersion.tbl_id == tbl_id))
-        conn.execute(sql.delete(schema.TableVersion.__table__).where(schema.TableVersion.tbl_id == tbl_id))
-        conn.execute(sql.delete(schema.Table.__table__).where(schema.Table.id == tbl_id))
-
     def drop(self) -> None:
+        from .catalog import Catalog
+
+        cat = Catalog.get()
         # delete this table and all associated data
         MediaStore.delete(self.id)
         FileCache.get().clear(tbl_id=self.id)
-        self.delete_md(self.id)
+        cat.delete_tbl_md(self.id)
         self.store_tbl.drop()
-
         # de-register table version from catalog
-        from .catalog import Catalog
-
-        cat = Catalog.get()
         cat.remove_tbl_version(self)
 
     def _init_schema(self, tbl_md: schema.TableMd, schema_version_md: schema.TableSchemaVersionMd) -> None:
@@ -381,7 +377,7 @@ class TableVersion:
 
             # make sure to traverse columns ordered by position = order in which cols were created;
             # this guarantees that references always point backwards
-            if col_md.value_expr is not None:
+            if not self.is_snapshot and col_md.value_expr is not None:
                 self._record_refd_columns(col)
 
     def _init_idxs(self, tbl_md: schema.TableMd) -> None:
@@ -437,29 +433,15 @@ class TableVersion:
                 specified preceding schema version
         """
         assert update_tbl_version or preceding_schema_version is None
+        from pixeltable.catalog import Catalog
 
-        conn = Env.get().conn
-        conn.execute(
-            sql.update(schema.Table.__table__)
-            .values({schema.Table.md: dataclasses.asdict(self._create_tbl_md())})
-            .where(schema.Table.id == self.id)
+        tbl_md = self._create_tbl_md()
+        version_md = self._create_version_md(timestamp) if update_tbl_version else None
+        schema_version_md = (
+            self._create_schema_version_md(preceding_schema_version) if preceding_schema_version is not None else None
         )
 
-        if update_tbl_version:
-            version_md = self._create_version_md(timestamp)
-            conn.execute(
-                sql.insert(schema.TableVersion.__table__).values(
-                    tbl_id=self.id, version=self.version, md=dataclasses.asdict(version_md)
-                )
-            )
-
-        if preceding_schema_version is not None:
-            schema_version_md = self._create_schema_version_md(preceding_schema_version)
-            conn.execute(
-                sql.insert(schema.TableSchemaVersion.__table__).values(
-                    tbl_id=self.id, schema_version=self.schema_version, md=dataclasses.asdict(schema_version_md)
-                )
-            )
+        Catalog.get().store_tbl_md(self.id, tbl_md, version_md, schema_version_md)
 
     def ensure_md_loaded(self) -> None:
         """Ensure that table metadata is loaded."""
@@ -480,33 +462,36 @@ class TableVersion:
         _logger.info(f'Added index {idx_name} on column {col.name} to table {self.name}')
         return status
 
-    def _add_default_index(self, col: Column) -> Optional[UpdateStatus]:
-        """Add a B-tree index on this column if it has a compatible type"""
+    def _is_btree_indexable(self, col: Column) -> bool:
         if not col.stored:
             # if the column is intentionally not stored, we want to avoid the overhead of an index
-            return None
+            return False
         # Skip index for stored media columns produced by an iterator
         if col.col_type.is_media_type() and self.is_iterator_column(col):
-            return None
+            return False
         if not col.col_type.is_scalar_type() and not (col.col_type.is_media_type() and not col.is_computed):
             # wrong type for a B-tree
-            return None
-        if col.col_type.is_bool_type():
+            return False
+        if col.col_type.is_bool_type():  # noqa : SIM103 Supress `Return the negated condition directly` check
             # B-trees on bools aren't useful
+            return False
+        return True
+
+    def _add_default_index(self, col: Column) -> Optional[UpdateStatus]:
+        """Add a B-tree index on this column if it has a compatible type"""
+        if not self._is_btree_indexable(col):
             return None
         status = self._add_index(col, idx_name=None, idx=index.BtreeIndex(col))
         return status
 
-    def _add_index(self, col: Column, idx_name: Optional[str], idx: index.IndexBase) -> UpdateStatus:
+    def _create_index_columns(self, idx: index.IndexBase) -> Tuple[Column, Column]:
+        """Create value and undo columns for the given index.
+        Args:
+            idx:  index for which columns will be created.
+        Returns:
+            A tuple containing the value column and the undo column.
+        """
         assert not self.is_snapshot
-        idx_id = self.next_idx_id
-        self.next_idx_id += 1
-        if idx_name is None:
-            idx_name = f'idx{idx_id}'
-        else:
-            assert is_valid_identifier(idx_name)
-            assert idx_name not in [i.name for i in self.idx_md.values()]
-
         # add the index value and undo columns (which need to be nullable)
         val_col = Column(
             col_id=self.next_col_id,
@@ -535,7 +520,19 @@ class TableVersion:
         undo_col.tbl = self.create_handle()
         undo_col.col_type = undo_col.col_type.copy(nullable=True)
         self.next_col_id += 1
+        return val_col, undo_col
 
+    def _create_index(
+        self, col: Column, val_col: Column, undo_col: Column, idx_name: Optional[str], idx: index.IndexBase
+    ) -> None:
+        """Create the given index along with index md"""
+        idx_id = self.next_idx_id
+        self.next_idx_id += 1
+        if idx_name is None:
+            idx_name = f'idx{idx_id}'
+        else:
+            assert is_valid_identifier(idx_name)
+            assert idx_name not in [i.name for i in self.idx_md.values()]
         # create and register the index metadata
         idx_cls = type(idx)
         idx_md = schema.IndexMd(
@@ -553,14 +550,27 @@ class TableVersion:
         idx_info = self.IndexInfo(id=idx_id, name=idx_name, idx=idx, col=col, val_col=val_col, undo_col=undo_col)
         self.idx_md[idx_id] = idx_md
         self.idxs_by_name[idx_name] = idx_info
+        try:
+            idx.create_index(self._store_idx_name(idx_id), val_col)
+        finally:
 
+            def cleanup_index() -> None:
+                """Delete the newly added in-memory index structure"""
+                del self.idxs_by_name[idx_name]
+                del self.idx_md[idx_id]
+                self.next_idx_id = idx_id
+
+            # Run cleanup only if there has been an exception; otherwise, skip cleanup.
+            run_cleanup_on_exception(cleanup_index)
+
+    def _add_index(self, col: Column, idx_name: Optional[str], idx: index.IndexBase) -> UpdateStatus:
+        val_col, undo_vol = self._create_index_columns(idx)
         # add the columns and update the metadata
         # TODO support on_error='abort' for indices; it's tricky because of the way metadata changes are entangled
         # with the database operations
-        status = self._add_columns([val_col, undo_col], print_stats=False, on_error='ignore')
+        status = self._add_columns([val_col, undo_vol], print_stats=False, on_error='ignore')
         # now create the index structure
-        idx.create_index(self._store_idx_name(idx_id), val_col)
-
+        self._create_index(col, val_col, undo_vol, idx_name, idx)
         return status
 
     def drop_index(self, idx_id: int) -> None:
@@ -601,9 +611,21 @@ class TableVersion:
         self.version += 1
         preceding_schema_version = self.schema_version
         self.schema_version = self.version
-        status = self._add_columns(cols, print_stats=print_stats, on_error=on_error)
+        index_cols: dict[Column, tuple[index.BtreeIndex, Column, Column]] = {}
+        all_cols: list[Column] = []
         for col in cols:
-            _ = self._add_default_index(col)
+            all_cols.append(col)
+            if self._is_btree_indexable(col):
+                idx = index.BtreeIndex(col)
+                val_col, undo_col = self._create_index_columns(idx)
+                index_cols[col] = (idx, val_col, undo_col)
+                all_cols.append(val_col)
+                all_cols.append(undo_col)
+        # Add all columns
+        status = self._add_columns(all_cols, print_stats=print_stats, on_error=on_error)
+        # Create indices and their mds
+        for col, (idx, val_col, undo_col) in index_cols.items():
+            self._create_index(col, val_col, undo_col, idx_name=None, idx=idx)
         self._update_md(time.time(), preceding_schema_version=preceding_schema_version)
         _logger.info(f'Added columns {[col.name for col in cols]} to table {self.name}, new version: {self.version}')
 
@@ -619,9 +641,9 @@ class TableVersion:
         self, cols: Iterable[Column], print_stats: bool, on_error: Literal['abort', 'ignore']
     ) -> UpdateStatus:
         """Add and populate columns within the current transaction"""
-        cols = list(cols)
+        cols_to_add = list(cols)
         row_count = self.store_tbl.count()
-        for col in cols:
+        for col in cols_to_add:
             if not col.col_type.nullable and not col.is_computed and row_count > 0:
                 raise excs.Error(
                     f'Cannot add non-nullable column {col.name!r} to table {self.name!r} with existing rows'
@@ -629,7 +651,8 @@ class TableVersion:
 
         num_excs = 0
         cols_with_excs: list[Column] = []
-        for col in cols:
+        for col in cols_to_add:
+            excs_per_col = 0
             col.schema_version_add = self.schema_version
             # add the column to the lookup structures now, rather than after the store changes executed successfully,
             # because it might be referenced by the next column's value_expr
@@ -652,29 +675,32 @@ class TableVersion:
 
             plan, value_expr_slot_idx = Planner.create_add_column_plan(self.path, col)
             plan.ctx.num_rows = row_count
-
             try:
                 plan.open()
                 try:
-                    num_excs = self.store_tbl.load_column(col, plan, value_expr_slot_idx, on_error)
+                    excs_per_col = self.store_tbl.load_column(col, plan, value_expr_slot_idx, on_error)
                 except sql.exc.DBAPIError as exc:
                     # Wrap the DBAPIError in an excs.Error to unify processing in the subsequent except block
                     raise excs.Error(f'SQL error during execution of computed column `{col.name}`:\n{exc}') from exc
-                if num_excs > 0:
+                if excs_per_col > 0:
                     cols_with_excs.append(col)
-            except excs.Error as exc:
-                self.cols.pop()
-                for c in cols:
-                    # remove columns that we already added
-                    if c.id not in self.cols_by_id:
-                        continue
-                    if c.name is not None:
-                        del self.cols_by_name[c.name]
-                    del self.cols_by_id[c.id]
-                # we need to re-initialize the sqlalchemy schema
-                self.store_tbl.create_sa_tbl()
-                raise exc
+                    num_excs += excs_per_col
             finally:
+                # Ensure cleanup occurs if an exception or keyboard interruption happens during `load_column()`.
+                def cleanup_on_error() -> None:
+                    """Delete columns that are added as part of current add_columns operation and re-initialize
+                    the sqlalchemy schema"""
+                    self.cols = [col for col in self.cols if col not in cols_to_add]
+                    for col in cols_to_add:
+                        # remove columns that we already added
+                        if col.id in self.cols_by_id:
+                            del self.cols_by_id[col.id]
+                        if col.name is not None and col.name in self.cols_by_name:
+                            del self.cols_by_name[col.name]
+                    self.store_tbl.create_sa_tbl()
+
+                # Run cleanup only if there has been an exception; otherwise, skip cleanup.
+                run_cleanup_on_exception(cleanup_on_error)
                 plan.close()
 
         if print_stats:
@@ -1298,6 +1324,7 @@ class TableVersion:
         column_md: dict[int, schema.ColumnMd] = {}
         for col in cols:
             value_expr_dict = col.value_expr.as_dict() if col.value_expr is not None else None
+            assert col.is_pk is not None
             column_md[col.id] = schema.ColumnMd(
                 id=col.id,
                 col_type=col.col_type.as_dict(),
@@ -1320,6 +1347,7 @@ class TableVersion:
             tbl_id=str(self.id),
             name=self.name,
             user=self.user,
+            is_replica=self.is_replica,
             current_version=self.version,
             current_schema_version=self.schema_version,
             next_col_id=self.next_col_id,

pixeltable/catalog/view.py

@@ -8,7 +8,7 @@ from uuid import UUID
 import pixeltable.exceptions as excs
 import pixeltable.metadata.schema as md_schema
 import pixeltable.type_system as ts
-from pixeltable import exprs, func
+from pixeltable import catalog, exprs, func
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
 
@@ -20,7 +20,7 @@ from .table_version_handle import TableVersionHandle
 from .table_version_path import TableVersionPath
 
 if TYPE_CHECKING:
-    import pixeltable as pxt
+    from pixeltable.globals import TableDataSource
 
 _logger = logging.getLogger('pixeltable')
 
@@ -65,7 +65,7 @@ class View(Table):
         base: TableVersionPath,
         select_list: Optional[list[tuple[exprs.Expr, Optional[str]]]],
         additional_columns: dict[str, Any],
-        predicate: Optional['pxt.exprs.Expr'],
+        predicate: Optional['exprs.Expr'],
         is_snapshot: bool,
         num_retained_versions: int,
         comment: str,
@@ -242,7 +242,7 @@ class View(Table):
             # there is not TableVersion to drop
             self._check_is_dropped()
             self.is_dropped = True
-            TableVersion.delete_md(self._id)
+            catalog.Catalog.get().delete_tbl_md(self._id)
         else:
             super()._drop()
 
@@ -252,11 +252,6 @@ class View(Table):
         md['is_snapshot'] = self._tbl_version_path.is_snapshot()
         return md
 
-    if TYPE_CHECKING:
-        import datasets  # type: ignore[import-untyped]
-
-        from pixeltable.globals import RowData, TableDataSource
-
     def insert(
         self,
         source: Optional[TableDataSource] = None,

pixeltable/exec/sql_node.py

@@ -103,7 +103,6 @@ class SqlNode(ExecNode):
         # create Select stmt
         self.sql_elements = sql_elements
         self.tbl = tbl
-        assert all(not isinstance(e, exprs.Literal) for e in select_list)  # we're never asked to materialize literals
         self.select_list = exprs.ExprSet(select_list)
         # unstored iter columns: we also need to retrieve whatever is needed to materialize the iter args
         for iter_arg in row_builder.unstored_iter_args.values():

pixeltable/exprs/json_path.py

@@ -5,7 +5,6 @@ from typing import Any, Optional, Union
 import jmespath
 import sqlalchemy as sql
 
-import pixeltable as pxt
 from pixeltable import catalog, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
@@ -19,10 +18,7 @@ from .sql_element_cache import SqlElementCache
 
 class JsonPath(Expr):
     def __init__(
-        self,
-        anchor: Optional['pxt.exprs.Expr'],
-        path_elements: Optional[list[Union[str, int, slice]]] = None,
-        scope_idx: int = 0,
+        self, anchor: Optional[Expr], path_elements: Optional[list[Union[str, int, slice]]] = None, scope_idx: int = 0
     ) -> None:
         """
         anchor can be None, in which case this is a relative JsonPath and the anchor is set later via set_anchor().

pixeltable/func/__init__.py

@@ -5,7 +5,7 @@ from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function, InvalidFunction
 from .function_registry import FunctionRegistry
-from .query_template_function import QueryTemplateFunction, query
+from .query_template_function import QueryTemplateFunction, query, retrieval_udf
 from .signature import Batch, Parameter, Signature
 from .tools import Tool, ToolChoice, Tools
 from .udf import expr_udf, make_function, udf

pixeltable/func/aggregate_function.py

@@ -159,7 +159,7 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.agg_classes[0])
 
     def help_str(self) -> str:

pixeltable/func/callable_function.py

@@ -60,7 +60,7 @@ class CallableFunction(Function):
     def is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.py_fn)
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return inspect.getdoc(self.py_fns[0])
 
     @property

pixeltable/func/expr_template_function.py

@@ -95,9 +95,9 @@ class ExprTemplateFunction(Function):
         )
         return substituted_expr.col_type
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
-            return self.templates[0].expr.fn._docstring()
+            return self.templates[0].expr.fn.comment()
         return None
 
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:

pixeltable/func/function.py

@@ -10,8 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 import sqlalchemy as sql
 from typing_extensions import Self
 
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 
 from .globals import resolve_symbol
 from .signature import Signature
@@ -106,11 +105,11 @@ class Function(ABC):
     @abstractmethod
     def is_async(self) -> bool: ...
 
-    def _docstring(self) -> Optional[str]:
+    def comment(self) -> Optional[str]:
         return None
 
     def help_str(self) -> str:
-        docstring = self._docstring()
+        docstring = self.comment()
         display = self.display_name + str(self.signatures[0])
         if docstring is None:
             return display

pixeltable/func/query_template_function.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional, overload
+from functools import reduce
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Union, overload
 
-from pixeltable import exprs, type_system as ts
+from pixeltable import catalog, exceptions as excs, exprs, func, type_system as ts
 
 from .function import Function
 from .signature import Signature
@@ -17,6 +18,7 @@ class QueryTemplateFunction(Function):
 
     template_df: Optional['DataFrame']
     self_name: Optional[str]
+    _comment: Optional[str]
 
     @classmethod
     def create(
@@ -34,15 +36,21 @@ class QueryTemplateFunction(Function):
         assert isinstance(template_df, DataFrame)
         # we take params and return json
         sig = Signature(return_type=ts.JsonType(), parameters=params)
-        return QueryTemplateFunction(template_df, sig, path=path, name=name)
+        return QueryTemplateFunction(template_df, sig, path=path, name=name, comment=inspect.getdoc(template_callable))
 
     def __init__(
-        self, template_df: Optional['DataFrame'], sig: Signature, path: Optional[str] = None, name: Optional[str] = None
+        self,
+        template_df: Optional['DataFrame'],
+        sig: Signature,
+        path: Optional[str] = None,
+        name: Optional[str] = None,
+        comment: Optional[str] = None,
     ):
         assert sig is not None
         super().__init__([sig], self_path=path)
         self.self_name = name
         self.template_df = template_df
+        self._comment = comment
 
     def _update_as_overload_resolution(self, signature_idx: int) -> None:
         pass  # only one signature supported for QueryTemplateFunction
@@ -74,6 +82,9 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
+    def comment(self) -> Optional[str]:
+        return self._comment
+
     def _as_dict(self) -> dict:
         return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
 
@@ -112,3 +123,75 @@ def query(*args: Any, **kwargs: Any) -> Any:
     else:
         assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
         return lambda py_fn: make_query_template(py_fn, kwargs['param_types'])
+
+
+def retrieval_udf(
+    table: catalog.Table,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    parameters: Optional[Iterable[Union[str, exprs.ColumnRef]]] = None,
+    limit: Optional[int] = 10,
+) -> func.QueryTemplateFunction:
+    """
+    Constructs a retrieval UDF for the given table. The retrieval UDF is a UDF whose parameters are
+    columns of the table and whose return value is a list of rows from the table. The return value of
+    ```python
+    f(col1=x, col2=y, ...)
+    ```
+    will be a list of all rows from the table that match the specified arguments.
+
+    Args:
+        table: The table to use as the dataset for the retrieval tool.
+        name: The name of the tool. If not specified, then the name of the table will be used by default.
+        description: The description of the tool. If not specified, then a default description will be generated.
+        parameters: The columns of the table to use as parameters. If not specified, all data columns
+            (non-computed columns) will be used as parameters.
+
+            All of the specified parameters will be required parameters of the tool, regardless of their status
+            as columns.
+        limit: The maximum number of rows to return. If not specified, then all matching rows will be returned.
+
+    Returns:
+        A list of dictionaries containing data from the table, one per row that matches the input arguments.
+        If there are no matching rows, an empty list will be returned.
+    """
+    # Argument validation
+    col_refs: list[exprs.ColumnRef]
+    if parameters is None:
+        col_refs = [table[col_name] for col_name in table.columns if not table[col_name].col.is_computed]
+    else:
+        for param in parameters:
+            if isinstance(param, str) and param not in table.columns:
+                raise excs.Error(f'The specified parameter {param!r} is not a column of the table {table._path!r}')
+        col_refs = [table[param] if isinstance(param, str) else param for param in parameters]
+
+    if len(col_refs) == 0:
+        raise excs.Error('Parameter list cannot be empty.')
+
+    # Construct the dataframe
+    predicates = [col_ref == exprs.Variable(col_ref.col.name, col_ref.col.col_type) for col_ref in col_refs]
+    where_clause = reduce(lambda c1, c2: c1 & c2, predicates)
+    df = table.select().where(where_clause)
+    if limit is not None:
+        df = df.limit(limit)
+
+    # Construct the signature
+    query_params = [
+        func.Parameter(col_ref.col.name, col_ref.col.col_type, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        for col_ref in col_refs
+    ]
+    query_signature = func.Signature(return_type=ts.JsonType(), parameters=query_params)
+
+    # Construct a name and/or description if not provided
+    if name is None:
+        name = table._name
+    if description is None:
+        description = (
+            f'Retrieves an entry from the dataset {name!r} that matches the given parameters.\n\nParameters:\n'
+        )
+        description += '\n'.join(
+            [f'    {col_ref.col.name}: of type `{col_ref.col.col_type._to_base_str()}`' for col_ref in col_refs]
+        )
+
+    fn = func.QueryTemplateFunction(df, query_signature, name=name, comment=description)
+    return fn

pixeltable/func/tools.py

@@ -39,7 +39,7 @@ class Tool(pydantic.BaseModel):
     def ser_model(self) -> dict[str, Any]:
         return {
             'name': self.name or self.fn.name,
-            'description': self.description or self.fn._docstring(),
+            'description': self.description or self.fn.comment(),
             'parameters': {
                 'type': 'object',
                 'properties': {param.name: param.col_type._to_json_schema() for param in self.parameters.values()},