feat(arguments): preserve union tag when decoding union-typed args

A DuckDB UNION / Arrow union argument is tagged — the discriminator lives in
UnionScalar.type_code, which plain Scalar.as_py() discards, so a union arg
arrived as just its member value with no way to tell which member was set.

Decode union scalars to a new TaggedUnion(tag, value) instead: tag is the
active member's field name (resolved from type_code via the union type), value
is its Python value. Routed through Arguments.get()/get_varargs(), so both
direct args.get() and table-function params.args.* benefit; every non-union
type is unchanged. TaggedUnion is exported from vgi and vgi.arguments.

Verified end-to-end over the RPC path (sparse union). Note: the C++ extension
still rejects dense unions; full union support needs a fix there too.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: rustyconover

tests/test_union_argument.py

@@ -0,0 +1,55 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Union-typed argument decoding.
+
+A DuckDB ``UNION`` / Arrow union argument is tagged: the discriminator lives in
+``UnionScalar.type_code``, which plain ``Scalar.as_py()`` discards. ``Arguments``
+decodes such scalars to ``TaggedUnion`` so the active member name is preserved.
+"""
+
+from __future__ import annotations
+
+import pyarrow as pa
+
+from vgi.arguments import Arguments, TaggedUnion
+
+_RF = pa.struct([pa.field("n_estimators", pa.list_(pa.int64())), pa.field("max_depth", pa.list_(pa.int64()))])
+_SVC = pa.struct([pa.field("C", pa.list_(pa.float64())), pa.field("kernel", pa.list_(pa.string()))])
+
+
+def _union_scalar(tag: str, code: int, members: dict) -> pa.UnionScalar:
+    """Build a one-element sparse-union scalar with named members, active = ``tag``."""
+    arr = pa.UnionArray.from_sparse(
+        pa.array([code], type=pa.int8()),
+        [
+            pa.array([members.get("random_forest_classifier")], type=_RF),
+            pa.array([members.get("svc")], type=_SVC),
+        ],
+        field_names=["random_forest_classifier", "svc"],
+        type_codes=[0, 1],
+    )
+    return arr[0]
+
+
+def test_union_arg_preserves_tag() -> None:
+    scalar = _union_scalar(
+        "random_forest_classifier", 0, {"random_forest_classifier": {"n_estimators": [100, 300], "max_depth": [3, 5]}}
+    )
+    got = Arguments(named={"config": scalar}).get("config")
+    assert isinstance(got, TaggedUnion)
+    assert got.tag == "random_forest_classifier"
+    assert got.value == {"n_estimators": [100, 300], "max_depth": [3, 5]}
+
+
+def test_union_arg_other_member() -> None:
+    scalar = _union_scalar("svc", 1, {"svc": {"C": [1.0, 10.0], "kernel": ["rbf", "linear"]}})
+    got = Arguments(named={"config": scalar}).get("config")
+    assert got.tag == "svc"
+    assert got.value == {"C": [1.0, 10.0], "kernel": ["rbf", "linear"]}
+
+
+def test_non_union_args_unchanged() -> None:
+    args = Arguments(named={"n": pa.scalar(5), "s": pa.scalar("hi")}, positional=(pa.scalar(1.5),))
+    assert args.get("n") == 5
+    assert args.get("s") == "hi"
+    assert args.get(0) == 1.5

vgi/__init__.py

@@ -39,6 +39,7 @@
     Param,
     Returns,
     TableInput,
+    TaggedUnion,
 )
 from vgi.auth import AuthContext, CallContext
 from vgi.metadata import (
@@ -143,6 +144,7 @@
     "TableInOutGenerator",
     "TableInput",
     "TableInputValidationError",
+    "TaggedUnion",
     "TypeMismatchError",
     "Worker",
     "functions_to_arrow",

vgi/arguments.py

@@ -159,6 +159,7 @@ def _accepts_none(annotated_inner_type: Any) -> bool:
     "PYTHON_TO_ARROW",
     "Returns",
     "TableInput",
+    "TaggedUnion",
     "TypeBoundPredicate",
     "OutputLength",
     "Setting",
@@ -168,6 +169,53 @@ def _accepts_none(annotated_inner_type: Any) -> bool:
 ]
 
 
+@dataclass(frozen=True, slots=True)
+class TaggedUnion:
+    """A decoded union-typed argument: which member is set (``tag``) and its ``value``.
+
+    DuckDB ``UNION`` / Arrow union arguments are *tagged*: the discriminator
+    (which member is present) lives in the Arrow ``UnionScalar.type_code``, not
+    in the member value. Plain ``Scalar.as_py()`` returns only the member value
+    and drops that tag, so union arguments are decoded into this wrapper
+    instead — ``tag`` is the active member's field name and ``value`` is its
+    Python value.
+
+    Example::
+
+        config: Annotated[TaggedUnion, Arg("config", arrow_type=pa.sparse_union([...]))]
+        ...
+        cfg = params.args.config            # TaggedUnion(tag=..., value=...)
+        if cfg.tag == "random_forest_classifier":
+            grid = cfg.value                # the member struct, as a dict
+
+    """
+
+    tag: str | None
+    value: Any
+
+
+def _scalar_to_py(scalar: "Scalar") -> Any:
+    """Convert an argument scalar to a Python value, preserving union tags.
+
+    Identical to ``scalar.as_py()`` for every type except unions: a
+    ``UnionScalar`` is decoded to a [`TaggedUnion`][] so the member
+    discriminator (which ``as_py()`` discards) is retained.
+    """
+    if isinstance(scalar, pa.UnionScalar):
+        union_type = scalar.type
+        tag = next(
+            (
+                union_type.field(i).name
+                for i in range(union_type.num_fields)
+                if union_type.type_codes[i] == scalar.type_code
+            ),
+            None,
+        )
+        inner = scalar.value
+        return TaggedUnion(tag=tag, value=inner.as_py() if inner is not None else None)
+    return scalar.as_py()
+
+
 class TableInput:
     """Sentinel type for table input parameters in table-in-out functions.
 
@@ -377,7 +425,7 @@ def get(
             else:
                 raise TypeError(f"Argument '{key}': expected {type}, got {scalar.type}")
 
-        return scalar.as_py()
+        return _scalar_to_py(scalar)
 
     def get_varargs(
         self,
@@ -410,7 +458,7 @@ def get_varargs(
             if type is not None and scalar.type != type:
                 raise TypeError(f"Argument {i}: expected {type}, got {scalar.type}")
 
-            values.append(scalar.as_py())
+            values.append(_scalar_to_py(scalar))
 
         return tuple(values)