mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-25 23:48:23 +00:00
fix: package ordinary @udf bodies and emit only the V1 type grammar (#4044)
Registering a real (embedding) Function failed on the client for three
reasons:
- `_package_source` treated `inspect.getclosurevars().unbound` as
"unresolved globals"; CPython puts attribute names there, so any body
with `np.linalg.norm(...)` or `body.split()` was rejected. Module-scope
references now come from Python's own scope analysis (`symtable`) over
the function source, recursively, and each is resolved the way the
interpreter would: the function's globals first (a module global may
shadow a builtin), then builtins. Free variables of nested scopes stay
lexical; postponed annotations are not runtime loads. A genuinely
missing global still fails.
- `_canonical_arrow_type` emitted spellings the server's frozen grammar
rejects (`fixed_size_list<T>[n]`, `timestamp[us]`, `struct<...>`,
zero-sized lists). It now emits exactly the grammar, with the server's
`fixed_size_list<item, size>` form, and the Rust declaration planner
parses that form too.
A shared golden
(`tests/fixtures/first_class_functions/v1/arrow_types.json`)
enumerates every grammar type, nested forms and rejected spellings; the
Python emitter and Rust parser are tested against it, and the same file
is under test in sophon. Packaging tests execute the shipped artifact in
a fresh namespace.
Contract changes (hence `breaking-change`):
- `@udf` now rejects namespace acquisition structurally
(`globals()`/`eval`/... by name, plus `import
sys`/`builtins`/`importlib`/`inspect` inside the body), requires the
function's captured `__builtins__` to be the standard mapping itself
(identity, so neither lookups nor implicit hooks such as `__import__`
can differ), rejects module globals that are namespace-bearing modules
(`builtins`, `sys`, ...), and treats the function's own name as
recursion only when the module binds it to the function or to the exact
`UdfDefinition` the decorator produced; it resolves module globals
through the function's real namespace (a module global may shadow a
builtin) and ships importable classes/functions as imports.
- List outputs must declare a non-nullable, metadata-free child named
`item` (`pa.list_(pa.field("item", t, nullable=False))`); that is what
the grammar means, and pyarrow's default nullable child was being
silently collapsed into it.
Contract, stated in the `udf` docstring: the artifact is a snapshot of
the function source plus exactly the module names it references.
Reaching the module namespace by another route is rejected where a
static packager can see it and is otherwise unsupported; there is no
dynamic-access detection beyond that.
This commit is contained in:
@@ -12,10 +12,13 @@ expression-backed refresh job.
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import builtins
|
||||
import base64
|
||||
import functools
|
||||
import hashlib
|
||||
import importlib
|
||||
import inspect
|
||||
import symtable
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
@@ -489,59 +492,58 @@ _FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
|
||||
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
_GRAMMAR_PRIMITIVES = (
|
||||
(pa.bool_(), "bool"),
|
||||
(pa.int8(), "int8"),
|
||||
(pa.int16(), "int16"),
|
||||
(pa.int32(), "int32"),
|
||||
(pa.int64(), "int64"),
|
||||
(pa.uint8(), "uint8"),
|
||||
(pa.uint16(), "uint16"),
|
||||
(pa.uint32(), "uint32"),
|
||||
(pa.uint64(), "uint64"),
|
||||
(pa.float16(), "float16"),
|
||||
(pa.float32(), "float32"),
|
||||
(pa.float64(), "float64"),
|
||||
(pa.string(), "utf8"),
|
||||
(pa.binary(), "binary"),
|
||||
(pa.date32(), "date32"),
|
||||
(pa.date64(), "date64"),
|
||||
)
|
||||
|
||||
|
||||
def _canonical_arrow_type(data_type: pa.DataType) -> str:
|
||||
primitive_types = (
|
||||
(pa.bool_(), "bool"),
|
||||
(pa.int8(), "int8"),
|
||||
(pa.int16(), "int16"),
|
||||
(pa.int32(), "int32"),
|
||||
(pa.int64(), "int64"),
|
||||
(pa.uint8(), "uint8"),
|
||||
(pa.uint16(), "uint16"),
|
||||
(pa.uint32(), "uint32"),
|
||||
(pa.uint64(), "uint64"),
|
||||
(pa.float16(), "float16"),
|
||||
(pa.float32(), "float32"),
|
||||
(pa.float64(), "float64"),
|
||||
(pa.string(), "utf8"),
|
||||
(pa.large_utf8(), "large_utf8"),
|
||||
(pa.binary(), "binary"),
|
||||
(pa.large_binary(), "large_binary"),
|
||||
(pa.date32(), "date32"),
|
||||
(pa.date64(), "date64"),
|
||||
)
|
||||
for candidate, name in primitive_types:
|
||||
"""The server's V1 Function type grammar. Anything outside it is rejected
|
||||
here rather than at registration."""
|
||||
for candidate, name in _GRAMMAR_PRIMITIVES:
|
||||
if data_type == candidate:
|
||||
return name
|
||||
if pa.types.is_fixed_size_binary(data_type):
|
||||
return f"fixed_size_binary[{data_type.byte_width}]"
|
||||
if pa.types.is_list(data_type):
|
||||
return f"list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
if pa.types.is_large_list(data_type):
|
||||
return f"large_list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
if pa.types.is_fixed_size_list(data_type):
|
||||
if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
|
||||
prefix = "list" if pa.types.is_list(data_type) else "large_list"
|
||||
return f"{prefix}<{_canonical_list_item(data_type)}>"
|
||||
if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
|
||||
return (
|
||||
f"fixed_size_list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
f"[{data_type.list_size}]"
|
||||
f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
|
||||
)
|
||||
if pa.types.is_struct(data_type):
|
||||
fields = ",".join(
|
||||
f"{field.name}:{_canonical_arrow_type(field.type)}" for field in data_type
|
||||
)
|
||||
return f"struct<{fields}>"
|
||||
if pa.types.is_timestamp(data_type):
|
||||
timezone = f",tz={data_type.tz}" if data_type.tz is not None else ""
|
||||
return f"timestamp[{data_type.unit}{timezone}]"
|
||||
if pa.types.is_time32(data_type) or pa.types.is_time64(data_type):
|
||||
return f"time[{data_type.unit}]"
|
||||
if pa.types.is_duration(data_type):
|
||||
return f"duration[{data_type.unit}]"
|
||||
if pa.types.is_decimal(data_type):
|
||||
bit_width = data_type.bit_width
|
||||
return f"decimal{bit_width}({data_type.precision},{data_type.scale})"
|
||||
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
|
||||
|
||||
|
||||
def _canonical_list_item(data_type: pa.DataType) -> str:
|
||||
"""The grammar names only the item type; it always means a non-nullable
|
||||
child called `item`, so any other child metadata cannot be represented."""
|
||||
child = data_type.value_field
|
||||
if child.name != "item" or child.nullable or child.metadata:
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: list items must be a "
|
||||
f"non-nullable field named 'item', got {child}"
|
||||
)
|
||||
return _canonical_arrow_type(child.type)
|
||||
|
||||
|
||||
def _list_of(item: pa.DataType) -> pa.DataType:
|
||||
return pa.list_(pa.field("item", item, nullable=False))
|
||||
|
||||
|
||||
def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
|
||||
nullable = False
|
||||
origin = get_origin(annotation)
|
||||
@@ -589,7 +591,7 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
|
||||
value_type, value_nullable = _annotation_type(arguments[0])
|
||||
if value_nullable:
|
||||
raise TypeError("nullable Function list elements are not supported")
|
||||
return pa.list_(value_type), nullable
|
||||
return _list_of(value_type), nullable
|
||||
raise TypeError(f"unsupported Function annotation: {annotation!r}")
|
||||
|
||||
|
||||
@@ -736,6 +738,104 @@ def _literal_source(value: Any) -> str:
|
||||
)
|
||||
|
||||
|
||||
_DYNAMIC_NAMESPACE_ACCESS = frozenset(
|
||||
{"globals", "locals", "vars", "eval", "exec", "compile", "__import__"}
|
||||
)
|
||||
# Modules that hand out namespaces (`sys.modules`, `builtins`, importers,
|
||||
# introspection). The artifact's module namespace holds only the names it was
|
||||
# packaged with, so reaching around it cannot be represented.
|
||||
_NAMESPACE_MODULES = frozenset(
|
||||
{"sys", "builtins", "importlib", "inspect", "gc", "ctypes", "types"}
|
||||
)
|
||||
|
||||
|
||||
def _namespace_acquisition(
|
||||
definition: ast.FunctionDef, references: set[str]
|
||||
) -> list[str]:
|
||||
found = set(references & _DYNAMIC_NAMESPACE_ACCESS)
|
||||
for node in ast.walk(definition):
|
||||
if isinstance(node, ast.Import):
|
||||
found.update(
|
||||
alias.name
|
||||
for alias in node.names
|
||||
if alias.name.split(".")[0] in _NAMESPACE_MODULES
|
||||
)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module:
|
||||
if node.module.split(".")[0] in _NAMESPACE_MODULES:
|
||||
found.add(node.module)
|
||||
return sorted(found)
|
||||
|
||||
|
||||
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."""
|
||||
|
||||
def visit(table: symtable.SymbolTable, found: set[str]) -> None:
|
||||
for symbol in table.get_symbols():
|
||||
if symbol.is_global() and (
|
||||
symbol.is_referenced() or symbol.is_declared_global()
|
||||
):
|
||||
found.add(symbol.get_name())
|
||||
for child in table.get_children():
|
||||
visit(child, found)
|
||||
|
||||
found: set[str] = set()
|
||||
for table in symtable.symtable(module_source, "<udf>", "exec").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}"
|
||||
)
|
||||
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)}"
|
||||
|
||||
|
||||
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")
|
||||
@@ -760,23 +860,46 @@ def _package_source(function: Callable[..., Any]) -> bytes:
|
||||
closure = inspect.getclosurevars(function)
|
||||
if closure.nonlocals:
|
||||
raise ValueError("@udf cannot package functions that capture closure values")
|
||||
if closure.unbound:
|
||||
raise ValueError(
|
||||
f"@udf source contains unresolved global names: {sorted(closure.unbound)!r}"
|
||||
)
|
||||
globals_source = []
|
||||
for name, value in sorted(closure.globals.items()):
|
||||
if isinstance(value, types.ModuleType):
|
||||
globals_source.append(f"import {value.__name__} as {name}")
|
||||
else:
|
||||
globals_source.append(f"{name} = {_literal_source(value)}")
|
||||
|
||||
function_source = ast.unparse(definition)
|
||||
parts = ["from __future__ import annotations"]
|
||||
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.
|
||||
# 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):
|
||||
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:
|
||||
raise ValueError(
|
||||
f"@udf source contains unresolved global names: {unresolved!r}"
|
||||
)
|
||||
|
||||
parts = [module_header]
|
||||
if globals_source:
|
||||
parts.extend(["", *globals_source])
|
||||
parts.extend(["", function_source, ""])
|
||||
return "\n".join(parts).encode("utf-8")
|
||||
packaged = "\n".join(parts)
|
||||
return packaged.encode("utf-8")
|
||||
|
||||
|
||||
class UdfDefinition:
|
||||
@@ -923,6 +1046,13 @@ def udf(
|
||||
python_version : str, optional
|
||||
Remote Python major/minor version. Defaults to the client 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.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UdfDefinition
|
||||
|
||||
@@ -3,7 +3,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import contextlib
|
||||
import functools
|
||||
import importlib.util
|
||||
import types
|
||||
from datetime import date
|
||||
import http.server
|
||||
import json
|
||||
from pathlib import Path
|
||||
@@ -16,6 +21,9 @@ import pytest
|
||||
import lancedb
|
||||
from lancedb.functions import UdfDefinition, udf
|
||||
|
||||
THRESHOLD = 20
|
||||
_CACHE = None
|
||||
|
||||
|
||||
FIXTURES = (
|
||||
Path(__file__).parents[3]
|
||||
@@ -71,9 +79,396 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
_assert_no_secret_values(request)
|
||||
|
||||
|
||||
def _run_packaged(definition, *args):
|
||||
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
|
||||
source = base64.b64decode(definition.registration_request.artifact.content.data)
|
||||
namespace: dict = {}
|
||||
exec(compile(source, "<udf>", "exec"), namespace)
|
||||
return namespace[definition.registration_request.artifact.entrypoint](*args)
|
||||
|
||||
|
||||
def test_udf_packages_attribute_access_and_body_imports():
|
||||
@udf
|
||||
def word_norm(body: str) -> float:
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
words = body.split()
|
||||
except AttributeError as error:
|
||||
raise ValueError(str(error)) from error
|
||||
return float(np.linalg.norm([len(w) for w in words]))
|
||||
|
||||
assert _run_packaged(word_norm, "aa bb") == pytest.approx(8**0.5)
|
||||
|
||||
|
||||
def test_udf_packages_module_globals_and_global_caches():
|
||||
@udf
|
||||
def label(value: int) -> str:
|
||||
return "big" if value >= THRESHOLD else "small"
|
||||
|
||||
assert _run_packaged(label, 21) == "big"
|
||||
|
||||
@udf
|
||||
def cached(value: int) -> int:
|
||||
global _CACHE
|
||||
if _CACHE is None:
|
||||
_CACHE = 40
|
||||
return _CACHE + value
|
||||
|
||||
assert _run_packaged(cached, 2) == 42
|
||||
|
||||
|
||||
def test_udf_annotations_are_not_runtime_names():
|
||||
@udf
|
||||
def identity(value: date) -> date:
|
||||
return value
|
||||
|
||||
assert _run_packaged(identity, date(2026, 8, 25)) == date(2026, 8, 25)
|
||||
|
||||
|
||||
def test_udf_nested_scopes_resolve_lexically():
|
||||
@udf
|
||||
def score(value: int) -> int:
|
||||
offset = 2
|
||||
|
||||
def add_offset() -> int:
|
||||
return value + offset
|
||||
|
||||
return add_offset() + sum(v for v in [0])
|
||||
|
||||
assert _run_packaged(score, 3) == 5
|
||||
|
||||
|
||||
def test_udf_resolves_module_globals_before_builtins(tmp_path):
|
||||
module_path = tmp_path / "shadowing_udfs.py"
|
||||
module_path.write_text(
|
||||
"max = 7\n"
|
||||
"len = lambda _: 99\n"
|
||||
"\n"
|
||||
"def uses_literal_shadow(value: int) -> int:\n"
|
||||
" def nested() -> int:\n"
|
||||
" return max\n"
|
||||
" return nested() + value\n"
|
||||
"\n"
|
||||
"def uses_callable_shadow(value: int) -> int:\n"
|
||||
" def nested() -> int:\n"
|
||||
" return len([1])\n"
|
||||
" return nested() + value\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("shadowing_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# The module's `max = 7` is what the interpreter would use, so it ships.
|
||||
assert _run_packaged(udf(module.uses_literal_shadow), 1) == 8
|
||||
# A callable global cannot ship; it must not be silently swapped for the builtin.
|
||||
with pytest.raises(TypeError, match="unsupported global value of type function"):
|
||||
udf(module.uses_callable_shadow)
|
||||
|
||||
|
||||
def test_canonical_arrow_type_is_exactly_the_grammar():
|
||||
from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
|
||||
|
||||
golden = json.loads(
|
||||
(
|
||||
Path(__file__).parents[3]
|
||||
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
|
||||
).read_text()
|
||||
)
|
||||
primitives = [
|
||||
case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
|
||||
]
|
||||
assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
|
||||
for outside in [
|
||||
pa.timestamp("us"),
|
||||
pa.decimal128(10, 2),
|
||||
pa.large_string(),
|
||||
pa.large_binary(),
|
||||
pa.binary(4),
|
||||
pa.duration("s"),
|
||||
pa.struct([pa.field("a", pa.int32())]),
|
||||
pa.list_(pa.float32(), 0),
|
||||
pa.list_(pa.timestamp("us")),
|
||||
]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(outside)
|
||||
|
||||
|
||||
def test_udf_nested_annotations_are_postponed_in_the_artifact():
|
||||
@udf
|
||||
def score(value: int) -> int:
|
||||
def identity(item: date) -> date:
|
||||
return item
|
||||
|
||||
identity(date(2026, 8, 25))
|
||||
return value
|
||||
|
||||
assert _run_packaged(score, 3) == 3
|
||||
|
||||
|
||||
def test_udf_ships_globals_the_body_deletes():
|
||||
@udf
|
||||
def clear(value: int) -> int:
|
||||
global _CACHE
|
||||
del _CACHE
|
||||
return value
|
||||
|
||||
assert _run_packaged(clear, 3) == 3
|
||||
|
||||
|
||||
def test_udf_rejects_a_module_global_that_does_not_import_as_itself(tmp_path):
|
||||
module_path = tmp_path / "fake_module_udfs.py"
|
||||
module_path.write_text(
|
||||
"import types\n"
|
||||
"np = types.ModuleType('numpy')\n"
|
||||
"np.sqrt = lambda x: 0\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return int(np.sqrt(value))\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("fake_module_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
with pytest.raises(TypeError, match="does not import as 'numpy'"):
|
||||
udf(module.score)
|
||||
|
||||
|
||||
def test_udf_rejects_a_module_level_namespace_alias(tmp_path):
|
||||
module_path = tmp_path / "aliasing_udfs.py"
|
||||
module_path.write_text(
|
||||
"import builtins as b\n"
|
||||
"THRESHOLD = 5\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return value + b.vars(b.__import__('aliasing_udfs'))['THRESHOLD']\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("aliasing_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
with pytest.raises(ValueError, match="dynamic namespace access"):
|
||||
udf(module.score)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"access",
|
||||
[
|
||||
"globals()['THRESHOLD']",
|
||||
"eval('THRESHOLD')",
|
||||
"(lambda g: g()['THRESHOLD'])(globals)",
|
||||
"__import__('sys').modules[__name__].THRESHOLD",
|
||||
"sys.modules[__name__].THRESHOLD",
|
||||
],
|
||||
)
|
||||
def test_udf_rejects_dynamic_namespace_access(access):
|
||||
namespace: dict = {}
|
||||
exec(
|
||||
f"def score(value: int) -> int:\n return value + {access}\n",
|
||||
{"THRESHOLD": 5},
|
||||
namespace,
|
||||
)
|
||||
with pytest.raises(ValueError, match="dynamic namespace access"):
|
||||
_package_from_text(
|
||||
"def score(value: int) -> int:\n"
|
||||
" import sys\n"
|
||||
f" return value + {access}\n"
|
||||
)
|
||||
|
||||
|
||||
def _package_from_text(source: str, module_globals: dict | None = None):
|
||||
"""Load `source` as a real module file so the packager can inspect it."""
|
||||
import tempfile
|
||||
|
||||
directory = tempfile.mkdtemp()
|
||||
path = Path(directory) / "generated_udf_module.py"
|
||||
path.write_text(source)
|
||||
spec = importlib.util.spec_from_file_location(f"generated_udf_{id(source)}", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
if module_globals:
|
||||
module.__dict__.update(module_globals)
|
||||
spec.loader.exec_module(module)
|
||||
functions = [
|
||||
value
|
||||
for value in vars(module).values()
|
||||
if callable(value) and getattr(value, "__module__", None) == module.__name__
|
||||
]
|
||||
return udf(functions[0])
|
||||
|
||||
|
||||
def test_udf_rejects_a_non_standard_builtins_environment():
|
||||
def score(value: int) -> int:
|
||||
return len([1]) + value
|
||||
|
||||
score.__globals__ # noqa: B018 -- real function, real globals
|
||||
import builtins
|
||||
|
||||
patched = types.FunctionType(
|
||||
score.__code__,
|
||||
{"__builtins__": {**vars(builtins), "len": lambda _: 99}},
|
||||
"score",
|
||||
)
|
||||
patched.__annotations__ = score.__annotations__
|
||||
assert patched(3) == 102
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(patched)
|
||||
|
||||
class ReportingDict(dict): # reports standard entries, resolves differently
|
||||
def __missing__(self, key):
|
||||
return vars(builtins)[key]
|
||||
|
||||
disguised = types.FunctionType(
|
||||
score.__code__, {"__builtins__": ReportingDict(len=lambda _: 99)}, "score"
|
||||
)
|
||||
disguised.__annotations__ = score.__annotations__
|
||||
assert disguised(3) == 102
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(disguised)
|
||||
|
||||
hooked = types.FunctionType(
|
||||
score.__code__,
|
||||
{"__builtins__": {**vars(builtins), "__import__": lambda *a, **k: None}},
|
||||
"score",
|
||||
)
|
||||
hooked.__annotations__ = score.__annotations__
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(hooked)
|
||||
|
||||
|
||||
def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
|
||||
module_path = tmp_path / "rebound_udfs.py"
|
||||
module_path.write_text(
|
||||
"def fact(value: int) -> int:\n"
|
||||
" return 1 if value <= 1 else value * fact(value - 1)\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return score + value\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("rebound_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
assert _run_packaged(udf(module.fact), 5) == 120
|
||||
raw = module.score
|
||||
module.score = 10
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw)
|
||||
# A wrapper that merely exposes __wrapped__ is not the function.
|
||||
module.score = functools.wraps(raw)(lambda value: 41)
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw)
|
||||
# The decorator's own result is; a subclass of it is not.
|
||||
module.fact = udf(module.fact)
|
||||
assert _run_packaged(module.fact, 4) == 24
|
||||
|
||||
class Twisted(UdfDefinition):
|
||||
def __call__(self, *args, **kwargs):
|
||||
return 41
|
||||
|
||||
raw_fact = module.fact._function
|
||||
module.fact = Twisted(
|
||||
raw_fact,
|
||||
name=None,
|
||||
input_schema=None,
|
||||
output_schema=None,
|
||||
pip=(),
|
||||
env={},
|
||||
secrets=(),
|
||||
python_version=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw_fact)
|
||||
|
||||
|
||||
def test_canonical_arrow_type_rejects_unrepresentable_list_children():
|
||||
from lancedb.functions import _canonical_arrow_type
|
||||
|
||||
for outside in [
|
||||
pa.list_(pa.float32()), # pyarrow default: nullable child
|
||||
pa.list_(pa.field("custom", pa.float32(), nullable=False)),
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
|
||||
]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(outside)
|
||||
assert (
|
||||
_canonical_arrow_type(
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 3)
|
||||
)
|
||||
== "fixed_size_list<float32, 3>"
|
||||
)
|
||||
|
||||
|
||||
def _calls_missing(value: int) -> int:
|
||||
return missing(value) # noqa: F821
|
||||
|
||||
|
||||
def _shadows_missing_in_a_comprehension(value: int) -> int:
|
||||
return missing(value) + sum(missing for missing in ()) # noqa: F821
|
||||
|
||||
|
||||
def _shadows_missing_in_a_lambda(value: int) -> int:
|
||||
return (lambda missing: missing)(value) + missing # noqa: F821
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function",
|
||||
[_calls_missing, _shadows_missing_in_a_comprehension, _shadows_missing_in_a_lambda],
|
||||
)
|
||||
def test_udf_rejects_a_truly_unresolved_global(function):
|
||||
with pytest.raises(ValueError, match=r"unresolved global names: \['missing'\]"):
|
||||
udf(function)
|
||||
|
||||
|
||||
def _arrow_type_from_golden(spec: dict) -> pa.DataType:
|
||||
kind = spec["type"]
|
||||
if kind in ("list", "large_list", "fixed_size_list"):
|
||||
item = _arrow_type_from_golden(spec["fields"][0]["type"])
|
||||
field = pa.field("item", item, nullable=False)
|
||||
if kind == "list":
|
||||
return pa.list_(field)
|
||||
if kind == "large_list":
|
||||
return pa.large_list(field)
|
||||
return pa.list_(field, spec["length"])
|
||||
return {
|
||||
"null": pa.null(),
|
||||
"bool": pa.bool_(),
|
||||
"utf8": pa.string(),
|
||||
"binary": pa.binary(),
|
||||
"float16": pa.float16(),
|
||||
"float32": pa.float32(),
|
||||
"float64": pa.float64(),
|
||||
"date32": pa.date32(),
|
||||
"date64": pa.date64(),
|
||||
}.get(kind) or getattr(pa, kind)()
|
||||
|
||||
|
||||
def test_arrow_type_grammar_matches_the_shared_golden():
|
||||
golden = json.loads(
|
||||
(
|
||||
Path(__file__).parents[3]
|
||||
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
|
||||
).read_text()
|
||||
)
|
||||
from lancedb.functions import _canonical_arrow_type
|
||||
|
||||
emitted = {
|
||||
case["arrow_type"]: _canonical_arrow_type(_arrow_type_from_golden(case["json"]))
|
||||
for case in golden["valid"]
|
||||
}
|
||||
assert emitted == {
|
||||
case["arrow_type"]: case["arrow_type"] for case in golden["valid"]
|
||||
}
|
||||
assert not set(emitted) & set(golden["invalid"])
|
||||
for case in golden["server_only"]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(_arrow_type_from_golden(case["json"]))
|
||||
|
||||
|
||||
def test_explicit_arrow_schema_is_deterministic():
|
||||
input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)])
|
||||
output_schema = pa.field("embedding", pa.list_(pa.float32(), 3), nullable=False)
|
||||
output_schema = pa.field(
|
||||
"embedding",
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 3),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
@udf(input_schema=input_schema, output_schema=output_schema)
|
||||
def explicit(value):
|
||||
@@ -82,7 +477,7 @@ def test_explicit_arrow_schema_is_deterministic():
|
||||
signature = explicit.registration_request.signature
|
||||
assert signature.inputs[0].arrow_type == "float32"
|
||||
assert signature.inputs[0].nullable is True
|
||||
assert signature.output.arrow_type == "fixed_size_list<float32>[3]"
|
||||
assert signature.output.arrow_type == "fixed_size_list<float32, 3>"
|
||||
assert signature.output.nullable is False
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user