Merge main and carry patched Lance Arrow through beta.12

python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lancedb-python"
-version = "0.40.0-beta.7"
+version = "0.40.0-beta.8"
 publish = false
 edition.workspace = true
 description = "Python bindings for LanceDB"

python/python/lancedb/functions.py

@@ -20,6 +20,7 @@ import importlib
 import inspect
 import symtable
 import json
+import linecache
 import math
 import re
 import sys
@@ -367,14 +368,20 @@ class FunctionVersion(_RemoteValue):
     metadata: Mapping[str, str]
     disabled: bool
 
-    def __call__(self, **inputs: Any) -> FunctionApplication:
+    def __call__(self, **arguments: Any) -> FunctionApplication:
         """Bind this exact version to named table columns.
 
         Every input must be a direct [lancedb.col][lancedb.expr.col]
-        reference. The returned application is immutable and retains a
-        named-struct output as one binding, so every row's sibling values
-        come from one logical Function evaluation. Map result fields to table
-        columns with
+        reference. A Function that declares initialization fields takes their
+        values as further keyword arguments; they are constants of this
+        binding, fixed for every row, and every instance the binding runs is
+        created with them. A non-nullable field without a default must be
+        given; an omitted nullable field is null, and the Function's own
+        default applies to a null value.
+
+        The returned application is immutable and retains a named-struct
+        output as one binding, so every row's sibling values come from one
+        logical Function evaluation. Map result fields to table columns with
         [FunctionApplication.rename][lancedb.functions.FunctionApplication.rename],
         then pass the application to
         [Table.add_columns][lancedb.table.Table.add_columns].
@@ -390,23 +397,34 @@ class FunctionVersion(_RemoteValue):
         ...     "token_count": "search_token_count",
         ... })
         >>> table.add_columns(application)  # doctest: +SKIP
+        >>> table.add_columns({  # doctest: +SKIP
+        ...     "embedding": embed(text=col("body"), model="small", dimensions=512)
+        ... })
         """
         from lancedb.expr import Expr
 
         parameters = tuple(parameter.name for parameter in self.signature.inputs)
-        missing = [parameter for parameter in parameters if parameter not in inputs]
-        unknown = sorted(set(inputs) - set(parameters))
-        if missing or unknown:
+        fields = {field.name: field for field in self.signature.initialization}
+        missing = [parameter for parameter in parameters if parameter not in arguments]
+        missing_initialization = [
+            name
+            for name, field in fields.items()
+            if not field.nullable and name not in arguments
+        ]
+        unknown = sorted(set(arguments) - set(parameters) - set(fields))
+        if missing or missing_initialization or unknown:
             details = []
             if missing:
                 details.append(f"missing inputs: {missing!r}")
+            if missing_initialization:
+                details.append(f"missing initialization: {missing_initialization!r}")
             if unknown:
-                details.append(f"unknown inputs: {unknown!r}")
-            raise TypeError("invalid Function inputs (" + "; ".join(details) + ")")
+                details.append(f"unknown arguments: {unknown!r}")
+            raise TypeError("invalid Function arguments (" + "; ".join(details) + ")")
 
         bindings = []
         for parameter in parameters:
-            value = inputs[parameter]
+            value = arguments[parameter]
             if not isinstance(value, Expr) or value._column_path is None:
                 raise TypeError(
                     f"Function input {parameter!r} must be a direct col(...) reference"
@@ -418,6 +436,10 @@ class FunctionVersion(_RemoteValue):
                     value={"path": value._column_path},
                 )
             )
+        initialization = {}
+        for name in fields:
+            if name in arguments:
+                initialization[name] = _initialization_value(name, arguments[name])
         return FunctionApplication(
             function=FunctionVersionRef(
                 name=self.name,
@@ -428,9 +450,40 @@ class FunctionVersion(_RemoteValue):
             ),
             inputs=tuple(bindings),
             output=self.signature.output,
+            initialization=initialization,
         )
 
 
