Skip to content

Commit eb7a878

Browse files
rustyconoverclaude
andcommitted
Let a required ANY argument be declared nullable, so it can validate itself
A function with a required `AnyArrow` argument had no way to produce a good error when a caller omitted it. The framework parses arguments before `on_bind()` runs, so the omission surfaced as KeyError: "Argument 'questions': not found" instead of the message the function had written for exactly that case, which names the argument and shows its shape. Two things blocked the obvious fix of `default=None`: 1. parameter_default_values_from_specs rejected ANY default on an ANY-typed parameter. Only a VALUE default is unrepresentable there: `None` is fine, because a null-typed array holds nulls and it is column PRESENCE, not the value, that marks a parameter as having a default — as that function's own docstring says. Narrowed the guard to non-None defaults. 2. The marker types (AnyArrow, AnyArrowValue, TableInput) are recognised by identity, so `AnyArrow | None` hid the marker: the parameter silently became NULL-typed, and DuckDB could then cast nothing into it — `Unimplemented type for cast (STRUCT(...) -> NULL)` at the first real call. Both detection sites now unwrap an optional union first, mirroring what `_accepts_none` already looks for. Together these let a function declare Annotated[AnyArrow | None, Arg("questions", default=None)] for an argument that is required in practice but validated by the function itself. Verified end to end against a real DuckDB ATTACH: the missing case now raises the function's own message, and the struct, JSON string and correlated LATERAL call shapes all still work with the argument present. Tests: an optional ANY argument keeps its ANY marker and null arrow type; a None default on an ANY parameter builds a one-row null column; a non-null default on one is still rejected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 39ad674 commit eb7a878

2 files changed

Lines changed: 101 additions & 10 deletions

File tree

‎tests/test_argument_spec.py‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from __future__ import annotations
66

7+
import dataclasses
78
from typing import Annotated, Any
89

