diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 38b45d8b6..2ab658293 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -1462,8 +1462,8 @@ class _SourcePackager: self._packaged: set[int] = set() self._names: dict[str, Any] = {} - def package(self, target: Any) -> str: - entry = self._definition(target, entry=True) + def package(self, target: Any, signature: FunctionSignature) -> str: + entry = self._definition(target, entry=True, signature=signature) parts = [_SOURCE_HEADER] if self._imports: parts.extend(["", *(self._imports[name] for name in sorted(self._imports))]) @@ -1472,12 +1472,20 @@ class _SourcePackager: parts.extend(["", entry, ""]) return "\n".join(parts) - def _definition(self, target: Any, *, entry: bool) -> str: + def _definition( + self, target: Any, *, entry: bool, signature: Optional[FunctionSignature] = None + ) -> str: definition = _definition_node(target, entry=entry) if not inspect.isclass(target) and inspect.getclosurevars(target).nonlocals: raise ValueError( "@udf cannot package functions that capture closure values" ) + namespace = _definition_namespace(target) + reader_name = ( + self._blob_annotations(definition, signature, namespace) + if signature is not None + else None + ) source = ast.unparse(definition) references = _module_references(f"{_SOURCE_HEADER}\n\n{source}\n") dynamic = _namespace_acquisition(definition, references) @@ -1485,9 +1493,10 @@ class _SourcePackager: raise ValueError( f"@udf cannot package dynamic namespace access: {dynamic!r}" ) - namespace = _definition_namespace(target) self._packaged.add(id(target)) self._claim(target.__name__, ("definition", id(target))) + if reader_name is not None: + self._bind(reader_name, BlobReader) unresolved = [] for name in sorted(references): if name == target.__name__: @@ -1508,6 +1517,44 @@ class _SourcePackager: ) return source + def _blob_annotations( + self, + definition: Union[ast.FunctionDef, ast.ClassDef], + signature: FunctionSignature, + namespace: Mapping[str, Any], + ) -> Optional[str]: + """Bind reader inputs to the worker type, including postponed aliases.""" + readers = { + parameter.name: parameter + for parameter in signature.inputs + if parameter.blob_access == "reader" + } + if not readers: + return None + call = definition + if isinstance(definition, ast.ClassDef): + call = next( + ( + member + for member in definition.body + if isinstance(member, ast.FunctionDef) and member.name == "__call__" + ), + None, + ) + if call is None: + return None + # Avoid shadowing globals used by the body or other packaged definitions. + name = "_lance_blob_reader" + while name in namespace or name in self._names or name == definition.name: + name += "_" + for argument in (*call.args.args, *call.args.kwonlyargs): + if argument.arg in readers: + annotation = name + if readers[argument.arg].nullable: + annotation += " | None" + argument.annotation = ast.parse(annotation, mode="eval").body + return name + def _claim(self, name: str, meaning: Any) -> None: claimed = self._names.setdefault(name, meaning) if claimed != meaning: @@ -1845,7 +1892,7 @@ class UdfDefinition: code_files, code_modules = _code_files(code) entry_source = _SourcePackager( code_modules, frozenset(_requirement_name(package) for package in packages) - ).package(function) + ).package(function, signature) if code_files: kind = _PYTHON_BUNDLE_ARTIFACT source = json.dumps( diff --git a/python/python/tests/test_function_packaging.py b/python/python/tests/test_function_packaging.py index 6e1e1cf43..fe86de14e 100644 --- a/python/python/tests/test_function_packaging.py +++ b/python/python/tests/test_function_packaging.py @@ -11,9 +11,10 @@ import json import subprocess import sys import textwrap +import types import warnings from pathlib import Path -from typing import Annotated, Optional +from typing import Annotated, Optional, get_type_hints import pyarrow as pa import pytest @@ -126,6 +127,95 @@ def test_function_without_initialization_keeps_its_wire_form(): assert request["artifact"]["kind"] == "python_callable" +@pytest.mark.parametrize("class_udf", [False, True]) +@pytest.mark.parametrize("postponed", [False, True]) +@pytest.mark.parametrize( + ("annotation", "nullable"), + [ + ("BlobReader", False), + ("Reader", False), + ('"Reader"', False), + ("Reader | None", True), + ("Optional[Reader]", True), + ('Optional["Reader"]', True), + ("MaybeReader", True), + ("lancedb.BlobReader", False), + ("lancedb.BlobReader | None", True), + ], +) +def test_blob_reader_annotations_resolve_in_the_artifact( + helper_package, monkeypatch, class_udf, postponed, annotation, nullable +): + definition = ( + f"class read_blob:\n" + f" def __call__(self, video: {annotation}) -> int:\n" + f" return 0\n" + if class_udf + else f"def read_blob(video: {annotation}) -> int:\n return 0\n" + ) + module = helper_package( + "fx_blob_annotations", + { + "__init__.py": ( + ("from __future__ import annotations\n" if postponed else "") + + "import lancedb\n" + "from typing import Optional\n" + "from lancedb import BlobReader, BlobReader as Reader\n" + "from lancedb.functions import udf\n" + "MaybeReader = Optional[Reader]\n\n" + f"@udf\n{definition}" + ) + }, + ) + request = module.read_blob.registration_request + parameter = request.signature.inputs[0] + assert parameter.arrow_type == "blob_v2" + assert parameter.blob_access == "reader" + assert parameter.nullable is nullable + + # Workers provide lance_blob rather than importing the client SDK's reader. + worker = types.ModuleType("lance_blob") + worker.BlobReader = type("BlobReader", (), {}) + monkeypatch.setitem(sys.modules, "lance_blob", worker) + namespace = {} + exec(compile(artifact(module.read_blob), "user_function.py", "exec"), namespace) + packaged = namespace["read_blob"] + if class_udf: + packaged = packaged.__call__ + assert get_type_hints(packaged) == { + "video": Optional[worker.BlobReader] if nullable else worker.BlobReader, + "return": int, + } + + +def test_blob_reader_annotation_does_not_shadow_body_globals( + helper_package, monkeypatch +): + module = helper_package( + "fx_blob_annotation_globals", + { + "__init__.py": """ + from lancedb import BlobReader + from lancedb.functions import udf + + _lance_blob_reader = 7 + + @udf + def read_blob(video: BlobReader) -> int: + return _lance_blob_reader + """ + }, + ) + worker = types.ModuleType("lance_blob") + worker.BlobReader = type("BlobReader", (), {}) + monkeypatch.setitem(sys.modules, "lance_blob", worker) + namespace = {} + exec(compile(artifact(module.read_blob), "user_function.py", "exec"), namespace) + packaged = namespace["read_blob"] + assert get_type_hints(packaged)["video"] is worker.BlobReader + assert packaged(worker.BlobReader()) == 7 + + @pytest.mark.parametrize( ("init", "message"), [