+def _initialization_value(name: str, value: Any) -> Any:
+    """One initialization value as plain JSON; the service checks its type."""
+    from lancedb.expr import Expr
+
+    def plain(value):
+        if isinstance(value, Expr):
+            raise TypeError(
+                f"initialization field {name!r} takes a constant, not a column "
+                "expression; every row of a binding shares one initialization"
+            )
+        if value is None or type(value) in (bool, int, str):
+            return value
+        if type(value) is float:
+            if not math.isfinite(value):
+                raise ValueError(
+                    f"initialization field {name!r} must be finite, got {value!r}"
+                )
+            return value
+        if isinstance(value, (list, tuple)):
+            return [plain(child) for child in value]
+        if isinstance(value, Mapping) and all(isinstance(key, str) for key in value):
+            return {key: plain(child) for key, child in value.items()}
+        raise TypeError(
+            f"initialization field {name!r} takes booleans, numbers, strings, "
+            f"lists, and string-keyed mappings, got {type(value).__name__}"
+        )
+
+    return plain(value)
+
+
 class FunctionRegistrationRequest(_RemoteValue):
     """Stable remote registration envelope produced by :func:`udf`.
 
@@ -919,7 +972,12 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
             raise TypeError(
                 "Annotated Function types require exactly one PyArrow DataType"
             )
-        _, base_nullable = _annotation_type(base)
+        # A struct value is written as a dict, which has no Arrow type of its
+        # own; the Annotated metadata supplies it.
+        if base is dict or get_origin(base) is dict:
+            base_nullable = False
+        else:
+            _, base_nullable = _annotation_type(base)
         return arrow_types[0], nullable or base_nullable
 
     if isinstance(annotation, pa.DataType):
@@ -1022,8 +1080,12 @@ def _infer_signature(
     function: Callable[..., Any],
     input_schema: Optional[pa.Schema],
     output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
+    *,
+    method: bool = False,
 ) -> FunctionSignature:
     parameters = _callable_parameters(function)
+    if method:
+        parameters = parameters[1:]
     if (input_schema is None) != (output_schema is None):
         raise ValueError("input_schema and output_schema must be provided together")
 
@@ -1135,10 +1197,11 @@ def _namespace_acquisition(
 
 
 def _module_references(module_source: str) -> set[str]:
-    """Names any scope in `module_source` binds or loads at module scope.
-    Python's own scope analysis on the exact text that ships: free variables
-    belong to an enclosing scope inside the function, and postponed
-    annotations are not runtime loads."""
+    """Names `module_source` loads from module scope, by Python's own scope
+    analysis on the exact text that ships: loads inside any definition, plus
+    what the definition statement itself evaluates at module scope (bases,
+    decorators, defaults). Free variables belong to an enclosing scope inside
+    the definition, and postponed annotations are not runtime loads."""
 
     def visit(table: symtable.SymbolTable, found: set[str]) -> None:
         for symbol in table.get_symbols():
@@ -1150,138 +1213,522 @@ def _module_references(module_source: str) -> set[str]:
             visit(child, found)
 
     found: set[str] = set()
-    for table in symtable.symtable(module_source, "<udf>", "exec").get_children():
+    module = symtable.symtable(module_source, "<udf>", "exec")
+    for symbol in module.get_symbols():
+        if (
+            symbol.is_referenced()
+            and not symbol.is_assigned()
+            and not symbol.is_imported()
+        ):
+            found.add(symbol.get_name())
+    for table in module.get_children():
         visit(table, found)
     return found
 
 
-def _global_source(name: str, value: Any) -> str:
-    """One module-level line that rebinds `name` to `value` in the artifact:
-    an import for modules and importable classes/functions, a literal otherwise."""
-    if isinstance(value, types.ModuleType):
-        if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
-            raise ValueError(
-                f"@udf cannot package dynamic namespace access: {value.__name__!r}"
-            )
+_SOURCE_HEADER = "from __future__ import annotations"
+_SOURCE_ENTRY_FILE = "user_function.py"
+_PYTHON_CALLABLE_ARTIFACT = "python_callable"
+_PYTHON_BUNDLE_ARTIFACT = "python_bundle"
+# Modules the Function image places beside user code.
+_RESERVED_SOURCE_MODULES = frozenset(
+    {"user_function", "lance_source", "lance_callable"}
+)
+# The service accepts bundle paths made of these components only.
+_SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+_REQUIREMENT_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*")
+
+
+def _class_source(cls: type) -> str:
+    """Source of a module-level class.
+
+    `inspect.getsource` finds a class through its module's file, which a
+    notebook cell or a doctest does not have. A method's code object records
+    where the class body really is, so that is consulted first.
+    """
+    for member in vars(cls).values():
+        if isinstance(member, (staticmethod, classmethod)):
+            member = member.__func__
+        code = (
+            getattr(inspect.unwrap(member), "__code__", None)
+            if callable(member)
+            else None
+        )
+        if code is None:
+            continue
+        lines = linecache.getlines(code.co_filename)
         try:
-            imported = importlib.import_module(value.__name__)
-        except ImportError:
-            imported = None
-        if imported is not value:
-            raise TypeError(
-                f"Function source references module {name!r} that does not import "
-                f"as {value.__name__!r}"
-            )
-        return f"import {value.__name__} as {name}"
-    module_name = getattr(value, "__module__", None)
-    qualname = getattr(value, "__qualname__", None)
-    if (
-        isinstance(module_name, str)
-        and isinstance(qualname, str)
-        and module_name != "__main__"
-        and "." not in qualname
-        and "<" not in qualname
-    ):
-        try:
-            imported = getattr(importlib.import_module(module_name), qualname)
-        except (ImportError, AttributeError):
-            imported = None
-        if imported is value:
-            return f"from {module_name} import {qualname} as {name}"
-    return f"{name} = {_literal_source(value)}"
+            module = ast.parse("".join(lines))
+        except SyntaxError:
+            continue
+        for node in module.body:
+            if (
+                isinstance(node, ast.ClassDef)
+                and node.name == cls.__name__
+                and node.lineno <= code.co_firstlineno <= node.end_lineno
+            ):
+                start = min(
+                    [node.lineno]
+                    + [decorator.lineno for decorator in node.decorator_list]
+                )
+                return "".join(lines[start - 1 : node.end_lineno])
+    return inspect.getsource(cls)
 
 
-def _is_recursive_reference(function: Callable[..., Any], name: str) -> bool:
-    """`name` inside the body means the function itself unless the module has
-    since bound it to something else."""
-    if name != function.__name__:
-        return False
-    bound = function.__globals__.get(name, function)
-    if bound is function:
-        return True
-    # The decorator's own result is the one wrapper known to call `function`
-    # unchanged; any other binding may behave differently from a self-call.
-    return type(bound) is UdfDefinition and bound._function is function
-
-
-def _package_source(function: Callable[..., Any]) -> bytes:
-    if not inspect.isfunction(function) or inspect.iscoroutinefunction(function):
-        raise TypeError("@udf requires a synchronous Python function")
+def _definition_node(
+    target: Any, *, entry: bool
+) -> Union[ast.FunctionDef, ast.ClassDef]:
     try:
-        source = textwrap.dedent(inspect.getsource(function))
+        source = textwrap.dedent(
+            _class_source(target)
+            if inspect.isclass(target)
+            else inspect.getsource(target)
+        )
     except (OSError, TypeError) as error:
-        raise ValueError("@udf requires inspectable Python source") from error
+        raise ValueError(
+            f"@udf requires inspectable Python source for {target.__qualname__!r}"
+        ) from error
     module = ast.parse(source)
+    is_class = inspect.isclass(target)
     definitions = [
         node
         for node in module.body
-        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
-        and node.name == function.__name__
+        if isinstance(
+            node,
+            (ast.ClassDef,) if is_class else (ast.FunctionDef, ast.AsyncFunctionDef),
+        )
+        and node.name == target.__name__
     ]
-    if len(definitions) != 1 or not isinstance(definitions[0], ast.FunctionDef):
-        raise ValueError("@udf source must contain exactly one synchronous function")
+    if len(definitions) != 1 or isinstance(definitions[0], ast.AsyncFunctionDef):
+        raise ValueError(
+            "@udf source must contain exactly one "
+            + ("class" if is_class else "synchronous function")
+        )
     definition = definitions[0]
-    if any(not _is_udf_decorator(decorator) for decorator in definition.decorator_list):
-        raise ValueError("@udf cannot package additional Python decorators")
-    definition.decorator_list = []
+    if entry:
+        if any(
+            not _is_udf_decorator(decorator) for decorator in definition.decorator_list
+        ):
+            raise ValueError("@udf cannot package additional Python decorators")
+        definition.decorator_list = []
+    return definition
 
-    closure = inspect.getclosurevars(function)
-    if closure.nonlocals:
-        raise ValueError("@udf cannot package functions that capture closure values")
-    function_source = ast.unparse(definition)
-    module_header = "from __future__ import annotations"
-    references = _module_references(f"{module_header}\n\n{function_source}\n")
-    dynamic = _namespace_acquisition(definition, references)
-    if dynamic:
-        raise ValueError(f"@udf cannot package dynamic namespace access: {dynamic!r}")
-    # Resolve every module-scope reference the way the interpreter would: the
-    # function's own globals first (a module global may shadow a builtin, and
-    # nested scopes are not visible to getclosurevars), then its builtins.
+
+def _definition_namespace(target: Any) -> Mapping[str, Any]:
+    """The module namespace a function or class resolves its globals in."""
+    if inspect.isclass(target):
+        module = sys.modules.get(target.__module__)
+        if module is None:
+            raise ValueError(
+                f"@udf cannot resolve the module of {target.__qualname__!r}"
+            )
+        namespace = vars(module)
+        environment = namespace.get("__builtins__", builtins)
+        environment = (
+            vars(environment)
+            if isinstance(environment, types.ModuleType)
+            else environment
+        )
+    else:
+        namespace = target.__globals__
+        environment = target.__builtins__
     # The artifact runs under the standard builtins; only the exact mapping is
     # provably equivalent (a subclass or copy can change lookups and hooks).
-    if function.__builtins__ is not vars(builtins):
+    if environment is not vars(builtins):
         raise ValueError("@udf cannot package a non-standard builtins environment")
-    globals_source = []
-    unresolved = []
-    for name in sorted(references):
-        if name == function.__name__:
-            if not _is_recursive_reference(function, name):
-                raise ValueError(
-                    f"@udf cannot package {name!r}: the module binds that name to "
-                    "another value, which the artifact's own definition would shadow"
-                )
-            continue
-        if name in function.__globals__:
-            globals_source.append(_global_source(name, function.__globals__[name]))
-        elif hasattr(builtins, name):
-            pass
-        else:
-            unresolved.append(name)
-    if unresolved:
+    return namespace
+
+
+def _is_self_reference(target: Any, namespace: Mapping[str, Any], name: str) -> bool:
+    """`name` inside the definition means the definition itself unless the
+    module has since bound it to something else."""
+    if name != target.__name__:
+        return False
+    bound = namespace.get(name, target)
+    if bound is target:
+        return True
+    # The decorator's own result is the one wrapper known to behave like
+    # `target`; any other binding may differ from the packaged definition.
+    return type(bound) is UdfDefinition and bound._function is target
+
+
+def _packaged_by_source(value: Any) -> bool:
+    """Functions and classes defined in `__main__` -- notebook cells and
+    scripts -- have no module a worker can import, so their source travels."""
+    if not (inspect.isfunction(value) or inspect.isclass(value)):
+        return False
+    if getattr(value, "__module__", None) != "__main__":
+        return False
+    qualname = value.__qualname__
+    if "<" in qualname or "." in qualname:
         raise ValueError(
-            f"@udf source contains unresolved global names: {unresolved!r}"
+            f"@udf cannot package {qualname!r}: only module-level functions and "
+            "classes travel with a Function, not lambdas, closures, or nested "
+            "definitions"
+        )
+    return True
+
+
+def _requirement_name(requirement: str) -> str:
+    requirement = requirement.rpartition("::")[2]
+    match = _REQUIREMENT_NAME.match(requirement.strip())
+    return _normalized_distribution(match.group(0) if match else requirement)
+
+
+def _normalized_distribution(name: str) -> str:
+    return re.sub(r"[-_.]+", "_", name).lower()
+
+
+def _installed_roots() -> tuple[str, ...]:
+    import os
+    import site
+    import sysconfig
+
+    roots = set(site.getsitepackages())
+    roots.add(site.getusersitepackages())
+    paths = sysconfig.get_paths()
+    for key in ("purelib", "platlib", "stdlib", "platstdlib"):
+        if key in paths:
+            roots.add(paths[key])
+    return tuple(os.path.realpath(root) + os.sep for root in roots)
+
+
+def _module_locations(module: types.ModuleType) -> tuple[str, ...]:
+    import os
+
+    file = getattr(module, "__file__", None)
+    if file:
+        return (os.path.realpath(file),)
+    return tuple(os.path.realpath(path) for path in getattr(module, "__path__", ()))
+
+
+class _SourcePackager:
+    """Builds the entry module of a Function artifact.
+
+    The module holds the definition plus exactly the module-level names it
+    references: modules as imports, importable classes and functions as
+    imports, literals inline, and functions and classes defined in
+    ``__main__`` by source, recursively, dependencies first.
+    """
+
+    def __init__(self, code_modules: frozenset[str], requirements: frozenset[str]):
+        self._code_modules = code_modules
+        self._requirements = requirements
+        self._imports: dict[str, str] = {}
+        self._definitions: list[str] = []
+        self._packaged: set[int] = set()
+        self._names: dict[str, Any] = {}
+
+    def package(self, target: Any) -> str:
+        entry = self._definition(target, entry=True)
+        parts = [_SOURCE_HEADER]
+        if self._imports:
+            parts.extend(["", *(self._imports[name] for name in sorted(self._imports))])
+        for definition in self._definitions:
+            parts.extend(["", definition])
+        parts.extend(["", entry, ""])
+        return "\n".join(parts)
+
+    def _definition(self, target: Any, *, entry: bool) -> 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"
+            )
+        source = ast.unparse(definition)
+        references = _module_references(f"{_SOURCE_HEADER}\n\n{source}\n")
+        dynamic = _namespace_acquisition(definition, references)
+        if dynamic:
+            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)))
+        unresolved = []
+        for name in sorted(references):
+            if name == target.__name__:
+                if not _is_self_reference(target, namespace, name):
+                    raise ValueError(
+                        f"@udf cannot package {name!r}: the module binds that name "
+                        "to another value, which the artifact's own definition "
+                        "would shadow"
+                    )
+                continue
+            if name in namespace:
+                self._bind(name, namespace[name])
+            elif not hasattr(builtins, name):
+                unresolved.append(name)
+        if unresolved:
+            raise ValueError(
+                f"@udf source contains unresolved global names: {unresolved!r}"
+            )
+        return source
+
+    def _claim(self, name: str, meaning: Any) -> None:
+        claimed = self._names.setdefault(name, meaning)
+        if claimed != meaning:
+            raise ValueError(
+                f"@udf cannot package {name!r}: packaged definitions bind that "
+                "name to different values"
+            )
+
+    def _bind(self, name: str, value: Any) -> None:
+        if _packaged_by_source(value):
+            if value.__name__ != name:
+                raise ValueError(
+                    f"@udf packages {value.__qualname__!r} by source, so it must be "
+                    f"referenced by its own name, not {name!r}"
+                )
+            if id(value) not in self._packaged:
+                self._definitions.append(self._definition(value, entry=False))
+            self._claim(name, ("definition", id(value)))
+            return
+        line = self._global_source(name, value)
+        self._claim(name, line)
+        self._imports[name] = line
+
+    def _global_source(self, name: str, value: Any) -> str:
+        """One module-level line that rebinds `name` to `value` in the artifact:
+        an import for modules and importable classes/functions, a literal
+        otherwise."""
+        if isinstance(value, types.ModuleType):
+            if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
+                raise ValueError(
+                    f"@udf cannot package dynamic namespace access: {value.__name__!r}"
+                )
+            try:
+                imported = importlib.import_module(value.__name__)
+            except ImportError:
+                imported = None
+            if imported is not value:
+                raise TypeError(
+                    f"Function source references module {name!r} that does not "
+                    f"import as {value.__name__!r}"
+                )
+            self._require_importable(value.__name__)
+            return f"import {value.__name__} as {name}"
+        module_name = getattr(value, "__module__", None)
+        qualname = getattr(value, "__qualname__", None)
+        if (
+            isinstance(module_name, str)
+            and isinstance(qualname, str)
+            and module_name != "__main__"
+            and "." not in qualname
+            and "<" not in qualname
+        ):
+            try:
+                imported = getattr(importlib.import_module(module_name), qualname)
+            except (ImportError, AttributeError):
+                imported = None
+            if imported is value:
+                self._require_importable(module_name)
+                return f"from {module_name} import {qualname} as {name}"
+        return f"{name} = {_literal_source(value)}"
+
+    def _require_importable(self, module_name: str) -> None:
+        """Refuse an import the Function's environment has no way to satisfy:
+        a module that lives in a local source tree, is not shipped with
+        ``code=``, and names no declared package."""
+        top = module_name.partition(".")[0]
+        if (
+            top in self._code_modules
+            or top in sys.stdlib_module_names
+            or top in sys.builtin_module_names
+            or _normalized_distribution(top) in self._requirements
+        ):
+            return
+        module = sys.modules.get(top)
+        locations = () if module is None else _module_locations(module)
+        roots = _installed_roots()
+        if not locations or any(location.startswith(roots) for location in locations):
+            return
+        raise ValueError(
+            f"@udf source imports {module_name!r} from {locations[0]}, a local "
+            "module the Function's environment cannot import; ship it with "
+            f"code=[{top}] or declare the package that provides it in pip/conda"
         )
 
-    parts = [module_header]
-    if globals_source:
-        parts.extend(["", *globals_source])
-    parts.extend(["", function_source, ""])
-    packaged = "\n".join(parts)
-    return packaged.encode("utf-8")
+
+def _code_files(modules: Sequence[Any]) -> tuple[dict[str, str], frozenset[str]]:
+    """The Python source files of the modules and packages in ``code=``,
+    keyed by their path on the Function's import path."""
+    import os
+    from pathlib import Path
+
+    files: dict[str, str] = {}
+    names: set[str] = set()
+    for module in modules:
+        if not isinstance(module, types.ModuleType):
+            raise TypeError(
+                f"code= takes imported modules or packages, got {type(module).__name__}"
+            )
+        name = module.__name__
+        if "." in name:
+            raise ValueError(
+                f"code= takes top-level modules or packages; pass "
+                f"{name.partition('.')[0]!r} instead of {name!r}"
+            )
+        if (
+            name in _RESERVED_SOURCE_MODULES
+            or name == "__main__"
+            or name in sys.stdlib_module_names
+            or not _SOURCE_PATH_COMPONENT.fullmatch(name)
+        ):
+            raise ValueError(f"code= cannot ship module {name!r}")
+        if name in names:
+            raise ValueError(f"code= lists module {name!r} more than once")
+        names.add(name)
+        paths = list(getattr(module, "__path__", ()))
+        if paths:
+            if len(paths) != 1 or not getattr(module, "__file__", None):
+                raise ValueError(
+                    f"code= cannot ship namespace package {name!r}; give it an "
+                    "__init__.py"
+                )
+            root = Path(paths[0])
+            sources = []
+            for directory, subdirectories, filenames in os.walk(root):
+                # Only importable modules travel: data files, caches, and
+                # directories that are not identifiers cannot be imported.
+                subdirectories[:] = sorted(
+                    child
+                    for child in subdirectories
+                    if _SOURCE_PATH_COMPONENT.fullmatch(child)
+                    and child != "__pycache__"
+                )
+                for filename in filenames:
+                    stem, suffix = os.path.splitext(filename)
+                    if suffix == ".py" and _SOURCE_PATH_COMPONENT.fullmatch(stem):
+                        sources.append(Path(directory) / filename)
+            for path in sources:
+                relative = path.relative_to(root).as_posix()
+                files[f"{name}/{relative}"] = _source_text(path)
+        else:
+            file = getattr(module, "__file__", None)
+            if not file or not file.endswith(".py"):
+                raise ValueError(f"code= ships Python source; module {name!r} has none")
+            files[f"{name}.py"] = _source_text(Path(file))
+    return files, frozenset(names)
+
+
+def _source_text(path) -> str:
+    try:
+        return path.read_text(encoding="utf-8")
+    except UnicodeDecodeError as error:
+        raise ValueError(f"code= ships UTF-8 source; {path} is not UTF-8") from error
+
+
+def _initialization_type(data_type: pa.DataType) -> bool:
+    """Types with one unambiguous JSON and SQL literal spelling."""
+    if data_type in (
+        pa.bool_(),
+        pa.int8(),
+        pa.int16(),
+        pa.int32(),
+        pa.int64(),
+        pa.uint8(),
+        pa.uint16(),
+        pa.uint32(),
+        pa.uint64(),
+        pa.float32(),
+        pa.float64(),
+        pa.string(),
+        pa.large_string(),
+    ):
+        return True
+    if (
+        pa.types.is_list(data_type)
+        or pa.types.is_large_list(data_type)
+        or pa.types.is_fixed_size_list(data_type)
+    ):
+        return _initialization_type(data_type.value_type)
+    if pa.types.is_struct(data_type):
+        return all(_initialization_type(field.type) for field in data_type)
+    return False
+
+
+def _initialization_signature(cls: type) -> tuple[FunctionParameter, ...]:
+    """The initialization fields a class Function's ``__init__`` declares."""
+    if cls.__init__ is object.__init__:
+        return ()
+    parameters = tuple(inspect.signature(cls.__init__).parameters.values())[1:]
+    for parameter in parameters:
+        if parameter.kind not in (
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        ):
+            raise TypeError(
+                "Function initialization parameters must be named and "
+                f"non-variadic: {parameter.name!r}"
+            )
+    try:
+        annotations = get_type_hints(cls.__init__, include_extras=True)
+    except Exception as error:
+        raise TypeError(
+            f"failed to resolve Function initialization annotations: {error}"
+        ) from error
+    missing = [
+        parameter.name for parameter in parameters if parameter.name not in annotations
+    ]
+    if missing:
+        raise TypeError(f"missing Function initialization annotations: {missing!r}")
+    fields = []
+    for parameter in parameters:
+        data_type, nullable = _annotation_type(annotations[parameter.name])
+        if not _initialization_type(data_type):
+            raise TypeError(
+                f"Function initialization parameter {parameter.name!r} has type "
+                f"{data_type}; initialization takes booleans, integers, floats, "
+                "strings, and lists or structs of them"
+            )
+        # A parameter with a default may be omitted from a binding; a null
+        # value then takes the default.
+        if parameter.default is not inspect.Parameter.empty:
+            nullable = True
+        fields.append(
+            FunctionParameter(
+                name=parameter.name,
+                arrow_type=_canonical_arrow_field(
+                    pa.field(parameter.name, data_type, nullable=nullable)
+                ),
+                nullable=nullable,
+            )
+        )
+    return tuple(fields)
+
+
+def _class_member(cls: type, name: str) -> Any:
+    for klass in cls.__mro__:
+        if klass is not object and name in vars(klass):
+            return vars(klass)[name]
+    return None
+
+
+def _class_call(cls: type) -> Callable[..., Any]:
+    call = _class_member(cls, "__call__")
+    if not inspect.isfunction(call) or inspect.iscoroutinefunction(call):
+        raise TypeError("a class Function must define a synchronous __call__ method")
+    close = _class_member(cls, "close")
+    if close is not None and (
+        not inspect.isfunction(close) or len(inspect.signature(close).parameters) != 1
+    ):
+        raise TypeError("a class Function's close must be a method taking no arguments")
+    return call
 
 
 class UdfDefinition:
