diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 15cbedbb4..ae2cef9e2 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -179,10 +179,17 @@ 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 -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 +206,15 @@ 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 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 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..e4c41b95e 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,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. + 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 ( @@ -521,6 +526,194 @@ 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 _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 _udf_probe_value(dtype): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. + + 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] + + 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. + """ + 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 + + 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 + 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 _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. + + 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. 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, + 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 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) + 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 + 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 can return None{field}, but {return_type} cannot hold None. " + 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) + 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)) + 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] + 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 ): @@ -536,15 +729,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] + arg_dtypes = [dtype] if include_indexes: wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + 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) + arg_dtypes.append(dtype2) yarg = ", y_ptr" yderef_expr = f", {yderef}" wrapper_sig = nt.void(*wrapper_args) + _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}):" @@ -586,6 +783,8 @@ def _get_udt_wrapper_indexbinary( UINT64.numba_type, tptr_type, ) + 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_indexbinary.py b/graphblas/tests/test_indexbinary.py index c56bb86f9..19b65d720 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 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 + 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..36152dfe0 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -2052,6 +2052,409 @@ def _add_corner(x, y): ) +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_shape_errors(): + """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 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") + + 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] + + # 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 +def test_udt_record_array_leaf_shape_errors(): + """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. 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") + + 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] + + # ``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] + + # 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"]) + + 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 + + # 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 +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)), + # 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)) + 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 +# _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. + + 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. 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 + + 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] + + # 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) + 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():