From d081cf3f0d278b8d748ace1d1d17dfa46ed3e30b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 31 Jul 2026 12:37:01 -0500 Subject: [PATCH 1/5] Reject UDT UDF returns that cannot fill the element A UDF over an array UDT whose return does not fit the element is not a type error. Numba's `Array` type carries `ndim` but not its extents, so returning `x[:2]` from a 9-element UDT types identically to returning `x`. The wrapper's `z[:] = ...` then raises inside the cfunc, where Numba prints the traceback and returns. GraphBLAS is handed no error, and the element keeps whatever was already in the buffer, so the caller gets a wrong answer and no exception. A record UDF that under-fills an array-typed leaf fails the same way and abandons the write part-way: every leaf after the short one, scalar leaves included, keeps what SuiteSparse had there. Extents can only be recovered by running the function, so `_get_udt_wrapper` and the IndexBinaryOp wrapper now call it once on stand-in operands and check that the result fits the UDT. Fitting is not the same as matching. Slice-assign broadcasts, so a `(1,)` return fills a `(6,)` element and a `(1, 3)` return fills both rows of a `(2, 3)` one; both work today and must keep working, and an equality check would reject them. `_fits_by_broadcast` applies numpy's assignment rule instead, and `test_udt_broadcast_matches_numba_slice_assign` pins it against Numba's own slice-assign in both directions, including the two ranks where broadcasting alone gives the wrong answer. The operands are ones rather than zeros for the same reason `dtypes._sample_values` avoids zeros: a UDF that divides by an operand would raise on a zero-filled probe, and a raising probe is treated as "cannot check" rather than as a failure. The check is best-effort, and this commit should not claim more. A UDF that raises on the probe values is not checked at all, so one that also returns a shape that does not fit still reaches the wrapper and still fails silently, exactly as it did before. The wrapper's slice-assign remains the only backstop that always runs. What the check buys is that the common case, a UDF that is simply wrong about the shape, becomes an error at registration instead of a wrong answer at apply time. Because the probe reports the shape for the values it was run on, the diagnostic says so. A UDF whose output shape depends on its operands' values rather than their types can otherwise be rejected on a shape it would never return for real data. This is a judgement call rather than a straight fix, and it is cheap to decline. Typing an op for a UDT now executes the user's function, where before it only compiled it. `OpBase.__contains__` is implemented as a typed lookup, so `udt in some_op` runs user code as a side effect of a membership test. Measured: registering an op runs no probe, since compilation stays lazy, but a single `udt in some_op` runs exactly one. That probe costs no extra compile, because it calls the function with an `ndarray` and the wrapper's `numba.carray` view needs that same specialization anyway. A/B with the check stubbed out: 93.4 ms median disabled against 92.1 ms enabled, two Numba specializations either way. Only array UDTs and records that actually have an array leaf are probed, so an ordinary record UDF is still typed without running user code. The array-UDT paragraph in `docs/user_guide/udt.rst`, added one commit earlier, described the rule as building "a new array of the same shape". It now states the rule the check enforces, and says the shape is learned by running the UDF. If this is declined, the array-UDT view change it sits on still stands. A return that does not fit stays memory-safe, because the slice-assign cannot overrun the element, and goes back to being silent. --- docs/user_guide/udt.rst | 9 +- graphblas/core/operator/base.py | 165 ++++++++++++++++++++++- graphblas/tests/test_indexbinary.py | 26 ++++ graphblas/tests/test_op.py | 201 ++++++++++++++++++++++++++++ 4 files changed, 399 insertions(+), 2 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 15cbedbb4..c8b4647b2 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -182,7 +182,9 @@ return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value For array UDTs each operand arrives as a numpy view of that element's values, in the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and -the UDF may return one of its operands or build a new array of the same shape: +the UDF may return one of its operands or build a new array that fills the +element, either at the element's own shape or at one that broadcasts to it (a +``(1,)`` return fills every slot of a ``(6,)`` element): .. code-block:: python @@ -199,6 +201,11 @@ the UDF may return one of its operands or build a new array of the same shape: c = a.ewise_mult(b, op[point3]).new() # c[0] = [5.0, 11.0, 17.0] +A return that cannot fill the element, such as ``x[:2]`` from a 3-element UDT, +is rejected when the op is typed for the UDT. The shape is learned by running +the UDF once on sample values, so a UDF that raises on those values is not +checked. + If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced from Numba's typing pass instead of a 200-line traceback. diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index 485857cfb..b081912cd 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -3,10 +3,13 @@ from operator import getitem from types import BuiltinFunctionType, ModuleType +import numpy as np + from ... import _STANDARD_OPERATOR_NAMES, backend, op from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, _supports_udfs, ffi, lib +from ..dtypes import _sample_values from ..expr import InfixExprBase from ..utils import output_type @@ -494,7 +497,15 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): the caller either: the resulting ``ValueError`` is raised inside a cfunc, which Numba prints and swallows, so the write stops there and every leaf from that point on keeps whatever SuiteSparse had in the - buffer. + buffer. ``_check_array_udf_shape`` and + ``_check_record_udf_leaf_shapes`` reject such a return up front, but + only when they can run the UDF, so this is the backstop that always + runs rather than dead code. + + Slice-assign also broadcasts, which is why those checks accept any + return that broadcasts to the element rather than requiring an exact + shape: filling a ``(6,)`` element from a ``(1,)`` return works here + and must keep working. """ if zkind == "array_elements": return ( @@ -521,6 +532,141 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): BL, BR, zname = zinfo return f"{signature_line}\n{body_setup} {zname} = {BL}{call_expr}{BR}\n" + def _udf_probe_value(dtype): + """Build a stand-in operand of ``dtype`` for a UDF probe. + + Ones rather than zeros, for the same reason ``dtypes._sample_values`` + avoids zeros: a UDF that divides by an operand raises on a zero-filled + probe, and :func:`_run_udf_probe` treats a raising probe as "cannot + check", which would quietly skip the check the probe exists to perform. + """ + if dtype._is_udt: + np_type = dtype.np_type + if np_type.subdtype is None: + return np.ones(1, dtype=np_type)[0] + base_np_type, shape = np_type.subdtype + return np.ones(shape, dtype=base_np_type) + return _sample_values[dtype] + + def _run_udf_probe(numba_func, operands): + """Run the UDF on stand-in operands. Returns ``(result,)``, or ``None``. + + Numba's ``Array`` type records ``ndim`` but not extents, so a UDF that + builds its result (``x + y``) rather than returning an operand can only + be shape-checked by running it. Doing that at registration turns a + wrong-shape return into an error the caller sees. Left to the wrapper + it raises inside a cfunc, where Numba prints the traceback (once per + element) and returns, handing back an uninitialized element and no + exception. + + The cost is that the first ``op[udt]`` lookup now executes the user's + function once, on ones, where before it only compiled it. Registering + the op does not, since compilation stays lazy, but + ``OpBase.__contains__`` is a typed lookup, so ``udt in some_op`` runs + the function as a side effect of a membership test. It costs no extra + compile: the probe passes an ``ndarray``, and the wrapper's + ``numba.carray`` view needs that same specialization anyway. + + The check is best-effort, not a guarantee. A UDF that raises on the + probe values returns ``None`` here and is not checked at all, so one + that also returns a wrong shape still reaches the wrapper and still + fails the way it always did: a ``ValueError`` inside the cfunc that + Numba prints and swallows, a caller that sees no exception, and an + element left as SuiteSparse found it. The wrapper's slice-assign + remains the only backstop that always runs. Only the call is guarded, + so mistakes in the probe itself still surface. + + One probe cannot settle a UDF whose output shape depends on operand + *values* rather than their types: it reports the shape for the probe + values, which is why the callers say so in their message. + """ + args = [_udf_probe_value(d) for d in operands] + try: + return (numba_func(*args),) + except NumbaError as exc: + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except Exception: + return None + + # Shared tail for the shape diagnostics: the probed shape is one sample, so + # a value-dependent UDF can be rejected on a shape it returns only here. + _SHAPE_HINT = ( + "A UDF whose output shape varies with its input values must still fit every element." + ) + + def _fits_by_broadcast(shape, expected): + """Whether a return of ``shape`` fills an ``expected``-shaped destination. + + The wrapper slice-assigns (``z[:] = ...``), and numpy broadcasts on + assignment, so an exact match is not the requirement: a ``(1,)`` return + legitimately fills a ``(6,)`` element, and a ``(1, 3)`` return fills + every row of a ``(2, 3)`` one. Rejecting those would refuse code that + works today. + + numpy's assignment rule is broadcasting plus a leading-``1`` strip when + the source has the higher rank, so ``(1, 6)`` fits ``(6,)`` while + ``(6, 1)`` does not. ``test_udt_broadcast_matches_numba_slice_assign`` + pins this against Numba's own slice-assign for both. + """ + shape = tuple(shape) + expected = tuple(expected) + while len(shape) > len(expected) and shape[0] == 1: + shape = shape[1:] + try: + return np.broadcast_shapes(shape, expected) == expected + except ValueError: + return False + + def _check_array_udf_shape(numba_func, return_type, operands): + """Reject a UDF whose built array cannot fill the array UDT's element.""" + probed = _run_udf_probe(numba_func, operands) + if probed is None: + return + expected = return_type.numba_type.shape + shape = getattr(probed[0], "shape", None) + if shape is not None and not _fits_by_broadcast(shape, expected): + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)} when run on sample " + f"values, but {return_type} elements are {tuple(expected)}. Return an " + f"array whose shape matches or broadcasts to that, or one of the " + f"operands. {_SHAPE_HINT}" + ) + + def _check_record_udf_leaf_shapes(numba_func, return_type, operands): + """Reject a record return whose array-typed leaf cannot fill its field. + + The array-leaf half of :func:`_check_array_udf_shape`. The wrapper + slice-assigns those leaves, so a wrong extent raises inside the cfunc + and abandons the write part-way: every leaf after it keeps whatever + SuiteSparse had in the buffer, scalar leaves included. + + Only records that actually have an array leaf are probed, so the + common record UDF still registers without running the user's code. + """ + from .udt_utils import _iter_record_leaves + + leaves = [(py, d) for py, _c, d in _iter_record_leaves(return_type.np_type)] + if not any(d.subdtype is not None for _py, d in leaves): + return + probed = _run_udf_probe(numba_func, operands) + if probed is None: + return + result = probed[0] + if not isinstance(result, tuple) or len(result) != len(leaves): + return + for (path, leaf_dtype), value in zip(leaves, result, strict=True): + if leaf_dtype.subdtype is None: + continue + expected = leaf_dtype.subdtype[1] + shape = getattr(value, "shape", None) + if shape is not None and not _fits_by_broadcast(shape, expected): + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)} for field " + f"{path} of {return_type} when run on sample values, which holds " + f"{tuple(expected)} there. Return an array whose shape matches or " + f"broadcasts to that. {_SHAPE_HINT}" + ) + def _get_udt_wrapper( numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None ): @@ -536,15 +682,22 @@ def _get_udt_wrapper( zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) xsetup, xderef, xptr_type = _input_operand(dtype, "x") wrapper_args = [zptr_type, xptr_type] + probe_operands = [dtype] if include_indexes: wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + probe_operands.extend([UINT64, UINT64]) ysetup, yderef_expr, yarg = "", "", "" if dtype2 is not None: ysetup, yderef, yptr_type = _input_operand(dtype2, "y") wrapper_args.append(yptr_type) + probe_operands.append(dtype2) yarg = ", y_ptr" yderef_expr = f", {yderef}" wrapper_sig = nt.void(*wrapper_args) + if zkind == "array_elements": + _check_array_udf_shape(numba_func, return_type, probe_operands) + elif zkind == "record_fields": + _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) rcidx = ", row, col" if include_indexes else "" signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" @@ -586,6 +739,16 @@ def _get_udt_wrapper_indexbinary( UINT64.numba_type, tptr_type, ) + # Same registration-time guards as :func:`_get_udt_wrapper`. Without + # them a wrong-shape return is memory-safe but silent: the wrapper's + # ``z[:] =`` raises inside the cfunc, Numba prints and swallows it, + # and the element is left as SuiteSparse found it. + if zkind in ("array_elements", "record_fields"): + probe_operands = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + if zkind == "array_elements": + _check_array_udf_shape(numba_func, return_type, probe_operands) + else: + _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index c56bb86f9..82ffbdbf0 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -300,6 +300,32 @@ def test_bind_raw_array_udt_theta(): delattr(indexbinary, "raw_array_udt_op") +def test_array_udt_udf_shape_is_checked(): + """An IBO whose UDF builds a wrong-shape array is rejected at registration. + + The same guard the binary and unary paths get. Numba's ``Array`` type + carries ``ndim`` but not extents, so the mismatch is not a type error; left + to the wrapper it raises inside the cfunc, where Numba prints the traceback + and returns, leaving the element as SuiteSparse found it. + """ + # A shape of its own: ``register_anonymous`` caches by numpy dtype, so + # sharing one with another test would hand back that test's DataType. + arr_udt = dtypes.register_anonymous(np.dtype((np.float64, (12,))), "_IboShapeArr12") + + def _truncate(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x + y)[:2] + + op = indexbinary.register_anonymous(_truncate, is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[arr_udt] + + def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + good = indexbinary.register_anonymous(_combine, is_udt=True) + assert good[arr_udt].return_type is arr_udt + + def test_bind_raw_udt_theta_without_dtype_errors(): """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 9bc3f2cdc..4381ade1f 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2052,6 +2052,207 @@ def _add_corner(x, y): ) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_shape_errors(): + """Array-UDT UDFs that can't fill the element are rejected at registration. + + Numba's ``Array`` type records ``ndim`` but not extents, so neither case + below is a type error. Both used to reach the cfunc, where the shape + mismatch raises in a context that swallows the exception, handing the + caller an uninitialized element and no error. + """ + # Shapes unique to this test: ``register_anonymous`` caches by dtype and + # freezes the JIT C name at first registration, so sharing a shape with + # another test makes both order-dependent. + udt9 = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_ShapeErr9") + udt10 = dtypes.register_anonymous(np.dtype((np.float64, (10,))), "_ShapeErr10") + + def _truncate(x): # pragma: no cover (numba) + return x[:2] + + op = UnaryOp.register_anonymous(_truncate, "_shape_err_trunc", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[udt9] + + # Two array UDTs sharing a base dtype and rank are indistinguishable once a + # UDF builds its result, so refuse to guess which one it meant. + def _built(x, y): # pragma: no cover (numba) + return y + 0.0 + + op2 = BinaryOp.register_anonymous(_built, "_shape_err_ambiguous", is_udt=True) + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op2[udt9, udt10] + + # Ambiguity is decided on the UDTs, not their Numba shapes: these two are + # separate DataTypes with separate GraphBLAS handles, but Numba collapses + # the layered dtype to the flat one's ``nestedarray(float64, (2, 3))``. + flat = dtypes.register_anonymous(np.dtype((np.float64, (3, 4))), "_ShapeErrFlat") + layered = dtypes.register_anonymous( + np.dtype((np.dtype((np.float64, (4,))), (3,))), "_ShapeErrLayered" + ) + assert flat.numba_type == layered.numba_type + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op2[flat, layered] + assert op2[flat, flat].return_type is flat # a same-type pair is not ambiguous + + # An array UDF whose result matches no input names the mismatch rather + # than telling the user to return an array, which is what they did. + def _recast(x): # pragma: no cover (numba) + return x.astype(np.float32) + + op3 = UnaryOp.register_anonymous(_recast, "_shape_err_recast", is_udt=True) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + op3[udt9] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_leaf_shape_errors(): + """A record UDF that under-fills an array-typed leaf is rejected at registration. + + The wrapper slice-assigns array leaves, so a short return raises inside + the cfunc and abandons the write part-way: leaves after it keep whatever + SuiteSparse had in the buffer, scalar leaves included. + """ + spec = np.dtype([("rl_vec", np.float64, (3,)), ("rl_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafShape") + + def _short(x, y): # pragma: no cover (numba) + return (x["rl_vec"][:2], x["rl_tag"]) + + op = BinaryOp.register_anonymous(_short, "_rec_leaf_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field .* holds \(3,\)"): + op[udt] + + def _full(x, y): # pragma: no cover (numba) + return (x["rl_vec"] + y["rl_vec"], x["rl_tag"] + y["rl_tag"]) + + op = BinaryOp.register_anonymous(_full, "_rec_leaf_full", is_udt=True) + v = Vector(udt, size=1) + v[0] = ([1.0, 2.0, 3.0], 7) + w = Vector(udt, size=1) + w[0] = ([4.0, 5.0, 6.0], 8) + got = v.ewise_mult(w, op).new()[0].new().value + np.testing.assert_array_equal(got["rl_vec"], [5.0, 7.0, 9.0]) + assert got["rl_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_broadcast_return(): + """A return that broadcasts to the element fills it, and is not rejected. + + The wrapper slice-assigns and numpy broadcasts on assignment, so a ``(1,)`` + return legitimately fills every slot of a ``(6,)`` element. Requiring an + exact shape would refuse this, which works. + """ + udt6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_BCast6") + + def _fill(x): # pragma: no cover (numba) + return x[:1] + 10.0 + + op1 = UnaryOp.register_anonymous(_fill, "_bcast_fill", is_udt=True) + assert op1[udt6].return_type is udt6 + v = Vector(udt6, size=1) + v[0] = np.arange(1.0, 7.0) + np.testing.assert_array_equal(v.apply(op1).new()[0].new().value, [11.0] * 6) + + # A row broadcast across a 2-D element: the same rule one rank up. + udt42 = dtypes.register_anonymous(np.dtype((np.float64, (4, 2))), "_BCast42") + + def _fill_rows(x): # pragma: no cover (numba) + return x[:1, :] + 100.0 + + op2 = UnaryOp.register_anonymous(_fill_rows, "_bcast_fill_rows", is_udt=True) + assert op2[udt42].return_type is udt42 + v2 = Vector(udt42, size=1) + v2[0] = np.arange(8.0).reshape(4, 2) + np.testing.assert_array_equal( + v2.apply(op2).new()[0].new().value, np.tile([100.0, 101.0], (4, 1)) + ) + + # The other side of the boundary: (2,) does not broadcast to (6,), Numba's + # slice-assign raises on it, and it stays rejected. + def _short(x): # pragma: no cover (numba) + return x[:2] + 10.0 + + op3 = UnaryOp.register_anonymous(_short, "_bcast_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op3[udt6] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_leaf_broadcast_return(): + """A broadcastable array leaf fills its field, and later leaves still land. + + Same boundary as the array case, and + ``test_udt_record_array_leaf_shape_errors`` holds the rejecting side. The + scalar leaf is worth asserting because a leaf that raises in the cfunc + abandons the write, leaving every leaf after it as SuiteSparse had it. + """ + spec = np.dtype([("bc_vec", np.float64, (11,)), ("bc_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafBCast") + + def _fill_leaf(x, y): # pragma: no cover (numba) + return (x["bc_vec"][:1] + y["bc_vec"][:1], x["bc_tag"] + y["bc_tag"]) + + op1 = BinaryOp.register_anonymous(_fill_leaf, "_rec_leaf_bcast", is_udt=True) + v = Vector(udt, size=1) + v[0] = (np.arange(11.0), 7) + w = Vector(udt, size=1) + w[0] = (np.arange(11.0) + 1.0, 8) + got = v.ewise_mult(w, op1).new()[0].new().value + np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) + assert got["bc_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_broadcast_matches_numba_slice_assign(): + """The shape check accepts exactly what the wrapper's slice-assign accepts. + + The check turns a silent cfunc failure into a registration error, so a + shape it rejects that Numba would have assigned is a false rejection, and + one it accepts that Numba raises on is the failure it exists to catch. Pin + both directions against Numba itself, including the two ranks where + broadcasting alone gives the wrong answer: ``(1, 6)`` fills a ``(6,)`` + destination because assignment drops leading ones, ``(6, 1)`` does not. + """ + import numba + + from graphblas.core.operator.base import _fits_by_broadcast + + @numba.njit + def _assign(z, src): # pragma: no cover (numba) + z[:] = src + + for dst, src in [ + ((6,), ()), + ((6,), (1,)), + ((6,), (6,)), + ((6,), (2,)), + ((6,), (12,)), + ((6,), (1, 6)), + ((6,), (6, 1)), + ((2, 3), (1, 3)), + ((2, 3), (2, 1)), + ((2, 3), (1, 1)), + ((2, 3), (3,)), + ((2, 3), (2, 3)), + ((2, 3), (6,)), + ((2, 3), (3, 2)), + ]: + try: + _assign(np.zeros(dst), np.ones(src)) + except ValueError: + numba_assigns = False + else: + numba_assigns = True + assert _fits_by_broadcast(src, dst) is numba_assigns, (src, dst) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_field_roundtrip(): From 4386f96e8cb35d987931e1c690f3542e48945473 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 23 Sep 2026 12:33:58 -0500 Subject: [PATCH 2/5] Drop the flat-vs-layered ambiguity case from the shape-error test test_udt_array_udf_shape_errors expected a flat FP64[3, 4] and a layered FP64[4] x 3 to be separate UDTs, so that a UDF building an array from the pair was ambiguous. "Flatten nested subarray dtypes so each UDT layout has one DataType" (#589) made them the same DataType, so nothing is ambiguous and the test failed with "DID NOT RAISE UdfParseError" on every CI job with Numba. Keep the same-type check that block ended with, now on udt9. The flattening itself is tested where it is introduced. --- graphblas/tests/test_op.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 4381ade1f..45e4736c7 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2084,17 +2084,7 @@ def _built(x, y): # pragma: no cover (numba) with pytest.raises(UdfParseError, match="matches more than one input array UDT"): op2[udt9, udt10] - # Ambiguity is decided on the UDTs, not their Numba shapes: these two are - # separate DataTypes with separate GraphBLAS handles, but Numba collapses - # the layered dtype to the flat one's ``nestedarray(float64, (2, 3))``. - flat = dtypes.register_anonymous(np.dtype((np.float64, (3, 4))), "_ShapeErrFlat") - layered = dtypes.register_anonymous( - np.dtype((np.dtype((np.float64, (4,))), (3,))), "_ShapeErrLayered" - ) - assert flat.numba_type == layered.numba_type - with pytest.raises(UdfParseError, match="matches more than one input array UDT"): - op2[flat, layered] - assert op2[flat, flat].return_type is flat # a same-type pair is not ambiguous + assert op2[udt9, udt9].return_type is udt9 # a same-type pair is not ambiguous # An array UDF whose result matches no input names the mismatch rather # than telling the user to return an array, which is what they did. From 72335600c66e7fc078facaee3d90a8562befcb48 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 24 Sep 2026 22:32:01 -0500 Subject: [PATCH 3/5] Run a UDT UDF for its shape only when Numba cannot tell it The shape check ran the UDF on sample values for every array return and every record return with an array leaf. An operand, or operand field, returned as-is keeps its extents in Numba's `NestedArray` type, so check those from the type and run the UDF only for an array it builds (a plain `Array`). Returning an operand no longer runs user code when the op is typed, and fields returned into the wrong slots are caught even if the UDF would raise on the sample values. The two checkers become one, `_check_udf_fills_output`. Also reject a record leaf the UDF can return as `None` (typed `Optional`, or always `None`), scalar leaves included. On that path the wrapper's assignment raises inside the cfunc and the element comes back holding another element's values. The probe docstring's reason for using ones was wrong: UDFs compile with `error_model="numpy"`, so dividing by zero does not raise. Drop two test cases that repeated #589's resolver tests, and add a test pinning when the UDF is run. --- docs/user_guide/udt.rst | 7 +- graphblas/core/operator/base.py | 241 +++++++++++++++----------------- graphblas/tests/test_op.py | 132 +++++++++++++---- 3 files changed, 220 insertions(+), 160 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index c8b4647b2..341ab6807 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -202,9 +202,10 @@ element, either at the element's own shape or at one that broadcasts to it (a # c[0] = [5.0, 11.0, 17.0] A return that cannot fill the element, such as ``x[:2]`` from a 3-element UDT, -is rejected when the op is typed for the UDT. The shape is learned by running -the UDF once on sample values, so a UDF that raises on those values is not -checked. +is rejected when the op is typed for the UDT (``op[point3]``, or the first +operation that uses it). Numba does not know the shape of an array the UDF +builds, so the check runs the UDF once on sample values; a UDF that raises on +those values is not checked. If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index b081912cd..a94fe2327 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -497,15 +497,9 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): the caller either: the resulting ``ValueError`` is raised inside a cfunc, which Numba prints and swallows, so the write stops there and every leaf from that point on keeps whatever SuiteSparse had in the - buffer. ``_check_array_udf_shape`` and - ``_check_record_udf_leaf_shapes`` reject such a return up front, but - only when they can run the UDF, so this is the backstop that always - runs rather than dead code. - - Slice-assign also broadcasts, which is why those checks accept any - return that broadcasts to the element rather than requiring an exact - shape: filling a ``(6,)`` element from a ``(1,)`` return works here - and must keep working. + buffer. :func:`_check_udf_fills_output` rejects such a return when the + op is typed, but it cannot always learn the shape, so this stays the + backstop. """ if zkind == "array_elements": return ( @@ -532,68 +526,6 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): BL, BR, zname = zinfo return f"{signature_line}\n{body_setup} {zname} = {BL}{call_expr}{BR}\n" - def _udf_probe_value(dtype): - """Build a stand-in operand of ``dtype`` for a UDF probe. - - Ones rather than zeros, for the same reason ``dtypes._sample_values`` - avoids zeros: a UDF that divides by an operand raises on a zero-filled - probe, and :func:`_run_udf_probe` treats a raising probe as "cannot - check", which would quietly skip the check the probe exists to perform. - """ - if dtype._is_udt: - np_type = dtype.np_type - if np_type.subdtype is None: - return np.ones(1, dtype=np_type)[0] - base_np_type, shape = np_type.subdtype - return np.ones(shape, dtype=base_np_type) - return _sample_values[dtype] - - def _run_udf_probe(numba_func, operands): - """Run the UDF on stand-in operands. Returns ``(result,)``, or ``None``. - - Numba's ``Array`` type records ``ndim`` but not extents, so a UDF that - builds its result (``x + y``) rather than returning an operand can only - be shape-checked by running it. Doing that at registration turns a - wrong-shape return into an error the caller sees. Left to the wrapper - it raises inside a cfunc, where Numba prints the traceback (once per - element) and returns, handing back an uninitialized element and no - exception. - - The cost is that the first ``op[udt]`` lookup now executes the user's - function once, on ones, where before it only compiled it. Registering - the op does not, since compilation stays lazy, but - ``OpBase.__contains__`` is a typed lookup, so ``udt in some_op`` runs - the function as a side effect of a membership test. It costs no extra - compile: the probe passes an ``ndarray``, and the wrapper's - ``numba.carray`` view needs that same specialization anyway. - - The check is best-effort, not a guarantee. A UDF that raises on the - probe values returns ``None`` here and is not checked at all, so one - that also returns a wrong shape still reaches the wrapper and still - fails the way it always did: a ``ValueError`` inside the cfunc that - Numba prints and swallows, a caller that sees no exception, and an - element left as SuiteSparse found it. The wrapper's slice-assign - remains the only backstop that always runs. Only the call is guarded, - so mistakes in the probe itself still surface. - - One probe cannot settle a UDF whose output shape depends on operand - *values* rather than their types: it reports the shape for the probe - values, which is why the callers say so in their message. - """ - args = [_udf_probe_value(d) for d in operands] - try: - return (numba_func(*args),) - except NumbaError as exc: - raise UdfParseError(_summarize_numba_typing_error(exc)) from exc - except Exception: - return None - - # Shared tail for the shape diagnostics: the probed shape is one sample, so - # a value-dependent UDF can be rejected on a shape it returns only here. - _SHAPE_HINT = ( - "A UDF whose output shape varies with its input values must still fit every element." - ) - def _fits_by_broadcast(shape, expected): """Whether a return of ``shape`` fills an ``expected``-shaped destination. @@ -617,55 +549,117 @@ def _fits_by_broadcast(shape, expected): except ValueError: return False - def _check_array_udf_shape(numba_func, return_type, operands): - """Reject a UDF whose built array cannot fill the array UDT's element.""" - probed = _run_udf_probe(numba_func, operands) - if probed is None: - return - expected = return_type.numba_type.shape - shape = getattr(probed[0], "shape", None) - if shape is not None and not _fits_by_broadcast(shape, expected): - raise UdfParseError( - f"UDT UDF returned an array of shape {tuple(shape)} when run on sample " - f"values, but {return_type} elements are {tuple(expected)}. Return an " - f"array whose shape matches or broadcasts to that, or one of the " - f"operands. {_SHAPE_HINT}" - ) - - def _check_record_udf_leaf_shapes(numba_func, return_type, operands): - """Reject a record return whose array-typed leaf cannot fill its field. + def _udf_probe_value(dtype): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. - The array-leaf half of :func:`_check_array_udf_shape`. The wrapper - slice-assigns those leaves, so a wrong extent raises inside the cfunc - and abandons the write part-way: every leaf after it keeps whatever - SuiteSparse had in the buffer, scalar leaves included. + Only the shape of the UDF's result is read, so any values it accepts + will do. + """ + if dtype._is_udt: + np_type = dtype.np_type + if np_type.subdtype is None: + return np.ones(1, dtype=np_type)[0] + base_np_type, shape = np_type.subdtype + return np.ones(shape, dtype=base_np_type) + return _sample_values[dtype] - Only records that actually have an array leaf are probed, so the - common record UDF still registers without running the user's code. + def _run_udf_probe(numba_func, arg_dtypes): + """Call the UDF on stand-in operands. Returns ``(result,)``, or ``None`` if it raised. + + A UDF that raises on the stand-in values is not checked, so one that + also returns a shape that does not fit is left to the wrapper, which + fails without telling the caller (see :func:`_compose_wrapper_body`). + Only the call is guarded, so a mistake in building the operands still + surfaces. + + This is where typing an op for a UDT runs the user's function rather + than only compiling it, and ``OpBase.__contains__`` is a typed lookup, + so ``udt in some_op`` can run it too. It costs no extra compile: Numba + types the stand-ins as it types the wrapper's operands (a C-contiguous + array, a record, a scalar), so the call compiles the specialization the + wrapper needs anyway. """ - from .udt_utils import _iter_record_leaves + args = [_udf_probe_value(d) for d in arg_dtypes] + try: + return (numba_func(*args),) + except NumbaError as exc: + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except Exception: + return None - leaves = [(py, d) for py, _c, d in _iter_record_leaves(return_type.np_type)] - if not any(d.subdtype is not None for _py, d in leaves): - return - probed = _run_udf_probe(numba_func, operands) - if probed is None: + def _raise_unless_fits(shape, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a return of ``shape`` fills an ``expected`` slot. + + ``path`` is the record field, or ``None`` for an array UDT's element. + """ + if _fits_by_broadcast(shape, expected): return - result = probed[0] - if not isinstance(result, tuple) or len(result) != len(leaves): + field = "" if path is None else f" for field {path} of {return_type}" + # A probed shape holds for the values the probe used. Saying so helps + # the author of a UDF whose output shape depends on its inputs' values. + how = " when run on sample values" if probed else "" + slot = f"{return_type} elements are" if path is None else "that field holds" + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)}{field}{how}, but {slot} " + f"{tuple(expected)}. Return an array of that shape, or one that broadcasts to it." + ) + + def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes): + """Reject a UDF whose return cannot fill the output element. + + The wrapper slice-assigns an array return, and each array-typed record + leaf, into the element, so a shape that does not fit raises inside the + cfunc, where the caller never sees it (see + :func:`_compose_wrapper_body`). So does a record leaf the UDF can + return as ``None`` (``v if cond else None``, typed ``Optional``), array + or scalar, on the path that returns it; a leaf that is always ``None`` + does not compile at all, and Numba's traceback says why at length. + Checking here makes each an error when the op is typed. + + An ``Optional`` leaf is rejected from its type, whichever branch the + sample values would take. Numba also keeps the extents of an operand, + or operand field, returned as-is (a ``NestedArray``), so those are + checked from the type alone. An array the UDF builds (``x + y``, + ``x[:2]``) is a plain ``Array``, which has a rank but no extents, so + only running the UDF, on stand-ins built from ``arg_dtypes``, can say + what shape it has. A UDF that returns no built array is never run. + """ + if zkind == "array_elements": + # (record field path, Numba type, slot shape, index into the tuple) + slots = [(None, numba_ret_type, return_type.np_type.shape, None)] + elif zkind == "record_fields": + from .udt_utils import _iter_record_leaves + + # The tuple is flat over the leaves, one element per leaf. + leaves = list(_iter_record_leaves(return_type.np_type)) + if len(numba_ret_type.types) != len(leaves): + return # a wrong field count, not a wrong shape + slots = [ + (path, numba_ret_type.types[i], leaf.shape, i) + for i, (path, _c, leaf) in enumerate(leaves) + ] + else: return - for (path, leaf_dtype), value in zip(leaves, result, strict=True): - if leaf_dtype.subdtype is None: - continue - expected = leaf_dtype.subdtype[1] - shape = getattr(value, "shape", None) - if shape is not None and not _fits_by_broadcast(shape, expected): + unknown = [] + for path, nb_type, expected, i in slots: + if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): + field = "" if path is None else f" for field {path}" raise UdfParseError( - f"UDT UDF returned an array of shape {tuple(shape)} for field " - f"{path} of {return_type} when run on sample values, which holds " - f"{tuple(expected)} there. Return an array whose shape matches or " - f"broadcasts to that. {_SHAPE_HINT}" + f"UDT UDF can return None{field}, but {return_type} cannot hold None. " + f"Return a value on every path." ) + if not expected: + continue # a scalar leaf: nothing to fit + if isinstance(nb_type, numba.core.types.NestedArray): + _raise_unless_fits(nb_type.shape, expected, return_type, path, probed=False) + elif isinstance(nb_type, numba.core.types.Array) or nb_type is None: + # No extents to read (or, with no return type given, no type + # at all), so learn the shape by running the UDF. + unknown.append((path, expected, i)) + if unknown and (ran := _run_udf_probe(numba_func, arg_dtypes)) is not None: + for path, expected, i in unknown: + value = ran[0] if i is None else ran[0][i] + _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) def _get_udt_wrapper( numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None @@ -682,22 +676,19 @@ def _get_udt_wrapper( zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) xsetup, xderef, xptr_type = _input_operand(dtype, "x") wrapper_args = [zptr_type, xptr_type] - probe_operands = [dtype] + arg_dtypes = [dtype] if include_indexes: wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) - probe_operands.extend([UINT64, UINT64]) + arg_dtypes.extend([UINT64, UINT64]) ysetup, yderef_expr, yarg = "", "", "" if dtype2 is not None: ysetup, yderef, yptr_type = _input_operand(dtype2, "y") wrapper_args.append(yptr_type) - probe_operands.append(dtype2) + arg_dtypes.append(dtype2) yarg = ", y_ptr" yderef_expr = f", {yderef}" wrapper_sig = nt.void(*wrapper_args) - if zkind == "array_elements": - _check_array_udf_shape(numba_func, return_type, probe_operands) - elif zkind == "record_fields": - _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) + _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) rcidx = ", row, col" if include_indexes else "" signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" @@ -739,16 +730,8 @@ def _get_udt_wrapper_indexbinary( UINT64.numba_type, tptr_type, ) - # Same registration-time guards as :func:`_get_udt_wrapper`. Without - # them a wrong-shape return is memory-safe but silent: the wrapper's - # ``z[:] =`` raises inside the cfunc, Numba prints and swallows it, - # and the element is left as SuiteSparse found it. - if zkind in ("array_elements", "record_fields"): - probe_operands = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] - if zkind == "array_elements": - _check_array_udf_shape(numba_func, return_type, probe_operands) - else: - _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) + arg_dtypes = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes) signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 45e4736c7..b69c1b844 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2055,18 +2055,17 @@ def _add_corner(x, y): @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_array_udf_shape_errors(): - """Array-UDT UDFs that can't fill the element are rejected at registration. + """An array-UDT UDF whose result can't fill the element is rejected when typed. - Numba's ``Array`` type records ``ndim`` but not extents, so neither case - below is a type error. Both used to reach the cfunc, where the shape - mismatch raises in a context that swallows the exception, handing the - caller an uninitialized element and no error. + Numba's ``Array`` type records ``ndim`` but not extents, so this is not a + type error. It used to reach the cfunc, where the shape mismatch raises in + a context that swallows the exception, handing the caller an uninitialized + element and no error. """ # Shapes unique to this test: ``register_anonymous`` caches by dtype and # freezes the JIT C name at first registration, so sharing a shape with # another test makes both order-dependent. udt9 = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_ShapeErr9") - udt10 = dtypes.register_anonymous(np.dtype((np.float64, (10,))), "_ShapeErr10") def _truncate(x): # pragma: no cover (numba) return x[:2] @@ -2075,35 +2074,16 @@ def _truncate(x): # pragma: no cover (numba) with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): op[udt9] - # Two array UDTs sharing a base dtype and rank are indistinguishable once a - # UDF builds its result, so refuse to guess which one it meant. - def _built(x, y): # pragma: no cover (numba) - return y + 0.0 - - op2 = BinaryOp.register_anonymous(_built, "_shape_err_ambiguous", is_udt=True) - with pytest.raises(UdfParseError, match="matches more than one input array UDT"): - op2[udt9, udt10] - - assert op2[udt9, udt9].return_type is udt9 # a same-type pair is not ambiguous - - # An array UDF whose result matches no input names the mismatch rather - # than telling the user to return an array, which is what they did. - def _recast(x): # pragma: no cover (numba) - return x.astype(np.float32) - - op3 = UnaryOp.register_anonymous(_recast, "_shape_err_recast", is_udt=True) - with pytest.raises(UdfParseError, match="matches no input array UDT"): - op3[udt9] - @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_leaf_shape_errors(): - """A record UDF that under-fills an array-typed leaf is rejected at registration. + """A record UDF that under-fills an array leaf, or can return None, is rejected. The wrapper slice-assigns array leaves, so a short return raises inside the cfunc and abandons the write part-way: leaves after it keep whatever - SuiteSparse had in the buffer, scalar leaves included. + SuiteSparse had in the buffer, scalar leaves included. A ``None`` leaf + raises there too. """ spec = np.dtype([("rl_vec", np.float64, (3,)), ("rl_tag", np.int64)], align=True) udt = dtypes.register_anonymous(spec, "_RecLeafShape") @@ -2115,6 +2095,33 @@ def _short(x, y): # pragma: no cover (numba) with pytest.raises(UdfParseError, match=r"shape \(2,\) for field .* holds \(3,\)"): op[udt] + # ``v if cond else None`` types as Optional. Numba unwraps it where the + # wrapper assigns, and a ``None`` raises there like a short array does, so + # the type alone rejects it: no probe, and no dependence on which branch + # the sample values take. A scalar leaf fails the same way. + def _maybe_none(x, y): # pragma: no cover (numba) + return (x["rl_vec"] + y["rl_vec"] if x["rl_tag"] > 0 else None, x["rl_tag"]) + + op = BinaryOp.register_anonymous(_maybe_none, "_rec_leaf_maybe_none", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_vec'\]"): + op[udt] + + def _maybe_no_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_tag"] if x["rl_tag"] > 0 else None) + + op = BinaryOp.register_anonymous(_maybe_no_tag, "_rec_leaf_maybe_no_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_tag'\]"): + op[udt] + + # A leaf that is always ``None`` would otherwise fail the wrapper compile + # with Numba's full traceback; it gets the same one-line diagnostic. + def _no_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], None) + + op = BinaryOp.register_anonymous(_no_tag, "_rec_leaf_no_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"can return None for field \['rl_tag'\]"): + op[udt] + def _full(x, y): # pragma: no cover (numba) return (x["rl_vec"] + y["rl_vec"], x["rl_tag"] + y["rl_tag"]) @@ -2243,6 +2250,75 @@ def _assign(z, src): # pragma: no cover (numba) assert _fits_by_broadcast(src, dst) is numba_assigns, (src, dst) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): + """The shape check runs the UDF only when Numba's type lacks the extents. + + An operand, or operand field, returned as-is keeps its extents in Numba's + ``NestedArray`` type, so it is checked without running user code. That + includes rejecting one returned into a field of another shape. Only an + array the UDF builds, typed as a plain ``Array``, needs the probe. + """ + from graphblas.core.operator import base as _base + + probed = [] + real_probe = _base._run_udf_probe + + def spy(numba_func, arg_dtypes): + probed.append(numba_func.py_func.__name__) + return real_probe(numba_func, arg_dtypes) + + monkeypatch.setattr(_base, "_run_udf_probe", spy) + udt = dtypes.register_anonymous(np.dtype((np.float32, (18,))), "_ProbeWhen18") + rec = dtypes.register_anonymous( + np.dtype( + [("pw_a", np.float64, (3,)), ("pw_b", np.float64, (2,)), ("pw_n", np.int64)], + align=True, + ), + "_ProbeWhenRec", + ) + + def _second(x, y): # pragma: no cover (numba) + return y + + op = BinaryOp.register_anonymous(_second, "_probe_when_second", is_udt=True) + assert op[udt].return_type is udt + + def _fields(x, y): # pragma: no cover (numba) + return (x["pw_a"], y["pw_b"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_fields, "_probe_when_fields", is_udt=True) + assert op[rec].return_type is rec + + def _swapped(x, y): # pragma: no cover (numba) + return (x["pw_b"], x["pw_a"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_swapped, "_probe_when_swapped", is_udt=True) + with pytest.raises( + UdfParseError, match=r"shape \(2,\) for field \['pw_a'\] of _ProbeWhenRec, but" + ): + op[rec] + assert probed == [] + + def _sum(x, y): # pragma: no cover (numba) + return x + y + + op = BinaryOp.register_anonymous(_sum, "_probe_when_sum", is_udt=True) + assert op[udt].return_type is udt + assert probed == ["_sum"] + + # A UDF that raises on the probe's values is not checked, and still types. + def _no_ones(x, y): # pragma: no cover (numba) + if x[0] == 1: + raise ValueError("the probe's values") + return x + y + + op = BinaryOp.register_anonymous(_no_ones, "_probe_when_raises", is_udt=True) + assert op[udt].return_type is udt + assert probed == ["_sum", "_no_ones"] + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_field_roundtrip(): From 955e30de995259c9f1d9b83f62959585e0573be9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Sun, 27 Sep 2026 10:03:21 -0500 Subject: [PATCH 4/5] Ignore the UDT repr size warning in the shape-check probe test Below SuiteSparse 9, register_anonymous stores a UDT's numpy repr in a type name capped at 128 characters, and warns when the repr is longer. _ProbeWhenRec's is 142, so the test failed on psg 8 draws that ran slow tests on the suitesparse backend. It is about the shape check, not serialization. --- graphblas/tests/test_op.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index b69c1b844..e7a7d23aa 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2252,6 +2252,8 @@ def _assign(z, src): # pragma: no cover (numba) @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow +# _ProbeWhenRec's numpy repr is 142 chars; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): """The shape check runs the UDF only when Numba's type lacks the extents. From 1c4bc0f101be90b5afa82e716a13deda0033be80 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 1 Oct 2026 15:03:17 -0700 Subject: [PATCH 5/5] Reject record returns that cannot be written to their fields Two kinds of record UDF return got past the shape check. A tuple or list of the wrong length for an array field failed silently, like a wrong-shape array. Numba copies a sequence into a field element by element, only into a 1-D field and only at its exact length. It does not broadcast, so even a 1-tuple for a 3-element field raises inside the cfunc, and the element keeps whatever was in the buffer. A tuple's length is in its type; a list's is not, so a list is run on sample values like a built array. An array, tuple or list returned for a scalar field failed only at the wrapper compile, with Numba's long TypingError. It is now rejected from its type with a UdfParseError, as a field that can be None already is. Also test the probe's IndexUnaryOp path, which nothing covered. If the row and column were dropped from the probe call, the check silently did nothing. Pin that a scalar or tuple for an array field is checked without running the UDF, and that every leading 1 is dropped in the broadcast rule. udt.rst names the error and documents the record-field rules. --- docs/user_guide/udt.rst | 16 +++- graphblas/core/operator/base.py | 61 +++++++++++- graphblas/tests/test_indexbinary.py | 2 +- graphblas/tests/test_op.py | 138 +++++++++++++++++++++++++++- 4 files changed, 206 insertions(+), 11 deletions(-) diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 341ab6807..ae2cef9e2 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -179,6 +179,11 @@ For nested record UDTs the tuple is *flat over the leaves*. Given return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value (e.g., one of the inputs) is also fine and preserves the nested shape. +Every field needs a value on every path through the UDF. A record cannot hold +``None``, so a UDF that can return ``None`` for a field +(``v if cond else None``) is rejected when the op is typed. So is an array, +tuple, or list returned for a scalar field. + For array UDTs each operand arrives as a numpy view of that element's values, in the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and @@ -202,10 +207,13 @@ element, either at the element's own shape or at one that broadcasts to it (a # c[0] = [5.0, 11.0, 17.0] A return that cannot fill the element, such as ``x[:2]`` from a 3-element UDT, -is rejected when the op is typed for the UDT (``op[point3]``, or the first -operation that uses it). Numba does not know the shape of an array the UDF -builds, so the check runs the UDF once on sample values; a UDF that raises on -those values is not checked. +is rejected with a ``UdfParseError`` when the op is typed for the UDT +(``op[point3]``, or the first operation that uses it). The same goes for an +array returned into an array-typed field of a record UDT. A tuple or list also +fills such a field, but only a one-dimensional one and only at its exact +length; unlike an array, it does not broadcast. Numba does not know the shape +of an array the UDF builds, or the length of a list, so the check runs the UDF +once on sample values; a UDF that raises on those values is not checked. If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index a94fe2327..e4c41b95e 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -604,6 +604,37 @@ def _raise_unless_fits(shape, expected, return_type, path, *, probed): f"{tuple(expected)}. Return an array of that shape, or one that broadcasts to it." ) + def _raise_unless_sequence_fits(kind, length, expected, return_type, path, *, probed): + """Raise ``UdfParseError`` unless a ``kind`` of ``length`` fills record field ``path``. + + Numba copies a tuple or list into an array field element by element, + which works only for a one-dimensional field and only at its exact + length. Unlike an array, a sequence does not broadcast, so a 1-tuple + does not fill a 3-element field. A wrong length raises inside the cfunc + just as a wrong-shape array does; a field of higher rank fails the + wrapper compile. + """ + if expected == (length,): + return + how = " when run on sample values" if probed else "" + if len(expected) == 1: + fix = f"Return {expected[0]} values, or an array that broadcasts to {expected}." + else: + fix = f"A {kind} only fills a one-dimensional field; return an array." + raise UdfParseError( + f"UDT UDF returned a {kind} of length {length} for field {path} of {return_type}" + f"{how}, but that field holds {expected}. {fix}" + ) + + # Numba types a record UDF might return for a scalar field, none of which the + # wrapper can write there, each with the kind the error names. ``Array`` + # covers ``NestedArray``. + _NON_SCALAR_KINDS = ( + ("an array", numba.core.types.Array), + ("a tuple", numba.core.types.BaseTuple), + ("a list", numba.core.types.List), + ) + def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_dtypes): """Reject a UDF whose return cannot fill the output element. @@ -613,8 +644,10 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ :func:`_compose_wrapper_body`). So does a record leaf the UDF can return as ``None`` (``v if cond else None``, typed ``Optional``), array or scalar, on the path that returns it; a leaf that is always ``None`` - does not compile at all, and Numba's traceback says why at length. - Checking here makes each an error when the op is typed. + does not compile at all, and Numba's traceback says why at length. Nor + does an array, tuple or list returned for a scalar leaf: Numba has no + record setitem for them. Checking here makes each an error when the op + is typed. An ``Optional`` leaf is rejected from its type, whichever branch the sample values would take. Numba also keeps the extents of an operand, @@ -622,7 +655,10 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ checked from the type alone. An array the UDF builds (``x + y``, ``x[:2]``) is a plain ``Array``, which has a rank but no extents, so only running the UDF, on stand-ins built from ``arg_dtypes``, can say - what shape it has. A UDF that returns no built array is never run. + what shape it has. A tuple or list for an array leaf must match its + length exactly (:func:`_raise_unless_sequence_fits`); a tuple's length + is in its type, and a list's is learned the same way as a built + array's. A UDF that returns no built array or list is never run. """ if zkind == "array_elements": # (record field path, Numba type, slot shape, index into the tuple) @@ -649,6 +685,12 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ f"Return a value on every path." ) if not expected: + for kind, cls in _NON_SCALAR_KINDS: + if isinstance(nb_type, cls): + raise UdfParseError( + f"UDT UDF returned {kind} for field {path} of {return_type}, but " + f"that field holds a scalar. Return a scalar for it." + ) continue # a scalar leaf: nothing to fit if isinstance(nb_type, numba.core.types.NestedArray): _raise_unless_fits(nb_type.shape, expected, return_type, path, probed=False) @@ -656,10 +698,21 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ # No extents to read (or, with no return type given, no type # at all), so learn the shape by running the UDF. unknown.append((path, expected, i)) + elif isinstance(nb_type, numba.core.types.BaseTuple): + _raise_unless_sequence_fits( + "tuple", len(nb_type.types), expected, return_type, path, probed=False + ) + elif isinstance(nb_type, numba.core.types.List): + unknown.append((path, expected, i)) # a list's length is not in its type if unknown and (ran := _run_udf_probe(numba_func, arg_dtypes)) is not None: for path, expected, i in unknown: value = ran[0] if i is None else ran[0][i] - _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) + if isinstance(value, list): + _raise_unless_sequence_fits( + "list", len(value), expected, return_type, path, probed=True + ) + else: + _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) def _get_udt_wrapper( numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index 82ffbdbf0..19b65d720 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -301,7 +301,7 @@ def test_bind_raw_array_udt_theta(): def test_array_udt_udf_shape_is_checked(): - """An IBO whose UDF builds a wrong-shape array is rejected at registration. + """An IBO whose UDF builds a wrong-shape array is rejected when typed. The same guard the binary and unary paths get. Numba's ``Array`` type carries ``ndim`` but not extents, so the mismatch is not a type error; left diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index e7a7d23aa..36152dfe0 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2074,6 +2074,16 @@ def _truncate(x): # pragma: no cover (numba) with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): op[udt9] + # IndexUnaryOp passes the row and column between x and the thunk, and so + # must the probe: with the thunk's array in an index slot, ``i + j`` would + # not type as a slice bound. The probe's indices are 1, so this is (2,). + def _head(x, i, j, thunk): # pragma: no cover (numba) + return (x + thunk)[: i + j] + + op = IndexUnaryOp.register_anonymous(_head, "_shape_err_head", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[udt9] + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow @@ -2122,6 +2132,23 @@ def _no_tag(x, y): # pragma: no cover (numba) with pytest.raises(UdfParseError, match=r"can return None for field \['rl_tag'\]"): op[udt] + # An array for a scalar leaf would fail the wrapper compile the same way + # (Numba has no record setitem for it), whether it is an operand's field + # or built. + def _arr_for_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_vec"]) + + op = BinaryOp.register_anonymous(_arr_for_tag, "_rec_leaf_arr_for_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned an array for field \['rl_tag'\]"): + op[udt] + + def _built_for_tag(x, y): # pragma: no cover (numba) + return (x["rl_vec"], x["rl_vec"] + y["rl_vec"]) + + op = BinaryOp.register_anonymous(_built_for_tag, "_rec_leaf_built_for_tag", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned an array for field \['rl_tag'\]"): + op[udt] + def _full(x, y): # pragma: no cover (numba) return (x["rl_vec"] + y["rl_vec"], x["rl_tag"] + y["rl_tag"]) @@ -2204,6 +2231,92 @@ def _fill_leaf(x, y): # pragma: no cover (numba) np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) assert got["bc_tag"] == 15 + # A scalar fills the leaf the same way. Its type says it is no array, so + # the check lets it through without running the UDF. + def _fill_leaf_scalar(x, y): # pragma: no cover (numba) + return (x["bc_vec"][0] + y["bc_vec"][0], x["bc_tag"] + y["bc_tag"]) + + op2 = BinaryOp.register_anonymous(_fill_leaf_scalar, "_rec_leaf_bcast_scalar", is_udt=True) + got = v.ewise_mult(w, op2).new()[0].new().value + np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) + assert got["bc_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_leaf_sequence_return(): + """A tuple or list fills an array leaf only at its exact length. + + Numba copies a sequence into the field element by element, and unlike an + array it does not broadcast one. A wrong length raises inside the cfunc, + and the element comes back holding whatever was in the buffer. A tuple's + length is in its type; a list's is learned by running the UDF. + """ + spec = np.dtype([("sq_v", np.float64, (3,)), ("sq_t", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafSeq") + v = Vector(udt, size=1) + v[0] = ([1.0, 2.0, 3.0], 7) + + def _tuple3(x, y): # pragma: no cover (numba) + return ((x["sq_v"][0], x["sq_v"][1], 9.0), x["sq_t"]) + + def _list3(x, y): # pragma: no cover (numba) + return ([x["sq_v"][0], x["sq_v"][1], 9.0], x["sq_t"]) + + for func in [_tuple3, _list3]: + op = BinaryOp.register_anonymous(func, f"_rec_leaf_seq{func.__name__}", is_udt=True) + got = v.ewise_mult(v, op).new()[0].new().value + np.testing.assert_array_equal(got["sq_v"], [1.0, 2.0, 9.0]) + assert got["sq_t"] == 7 + + def _tuple2(x, y): # pragma: no cover (numba) + return ((x["sq_v"][0], 9.0), x["sq_t"]) + + op = BinaryOp.register_anonymous(_tuple2, "_rec_leaf_seq_tuple2", is_udt=True) + with pytest.raises(UdfParseError, match=r"tuple of length 2 for field \['sq_v'\] of _RecL"): + op[udt] + + # As a (1,) array this would broadcast; as a tuple it does not. + def _tuple1(x, y): # pragma: no cover (numba) + return ((9.0,), x["sq_t"]) + + op = BinaryOp.register_anonymous(_tuple1, "_rec_leaf_seq_tuple1", is_udt=True) + with pytest.raises(UdfParseError, match=r"tuple of length 1 .* holds \(3,\)\. Return 3 values"): + op[udt] + + def _list2(x, y): # pragma: no cover (numba) + return ([x["sq_v"][0], 9.0], x["sq_t"]) + + op = BinaryOp.register_anonymous(_list2, "_rec_leaf_seq_list2", is_udt=True) + with pytest.raises(UdfParseError, match=r"list of length 2 .* when run on sample values"): + op[udt] + + # A sequence never fills a field of higher rank; Numba will not compile it. + spec2 = np.dtype([("sq_m", np.float64, (2, 2)), ("sq_n", np.int64)], align=True) + udt2 = dtypes.register_anonymous(spec2, "_RecLeafSeq2D") + + def _flat4(x, y): # pragma: no cover (numba) + return ((1.0, 2.0, 3.0, 4.0), x["sq_n"]) + + op = BinaryOp.register_anonymous(_flat4, "_rec_leaf_seq_flat4", is_udt=True) + with pytest.raises(UdfParseError, match=r"only fills a one-dimensional field"): + op[udt2] + + # Nor a scalar field, like an array (test_udt_record_array_leaf_shape_errors). + def _tuple_for_t(x, y): # pragma: no cover (numba) + return (x["sq_v"], (x["sq_t"], 1)) + + op = BinaryOp.register_anonymous(_tuple_for_t, "_rec_leaf_seq_tuple_for_t", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned a tuple for field \['sq_t'\] .* a scalar"): + op[udt] + + def _list_for_t(x, y): # pragma: no cover (numba) + return (x["sq_v"], [x["sq_t"], 1]) + + op = BinaryOp.register_anonymous(_list_for_t, "_rec_leaf_seq_list_for_t", is_udt=True) + with pytest.raises(UdfParseError, match=r"returned a list for field \['sq_t'\] .* a scalar"): + op[udt] + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow @@ -2240,6 +2353,10 @@ def _assign(z, src): # pragma: no cover (numba) ((2, 3), (2, 3)), ((2, 3), (6,)), ((2, 3), (3, 2)), + # Every leading one is dropped, not just the first. + ((6,), (1, 1, 6)), + ((2, 3), (1, 1, 3)), + ((2, 3), (2, 1, 3)), ]: try: _assign(np.zeros(dst), np.ones(src)) @@ -2259,8 +2376,10 @@ def test_udt_udf_shape_check_runs_udf_only_for_built_arrays(monkeypatch): An operand, or operand field, returned as-is keeps its extents in Numba's ``NestedArray`` type, so it is checked without running user code. That - includes rejecting one returned into a field of another shape. Only an - array the UDF builds, typed as a plain ``Array``, needs the probe. + includes rejecting one returned into a field of another shape. A tuple is + checked the same way, since its length is in its type. Only an array the + UDF builds, typed as a plain ``Array``, needs the probe; so does a list, + whose type has no length (``test_udt_record_leaf_sequence_return``). """ from graphblas.core.operator import base as _base @@ -2301,6 +2420,21 @@ def _swapped(x, y): # pragma: no cover (numba) UdfParseError, match=r"shape \(2,\) for field \['pw_a'\] of _ProbeWhenRec, but" ): op[rec] + + # A scalar for an array leaf broadcasts into it; its type says it is no + # array, so there is nothing to probe. + def _scalars(x, y): # pragma: no cover (numba) + return (x["pw_n"], y["pw_n"], x["pw_n"] + y["pw_n"]) + + op = BinaryOp.register_anonymous(_scalars, "_probe_when_scalars", is_udt=True) + assert op[rec].return_type is rec + + # A tuple carries its length in its type, so it is checked without a run. + def _tuple_leaf(x, y): # pragma: no cover (numba) + return ((x["pw_a"][0], x["pw_a"][1], 9.0), y["pw_b"], x["pw_n"]) + + op = BinaryOp.register_anonymous(_tuple_leaf, "_probe_when_tuple_leaf", is_udt=True) + assert op[rec].return_type is rec assert probed == [] def _sum(x, y): # pragma: no cover (numba)