-    """A scalar Python callable prepared for remote Function registration.
+    """A Python callable prepared for remote Function registration.
 
     Instances are created with :func:`udf`. Calling an instance executes the
-    original scalar Python function, which keeps local unit testing ordinary.
-    Remote execution adapts that scalar callable to the internal Arrow batch
-    ABI described by the registration artifact.
+    original Python function, or constructs the original class, which keeps
+    local unit testing ordinary. Remote execution adapts the callable to the
+    internal Arrow batch ABI described by the registration artifact.
     """
 
     def __init__(
         self,
-        function: Callable[..., Any],
+        function: Any,
         *,
         name: Optional[str],
         input_schema: Optional[pa.Schema],
@@ -1292,7 +1739,13 @@ class UdfDefinition:
         gpu: bool = False,
         conda: tuple[str, ...] = (),
         conda_channels: tuple[str, ...] = (),
+        code: tuple[types.ModuleType, ...] = (),
     ):
+        is_class = inspect.isclass(function)
+        if not is_class and (
+            not inspect.isfunction(function) or inspect.iscoroutinefunction(function)
+        ):
+            raise TypeError("@udf requires a synchronous Python function or a class")
         function_name = name or function.__name__
         if not _FUNCTION_NAME.fullmatch(function_name):
             raise ValueError(f"invalid Function name: {function_name!r}")
@@ -1315,8 +1768,43 @@ class UdfDefinition:
             for key, value in environment.items()
         ):
             raise TypeError("Function env keys and values must be strings")
-        signature = _infer_signature(function, input_schema, output_schema)
-        source = _package_source(function)
+        if is_class:
+            if "<" in function.__qualname__ or "." in function.__qualname__:
+                raise ValueError(
+                    "@udf classes must be defined at module level; nested classes "
+                    "would capture their enclosing scope"
+                )
+            signature = _infer_signature(
+                _class_call(function), input_schema, output_schema, method=True
+            )
+            initialization = _initialization_signature(function)
+            shared = sorted(
+                {field.name for field in initialization}
+                & {parameter.name for parameter in signature.inputs}
+            )
+            if shared:
+                raise ValueError(
+                    "Function initialization parameters and inputs share names "
+                    f"{shared!r}; a binding passes both as keyword arguments"
+                )
+            signature = signature._copy(update={"initialization": initialization})
+        else:
+            signature = _infer_signature(function, input_schema, output_schema)
+        code_files, code_modules = _code_files(code)
+        entry_source = _SourcePackager(
+            code_modules, frozenset(_requirement_name(package) for package in packages)
+        ).package(function)
+        if code_files:
+            kind = _PYTHON_BUNDLE_ARTIFACT
+            source = json.dumps(
+                {"files": {**code_files, _SOURCE_ENTRY_FILE: entry_source}},
+                ensure_ascii=False,
+                sort_keys=True,
+                separators=(",", ":"),
+            ).encode("utf-8")
+        else:
+            kind = _PYTHON_CALLABLE_ARTIFACT
+            source = entry_source.encode("utf-8")
         digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
         gpu_marker = _normalize_gpu_marker(gpu)
         runtime = PythonRuntimeSpec(
@@ -1331,7 +1819,7 @@ class UdfDefinition:
         self._request = FunctionRegistrationRequest(
             name=function_name,
             artifact=FunctionArtifactRequest(
-                kind="python_callable",
+                kind=kind,
                 digest=digest,
                 entrypoint=function.__name__,
                 content=FunctionArtifactContent(
@@ -1346,7 +1834,11 @@ class UdfDefinition:
             signature=signature,
             runtime=runtime,
         )
-        functools.update_wrapper(self, function)
+        # A class's own __dict__ holds its methods; copying it would shadow
+        # this wrapper's attributes.
+        functools.update_wrapper(
+            self, function, updated=() if is_class else functools.WRAPPER_UPDATES
+        )
 
     @property
     def registration_request(self) -> FunctionRegistrationRequest:
@@ -1438,6 +1930,7 @@ def udf(
     gpu: bool = False,
     conda: tuple[str, ...] | list[str] = (),
     conda_channels: tuple[str, ...] | list[str] = (),
+    code: Sequence[types.ModuleType] = (),
 ) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
 
 
@@ -1453,8 +1946,9 @@ def udf(
     gpu: bool = False,
     conda: tuple[str, ...] | list[str] = (),
     conda_channels: tuple[str, ...] | list[str] = (),
+    code: Sequence[types.ModuleType] = (),
 ):
-    """Prepare a scalar Python callable for remote Function registration.
+    """Prepare a Python function or class for remote Function registration.
 
     Input and output signatures are inferred from supported annotations. For
     Arrow types annotations cannot express precisely, pass ``input_schema``
@@ -1462,10 +1956,54 @@ def udf(
     named-struct field may be nullable; Enterprise preserves the struct's
     validity when the result is expanded into sibling columns.
 
+    **Functions and classes.** A decorated function is called once per row,
+    or once per batch when its parameters are annotated ``pyarrow.Array`` or
+    it takes one ``pyarrow.RecordBatch``. A decorated class is the same
+    callable with state: each remote instance calls ``__init__`` once, calls
+    ``__call__`` for every row or batch of every input it processes, and calls
+    ``close()``, if defined, once when it retires. Load models, open clients,
+    and build rate limiters in ``__init__``; every batch of the instance then
+    reuses them. An instance serves many batches of one binding, so state is
+    shared across them, but not across instances or workers.
+
+    **Initialization.** The annotated parameters of a class's ``__init__``
+    are the Function's initialization fields. Their values belong to each
+    binding, not to the Function: ``function(text=col("body"), model="small")``
+    binds column ``body`` and initializes every instance with ``model="small"``,
+    and another column can bind the same Function version with other values.
+    Initialization fields take booleans, integers, floats, strings, and lists
+    or structs of them (``Annotated[dict, pa.struct(...)]``). A parameter with
+    a default may be omitted from a binding; a null value then takes the
+    default. Use Secrets, not initialization, for credentials.
+
+    **What travels with the Function.** The artifact is a snapshot of the
+    decorated source plus exactly the module-level names it references:
+
+    - modules, and classes and functions importable from an installed
+      package, become imports; the package must be declared in ``pip`` or
+      ``conda``;
+    - literals (``None``, booleans, numbers, strings, bytes, tuples of them)
+      are inlined;
+    - functions and classes defined in ``__main__``, such as notebook cells,
+      are packaged by source, together with what they reference;
+    - modules and packages listed in ``code`` are shipped as Python source and
+      imported normally, so helpers shared by several Functions live in one
+      place.
+
+    A reference to a module in a local source tree that is neither shipped
+    with ``code`` nor provided by a declared package is rejected at
+    registration. Closures, lambdas, nested definitions, mutable module-level
+    objects, code that reaches the module namespace another way
+    (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a
+    non-standard ``__builtins__`` are rejected where they can be seen and
+    otherwise unsupported. Changing any packaged source, including a module in
+    ``code``, produces a new Function version; the environment built from
+    ``pip`` or ``conda`` is reused.
+
     Parameters
     ----------
-    function : Callable, optional
-        The synchronous scalar callable to package.
+    function : function or class, optional
+        The synchronous callable, or the class whose instances are called.
     name : str, optional
         The remote Function name. Defaults to the callable name.
     input_schema : pyarrow.Schema, optional
@@ -1483,20 +2021,17 @@ def udf(
     env : mapping of str to str, optional
         Environment variables included in the Function definition. Not for
         credentials -- these are ordinary configuration, stored with the
-        Function and visible wherever it is.
+        Function and visible wherever it is. Prefer initialization for
+        values that differ between bindings.
     python_version : str, optional
         Remote Python major/minor version. Defaults to the client version.
     gpu : bool, default False
         Whether every remote execution requires a GPU. The execution platform
         selects one compatible GPU for each worker. The requirement is part of
         the immutable Function version.
-
-    The packaged artifact is a snapshot: the function source plus exactly
-    the module-level names it references (modules as imports, importable
-    classes and functions as imports, literals inline). Code that reaches the
-    module namespace another way -- ``globals()``/``eval``, ``sys.modules``,
-    ``builtins`` -- is rejected where it can be seen and otherwise
-    unsupported; closures and a non-standard ``__builtins__`` are rejected.
+    code : sequence of module, optional
+        Top-level modules or packages whose importable ``.py`` files ship
+        with the Function and are importable by their own names.
 
     Returns
     -------
@@ -1519,9 +2054,19 @@ def udf(
     ...     return value * 2
     >>> gpu_score.registration_request.runtime.gpu
     True
+    >>> @udf
+    ... class scale:
+    ...     def __init__(self, factor: float, offset: float = 0.0):
+    ...         self.factor, self.offset = factor, offset
+    ...     def __call__(self, value: float) -> float:
+    ...         return value * self.factor + self.offset
+    >>> scale(factor=2.0)(1.5)
+    3.0
+    >>> [field.name for field in scale.registration_request.signature.initialization]
+    ['factor', 'offset']
     """
 
