pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/exprs/type_cast.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import sqlalchemy as sql
 
 from pixeltable import type_system as ts
@@ -19,7 +17,7 @@ class TypeCast(Expr):
     def __init__(self, underlying: Expr, new_type: ts.ColumnType):
         super().__init__(new_type)
         self.components: list[Expr] = [underlying]
-        self.id: Optional[int] = self._create_id()
+        self.id: int | None = self._create_id()
 
     def _equals(self, other: 'TypeCast') -> bool:
         # `TypeCast` has no properties beyond those captured by `Expr`.
@@ -29,7 +27,7 @@ class TypeCast(Expr):
     def _op1(self) -> Expr:
         return self.components[0]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         """
         sql_expr() is unimplemented for now, in order to sidestep potentially thorny
         questions about consistency of doing type conversions in both Python and Postgres.
@@ -40,7 +38,7 @@ class TypeCast(Expr):
         original_val = data_row[self._op1.slot_idx]
         data_row[self.slot_idx] = self.col_type.create_literal(original_val)
 
-    def as_literal(self) -> Optional[Literal]:
+    def as_literal(self) -> Literal | None:
         op1_lit = self._op1.as_literal()
         if op1_lit is None:
             return None

pixeltable/func/__init__.py

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

pixeltable/func/aggregate_function.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import abc
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Sequence, overload
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Sequence, overload
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
@@ -42,7 +42,7 @@ class AggregateFunction(Function):
     def __init__(
         self,
         agg_class: type[Aggregator],
-        type_substitutions: Optional[Sequence[dict]],
+        type_substitutions: Sequence[dict] | None,
         self_path: str,
         requires_order_by: bool,
         allows_std_agg: bool,
@@ -75,7 +75,7 @@ class AggregateFunction(Function):
         self.init_param_names = [self.init_param_names[signature_idx]]
 
     def __cls_to_signature(
-        self, cls: type[Aggregator], type_substitutions: Optional[dict] = None
+        self, cls: type[Aggregator], type_substitutions: dict | None = None
     ) -> tuple[Signature, list[str]]:
         """Inspects the Aggregator class to infer the corresponding function signature. Returns the
         inferred signature along with the list of init_param_names (for downstream error handling).
@@ -159,7 +159,7 @@ class AggregateFunction(Function):
         self.init_param_names.append(init_param_names)
         return self
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return inspect.getdoc(self.agg_classes[0])
 
     def help_str(self) -> str:
@@ -173,7 +173,7 @@ class AggregateFunction(Function):
         from pixeltable import exprs
 
         # perform semantic analysis of special parameters 'order_by' and 'group_by'