910
import pyarrow as pa
@@ -856,7 +857,7 @@ def compute(
856857
assert defaults.column("label")[0].as_py() is None
857858

858859
def test_typed_default_requires_concrete_parameter_type(self) -> None:
859-
"""An ANY parameter cannot carry an untyped protocol default."""
860+
"""An ANY parameter cannot carry an untyped protocol default VALUE."""
860861
spec = ArgumentSpec(
861862
name="value",
862863
position=0,
@@ -865,9 +866,65 @@ def test_typed_default_requires_concrete_parameter_type(self) -> None:
865866
default_value=1,
866867
)
867868

868-
with pytest.raises(ValueError, match="typed defaults require an explicit parameter type"):
869+
with pytest.raises(ValueError, match="a typed default requires an explicit parameter type"):
869870
parameter_default_values_from_specs([spec], require_trailing=True)
870871

872+
def test_an_any_parameter_may_default_to_none(self) -> None:
873+
"""`None` IS representable on an ANY parameter — a null array holds nulls.
874+
875+
It is column presence, not the value, that marks a parameter as having a
876+
default, so this costs the protocol nothing. It buys a function the
877+
ability to declare a *required* ANY argument with ``default=None`` so
878+
that omitting it reaches ``on_bind()``, where the function can name the
879+
argument and show its shape, instead of raising a bare ``KeyError``
880+
during argument parsing before any user code runs.
881+
"""
882+
spec = ArgumentSpec(
883+
name="questions",
884+
position="questions",
885+
arrow_type=pa.null(),
886+
is_any_type=True,
887+
default_value=None,
888+
)
889+
890+
defaults = parameter_default_values_from_specs([spec], require_trailing=True)
891+
892+
assert defaults is not None
893+
assert defaults.schema.names == ["questions"]
894+
assert defaults.column("questions").type == pa.null()
895+
assert defaults.column("questions")[0].as_py() is None
896+
897+
def test_an_optional_any_argument_is_still_an_any_argument(self) -> None:
898+
"""``AnyArrow | None`` must stay ANY-typed, not collapse to a NULL column.
899+
900+
A marker type is recognised by identity, so an Optional wrapper used to
901+
hide it: the parameter silently became NULL-typed and DuckDB could then
902+
cast nothing into it (`Unimplemented type for cast (STRUCT(...) -> NULL)`).
903+
Declaring a REQUIRED ANY argument nullable, with ``default=None``, is how
904+
a function gets to reject the missing case in its own ``on_bind()`` with
905+
a message naming the argument.
906+
"""
907+
908+
@dataclasses.dataclass(slots=True, frozen=True, kw_only=True)
909+
class _Args:
910+
"""Args with a required-but-nullable ANY parameter."""
911+
912+
payload: Annotated[AnyArrow | None, Arg("payload", default=None)] = None
913+
914+
class _Fn(TableInOutFunction[_Args]):
915+
"""A function declaring one optional ANY argument."""
916+
917+
FunctionArguments = _Args
918+
919+
class Meta:
920+
"""Metadata."""
921+
922+
name = "optional_any"
923+
924+
(spec,) = [sp for sp in extract_argument_specs(_Fn) if sp.name == "payload"]
925+
assert spec.is_any_type, "an optional ANY argument lost its ANY marker"
926+
assert pa.types.is_null(spec.arrow_type)
927+
871928
def test_typed_defaults_reject_case_insensitive_duplicate_names(self) -> None:
872929
"""Default-batch field names must identify one signature slot."""
873930
specs = [

‎vgi/argument_spec.py‎

Lines changed: 42 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@
1919
from collections.abc import Sequence
2020
from dataclasses import dataclass
2121
from functools import cache
22-
from typing import Annotated, Any, get_args, get_origin, get_type_hints
22+
from types import UnionType
23+
from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints
2324

2425
import pyarrow as pa
2526

@@ -395,14 +396,22 @@ def parameter_default_values_from_specs(
395396
arrays: list[pa.Array[Any]] = []
396397
names: list[str] = []
397398
for spec in default_specs:
398-
if pa.types.is_null(spec.arrow_type):
399-
raise ValueError(
400-
f"Parameter {spec.name!r} has a default but no concrete Arrow type; "
401-
"typed defaults require an explicit parameter type"
402-
)
403399
value = spec.default_value
404400
if isinstance(value, pa.Scalar):
405401
value = value.as_py()
402+
# An ANY-typed parameter (AnyArrow) has no concrete Arrow type, so a
403+
# VALUE default cannot be represented. `None` still can: a null-typed
404+
# array holds nulls, and it is column PRESENCE, not the value, that
405+
# marks a parameter as having a default. Allowing it lets a required
406+
# ANY parameter be declared `default=None` so that omitting it reaches
407+
# the function's own on_bind() validation — which can name the argument
408+
# and show its shape — instead of raising a bare KeyError during
409+
# argument parsing, before any user code runs.
410+
if pa.types.is_null(spec.arrow_type) and value is not None:
411+
raise ValueError(
412+
f"Parameter {spec.name!r} has a non-null default but no concrete Arrow type; "
413+
"a typed default requires an explicit parameter type (None is allowed)"
414+
)
406415
try:
407416
arrays.append(pa.array([value], type=spec.arrow_type))
408417
except (pa.ArrowException, TypeError, ValueError) as exc:
@@ -538,6 +547,30 @@ def _format_range(
538547
return f"{low}, {high}"
539548

540549

550+
def _unwrap_optional(annotation: Any) -> Any:
551+
"""``T | None`` -> ``T``; anything else unchanged.
552+
553+
A special marker type (``AnyArrow``, ``TableInput``) is recognised by
554+
identity, so an Optional wrapper would hide it — an ``AnyArrow | None``
555+
parameter would silently become a NULL-typed one, and DuckDB then cannot
556+
cast anything into it. Unwrapping first lets a marker be declared nullable,
557+
which is how a REQUIRED ANY argument opts into reaching its function's own
558+
``on_bind()`` validation (with ``default=None``) instead of being rejected
559+
during parsing. Mirrors what ``_accepts_none`` already looks for.
560+
561+
Args:
562+
annotation: The declared type of an ``Annotated[T, Arg(...)]`` field.
563+
564+
Returns:
565+
The single non-``None`` member of an optional union, else ``annotation``.
566+
"""
567+
origin = get_origin(annotation)
568+
if origin is not Union and origin is not UnionType:
569+
return annotation
570+
members = [a for a in get_args(annotation) if a is not type(None)]
571+
return members[0] if len(members) == 1 else annotation
572+
573+
541574
def _constraint_kwargs(arg: Arg[Any]) -> dict[str, Any]:
542575
"""Extract discovery-facing constraint metadata from an [`Arg`][] descriptor.
543576
@@ -744,6 +777,7 @@ def _extract_argument_specs(
744777
if type_args:
745778
infer_type = type_args[0]
746779

780+
infer_type = _unwrap_optional(infer_type)
747781
is_table_input = infer_type is TableInput
748782
is_any_type = infer_type is AnyArrow or infer_type is AnyArrowValue
749783

@@ -790,8 +824,8 @@ def _extract_argument_specs(
790824
# Check for special types (AnyArrow, TableInput)
791825
# Priority: Arg subscript type (Arg[AnyArrow]) > class type hint
792826
# Also check _returns_any_arrow_value for Annotated[AnyArrowValue, ...]
793-
hint = hints.get(attr_name)
794-
type_param = getattr(arg_legacy, "_type_param", None)
827+
hint = _unwrap_optional(hints.get(attr_name))
828+
type_param = _unwrap_optional(getattr(arg_legacy, "_type_param", None))
795829
is_table_input = type_param is TableInput or hint is TableInput
796830
is_any_type = (
797831
type_param is AnyArrow

0 commit comments

Comments
 (0)