-    def decorate(target: Callable[..., Any]) -> UdfDefinition:
+    def decorate(target: Any) -> UdfDefinition:
         return UdfDefinition(
             target,
             name=name,
@@ -1533,6 +2078,7 @@ def udf(
             gpu=gpu,
             conda=tuple(conda),
             conda_channels=tuple(conda_channels),
+            code=tuple(code),
         )
 
     if function is None:

python/python/lancedb/table.py

@@ -771,12 +771,13 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
         if json_storage is not None:
             # Labelled through metadata rather than pa.json_(), which only exists on
             # newer PyArrow; Lance reads the extension name off the field either way.
-            return pa.field(
-                field.name,
-                json_storage,
-                field.nullable,
-                {"ARROW:extension:name": "arrow.json"},
-            )
+            # The other metadata keys mirror the table field's: Lance swaps the name
+            # back to lance.json on write and then requires the field to match the
+            # stored one exactly, including the empty ``ARROW:extension:metadata``
+            # that pyarrow records for a ``pa.json_()`` column.
+            metadata = dict(target_field.metadata or {})
+            metadata[b"ARROW:extension:name"] = b"arrow.json"
+            return pa.field(field.name, json_storage, field.nullable, metadata)
     if pa.types.is_struct(target_field.type):
         if pa.types.is_struct(field.type):
             new_type = pa.struct(

python/python/tests/test_blob.py

@@ -921,6 +921,70 @@ def test_fetch_blob_files_lazy_read():
     handles = table.fetch_blob_files("image", [by_id[1]])
     assert len(handles) == 1
     assert handles[0].read() == payload
+    handles[0].close()
+    assert handles[0].closed
+    with pytest.raises(RuntimeError, match="already closed"):
+        handles[0].read()
+
+
+def test_blob_file_finalizer_and_close_on_runtime_worker():
+    # A tokio panic in IOBase.__del__ is printed but does not fail the add call.
+    # Run in a child process so the assertion catches that panic on stderr.
+    script = textwrap.dedent(
+        """\
+        import gc
+        import threading
+
+        import lancedb
+        import pyarrow as pa
+
+        db = lancedb.connect("memory:///")
+        schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+        table = db.create_table("finalize", schema=schema)
+        table.add([{"id": 1, "image": b"first"}])
+        hits = table.search().with_row_id(True).limit(1).to_arrow()
+        row_id = hits["_rowid"][0].as_py()
+        gc_handle, explicit_handle = table.fetch_blob_files("image", [row_id, row_id])
+
+        gc.disable()
+        cycle = [gc_handle]
+        cycle.append(cycle)
+        del gc_handle, cycle
+
+        observed = []
+        def on_progress(_):
+            if not observed:
+                observed.append((
+                    threading.current_thread() is not threading.main_thread(),
+                    explicit_handle.closed,
+                    gc.collect(),
+                ))
+                explicit_handle.close()
+
+        table.add([{"id": 2, "image": b"second"}], progress=on_progress)
+        gc.enable()
+
+        assert len(observed) == 1, observed
+        assert observed[0][0] is True, observed
+        assert observed[0][1] is False, observed
+        assert observed[0][2] > 0, observed
+        assert explicit_handle.closed
+        explicit_handle.close()  # Close remains idempotent after the worker call.
+        try:
+            explicit_handle.read()
+        except RuntimeError as error:
+            assert "already closed" in str(error)
+        else:
+            raise AssertionError("a closed blob file was readable")
+        print("worker close completed")
+        """
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, check=False
+    )
+    assert result.returncode == 0, result.stderr
+    assert "Cannot start a runtime from within a runtime" not in result.stderr
+    assert "worker close completed" in result.stdout
 
 
 def test_fetch_blob_files_null_alignment():

python/python/tests/test_first_class_function_slice1.py

@@ -169,7 +169,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
 
     with pytest.raises(TypeError, match=r"missing inputs: \['text'\]"):
         version()
-    with pytest.raises(TypeError, match=r"unknown inputs: \['body'\]"):
+    with pytest.raises(TypeError, match=r"unknown arguments: \['body'\]"):
         version(text=col("text"), body=col("body"))
     with pytest.raises(TypeError, match="direct col"):
         version(text=col("text").lower())

python/python/tests/test_function_packaging.py

@@ -0,0 +1,521 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+"""Class Functions, initialization, and the source that travels with a Function."""
+
+from __future__ import annotations
+
+import base64
+import importlib
+import json
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+from typing import Annotated, Optional
+
+import pyarrow as pa
+import pytest
+
+from lancedb import col
+from lancedb.functions import FunctionVersion, udf
+
+FIXTURES = (
+    Path(__file__).parents[3]
+    / "rust"
+    / "lancedb"
+    / "tests"
+    / "fixtures"
+    / "first_class_functions"
+    / "v1"
+)
+
+
+def artifact(definition) -> bytes:
+    return base64.b64decode(definition.registration_request.artifact.content.data)
+
+
+def bundle_files(definition) -> dict[str, str]:
+    assert definition.registration_request.artifact.kind == "python_bundle"
+    return json.loads(artifact(definition))["files"]
+
+
+@pytest.fixture
+def helper_package(tmp_path, monkeypatch):
+    """A local package on sys.path, the way a project's own helpers are."""
+
+    def write(name: str, files: dict[str, str]):
+        root = tmp_path / name
+        for path, source in files.items():
+            target = root / path
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_text(textwrap.dedent(source))
+        monkeypatch.syspath_prepend(str(tmp_path))
+        for module in [
+            module for module in sys.modules if module.split(".")[0] == name
+        ]:
+            del sys.modules[module]
+        importlib.invalidate_caches()
+        return importlib.import_module(name)
+
+    return write
+
+
+@udf
+class scale:
+    def __init__(
+        self,
+        factor: float,
+        label: str,
+        offset: float = 0.0,
+        tags: Optional[list[str]] = None,
+        retry: Annotated[dict, pa.struct([("attempts", pa.int32())])] = None,
+    ):
+        self.factor, self.offset = factor, offset
+
+    def __call__(self, value: float) -> float:
+        return value * self.factor + self.offset
+
+
+class NoCall:
+    pass
+
+
+class BadClose:
+    def __call__(self, value: int) -> int:
+        return value
+
+    def close(self, force):
+        pass
+
+
+def test_class_function_declares_inputs_output_and_initialization():
+    signature = scale.registration_request.signature
+    assert [(p.name, p.arrow_type) for p in signature.inputs] == [("value", "float64")]
+    assert [
+        (field.name, field.arrow_type, field.nullable)
+        for field in signature.initialization
+    ] == [
+        ("factor", "float64", False),
+        ("label", "utf8", False),
+        ("offset", "float64", True),
+        ("tags", "list<utf8>", True),
+        (
+            "retry",
+            '{"fields":[{"name":"attempts","nullable":true,"type":{"type":"int32"}}],'
+            '"type":"struct"}',
+            True,
+        ),
+    ]
+    # Calling the definition constructs the class, so local tests are ordinary.
+    assert scale(factor=2.0, label="x")(1.5) == 3.0
+    assert scale.registration_request.artifact.entrypoint == "scale"
+    source = artifact(scale).decode()
+    assert source.startswith("from __future__ import annotations\n")
+    assert "class scale:" in source and "@udf" not in source
+
+
+def test_function_without_initialization_keeps_its_wire_form():
+    @udf
+    def double(value: int) -> int:
+        return value * 2
+
+    request = json.loads(double.registration_request.to_canonical_json())
+    assert "initialization" not in request["signature"]
+    assert request["artifact"]["kind"] == "python_callable"
+
+
+@pytest.mark.parametrize(
+    ("init", "message"),
+    [
+        ("def __init__(self, blob: bytes): pass", "initialization takes booleans"),
+        ("def __init__(self, *args: int): pass", "named and non-variadic"),
+        ("def __init__(self, factor): pass", "missing Function initialization"),
+        ("def __init__(self, value: int): pass", "share names \\['value'\\]"),
+    ],
+)
+def test_class_function_rejects_unsupported_initialization(init, message):
+    namespace = {}
+    exec(
+        textwrap.dedent(
+            f"""
+            class bad:
+                {init}
+                def __call__(self, value: int) -> int:
+                    return value
+            """
+        ),
+        namespace,
+    )
+    with pytest.raises((TypeError, ValueError), match=message):
+        udf(namespace["bad"])
+
+
+def test_class_function_requires_call_and_a_plain_close():
+    with pytest.raises(TypeError, match="synchronous __call__"):
+        udf(NoCall)
+    with pytest.raises(TypeError, match="close must be a method taking no"):
+        udf(BadClose)
+
+
+def test_nested_class_and_closure_are_rejected_with_the_reason():
+    def build():
+        class nested:
+            def __call__(self, value: int) -> int:
+                return value
+
+        return nested
+
+    with pytest.raises(ValueError, match="defined at module level"):
+        udf(build())
+
+    factor = 3
+
+    with pytest.raises(ValueError, match="capture closure values"):
+
+        @udf
+        def closes_over(value: int) -> int:
+            return value * factor
+
+
+def test_binding_passes_initialization_as_constants():
+    version = FunctionVersion.from_json(
+        json.dumps(
+            {
+                **json.loads(
+                    (FIXTURES / "remote_function_version.canonical.json").read_text()
+                ),
+                "signature": {
+                    "inputs": [
+                        {"name": "text", "arrow_type": "utf8", "nullable": True}
+                    ],
+                    "output": {
+                        "kind": "scalar",
+                        "arrow_type": "list<float32>",
+                        "nullable": False,
+                    },
+                    "initialization": [
+                        {"name": "model", "arrow_type": "utf8", "nullable": False},
+                        {"name": "dimensions", "arrow_type": "int32", "nullable": True},
+                    ],
+                },
+            }
+        )
+    )
+    application = version(text=col("body"), model="small", dimensions=(512))
+    assert dict(application.initialization) == {"model": "small", "dimensions": 512}
+    assert json.loads(application.to_canonical_json())["initialization"] == {
+        "dimensions": 512,
+        "model": "small",
+    }
+    assert version(text=col("body"), model="small").initialization == {"model": "small"}
+
+    with pytest.raises(TypeError, match=r"missing initialization: \['model'\]"):
+        version(text=col("body"))
+    with pytest.raises(TypeError, match="takes a constant"):
+        version(text=col("body"), model=col("model"))
+    with pytest.raises(ValueError, match="must be finite"):
+        version(text=col("body"), model="small", dimensions=float("nan"))
+    with pytest.raises(TypeError, match="takes booleans"):
+        version(text=col("body"), model=object())
+
+
+def test_code_packages_ship_as_a_bundle_and_version_with_their_source(
+    tmp_path, helper_package
+):
+    files = {
+        "__init__.py": "",
+        "rate.py": """
+            class TokenBucket:
+                def __init__(self, tokens_per_minute):
+                    self.tokens_per_minute = tokens_per_minute
+        """,
+        "__pycache__/rate.cpython-312.pyc": "",
+        "fixtures/data.json": "{}",
+        "not-a-module.py": "",
+    }
+    helper_package("fx_helpers", files)
+    (tmp_path / "fx_helpers_functions.py").write_text(
+        textwrap.dedent(
+            """
+            import fx_helpers
+            from fx_helpers.rate import TokenBucket
+            from lancedb import udf
+
+
+            @udf(code=[fx_helpers])
+            class limited:
+                def __init__(self, tpm: int):
+                    self.limiter = TokenBucket(tpm)
+
+                def __call__(self, value: int) -> int:
+                    return value
+            """
+        )
+    )
+
+    def define():
+        sys.modules.pop("fx_helpers_functions", None)
+        return importlib.import_module("fx_helpers_functions").limited
+
+    first = define()
+    shipped = bundle_files(first)
+    assert sorted(shipped) == [
+        "fx_helpers/__init__.py",
+        "fx_helpers/rate.py",
+        "user_function.py",
+    ]
+    assert (
+        "from fx_helpers.rate import TokenBucket as TokenBucket"
+        in (shipped["user_function.py"])
+    )
+
+    # Only the helper changes: a new artifact digest, the same environment.
+    helper_package(
+        "fx_helpers", {**files, "rate.py": textwrap.dedent(files["rate.py"]) + "# v2\n"}
+    )
+    second = define()
+    assert (
+        first.registration_request.artifact.digest
+        != second.registration_request.artifact.digest
+    )
+    assert first.registration_request.runtime == second.registration_request.runtime
+
+
+def test_code_rejects_submodules_duplicates_and_non_modules(helper_package):
+    helpers = helper_package("fx_rejects", {"__init__.py": "", "sub.py": ""})
+    import fx_rejects.sub
+
+    for code, message in [
+        ([fx_rejects.sub], "top-level modules"),
+        ([helpers, helpers], "more than once"),
+        (["fx_rejects"], "imported modules"),
+        ([json], "cannot ship module 'json'"),
+    ]:
+        with pytest.raises((TypeError, ValueError), match=message):
+
+            @udf(code=code)
+            def f(value: int) -> int:
+                return value
+
+
+def test_local_module_outside_code_is_rejected_at_registration(
+    tmp_path, helper_package
+):
+    helper_package(
+        "fx_local", {"__init__.py": "def normalize(value):\n    return value\n"}
+    )
+
+    def define(name, options):
+        (tmp_path / f"{name}.py").write_text(
+            textwrap.dedent(
+                f"""
+                from fx_local import normalize
+                from lancedb import udf
+
+
+                @udf({options})
+                def {name}(value: int) -> int:
+                    return normalize(value)
+                """
+            )
+        )
+        return getattr(importlib.import_module(name), name)
+
+    with pytest.raises(ValueError, match=r"local module.*code=\[fx_local\]"):
+        define("fx_uses_local", "")
+
+    # A declared package that provides the module is trusted to install it.
+    installed = define("fx_uses_installed", "pip=['fx-local==1.0']")
+    assert "from fx_local import normalize" in artifact(installed).decode()
+
+
+MMLB_HELPERS = {
+    "__init__.py": "",
+    "rate.py": """
+        class TokenBucket:
+            constructed = 0
+
+            def __init__(self, tokens_per_minute):
+                TokenBucket.constructed += 1
+                self.tokens_per_minute = tokens_per_minute
+                self.acquired = 0
+
+            def acquire(self, tokens):
+                self.acquired += tokens
+    """,
+    "embedding.py": """
+        from fx_mmlb.rate import TokenBucket
+
+
+        class OpenAIEmbedder:
+            def __init__(self, model: str, tpm_limit: int, max_attempts: int = 6):
+                self.model = model
+                self.limiter = TokenBucket(tpm_limit)
+                self.max_attempts = max_attempts
+
+            def embed(self, texts):
+                self.limiter.acquire(sum(len(text) for text in texts))
+                return [[float(len(text))] for text in texts]
+    """,
+}
+
+MMLB_FUNCTIONS = """
+import pyarrow as pa
+
+import fx_mmlb
+from fx_mmlb.embedding import OpenAIEmbedder
+from lancedb import udf
+
+TEXT = pa.schema([pa.field("text", pa.string())])
+EMBEDDING = pa.list_(pa.field("item", pa.float64(), nullable=False))
+
+
+@udf(code=[fx_mmlb], input_schema=TEXT, output_schema=EMBEDDING)
+class openai_embedding_batch(OpenAIEmbedder):
+    def __call__(self, text: pa.Array) -> pa.Array:
+        item = pa.field("item", pa.float64(), nullable=False)
+        return pa.array(self.embed(text.to_pylist()), pa.list_(item))
+
+
+@udf(code=[fx_mmlb])
+class openai_embedding(OpenAIEmbedder):
+    def __call__(self, text: str) -> list[float]:
+        return self.embed([text])[0]
+"""
+
+
+def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package):
+    """The MMLB rate limiter, shared by a row and a batch Function: both are
+    initialized through the same helper, so neither path can skip it."""
+    helper_package("fx_mmlb", MMLB_HELPERS)
+    (tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS)
+    functions = importlib.import_module("fx_mmlb_functions")
+
+    signatures = [
+        functions.openai_embedding_batch.registration_request.signature,
+        functions.openai_embedding.registration_request.signature,
+    ]
+    for signature in signatures:
+        assert [field.name for field in signature.initialization] == [
+            "model",
+            "tpm_limit",
+            "max_attempts",
+        ]
+
+    # Each bundle imports on its own, in a process that has never seen the
+    # client's modules, the way the worker's image code layer does.
+    for definition, call in [
+        (
+            functions.openai_embedding_batch,
+            "instance(pa.array(['ab', 'abc'])).to_pylist()",
+        ),
+        (functions.openai_embedding, "[instance('ab'), instance('abc')]"),
+    ]:
+        worker = tmp_path / definition.__name__
+        for path, source in bundle_files(definition).items():
+            (worker / path).parent.mkdir(parents=True, exist_ok=True)
+            (worker / path).write_text(source)
+        probe = textwrap.dedent(
+            f"""
+            import sys
+            sys.path.insert(0, {str(worker)!r})
+            import pyarrow as pa
+            import user_function
+            from fx_mmlb.rate import TokenBucket
+            cls = getattr(user_function, {definition.__name__!r})
+            instance = cls(model="small", tpm_limit=1000)
+            print(repr(({call}, TokenBucket.constructed, instance.limiter.acquired)))
+            """
+        )
+        result = subprocess.run(
+            [sys.executable, "-I", "-c", probe],
+            capture_output=True,
+            text=True,
+            cwd=worker,
+        )
+        assert result.returncode == 0, result.stderr
+        assert result.stdout.strip() == "([[2.0], [3.0]], 1, 5)"
+
+
+def test_main_definitions_travel_by_source(tmp_path):
+    """Notebook cells and scripts define helpers in __main__, which a worker
+    cannot import; they are packaged by source, dependencies first."""
+    script = tmp_path / "notebook.py"
+    script.write_text(
+        textwrap.dedent(
+            """
+            import json
+            import math
+            from dataclasses import dataclass
+
+            from lancedb import udf
+
+            SCALE = 2
+
+
+            @dataclass
+            class Settings:
+                factor: int
+
+
+            def helper(value):
+                return math.floor(value) * SCALE
+
+
+            class Base:
+                def run(self, value):
+                    return helper(value) * self.settings.factor
+
+
+            @udf
+            class scaled(Base):
+                def __init__(self, factor: int):
+                    self.settings = Settings(factor)
+
+                def __call__(self, value: float) -> int:
+                    return self.run(value)
+
+
+            print(json.dumps(scaled.registration_request.to_canonical_json()))
+            """
+        )
+    )
+    result = subprocess.run(
+        [sys.executable, str(script)], capture_output=True, text=True
+    )
+    assert result.returncode == 0, result.stderr
+    request = json.loads(json.loads(result.stdout))
+    assert request["artifact"]["kind"] == "python_callable"
+    source = base64.b64decode(request["artifact"]["content"]["data"]).decode()
+    assert source.index("def helper") < source.index("class Base")
+    assert source.index("class Base") < source.index("class scaled")
+    assert source.index("class Settings") < source.index("class scaled")
+    assert "@dataclass" in source
+    assert "SCALE = 2" in source
+
+    namespace = {}
+    exec(compile(source, "user_function.py", "exec"), namespace)
+    assert namespace["scaled"](factor=3)(2.7) == 12
+
+
+def test_main_lambdas_are_rejected(tmp_path):
+    script = tmp_path / "lambda.py"
+    script.write_text(
+        textwrap.dedent(
+            """
+            from lancedb import udf
+
+            double = lambda value: value * 2
+
+            @udf
+            def uses_lambda(value: int) -> int:
+                return double(value)
+            """
+        )
+    )
+    result = subprocess.run(
+        [sys.executable, str(script)], capture_output=True, text=True
+    )
+    assert result.returncode != 0
+    assert "not lambdas, closures, or nested definitions" in result.stderr

python/python/tests/test_remote_db.py

@@ -2470,6 +2470,10 @@ def test_remote_blob_files_are_lazy_seekable_handles():
         assert alpha.read_range(1, 3) == b"lph"
         gamma.seek(2)
         assert gamma.read() == b"mma"
+        alpha.close()
+        assert alpha.closed
+        with pytest.raises(RuntimeError, match="already closed"):
+            alpha.read_range(0, 1)
 
 
 def test_remote_blob_fetch_accepts_query_table():

python/python/tests/test_util.py

@@ -948,6 +948,55 @@ def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
     ]
 
 