-        order_by_clause: Optional[Any] = None
+        order_by_clause: Any | None = None
         if self.ORDER_BY_PARAM in kwargs:
             if self.requires_order_by:
                 raise excs.Error(
@@ -198,7 +198,7 @@ class AggregateFunction(Function):
             # don't pass the first parameter on, the Function doesn't get to see it
             args = args[1:]
 
-        group_by_clause: Optional[Any] = None
+        group_by_clause: Any | None = None
         if self.GROUP_BY_PARAM in kwargs:
             if not self.allows_window:
                 raise excs.Error(
@@ -248,7 +248,7 @@ def uda(
     requires_order_by: bool = False,
     allows_std_agg: bool = True,
     allows_window: bool = False,
-    type_substitutions: Optional[Sequence[dict]] = None,
+    type_substitutions: Sequence[dict] | None = None,
 ) -> Callable[[type[Aggregator]], AggregateFunction]: ...
 
 
@@ -302,7 +302,7 @@ def make_aggregator(
     requires_order_by: bool = False,
     allows_std_agg: bool = True,
     allows_window: bool = False,
-    type_substitutions: Optional[Sequence[dict]] = None,
+    type_substitutions: Sequence[dict] | None = None,
 ) -> AggregateFunction:
     class_path = f'{cls.__module__}.{cls.__qualname__}'
     instance = AggregateFunction(cls, type_substitutions, class_path, requires_order_by, allows_std_agg, allows_window)

pixeltable/func/callable_function.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Sequence
 from uuid import UUID
 
 import cloudpickle  # type: ignore[import-untyped]
@@ -25,16 +25,16 @@ class CallableFunction(Function):
     """
 
     py_fns: list[Callable]
-    self_name: Optional[str]
-    batch_size: Optional[int]
+    self_name: str | None
+    batch_size: int | None
 
     def __init__(
         self,
         signatures: list[Signature],
         py_fns: list[Callable],
-        self_path: Optional[str] = None,
-        self_name: Optional[str] = None,
-        batch_size: Optional[int] = None,
+        self_path: str | None = None,
+        self_name: str | None = None,
+        batch_size: int | None = None,
         is_method: bool = False,
         is_property: bool = False,
     ):
@@ -60,7 +60,7 @@ class CallableFunction(Function):
     def is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.py_fn)
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return inspect.getdoc(self.py_fns[0])
 
     @property
@@ -138,7 +138,7 @@ class CallableFunction(Function):
         batched_kwargs = {k: v for k, v in kwargs.items() if k not in constant_param_names}
         return constant_kwargs, batched_kwargs
 
-    def get_batch_size(self, *args: Any, **kwargs: Any) -> Optional[int]:
+    def get_batch_size(self, *args: Any, **kwargs: Any) -> int | None:
         return self.batch_size
 
     @property
@@ -187,7 +187,7 @@ class CallableFunction(Function):
         return md, cloudpickle.dumps(self.py_fn)
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: str | None, md: dict, binary_obj: bytes) -> Function:
         py_fn = cloudpickle.loads(binary_obj)
         assert callable(py_fn)
         sig = Signature.from_dict(md['signature'])

pixeltable/func/expr_template_function.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Sequence
+from typing import Any, Sequence
 
 from pixeltable import exceptions as excs, exprs, type_system as ts
 
@@ -41,7 +41,7 @@ class ExprTemplateFunction(Function):
     templates: list[ExprTemplate]
     self_name: str
 
-    def __init__(self, templates: list[ExprTemplate], self_path: Optional[str] = None, name: Optional[str] = None):
+    def __init__(self, templates: list[ExprTemplate], self_path: str | None = None, name: str | None = None):
         self.templates = templates
         self.self_name = name
 
@@ -85,29 +85,29 @@ class ExprTemplateFunction(Function):
         conditional_return_type).
         """
         assert not self.is_polymorphic
-        template = self.template
         with_defaults = bound_args.copy()
         with_defaults.update(
-            {param_name: default for param_name, default in template.defaults.items() if param_name not in bound_args}
+            {
+                param_name: default
+                for param_name, default in self.template.defaults.items()
+                if param_name not in bound_args
+            }
         )
         substituted_expr = self.template.expr.copy().substitute(
-            {template.param_exprs[name]: expr for name, expr in with_defaults.items()}
+            {self.template.param_exprs[name]: expr for name, expr in with_defaults.items()}
         )
         return substituted_expr.col_type
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
             return self.templates[0].expr.fn.comment()
         return None
 
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
-        from pixeltable import exec
-
         assert not self.is_polymorphic
         expr = self.instantiate(args, kwargs)
         row_builder = exprs.RowBuilder(output_exprs=[expr], columns=[], input_exprs=[])
-        row_batch = exec.DataRowBatch(tbl=None, row_builder=row_builder, num_rows=1)
-        row = row_batch[0]
+        row = row_builder.make_row()
         row_builder.eval(row, ctx=row_builder.default_eval_ctx)
         return row[row_builder.get_output_exprs()[0].slot_idx]
 

pixeltable/func/function.py

@@ -5,7 +5,7 @@ import inspect
 import typing
 from abc import ABC, abstractmethod
 from copy import copy
-from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, Sequence, cast
 
 import sqlalchemy as sql
 from typing_extensions import Self
@@ -30,10 +30,10 @@ class Function(ABC):
     """
 
     signatures: list[Signature]
-    self_path: Optional[str]
+    self_path: str | None
     is_method: bool
     is_property: bool
-    _conditional_return_type: Optional[Callable[..., ts.ColumnType]]
+    _conditional_return_type: Callable[..., ts.ColumnType] | None
 
     # We cache the overload resolutions in self._resolutions. This ensures that each resolution is represented
     # globally by a single Python object. We do this dynamically rather than pre-constructing them in order to
@@ -43,17 +43,17 @@ class Function(ABC):
     # Translates a call to this function with the given arguments to its SQLAlchemy equivalent.
     # Overriden for specific Function instances via the to_sql() decorator. The override must accept the same
     # parameter names as the original function. Each parameter is going to be of type sql.ColumnElement.
-    _to_sql: Callable[..., Optional[sql.ColumnElement]]
+    _to_sql: Callable[..., sql.ColumnElement | None]
 
     # Returns the resource pool to use for calling this function with the given arguments.
     # Overriden for specific Function instances via the resource_pool() decorator. The override must accept a subset
     # of the parameters of the original function, with the same type.
-    _resource_pool: Callable[..., Optional[str]]
+    _resource_pool: Callable[..., str | None]
 
     def __init__(
         self,
         signatures: list[Signature],
-        self_path: Optional[str] = None,
+        self_path: str | None = None,
         is_method: bool = False,
         is_property: bool = False,
     ):
@@ -105,7 +105,7 @@ class Function(ABC):
     @abstractmethod
     def is_async(self) -> bool: ...
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return None
 
     def help_str(self) -> str:
@@ -175,7 +175,7 @@ class Function(ABC):
 
     def _bind_to_matching_signature(self, args: Sequence[Any], kwargs: dict[str, Any]) -> tuple[Self, dict[str, Any]]:
         result: int = -1
-        bound_args: Optional[dict[str, Any]] = None
+        bound_args: dict[str, Any] | None = None
         assert len(self.signatures) > 0
         if len(self.signatures) == 1:
             # Only one signature: call _bind_to_signature() and surface any errors directly
@@ -206,7 +206,7 @@ class Function(ABC):
         self._resolved_fns[signature_idx].validate_call(normalized_args)
         return normalized_args
 
-    def validate_call(self, bound_args: dict[str, Optional['exprs.Expr']]) -> None:
+    def validate_call(self, bound_args: dict[str, 'exprs.Expr' | None]) -> None:
         """Override this to do custom validation of the arguments"""
         assert not self.is_polymorphic
         self.signature.validate_args(bound_args, context=f'in function {self.name!r}')
@@ -252,9 +252,7 @@ class Function(ABC):
 
         return return_type
 
-    def _assemble_callable_args(
-        self, callable: Callable, bound_args: dict[str, 'exprs.Expr']
-    ) -> Optional[dict[str, Any]]:
+    def _assemble_callable_args(self, callable: Callable, bound_args: dict[str, 'exprs.Expr']) -> dict[str, Any] | None:
         """
         Return the kwargs to pass to callable, given bound_args passed to this function.
 
@@ -286,7 +284,7 @@ class Function(ABC):
                 return None
             assert isinstance(arg, exprs.Expr)
 
-            expects_expr: Optional[type[exprs.Expr]] = None
+            expects_expr: type[exprs.Expr] | None = None
             type_hint = callable_type_hints.get(param.name)
             if typing.get_origin(type_hint) is not None:
                 type_hint = typing.get_origin(type_hint)  # Remove type subscript if one exists
@@ -410,12 +408,12 @@ class Function(ABC):
         """Execute the function with the given arguments and return the result."""
         raise NotImplementedError()
 
-    def to_sql(self, fn: Callable[..., Optional[sql.ColumnElement]]) -> Callable[..., Optional[sql.ColumnElement]]:
+    def to_sql(self, fn: Callable[..., sql.ColumnElement | None]) -> Callable[..., sql.ColumnElement | None]:
         """Instance decorator for specifying the SQL translation of this function"""
         self._to_sql = fn
         return fn
 
-    def __default_to_sql(self, *args: Any, **kwargs: Any) -> Optional[sql.ColumnElement]:
+    def __default_to_sql(self, *args: Any, **kwargs: Any) -> sql.ColumnElement | None:
         """The default implementation of SQL translation, which provides no translation"""
         return None
 
@@ -425,7 +423,7 @@ class Function(ABC):
         self._resource_pool = fn
         return fn
 
-    def __default_resource_pool(self) -> Optional[str]:
+    def __default_resource_pool(self) -> str | None:
         return None
 
     def __eq__(self, other: object) -> bool:
@@ -495,7 +493,7 @@ class Function(ABC):
         raise NotImplementedError()
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: str | None, md: dict, binary_obj: bytes) -> Function:
         """
         Create a Function instance from the serialized representation returned by to_store()
         """
@@ -504,12 +502,12 @@ class Function(ABC):
 
 class InvalidFunction(Function):
     fn_dict: dict[str, Any]
-    errormsg: str
+    error_msg: str
 
-    def __init__(self, self_path: str, fn_dict: dict[str, Any], errormsg: str):
+    def __init__(self, self_path: str, fn_dict: dict[str, Any], error_msg: str):
         super().__init__([], self_path)
         self.fn_dict = fn_dict
-        self.errormsg = errormsg
+        self.error_msg = error_msg
 
     def _as_dict(self) -> dict:
         """

pixeltable/func/function_registry.py

@@ -4,7 +4,6 @@ import dataclasses
 import importlib
 import logging
 import sys
-from typing import Optional
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -23,7 +22,7 @@ class FunctionRegistry:
     Function are loaded from the store on demand.
     """
 
-    _instance: Optional[FunctionRegistry] = None
+    _instance: FunctionRegistry | None = None
 
     @classmethod
     def get(cls) -> FunctionRegistry:
@@ -94,7 +93,7 @@ class FunctionRegistry:
         #         stored_fn_md.append(md)
         return list(self.module_fns.values())
 
-    # def get_function(self, *, id: Optional[UUID] = None, fqn: Optional[str] = None) -> Function:
+    # def get_function(self, *, id: UUID | None = None, fqn: str | None = None) -> Function:
     #     assert (id is not None) != (fqn is not None)
     #     if id is not None:
     #         if id not in self.stored_fns_by_id:
@@ -143,7 +142,7 @@ class FunctionRegistry:
             return list(self.type_methods[base_type].values())
         return []
 
-    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Optional[Function]:
+    def lookup_type_method(self, base_type: ts.ColumnType.Type, name: str) -> Function | None:
         """
         Look up a method (or property) by name for a given base type. If no such method is registered, return None.
         """
@@ -151,8 +150,8 @@ class FunctionRegistry:
             return self.type_methods[base_type][name]
         return None
 
-    # def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: Optional[UUID] = None) -> UUID:
-    def create_stored_function(self, pxt_fn: Function, dir_id: Optional[UUID] = None) -> UUID:
+    # def create_function(self, md: schema.FunctionMd, binary_obj: bytes, dir_id: UUID | None = None) -> UUID:
+    def create_stored_function(self, pxt_fn: Function, dir_id: UUID | None = None) -> UUID:
         fn_md, binary_obj = pxt_fn.to_store()
         md = schema.FunctionMd(name=pxt_fn.name, md=fn_md, py_version=sys.version, class_name=pxt_fn.__class__.__name__)
         with env.Env.get().engine.begin() as conn:
@@ -184,7 +183,7 @@ class FunctionRegistry:
             self.stored_fns_by_id[id] = instance
             return instance
 
-    # def create_function(self, fn: Function, dir_id: Optional[UUID] = None, name: Optional[str] = None) -> None:
+    # def create_function(self, fn: Function, dir_id: UUID | None = None, name: str | None = None) -> None:
     #     with env.Env.get().engine.begin() as conn:
     #         _logger.debug(f'Pickling function {name}')
     #         eval_fn_str = cloudpickle.dumps(fn.eval_fn) if fn.eval_fn is not None else None

pixeltable/func/globals.py

@@ -1,14 +1,13 @@
 import importlib
 import inspect
 from types import ModuleType
-from typing import Optional
 
 import pixeltable.exceptions as excs
 
 
-def resolve_symbol(symbol_path: str) -> Optional[object]:
+def resolve_symbol(symbol_path: str) -> object | None:
     path_elems = symbol_path.split('.')
-    module: Optional[ModuleType] = None
+    module: ModuleType | None = None
     i = len(path_elems) - 1
     while i > 0 and module is None:
         try:

pixeltable/func/mcp.py

@@ -0,0 +1,74 @@
+import asyncio
+import inspect
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable import exceptions as excs, type_system as ts
+from pixeltable.func.signature import Parameter
+
+if TYPE_CHECKING:
+    import mcp
+
+
+def mcp_udfs(url: str) -> list['pxt.func.Function']:
+    return asyncio.run(mcp_udfs_async(url))
+
+
+async def mcp_udfs_async(url: str) -> list['pxt.func.Function']:
+    import mcp
+    from mcp.client.streamable_http import streamablehttp_client
+
+    list_tools_result: mcp.types.ListToolsResult | None = None
+    async with (
+        streamablehttp_client(url) as (read_stream, write_stream, _),
+        mcp.ClientSession(read_stream, write_stream) as session,
+    ):
+        await session.initialize()
+        list_tools_result = await session.list_tools()
+    assert list_tools_result is not None
+
+    return [mcp_tool_to_udf(url, tool) for tool in list_tools_result.tools]
+
+
+def mcp_tool_to_udf(url: str, mcp_tool: 'mcp.types.Tool') -> 'pxt.func.Function':
+    import mcp
+    from mcp.client.streamable_http import streamablehttp_client
+
+    async def invoke(**kwargs: Any) -> str:
+        # TODO: Cache session objects rather than creating a new one each time?
+        async with (
+            streamablehttp_client(url) as (read_stream, write_stream, _),
+            mcp.ClientSession(read_stream, write_stream) as session,
+        ):
+            await session.initialize()
+            res = await session.call_tool(name=mcp_tool.name, arguments=kwargs)
+            # TODO Handle image/audio responses?
+            return res.content[0].text  # type: ignore[union-attr]
+
+    if mcp_tool.description is not None:
+        invoke.__doc__ = mcp_tool.description
+
+    input_schema = mcp_tool.inputSchema
+    params = {
+        name: __mcp_param_to_pxt_type(mcp_tool.name, name, param) for name, param in input_schema['properties'].items()
+    }
+    required = input_schema.get('required', [])
+
+    # Ensure that any params not appearing in `required` are nullable.
+    # (A required param might or might not be nullable, since its type might be an 'anyOf' containing a null.)
+    for name in params.keys() - required:
+        params[name] = params[name].copy(nullable=True)
+
+    signature = pxt.func.Signature(
+        return_type=ts.StringType(),  # Return type is always string
+        parameters=[Parameter(name, col_type, inspect.Parameter.KEYWORD_ONLY) for name, col_type in params.items()],
+    )
+
+    return pxt.func.CallableFunction(signatures=[signature], py_fns=[invoke], self_name=mcp_tool.name)
+
+
+def __mcp_param_to_pxt_type(tool_name: str, name: str, param: dict[str, Any]) -> ts.ColumnType:
+    pxt_type = ts.ColumnType.from_json_schema(param)
+    if pxt_type is None:
+        raise excs.Error(f'Unknown type schema for MCP parameter {name!r} of tool {tool_name!r}: {param}')
+    return pxt_type

pixeltable/func/query_template_function.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import inspect
 from functools import reduce
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, overload
 
 from pixeltable import catalog, exceptions as excs, exprs, func, type_system as ts
 
@@ -16,13 +16,13 @@ if TYPE_CHECKING:
 class QueryTemplateFunction(Function):
     """A parameterized query/DataFrame from which an executable DataFrame is created with a function call."""
 
-    template_df: Optional['DataFrame']
-    self_name: Optional[str]
-    _comment: Optional[str]
+    template_df: 'DataFrame' | None
+    self_name: str | None
+    _comment: str | None
 
     @classmethod
     def create(
-        cls, template_callable: Callable, param_types: Optional[list[ts.ColumnType]], path: str, name: str
+        cls, template_callable: Callable, param_types: list[ts.ColumnType] | None, path: str, name: str
     ) -> QueryTemplateFunction:
         # we need to construct a template df and a signature
         py_sig = inspect.signature(template_callable)
@@ -40,11 +40,11 @@ class QueryTemplateFunction(Function):
 
     def __init__(
         self,
-        template_df: Optional['DataFrame'],
+        template_df: 'DataFrame' | None,
         sig: Signature,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        comment: Optional[str] = None,
+        path: str | None = None,
+        name: str | None = None,
+        comment: str | None = None,
     ):
         assert sig is not None
         super().__init__([sig], self_path=path)
@@ -82,7 +82,7 @@ class QueryTemplateFunction(Function):
     def name(self) -> str:
         return self.self_name
 
-    def comment(self) -> Optional[str]:
+    def comment(self) -> str | None:
         return self._comment
 
     def _as_dict(self) -> dict:
@@ -100,11 +100,11 @@ def query(py_fn: Callable) -> QueryTemplateFunction: ...
 
 
 @overload
-def query(*, param_types: Optional[list[ts.ColumnType]] = None) -> Callable[[Callable], QueryTemplateFunction]: ...
+def query(*, param_types: list[ts.ColumnType] | None = None) -> Callable[[Callable], QueryTemplateFunction]: ...
 
 
 def query(*args: Any, **kwargs: Any) -> Any:
-    def make_query_template(py_fn: Callable, param_types: Optional[list[ts.ColumnType]]) -> QueryTemplateFunction:
+    def make_query_template(py_fn: Callable, param_types: list[ts.ColumnType] | None) -> QueryTemplateFunction:
         if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
             # this is a named function in a module
             function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
@@ -127,10 +127,10 @@ def query(*args: Any, **kwargs: Any) -> Any:
 
 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,
+    name: str | None = None,
+    description: str | None = None,
+    parameters: Iterable[str | exprs.ColumnRef] | None = None,
+    limit: int | None = 10,
 ) -> func.QueryTemplateFunction:
     """
     Constructs a retrieval UDF for the given table. The retrieval UDF is a UDF whose parameters are
@@ -157,11 +157,13 @@ def retrieval_udf(
     """
     # Argument validation
     col_refs: list[exprs.ColumnRef]
+    # TODO: get rid of references to ColumnRef internals and replace instead with a public interface
+    col_names = table.columns()
     if parameters is None:
-        col_refs = [table[col_name] for col_name in table.columns if not table[col_name].col.is_computed]
+        col_refs = [table[col_name] for col_name in col_names if not table[col_name].col.is_computed]
     else:
         for param in parameters:
-            if isinstance(param, str) and param not in table.columns:
+            if isinstance(param, str) and param not in col_names:
                 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]