+def test_cast_to_target_schema_labels_json_text_with_table_extension_metadata():
+    # A table created from pa.json_() stores the empty ARROW:extension:metadata that
+    # pyarrow exports next to the name. Lance keeps that key when it turns the
+    # arrow.json label back into lance.json on write and then requires the field to
+    # match the stored one exactly, so the label has to mirror the table's metadata.
+    target = pa.schema(
+        [
+            pa.field(
+                "payload",
+                pa.large_binary(),
+                metadata={
+                    b"ARROW:extension:name": b"lance.json",
+                    b"ARROW:extension:metadata": b"",
+                },
+            ),
+            pa.field(
+                "docs",
+                pa.list_(
+                    pa.field(
+                        "item",
+                        pa.large_binary(),
+                        metadata={b"ARROW:extension:name": b"lance.json"},
+                    )
+                ),
+            ),
+        ]
+    )
+    data = pa.table(
+        {
+            "payload": pa.array(['{"a": 1}', None], type=pa.string()),
+            "docs": pa.array([['{"b": 2}'], []], type=pa.list_(pa.string())),
+        }
+    )
+
+    output = _cast_to_target_schema(data.to_reader(), target).read_all()
+
+    payload = output.schema.field("payload")
+    assert payload.type == pa.string()
+    assert payload.metadata == {
+        b"ARROW:extension:name": b"arrow.json",
+        b"ARROW:extension:metadata": b"",
+    }
+    docs_item = output.schema.field("docs").type.value_field
+    assert docs_item.type == pa.string()
+    assert docs_item.metadata == {b"ARROW:extension:name": b"arrow.json"}
+    assert output["payload"].to_pylist() == ['{"a": 1}', None]
+    assert output["docs"].to_pylist() == [['{"b": 2}'], []]
+
+
 def test_cast_to_target_schema_coerces_nested_binary_blob():
     data = pa.table(
         {

python/src/runtime.rs

@@ -197,6 +197,27 @@ pub fn block_on<F: std::future::Future>(fut: F) -> F::Output {
     get_runtime().block_on(fut)
 }
 
+/// Run a detached task, including when the caller is already on a runtime worker.
+/// Keep it visible to [`shutdown`] until it finishes.
+pub fn spawn_background<F>(fut: F)
+where
+    F: Future<Output = ()> + Send + 'static,
+{
+    let guard = OutstandingGuard::new();
+    let task = async move {
+        let _guard = guard;
+        fut.await;
+    };
+    match runtime::Handle::try_current() {
+        Ok(handle) => {
+            let _ = handle.spawn(task);
+        }
+        Err(_) => {
+            let _ = get_runtime().spawn(task);
+        }
+    }
+}
+
 /// Gracefully quiesce the shared runtime, meant to run at normal process exit.
 ///
 /// Waits (bounded by `timeout`) for [`OUTSTANDING`] to reach zero -- i.e.

python/src/table.rs

@@ -1,8 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+};
 
-use crate::runtime::{block_on, future_into_py};
+use crate::runtime::{block_on, future_into_py, spawn_background};
 use crate::{
     connection::Connection,
     error::PythonErrorExt,
@@ -680,11 +686,23 @@ impl From<lancedb::table::DropColumnsResult> for DropColumnsResult {
 #[pyclass(name = "BlobFile")]
 pub struct PyBlobFile {
     inner: Arc<BlobFile>,
+    closed: AtomicBool,
+}
+
+impl PyBlobFile {
+    fn ensure_open(&self) -> PyResult<()> {
+        if self.closed.load(Ordering::Acquire) {
+            Err(PyRuntimeError::new_err("blob file is already closed"))
+        } else {
+            Ok(())
+        }
+    }
 }
 
 #[pymethods]
 impl PyBlobFile {
     fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
+        self_.ensure_open()?;
         let inner = self_.inner.clone();
         let py = self_.py();
         let bytes = py
@@ -694,6 +712,7 @@ impl PyBlobFile {
     }
 
     pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+        self_.ensure_open()?;
         let inner = self_.inner.clone();
         future_into_py(self_.py(), async move {
             let bytes = inner
@@ -705,7 +724,20 @@ impl PyBlobFile {
     }
 
     fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
+        if self_.closed.swap(true, Ordering::AcqRel) {
+            return Ok(());
+        }
         let inner = self_.inner.clone();
+        if tokio::runtime::Handle::try_current().is_ok() {
+            // IOBase.__del__ can call close while cyclic GC runs on a worker.
+            // The status changes immediately; release Lance's async state there.
+            spawn_background(async move {
+                if let Err(error) = inner.close().await {
+                    log::warn!("blob close failed: {error}");
+                }
+            });
+            return Ok(());
+        }
         self_
             .py()
             .detach(move || block_on(async move { inner.close().await }))
@@ -713,13 +745,11 @@ impl PyBlobFile {
     }
 
     fn is_closed(self_: PyRef<'_, Self>) -> bool {
-        let inner = self_.inner.clone();
-        self_
-            .py()
-            .detach(move || block_on(async move { inner.is_closed().await }))
+        self_.closed.load(Ordering::Acquire)
     }
 
     fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
+        self_.ensure_open()?;
         let inner = self_.inner.clone();
         self_
             .py()
@@ -728,6 +758,7 @@ impl PyBlobFile {
     }
 
     fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
+        self_.ensure_open()?;
         let inner = self_.inner.clone();
         self_
             .py()
@@ -741,6 +772,7 @@ impl PyBlobFile {
 
     /// Read a blob-local byte range without moving the cursor.
     fn read_range(self_: PyRef<'_, Self>, offset: u64, length: usize) -> PyResult<Py<PyBytes>> {
+        self_.ensure_open()?;
         let end = offset
             .checked_add(length as u64)
             .ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
@@ -753,6 +785,7 @@ impl PyBlobFile {
     }
 
     fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
+        self_.ensure_open()?;
         let inner = self_.inner.clone();
         let py = self_.py();
         let bytes = py
@@ -1455,6 +1488,7 @@ impl Table {
                 .map(|handle| {
                     handle.map(|file| PyBlobFile {
                         inner: Arc::new(file),
+                        closed: AtomicBool::new(false),
                     })
                 })
                 .collect::<Vec<_>>())

python/vendor/lance-arrow/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lance-arrow"
-version = "13.0.0-beta.8"
+version = "13.0.0-beta.12"
 edition = "2024"
 authors = ["Lance Devs <dev@lance.org>"]
 license = "Apache-2.0"

python/vendor/lance-arrow/README.lancedb.md

@@ -1,5 +1,5 @@
-This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.8`
-(commit `3c20b7875f621c929fdcecabe74dbca507ec0f0d`), licensed under Apache-2.0.
+This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.12`
+(commit `5d2946fa452c61cb018144f3200892b6dd4a8b4f`), licensed under Apache-2.0.
 
 LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps
 until a Lance release includes the fix. The patch is selected in the root

python/vendor/lance-arrow/src/json.rs

@@ -7,13 +7,15 @@
 use std::convert::TryFrom;
 use std::sync::Arc;
 
-use arrow_array::builder::{LargeBinaryBuilder, StringBuilder};
+use arrow_array::builder::{GenericStringBuilder, LargeBinaryBuilder};
 use arrow_array::cast::AsArray;
 use arrow_array::{
-    Array, ArrayRef, FixedSizeListArray, LargeBinaryArray, LargeListArray, LargeStringArray,
-    ListArray, MapArray, RecordBatch, StringArray, StructArray,
+    Array, ArrayRef, FixedSizeListArray, GenericStringArray, LargeBinaryArray, LargeListArray,
+    LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray,
 };
 use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema};
+use jsonb::OwnedJsonb;
+use jsonb::jsonpath::{JsonPath, Selector};
 
 use crate::ARROW_EXT_NAME_KEY;
 
@@ -76,6 +78,72 @@ pub fn has_arrow_json_fields(field: &ArrowField) -> bool {
     }
 }
 
+/// The Arrow representation holding a JSON field's values.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum JsonEncoding {
+    /// JSONB bytes: [`JSON_EXT_NAME`] over `LargeBinary`. Lance stores JSON
+    /// this way.
+    Jsonb,
+    /// JSON text: [`ARROW_JSON_EXT_NAME`] over `Utf8` or `LargeUtf8`. Callers
+    /// write JSON this way, and reads return it this way by default.
+    Text,
+}
+
+impl JsonEncoding {
+    /// The encoding of `field`'s values, or `None` when `field` does not hold
+    /// JSON.
+    pub fn of_field(field: &ArrowField) -> Option<Self> {
+        if is_json_field(field) {
+            Some(Self::Jsonb)
+        } else if is_arrow_json_field(field) {
+            Some(Self::Text)
+        } else {
+            None
+        }
+    }
+}
+
+/// The values of a JSON column, readable as JSONB or as text whichever
+/// encoding they arrived in.
+///
+/// Code that computes on JSON, such as tokenizers and indices, reads values
+/// through this type, so its behavior does not depend on whether they come
+/// from storage or from a caller.
+#[derive(Debug, Clone, Copy)]
+pub struct JsonValues<'a> {
+    encoding: JsonEncoding,
+    array: &'a ArrayRef,
+}
+
+impl<'a> JsonValues<'a> {
+    /// Read `array` as the values of `field`, or `None` when `field` does not
+    /// hold JSON.
+    pub fn try_new(field: &ArrowField, array: &'a ArrayRef) -> Option<Self> {
+        JsonEncoding::of_field(field).map(|encoding| Self { encoding, array })
+    }
+
+    pub fn encoding(&self) -> JsonEncoding {
+        self.encoding
+    }
+
+    /// The values as JSONB. Text values are parsed, so invalid JSON fails.
+    pub fn to_jsonb(&self) -> Result<LargeBinaryArray, ArrowError> {
+        match self.encoding {
+            JsonEncoding::Jsonb => Ok(self.array.as_binary::<i64>().clone()),
+            JsonEncoding::Text => Ok(JsonArray::try_from(self.array.clone())?.into_inner()),
+        }
+    }
+
+    /// The values as JSON text: the input array when it already holds text,
+    /// otherwise `LargeUtf8` decoded from JSONB.
+    pub fn to_text(&self) -> ArrayRef {
+        match self.encoding {
+            JsonEncoding::Jsonb => Arc::new(decode_jsonb_array::<i64>(self.array.as_binary())),
+            JsonEncoding::Text => self.array.clone(),
+        }
+    }
+}
+
 /// Create a JSON field with the appropriate extension metadata
 pub fn json_field(name: &str, nullable: bool) -> ArrowField {
     let mut field = ArrowField::new(name, DataType::LargeBinary, nullable);
@@ -150,28 +218,13 @@ impl JsonArray {
             return Ok(None);
         }
 
-        let jsonb_bytes = self.inner.value(i);
-        get_json_path(jsonb_bytes, path).map_err(|e| {
-            ArrowError::InvalidArgumentError(format!("Failed to extract JSONPath: {}", e))
-        })
+        let path = parse_json_path(path)?;
+        Ok(select_json_path(self.inner.value(i), &path)?.map(|value| value.to_string()))
     }
 
     /// Convert to Arrow string array (JSON as UTF-8)
     pub fn to_arrow_json(&self) -> ArrayRef {
-        let mut builder = arrow_array::builder::StringBuilder::new();
-
-        for i in 0..self.inner.len() {
-            if self.inner.is_null(i) {
-                builder.append_null();
-            } else {
-                let jsonb_bytes = self.inner.value(i);
-                let json_str = decode_json(jsonb_bytes);
-                builder.append_value(&json_str);
-            }
-        }
-
-        // Return as UTF-8 string array (Arrow represents JSON as strings)
-        Arc::new(builder.finish())
+        Arc::new(decode_jsonb_array::<i32>(&self.inner))
     }
 
     pub fn len(&self) -> usize {
@@ -200,23 +253,7 @@ impl TryFrom<&StringArray> for JsonArray {
     type Error = ArrowError;
 
     fn try_from(array: &StringArray) -> Result<Self, Self::Error> {
-        let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len());
-
-        for i in 0..array.len() {
-            if array.is_null(i) {
-                builder.append_null();
-            } else {
-                let json_str = array.value(i);
-                let encoded = encode_json(json_str).map_err(|e| {
-                    ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e))
-                })?;
-                builder.append_value(&encoded);
-            }
-        }
-
-        Ok(Self {
-            inner: builder.finish(),
-        })
+        encode_json_text_array(array)
     }
 }
 
@@ -232,24 +269,40 @@ impl TryFrom<&LargeStringArray> for JsonArray {
     type Error = ArrowError;
 
     fn try_from(array: &LargeStringArray) -> Result<Self, Self::Error> {
-        let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len());
+        encode_json_text_array(array)
+    }
+}
 
-        for i in 0..array.len() {
-            if array.is_null(i) {
-                builder.append_null();
-            } else {
-                let json_str = array.value(i);
+fn encode_json_text_array<O: OffsetSizeTrait>(
+    array: &GenericStringArray<O>,
+) -> Result<JsonArray, ArrowError> {
+    let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len());
+    for value in array {
+        match value {
+            Some(json_str) => {
                 let encoded = encode_json(json_str).map_err(|e| {
                     ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e))
                 })?;
                 builder.append_value(&encoded);
             }
+            None => builder.append_null(),
         }
-
-        Ok(Self {
-            inner: builder.finish(),
-        })
     }
+    Ok(JsonArray {
+        inner: builder.finish(),
+    })
+}
+
+fn decode_jsonb_array<O: OffsetSizeTrait>(array: &LargeBinaryArray) -> GenericStringArray<O> {
+    let mut builder =
+        GenericStringBuilder::<O>::with_capacity(array.len(), array.value_data().len());
+    for value in array {
+        match value {
+            Some(jsonb_bytes) => builder.append_value(decode_json(jsonb_bytes)),
+            None => builder.append_null(),
+        }
+    }
+    builder.finish()
 }
 
 impl TryFrom<ArrayRef> for JsonArray {
@@ -293,21 +346,46 @@ pub fn decode_json(jsonb_bytes: &[u8]) -> String {
     raw_jsonb.to_string()
 }
 
-/// Extract JSONPath value from JSONB
-fn get_json_path(
-    jsonb_bytes: &[u8],
-    path: &str,
-) -> Result<Option<String>, Box<dyn std::error::Error>> {
-    let json_path = jsonb::jsonpath::parse_json_path(path.as_bytes())?;
-    let raw_jsonb = jsonb::RawJsonb::new(jsonb_bytes);
-    let mut selector = jsonb::jsonpath::Selector::new(raw_jsonb);
+/// Parse a JSONPath expression.
+pub fn parse_json_path(path: &str) -> Result<JsonPath<'_>, ArrowError> {
+    jsonb::jsonpath::parse_json_path(path.as_bytes())
+        .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}")))
+}
 
-    let values = selector.select_values(&json_path)?;
-    if values.is_empty() {
-        Ok(None)
-    } else {
-        Ok(Some(values[0].to_string()))
-    }
+/// Select the value at `path` in a JSONB value.
+///
+/// A path matching several values selects them together as one JSON array,
+/// so no match is silently dropped.
+pub fn select_json_path(
+    jsonb_bytes: &[u8],
+    path: &JsonPath<'_>,
+) -> Result<Option<OwnedJsonb>, ArrowError> {
+    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
+        .select_value(path)
+        .map_err(|e| {
+            ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}"))
+        })
+}
+
+/// Whether `path` matches any value in a JSONB value.
+pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result<bool, ArrowError> {
+    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
+        .exists(path)
+        .map_err(|e| {
+            ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}"))
+        })
+}
+
+/// Select every value matched by `path` in a JSONB value.
+pub fn select_json_path_values(
+    jsonb_bytes: &[u8],
+    path: &JsonPath<'_>,
+) -> Result<Vec<OwnedJsonb>, ArrowError> {
+    Selector::new(jsonb::RawJsonb::new(jsonb_bytes))
+        .select_values(path)
+        .map_err(|e| {
+            ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}"))
+        })
 }
 
 /// Convert an Arrow JSON field to Lance JSON field (with JSONB storage)
@@ -532,25 +610,9 @@ fn convert_lance_json_array(
 ) -> Result<(ArrowField, ArrayRef, bool), ArrowError> {
     convert_json_array(field, array, &|field, array| {
         if is_json_field(field) {
-            let binary_array = array
-                .as_any()
-                .downcast_ref::<LargeBinaryArray>()
-                .expect("Lance JSON field must be LargeBinaryArray");
-            let mut builder = StringBuilder::new();
-
-            for i in 0..binary_array.len() {
-                if binary_array.is_null(i) {
-                    builder.append_null();
-                } else {
-                    let jsonb_bytes = binary_array.value(i);
-                    let json_str = decode_json(jsonb_bytes);
-                    builder.append_value(&json_str);
-                }
-            }
-
             Ok(Some((
                 lance_json_to_arrow_json(field),
-                Arc::new(builder.finish()) as ArrayRef,
+                Arc::new(decode_jsonb_array::<i32>(array.as_binary())) as ArrayRef,
             )))
         } else {
             Ok(None)
@@ -677,6 +739,41 @@ mod tests {
 
         let age = json_array.json_path(1, "$.user.age").unwrap();
         assert_eq!(age, None);
+
+        // Several matches are selected together, as the SQL JSON functions do.
+        let json_array =
+            JsonArray::try_from_iter(vec![Some(r#"{"a": [{"b": 1}, {"b": 2}]}"#)]).unwrap();
+        let bs = json_array.json_path(0, "$.a[*].b").unwrap();
+        assert_eq!(bs, Some("[1,2]".to_string()));
+    }
+
+    #[test]
+    fn test_json_values_read_either_encoding() {
+        let docs = vec![Some(r#"{"a":1}"#), None];
+        let text: ArrayRef = Arc::new(StringArray::from(docs.clone()));
+        let jsonb: ArrayRef = Arc::new(JsonArray::try_from_iter(docs).unwrap().into_inner());
+        let mut text_field = ArrowField::new("j", DataType::Utf8, true);
+        text_field.set_metadata(std::collections::HashMap::from([(
+            ARROW_EXT_NAME_KEY.to_string(),
+            ARROW_JSON_EXT_NAME.to_string(),
+        )]));
+        let jsonb_field = json_field("j", true);
+
+        let from_text = JsonValues::try_new(&text_field, &text).unwrap();
+        let from_jsonb = JsonValues::try_new(&jsonb_field, &jsonb).unwrap();
+        assert_eq!(from_text.encoding(), JsonEncoding::Text);
+        assert_eq!(from_jsonb.encoding(), JsonEncoding::Jsonb);
+
+        assert_eq!(&from_text.to_jsonb().unwrap(), jsonb.as_binary::<i64>());
+        assert_eq!(&from_jsonb.to_jsonb().unwrap(), jsonb.as_binary::<i64>());
+        assert_eq!(&from_text.to_text(), &text);
+        let decoded = from_jsonb.to_text();
+        let decoded = decoded.as_string::<i64>();
+        assert_eq!(decoded.value(0), r#"{"a":1}"#);
+        assert!(decoded.is_null(1));
+
+        let plain = ArrowField::new("j", DataType::Utf8, true);
+        assert!(JsonValues::try_new(&plain, &text).is_none());
     }
 
     #[test]
@@ -1301,12 +1398,7 @@ mod tests {
         // Invalid JSONPath syntax should return error
         let result = json_array.json_path(0, "invalid path without $");
         assert!(result.is_err());
-        assert!(
-            result
-                .unwrap_err()
-                .to_string()
-                .contains("Failed to extract JSONPath")
-        );
+        assert!(result.unwrap_err().to_string().contains("Invalid JSONPath"));
     }
 
     #[test]

python/vendor/lance-arrow/src/lib.rs

@@ -472,7 +472,11 @@ pub fn iter_str_array(arr: &dyn Array) -> Box<dyn Iterator<Item = Option<&str>>
     match arr.data_type() {
         DataType::Utf8 => Box::new(arr.as_string::<i32>().iter()),
         DataType::LargeUtf8 => Box::new(arr.as_string::<i64>().iter()),
-        _ => panic!("Expecting Utf8 or LargeUtf8, found {:?}", arr.data_type()),
+        DataType::Utf8View => Box::new(arr.as_string_view().iter()),
+        _ => panic!(
+            "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}",
+            arr.data_type()
+        ),
     }
 }