diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 88b77086a..9bb5e803a 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -209,7 +209,10 @@ dtype, so it never narrows. Below, ``v`` is a Vector of an ``INT8[3]`` UDT. but converts exactly, such as ``0`` for a bool field or ``2.0`` for an int field, still converts with a ``DeprecationWarning`` and will raise in a future version; any other raises ``ValueError``, and a float too - large for a float field raises ``OverflowError``. + large for a float field raises ``OverflowError``. A Python or numpy + number beside a user-defined op that cannot be typed with an element + (``x[:2] * t`` with a 6-element ``t``) is typed as itself instead, as a + ``Scalar`` of it would be. - ``v.ewise_union(w, binary.plus, 0.5, 0)`` raises A Monoid used element-wise (``apply``, ``ewise_mult``, ``ewise_union``) uses @@ -365,12 +368,55 @@ 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. +on sample values; a UDF that raises on those values is not checked, since it +may raise only for some values. An exception the UDF raises when the op runs +is printed and goes no further (Numba cannot pass it to GraphBLAS), so that +element is not written; a built-in output, such as a select's bool, is 0 or +``False`` there. 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. +Naming the output type +~~~~~~~~~~~~~~~~~~~~~~ + +By default the return type is worked out from what the UDF returns, matched +against the input dtypes. That can only name a type the operator already has +in hand, so an output UDT that is not one of the operands is out of reach. +The ``x[:2]`` rejection above is a case of this: the output really is a +2-element UDT, but nothing says so. + +Pass ``ret_dtype`` to say it outright. It takes anything ``lookup_dtype`` +accepts and requires ``is_udt=True``: + +.. code-block:: python + + nine = gb.dtypes.register_anonymous(np.dtype((np.float64, (9,))), "Nine") + three = gb.dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Three") + + head = gb.unary.register_anonymous(lambda x: x[:3], "head", is_udt=True, ret_dtype=three) + head[nine].return_type # Three + +``ret_dtype`` is available on ``register_anonymous`` and ``register_new`` for +``UnaryOp``, ``BinaryOp``, ``IndexUnaryOp``, and ``IndexBinaryOp``. It is a +property of the operator, not of a particular input dtype: the same output +type applies to every dtype the operator is typed for. An operator whose +output type should vary with its inputs still needs one registration per +output type. + +Declaring the type does not switch off the shape check described above; it +points the check at the declared type instead. The result must also be +something the declared type can hold: one value per field for a record, an +array or a scalar for an array UDT, and a scalar for a scalar type. A UDF that +breaks either rule is rejected when the op is typed (``op[dtype]``, or its +first use), so a wrong ``ret_dtype`` is an error rather than a silently +mis-typed result. + +It does not apply to ``SelectOp``, whose return type GraphBLAS fixes at +``BOOL``, nor to an op registered without ``is_udt=True``, whose return type +comes from compiling the function against each input type in turn. + .. _udt_jit_introspection: JIT and introspection diff --git a/graphblas/core/literals.py b/graphblas/core/literals.py index d43c97aa8..b9255093e 100644 --- a/graphblas/core/literals.py +++ b/graphblas/core/literals.py @@ -127,10 +127,52 @@ def _literal_operand(dtype, value, op, *, left=False): literal_type = _literal_type(dtype, value, exact=exact) if literal_type is not None: return value, literal_type + if (own := _own_scalar_type_for_user_op(dtype, value, op, left)) is not None: + return value, own _check_literal_fits(dtype, value) return value, dtype +def _own_scalar_type_for_user_op(dtype, value, op, left): + """Return a number's own type, for a user's op that cannot take it as an element, or None. + + Beside a user's ``is_udt`` op (an untyped BinaryOp, IndexUnaryOp or + SelectOp), a literal becomes an element of the UDT. A Python or numpy + number for an op whose function cannot be typed with that element (it + fails to compile, or raises on every set of sample values: ``x[:2] * t`` + with a 6-element ``t``) is typed as itself instead, as a ``Scalar`` of it + would be, when the op can be typed so. Otherwise it stays an element, + and the op reports why it cannot take it. + """ + from ..exceptions import UdfParseError + + if ( + not getattr(op, "_is_udt", False) + or isinstance(op, (TypedOpBase, Monoid)) + or not ( + type(value) in {bool, int, float, complex} or isinstance(value, (np.number, np.bool_)) + ) + or (type(value) is int and not -(2**63) <= value < 2**63) # no INT64 holds it + or (own := lookup_dtype(type(value), value))._is_udt + ): + return None + + from .operator.udf import _RAISES_ON_SAMPLES + + def types(dtype2): + try: + return op[(dtype2, dtype) if left else (dtype, dtype2)] + except UdfParseError: + return None + + # A UDF that raised on every stand-in element types, but cannot take an element. + # The number's own type need only type: its stand-in is a fixed sample, not the number. + element = types(dtype) + if element is not None and element not in _RAISES_ON_SAMPLES or types(own) is None: + return None + return own + + # Built-in ops that compare their operands, and so keep no result type. _COMPARISONS = {"eq", "ne", "lt", "le", "gt", "ge"} _COMPARISONS.update([f"value{name}" for name in list(_COMPARISONS)]) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index f895c7a87..0b539744d 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -93,6 +93,39 @@ def _hasop(module, name): ) +def _validate_ret_dtype(ret_dtype, opclass, *, is_udt, parameterized): + """Normalize a user-supplied ``ret_dtype`` to a DataType, or raise. + + ``ret_dtype`` names the operator's output type outright instead of letting + it be inferred from what the UDF returns. Inference can only name a type it + can see, which is why this is limited to the UDT path: the builtin path + derives its output from Numba's typing of each sample input, and forcing a + single type across all of them would silently recast results. + """ + if ret_dtype is None: + return None + if not is_udt: + raise ValueError( + f"{opclass}: ret_dtype requires is_udt=True. The return type for builtin " + f"dtypes comes from compiling the function for each input type, so a single " + f"fixed type cannot describe it." + ) + if parameterized: + raise ValueError( + f"{opclass}: ret_dtype does not work with parameterized=True. " + f"A parameterized operator builds and registers its function when called, " + f"and that inner registration does not accept a return dtype; register the " + f"built function without parameterized=True to declare one." + ) + try: + return lookup_dtype(ret_dtype) + except (ValueError, TypeError) as exc: + raise ValueError( + f"{opclass}: ret_dtype={ret_dtype!r} is not a recognized dtype. " + f"Pass a DataType, a numpy dtype, or a name such as 'FP64'." + ) from exc + + class OpPath: def __init__(self, parent, name): self._parent = parent @@ -576,21 +609,25 @@ def _deserialize(cls, name, *args): return cls.register_new(name, *args) @classmethod - def _deserialize_udf(cls, name, orig_func, is_udt): + def _deserialize_udf(cls, name, orig_func, is_udt, ret_dtype=None): """Re-register a named UDF on unpickle, or reuse if already present. Shared by the five UDF-capable subclasses (UnaryOp, BinaryOp, IndexUnaryOp, SelectOp, IndexBinaryOp), all of which use the - default ``__reduce__`` below. + default ``__reduce__`` below. ``ret_dtype`` is passed on only when set: + SelectOp shares this path and takes no ret_dtype, and an op without one + pickles as the 3-tuple it always did. """ if (rv := cls._find(name)) is not None: return rv - return cls.register_new(name, orig_func, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_new(name, orig_func, is_udt=is_udt, **kwargs) @classmethod - def _deserialize_anon_udf(cls, func, name, is_udt): + def _deserialize_anon_udf(cls, func, name, is_udt, ret_dtype=None): """Re-register an anonymous UDF on unpickle.""" - return cls.register_anonymous(func, name, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_anonymous(func, name, is_udt=is_udt, **kwargs) def __reduce__(self): """Default ``__reduce__`` for UDF-capable subclasses. @@ -600,13 +637,19 @@ def __reduce__(self): ``Aggregator`` define their own ``__reduce__`` because their pickle shape differs (they hold a binary op + identity, etc.). """ + # ``ret_dtype`` rides along only when set, so an op without one pickles + # exactly as before and older versions can still read it. + ret = () if (ret_dtype := getattr(self, "_ret_dtype", None)) is None else (ret_dtype,) if self._anonymous: if hasattr(self.orig_func, "_parameterized_info"): return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (type(self)._deserialize_anon_udf, (self.orig_func, self.name, self._is_udt)) + return ( + type(self)._deserialize_anon_udf, + (self.orig_func, self.name, self._is_udt, *ret), + ) if (name := f"{self._modname}.{self.name}") in _STANDARD_OPERATOR_NAMES: return name - return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt)) + return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt, *ret)) @classmethod def _check_supports_udf(cls, method_name): diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 517ba72e1..cb8abde73 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -37,6 +37,7 @@ TypedOpBase, _call_op, _hasop, + _validate_ret_dtype, ) if _has_numba: @@ -48,7 +49,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) try: @@ -504,26 +505,38 @@ def _finalize_typed_binaryop(parent, binary_udf, name, type_, ret_type): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(binary_udf(bool(x0), bool(y0))) else: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(bool(x[0]), bool(y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = binary_udf(bool(x0), bool(y0)) elif ret_type == BOOL: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(x[0], y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(binary_udf(x0, y0)) else: def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(x[0], y[0]) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = binary_udf(x0, y0) binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(binary_wrapper) new_binary = ffi_new("GrB_BinaryOp*") @@ -608,6 +621,7 @@ class BinaryOp(OpBase): "_numba_func", "_custom_dtype", "_defer_builds", + "_ret_dtype", ) _module = binary _modname = "binary" @@ -713,11 +727,17 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + # This rejects ret_dtype unless is_udt, which keeps it disjoint from the + # deferred builtin path below: that one only runs under ``not is_udt``, + # so an op can never be both deferred and carrying a declared return + # type. ``_build_deferred`` reads its ret_type back from ``.types`` and + # never consults ``_ret_dtype``. + ret_dtype = _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=False) success = False # The error model has to be set here as well as on the cfunc wrapper in # ``_finalize_typed_binaryop``. A Dispatcher keeps one compilation per @@ -741,7 +761,14 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): binary_udf = numba.njit(func, error_model="numpy") else: binary_udf = numba.njit(func, error_model="numpy") - new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=binary_udf, + ret_dtype=ret_dtype, + ) return_types = {} if not is_udt: # ``cache=True`` marks the module-level built-in UDFs (floordiv and @@ -827,7 +854,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="binary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) binary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) @@ -886,7 +913,9 @@ def _compile_udt(self, dtype, dtype2): return op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a BinaryOp without registering it in the ``graphblas.binary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -915,6 +944,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- BinaryOp or ParameterizedBinaryOp @@ -922,12 +958,21 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) return ParameterizedBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod def register_new( - cls, name, func, *, parameterized=False, is_udt=False, lazy=False, _cache=False + cls, + name, + func, + *, + parameterized=False, + is_udt=False, + lazy=False, + ret_dtype=None, + _cache=False, ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. @@ -963,6 +1008,10 @@ def register_new( delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> def max_zero(x, y): @@ -991,6 +1040,9 @@ def register_new( """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( @@ -1000,6 +1052,7 @@ def register_new( "func": func, "parameterized": parameterized, "is_udt": is_udt, + "ret_dtype": ret_dtype, "_cache": _cache, }, ) @@ -1007,7 +1060,7 @@ def register_new( binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache, ret_dtype=ret_dtype) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -1232,8 +1285,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self._monoid = None self._commutes_to = None self._semiring_commutes_to = None diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index 0e0eff04d..d9c427d65 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -9,7 +9,7 @@ from ...exceptions import EmptyObject, UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values, unify -from .base import OpBase, ParameterizedUdf, TypedOpBase +from .base import OpBase, ParameterizedUdf, TypedOpBase, _validate_ret_dtype from .binary import BinaryOp _has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") @@ -22,7 +22,7 @@ _compile_for_dtype, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -330,7 +330,7 @@ class IndexBinaryOp(OpBase): no built-ins; all IndexBinaryOps are user-defined. """ - __slots__ = "orig_func", "_is_udt", "_numba_func" + __slots__ = "orig_func", "_is_udt", "_numba_func", "_ret_dtype" _module = indexbinary _modname = "indexbinary" _custom_dtype = None @@ -349,7 +349,7 @@ class IndexBinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not _has_idxbinop: raise RuntimeError( "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " @@ -359,12 +359,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype( + ret_dtype, "indexbinary", is_udt=is_udt, parameterized=False + ) success = False # Set on the Dispatcher, not just the cfunc wrapper; see the note in # ``BinaryOp._build``. indexbinary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( - name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexbinary_udf + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexbinary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -412,16 +420,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 z[0] = bool( - indexbinary_udf( - bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) - ) + indexbinary_udf(bool(x0), ix, jx, bool(y0), iy, jy, bool(theta0)) ) else: @@ -429,23 +441,27 @@ def indexbinary_wrapper( def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = indexbinary_udf( - bool(x[0]), ix, jx, bool(y[0]), iy, jy, bool(theta[0]) - ) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = indexbinary_udf(bool(x0), ix, jx, bool(y0), iy, jy, bool(theta0)) elif ret_type == BOOL: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = bool(indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0])) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = bool(indexbinary_udf(x0, ix, jx, y0, iy, jy, theta0)) else: def indexbinary_wrapper( z, x, ix, jx, y, iy, jy, theta ): # pragma: no cover (numba) - z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) + x0, y0, theta0 = x[0], y[0], theta[0] + z[0] = 0 + z[0] = indexbinary_udf(x0, ix, jx, y0, iy, jy, theta0) indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( indexbinary_wrapper @@ -509,7 +525,7 @@ def _compile_udt(self, dtype, dtype2, theta_type=None): numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=dtypes ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, *dtypes) + ret_type = _udt_ret_type(self, numba_ret_type, *dtypes) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( numba_func, ret_type, dtype, dtype2, theta_type, numba_ret_type=numba_ret_type ) @@ -572,7 +588,9 @@ def _add(self, op, *, is_jit=False): self._udt_ops[dtypes] = op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register an IndexBinaryOp without adding it to the ``indexbinary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -593,6 +611,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals is_udt : bool, default False Whether the operator is intended to operate on user-defined types. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- IndexBinaryOp or ParameterizedIndexBinaryOp @@ -600,11 +625,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) return ParameterizedIndexBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexBinaryOp under the ``graphblas.indexbinary`` namespace. Parameters @@ -624,23 +652,36 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal lazy : bool, default False When True, defer compilation until the operator is first used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexbinary.register_new("index_dist", lambda x, ix, jx, y, iy, jy, t: abs(ix - iy)) """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, idxbinop) else: - idxbinop = cls._build(name, func, is_udt=is_udt) + idxbinop = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, idxbinop) if not cls._initialized: # pragma: no cover (safety) @@ -656,8 +697,11 @@ def _initialize(cls): # No built-in IndexBinaryOps to register. cls._initialized = True - def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None): + def __init__( + self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None, ret_dtype=None + ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self._is_udt = is_udt diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index 7d187bfe7..7f1a2e310 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -7,7 +7,7 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op +from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op, _validate_ret_dtype if _has_numba: import numba @@ -18,7 +18,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -86,7 +86,7 @@ class IndexUnaryOp(OpBase): Built-in and registered IndexUnaryOps are located in the ``graphblas.indexunary`` namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _module = indexunary _modname = "indexunary" _custom_dtype = None @@ -111,17 +111,23 @@ class IndexUnaryOp(OpBase): "rowindex", "colindex"} # fmt: skip @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=False) success = False # Set on the Dispatcher, not just the cfunc wrapper; see the note in # ``BinaryOp._build``. indexunary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( - name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexunary_udf + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexunary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -156,26 +162,38 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = bool(indexunary_udf(bool(x[0]), row, col, bool(y[0]))) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(indexunary_udf(bool(x0), row, col, bool(y0))) else: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = indexunary_udf(bool(x[0]), row, col, bool(y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = indexunary_udf(bool(x0), row, col, bool(y0)) elif ret_type == BOOL: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = bool(indexunary_udf(x[0], row, col, y[0])) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = bool(indexunary_udf(x0, row, col, y0)) else: def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) - z[0] = indexunary_udf(x[0], row, col, y[0]) + x0, y0 = x[0], y[0] + z[0] = 0 + z[0] = indexunary_udf(x0, row, col, y0) indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( indexunary_wrapper @@ -218,7 +236,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="indexunary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) indexunary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) @@ -227,7 +245,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a IndexUnary without registering it in the ``graphblas.indexunary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -261,6 +281,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- return IndexUnaryOp or ParameterizedIndexUnaryOp @@ -268,11 +295,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) return ParameterizedIndexUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexUnaryOp and save it to ``graphblas.indexunary`` namespace. If the return type is Boolean, the function will also be registered as a SelectOp @@ -315,6 +345,10 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexunary.register_new("row_mod", lambda x, i, j, thunk: i % max(thunk, 2)) @@ -323,20 +357,34 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: indexunary_op = ParameterizedIndexUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, indexunary_op) else: - indexunary_op = cls._build(name, func, is_udt=is_udt) + indexunary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, indexunary_op) - # If return type is BOOL, register additionally as a SelectOp - if all(x == BOOL for x in indexunary_op.types.values()): + # If return type is BOOL, register additionally as a SelectOp. A UDT + # op has no types until it is compiled, so only a declared + # ``ret_dtype`` can rule it out here. + declared = indexunary_op._ret_dtype + if (declared is None or declared == BOOL) and all( + x == BOOL for x in indexunary_op.types.values() + ): from .select import SelectOp select_module, funcname = SelectOp._remove_nesting(name, strict=False) @@ -397,8 +445,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 1c1138ae5..a78b391b1 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -272,8 +272,10 @@ def _compile_udt(self, dtype, dtype2): _compile_udf_for_udt( numba_func, sig, op_kind="select", op_name=self.name, dtypes=(dtype, dtype2) ) + # Its return type lets the wrapper see an array returned for the bool + numba_ret_type = numba_func.overloads[sig].signature.return_type select_wrapper, wrapper_sig = _get_udt_wrapper( - numba_func, BOOL, dtype, dtype2, include_indexes=True + numba_func, BOOL, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) return _finalize_udt_op( self, dtype, dtype2, BOOL, select_wrapper, wrapper_sig, TypedUserSelectOp diff --git a/graphblas/core/operator/udf.py b/graphblas/core/operator/udf.py index ab5f0b8a2..1fb644b73 100644 --- a/graphblas/core/operator/udf.py +++ b/graphblas/core/operator/udf.py @@ -19,6 +19,7 @@ import re import warnings +import weakref import numba import numpy as np @@ -100,6 +101,7 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ + raises_on_samples = getattr(wrapper, "_raises_on_samples", None) try: wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) except NumbaError as exc: @@ -130,6 +132,8 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t key = (dtype, dtype2) parent_op._udt_types[key] = ret_type parent_op._udt_ops[key] = op + if raises_on_samples is not None: + _RAISES_ON_SAMPLES[op] = raises_on_samples return op @@ -190,6 +194,92 @@ def _summarize_numba_typing_error(exc): return "Numba could not compile the function for these input types" +def _udt_ret_type(parent_op, numba_ret_type, *dtypes): + """Return the operator's declared ``ret_dtype``, else infer one from the UDF. + + Inference names a built-in dtype or an operand's type, so an output UDT + that appears nowhere in the inputs is unreachable without this. A + declared type is checked against the UDF's return first. + """ + if (ret_dtype := parent_op._ret_dtype) is not None: + _check_declared_ret_type(ret_dtype, numba_ret_type) + return ret_dtype + return _resolve_udt_return_type(numba_ret_type, *dtypes) + + +def _check_declared_ret_type(ret_dtype, numba_ret_type): + """Reject a UDF whose return cannot be written as the declared ``ret_dtype``. + + Inference only names a type the return can be written as; a declared + type skips it. Left to the wrapper, a mismatch is a raw Numba error when + the wrapper compiles, or a silent wrong answer when it runs: a longer + tuple loses its extra values, an array written to a scalar keeps its + first element, and a tuple written to an array element leaves it + uninitialized. Shapes, and record leaves that can be ``None``, are + checked afterwards by :func:`_check_udf_fills_output`. + """ + nt = numba.core.types + np_type = ret_dtype.np_type + if ret_dtype._is_udt and np_type.names is not None: + nleaves = len(list(_iter_record_leaves(np_type))) + if isinstance(numba_ret_type, nt.BaseTuple): + if len(numba_ret_type.types) == nleaves: + return + # The tuple is flat over the leaves, so a nested record counts + # the fields of its inner records, not the inner records. + nested = "" if nleaves == len(np_type.names) else " once nested records are flattened" + raise UdfParseError( + f"UDT UDF returned a tuple of length {len(numba_ret_type.types)}, but " + f"ret_dtype={ret_dtype} has {nleaves} fields{nested}. " + f"Return one value per field." + ) + if numba_ret_type == ret_dtype.numba_type: + return # a ``ret_dtype`` value, such as an operand of that type + wanted = f"a tuple of its {nleaves} fields, or a {ret_dtype} value" + elif ret_dtype._is_udt and np_type.subdtype is not None: + _refuse_complex_into_real(ret_dtype, np_type.subdtype[0], numba_ret_type) + if isinstance(numba_ret_type, (nt.Array, nt.Number, nt.Boolean)): + return # a scalar fills every slot; the shape is checked later + wanted = "an array, or a scalar to fill every slot" + else: + _refuse_complex_into_real(ret_dtype, np_type, numba_ret_type) + if isinstance(numba_ret_type, (nt.Number, nt.Boolean)): + return + if ret_dtype._is_udt and numba_ret_type == ret_dtype.numba_type: + return # a value of a date, duration or text UDT, such as an operand returned as is + wanted = "a scalar" + if isinstance(numba_ret_type, nt.BaseTuple): + returned = f"a tuple of length {len(numba_ret_type.types)}" + elif isinstance(numba_ret_type, nt.Array): + returned = "an array" + elif isinstance(numba_ret_type, nt.Record): + returned = f"a record with fields {list(numba_ret_type.fields)}" + elif isinstance(numba_ret_type, (nt.Optional, nt.NoneType)): + returned = "a value that can be None" + else: + returned = f"{numba_ret_type}" + raise UdfParseError( + f"UDT UDF returned {returned}, which cannot be written as " + f"ret_dtype={ret_dtype}. Return {wanted}." + ) + + +def _refuse_complex_into_real(ret_dtype, element, numba_ret_type): + """Raise unless a complex return goes into complex, or bool, ``element``s of ``ret_dtype``. + + Numba has no cast from complex to real, so the wrapper fails to compile, + with a message that loses why (``Failed in nopython mode pipeline``). A + bool output is written as ``bool(...)``, which takes a complex value. + """ + nt = numba.core.types + returned = numba_ret_type.dtype if isinstance(numba_ret_type, nt.Array) else numba_ret_type + if isinstance(returned, nt.Complex) and element.kind not in "cb": + raise UdfParseError( + f"UDT UDF returned {numba_ret_type}, which cannot be written as " + f"ret_dtype={ret_dtype}. Return a real value, or declare a complex ret_dtype." + ) + + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. @@ -353,12 +443,15 @@ def _input_operand(dtype, var): var, nt.CPointer(base_numba), ) + # A built-in operand is read before the wrapper zeros a built-in output, which + # GraphBLAS can pass as an operand too (z = f(z, y) in a reduce). + setup = f" {var}0 = {var}_ptr[0]\n" if dtype == BOOL: # Numba can't compile bool ptrs (numba/numba#5395); expose them # as int8 and cast on deref. # MAINT 2026-10-18: still hits on Numba 0.68; re-test periodically. - return "", f"bool({var}_ptr[0])", nt.CPointer(INT8.numba_type) - return "", f"{var}_ptr[0]", nt.CPointer(dtype.numba_type) + return setup, f"bool({var}0)", nt.CPointer(INT8.numba_type) + return setup, f"{var}0", nt.CPointer(dtype.numba_type) def _output_handler(return_type, numba_ret_type): @@ -421,13 +514,22 @@ def _compose_wrapper_body(zkind, zinfo, signature_line, body_setup, call_expr): compile a nested-tuple assignment to an outer record field. Array returns and array-typed record leaves slice-assign into the caller's buffer, so a wrong-shape return cannot overrun it. The error does not reach - 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. :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. + the caller either: the resulting ``ValueError``, like any exception the + UDF raises, is raised inside a cfunc, which Numba prints and swallows, + and GraphBLAS takes the output as the wrapper left it. So a built-in + output (a select's bool, say) is set to 0 before the UDF is called: an + element the UDF raises for holds 0, or False, which a select reads as + "do not keep". A select's output is an uninitialized bool in SuiteSparse, + so it could keep the entry and corrupt the entry count. A UDT output is + not zeroed: it can be an operand too (z = f(z, y) in a reduce), and + copying a UDT operand first would cost a copy of every element. Numba + cannot re-raise, so catching the exception would silence it. + :func:`_check_udf_fills_output` rejects a return that cannot fill the + output when the op is typed, but it cannot always learn the shape, so + this stays the backstop. """ + if zkind == "direct" and zinfo[2] == "z_ptr[0]": + body_setup = f"{body_setup} z_ptr[0] = 0\n" if zkind == "array_elements": return ( f"{signature_line}\n" @@ -473,8 +575,8 @@ def _fits_by_broadcast(shape, expected): return False -def _udf_probe_value(dtype): - """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`. +def _udf_probe_value(dtype, fill=1): + """Build a stand-in operand of ``dtype`` for :func:`_run_udf_probe`, ``fill`` in every leaf. Only the shape of the UDF's result is read, so any values it accepts will do. @@ -482,20 +584,25 @@ def _udf_probe_value(dtype): if dtype._is_udt: np_type = dtype.np_type if np_type.subdtype is None: - return np.ones(1, dtype=np_type)[0] + return np.full(1, fill, dtype=np_type)[0] base_np_type, shape = np_type.subdtype - return np.ones(shape, dtype=base_np_type) + return np.full(shape, fill, 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. + """Call the UDF on stand-in operands. Returns ``(ran, error)``. + + ``ran`` is ``(result,)``, or ``None`` if the UDF raised on the stand-ins + of ones: then its result is not checked, and if it also returns a shape + that does not fit, that is left to the wrapper, whose output then keeps + its zeros (see :func:`_compose_wrapper_body`). Such a UDF may raise only + for some values, so it still types (``np.linalg.inv`` raises on a matrix + of ones). ``error`` names the exception (as text) when it raised on twos + as well, and it was no ``ArithmeticError``; only the literal rule reads it + (see :data:`_RAISES_ON_SAMPLES`). The exception itself is not kept: its + traceback holds the frames that hold the op. 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, @@ -504,14 +611,28 @@ def _run_udf_probe(numba_func, arg_dtypes): 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: # pragma: no cover (safety) - # The wrapper's operands have these types, so it would fail to compile too. - raise UdfParseError(_summarize_numba_typing_error(exc)) from exc - except Exception: - return None + error = None + for fill in (1, 2): + args = [_udf_probe_value(d, fill) for d in arg_dtypes] + try: + result = numba_func(*args) + except NumbaError as exc: # pragma: no cover (safety) + # The wrapper's operands have these types, so it would fail to compile too. + raise UdfParseError(_summarize_numba_typing_error(exc)) from exc + except ArithmeticError: + return None, None + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + continue + return ((result,) if fill == 1 else None), None + return None, error + + +# The UDT ops whose UDF raised, other than an ArithmeticError, on both sets of stand-in +# operands, with the exception as text. Such an op still types, since the UDF may raise only +# for some values, but beside it a number is typed as itself rather than as an element +# (``literals._own_scalar_type_for_user_op``), as a Scalar of it would be. +_RAISES_ON_SAMPLES = weakref.WeakKeyDictionary() def _raise_unless_fits(shape, expected, return_type, path, *, probed): @@ -589,6 +710,8 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ 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. + Returns what the UDF raised on every stand-in, as text (see + :func:`_run_udf_probe`), or None. """ if zkind == "array_elements": # (record field path, Numba type, slot shape, index into the tuple) @@ -604,8 +727,16 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ (path, numba_ret_type.types[i], leaf.shape, i) for i, (path, _c, leaf) in enumerate(leaves) ] + elif isinstance(numba_ret_type, numba.core.types.Array): + # An array returned for a scalar (a select's bool, say) fills it only with one + # value, which running the UDF shows; one that cannot is recorded, not refused, + # since its shape may depend on the values (``x[x > 0]``). + ran, error = _run_udf_probe(numba_func, arg_dtypes) + if ran is not None and np.size(ran[0]) != 1: + error = f"returned an array of shape {np.shape(ran[0])} for a scalar" + return error else: - return + return None unknown = [] for path, nb_type, expected, i in slots: if isinstance(nb_type, (numba.core.types.Optional, numba.core.types.NoneType)): @@ -634,7 +765,10 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ ) 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: + if not unknown: + return None + ran, error = _run_udf_probe(numba_func, arg_dtypes) + if ran is not None: for path, expected, i in unknown: value = ran[0] if i is None else ran[0][i] if isinstance(value, list): @@ -643,6 +777,7 @@ def _check_udf_fills_output(numba_func, numba_ret_type, return_type, zkind, arg_ ) else: _raise_unless_fits(np.shape(value), expected, return_type, path, probed=True) + return error def _get_udt_wrapper( @@ -672,7 +807,7 @@ def _get_udt_wrapper( 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) + error = _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}):" @@ -687,6 +822,7 @@ def _get_udt_wrapper( source_label=f"", extra_ns={"numba_func": numba_func}, ) + wrapper._raises_on_samples = error # for _finalize_udt_op return wrapper, wrapper_sig diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index b6e615ca5..23438328a 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -28,6 +28,7 @@ ParameterizedUdf, TypedOpBase, _hasop, + _validate_ret_dtype, ) if _supports_complex: @@ -41,7 +42,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -138,7 +139,7 @@ class UnaryOp(OpBase): as well as in the ``graphblas.op`` combined namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _custom_dtype = None _module = unary _modname = "unary" @@ -172,16 +173,24 @@ class UnaryOp(OpBase): _positional = {"positioni", "positioni1", "positionj", "positionj1"} @classmethod - def _build(cls, name, func, *, anonymous=False, is_udt=False): + def _build(cls, name, func, *, anonymous=False, is_udt=False, ret_dtype=None): if type(func) is not FunctionType: raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=False) success = False # Set on the Dispatcher, not just the cfunc wrapper; see the note in # ``BinaryOp._build``. unary_udf = numba.njit(func, error_model="numpy") - new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=unary_udf, + ret_dtype=ret_dtype, + ) return_types = {} nt = numba.types if not is_udt: @@ -212,26 +221,38 @@ def _build(cls, name, func, *, anonymous=False, is_udt=False): nt.CPointer(input_type.numba_type), ) + # Each wrapper writes 0 (False) before it calls the UDF: if the UDF raises, Numba + # prints the error and returns, and GraphBLAS uses the output as it is (for a + # select, an unwritten bool decides whether to keep the entry). The operands are + # read first, since the output can be one of them (z = f(z, y) in a reduce). if type_ == BOOL: if ret_type == BOOL: def unary_wrapper(z, x): - z[0] = bool(unary_udf(bool(x[0]))) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = bool(unary_udf(bool(x0))) # pragma: no cover (numba) else: def unary_wrapper(z, x): - z[0] = unary_udf(bool(x[0])) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = unary_udf(bool(x0)) # pragma: no cover (numba) elif ret_type == BOOL: def unary_wrapper(z, x): - z[0] = bool(unary_udf(x[0])) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = bool(unary_udf(x0)) # pragma: no cover (numba) else: def unary_wrapper(z, x): - z[0] = unary_udf(x[0]) # pragma: no cover (numba) + x0 = x[0] # pragma: no cover (numba) + z[0] = 0 # pragma: no cover (numba) + z[0] = unary_udf(x0) # pragma: no cover (numba) unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( unary_wrapper @@ -288,7 +309,7 @@ def _compile_udt(self, dtype, dtype2): sig = (dtype.numba_type,) _compile_udf_for_udt(numba_func, sig, op_kind="unary", op_name=self.name, dtypes=(dtype,)) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + ret_type = _udt_ret_type(self, numba_ret_type, dtype) unary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, numba_ret_type=numba_ret_type ) @@ -297,7 +318,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a UnaryOp without registering it in the ``graphblas.unary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -324,6 +347,12 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Whether the operator is intended to operate on user-defined types. If True, then the function will not be automatically compiled for builtin types, and it will be compiled "just in time" when used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which names a built-in dtype or one of the input UDTs, + so an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. Returns ------- @@ -332,11 +361,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) return ParameterizedUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new UnaryOp and save it to ``graphblas.unary`` namespace. Parameters @@ -369,6 +401,9 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal Compiling functions can be slow, however, so you may want to delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. Examples -------- @@ -378,17 +413,26 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) else: - unary_op = cls._build(name, func, is_udt=is_udt) + unary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, unary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -485,12 +529,14 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional self._is_udt = is_udt + self._ret_dtype = ret_dtype if is_udt: self._udt_types = {} # {dtype: DataType} self._udt_ops = {} # {dtype: TypedUserUnaryOp} diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index ea39ac5ac..32bf229c7 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,5 +1,6 @@ import itertools import os +import pickle import subprocess import sys from pathlib import Path @@ -3201,6 +3202,361 @@ def _no_ones(x, y): # pragma: no cover (numba) assert probed == ["_sum", "_no_ones"] +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udf_raising_for_some_elements_gives_zeros_there(): + """An element the UDF raises for holds 0, and a select does not keep it. + + Numba prints and swallows an exception raised in a cfunc, so the wrapper + wrote nothing there. A select then read SuiteSparse's ``bool keep``, an + uninitialized local, which could keep the entry and add a stray byte to + the entry count (r612 saw nvals from 6 to 256 in a Vector of size 2). The + wrappers now write 0, or False, into a built-in output before they call + the UDF. What this checks held by chance before (the stray byte was 0); + test_udt_wrapper_zeros_a_builtin_output_before_the_udf checks the wrapper + itself. + """ + + def keep_above(x, i, j, t): # pragma: no cover (numba) + if x < 0: + raise ValueError("negative") + return x > t + + def times_ten(x): # pragma: no cover (numba) + if x < 0: + raise ValueError("negative") + return x * 10 + + v = Vector.from_coo([0, 1, 2, 3, 4], [-1, 5, 2, -3, 7]) + # -1 and -3 raise, so they are not kept; of the others, 5 and 7 are above 3 + kept = v.select(SelectOp.register_anonymous(keep_above, "_raise_keep_above"), 3).new() + assert kept.nvals == 2 + assert kept.to_coo()[0].tolist() == [1, 4] + assert kept.to_coo()[1].tolist() == [5, 7] + # -1 and -3 raise, so they hold 0 + tens = v.apply(UnaryOp.register_anonymous(times_ten, "_raise_times_ten")).new() + assert tens.to_coo()[1].tolist() == [0, 50, 20, 0, 70] + + rec = dtypes.register_anonymous(np.dtype([("rz_a", np.int64), ("rz_b", np.float64)])) + + def keep_rec(x, i, j, t): # pragma: no cover (numba) + if x.rz_a < 0: + raise ValueError("negative") + return x.rz_b > t + + r = Vector(rec, size=4) + r[0] = (-1, 9.0) + r[1] = (1, 9.0) + r[2] = (1, 0.0) + r[3] = (-2, 9.0) + # Entries 0 and 3 raise; of the others, only 1 has rz_b above 1 + sel = SelectOp.register_anonymous(keep_rec, "_raise_keep_rec", is_udt=True) + kept = r.select(sel, gb.Scalar.from_value(1.0)).new() + assert kept.nvals == 1 + assert kept.to_coo()[0].tolist() == [1] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udt_wrapper_zeros_a_builtin_output_before_the_udf(): + """A UDT op's wrapper sets a built-in output to 0 (False) before it calls the UDF. + + Called on an output full of 0xAB with an input the UDF raises for, it + leaves 0: a select's bool, and an FP64. It left the 0xAB bytes, which a + select read as "keep". (A UDT output is not zeroed; see + ``udf._compose_wrapper_body``.) + """ + import ctypes + + import numba + + from graphblas.core.operator.udf import _get_udt_wrapper + + rec = dtypes.register_anonymous( + np.dtype([("zo_n", np.int8), ("zo_s", "S3"), ("zo_d", "M8[s]")], align=True) + ) + + def keep(x, i, j, t): # pragma: no cover (numba) + if x.zo_n < 0: + raise ValueError("negative") + return True + + def half(x): # pragma: no cover (numba) + if x.zo_n < 0: + raise ValueError("negative") + return 0.5 + + x = np.zeros(1, dtype=rec.np_type) + x["zo_n"] = -1 + thunk = np.zeros(1, np.float64) + index = np.zeros(1, np.uint64) + for func, out, kwargs in [ + (keep, dtypes.BOOL, {"include_indexes": True, "dtype2": FP64}), + (half, FP64, {}), + ]: + numba_func = numba.njit(func) + arg_types = (rec.numba_type,) + if kwargs: + arg_types += (numba.types.uint64, numba.types.uint64, numba.types.float64) + numba_func.compile(arg_types) + ret = numba_func.overloads[arg_types].signature.return_type + wrapper, sig = _get_udt_wrapper(numba_func, out, rec, numba_ret_type=ret, **kwargs) + cfunc = numba.cfunc(sig, nopython=True)(wrapper) + z = np.full( + np.dtype(np.int8 if out is dtypes.BOOL else out.np_type).itemsize, 0xAB, np.uint8 + ) + args = [ctypes.cast(z.ctypes.data, cfunc.ctypes.argtypes[0])] + args.append(ctypes.cast(x.ctypes.data, cfunc.ctypes.argtypes[1])) + if kwargs: + args += [index[0], index[0], ctypes.cast(thunk.ctypes.data, cfunc.ctypes.argtypes[4])] + cfunc.ctypes(*args) + assert not z.any(), out + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_udf_raising_on_stand_ins_still_types(): + """A UDF that raises on the probe's stand-in operands still types, and runs on real data. + + Stand-ins of ones and twos are degenerate: a matrix of ones is singular, + and an interval check fails on equal fields. The index and thunk + stand-ins do not vary at all. So no UDF is refused for raising on them; + "Refuse a UDT UDF that raises for its operand types when the op is typed" + refused each of these. Only the literal rule reads that a UDF raised on + every stand-in: beside such an op a number is typed as itself. + """ + m22 = dtypes.register_anonymous(np.dtype((np.float64, (2, 2))), "_StandIn22") + p2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_StandIn2") + p3 = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_StandIn3") + p6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_StandIn6") + + def inverse(x): # pragma: no cover (numba) + return np.linalg.inv(x) + + def widths(x): # pragma: no cover (numba) + if x[0] >= x[1]: + raise ValueError("empty interval") + return x * 2.0 + + def rows(x, i, j, t): # pragma: no cover (numba) + return x.reshape((int(t), -1)).sum(axis=0) + + def need_t(x, i, j, t): # pragma: no cover (numba) + if t < 1.0: + raise ValueError("t must be at least 1") + return x[:2] * t + + m = Vector(m22, size=1) + m[0] = [[2.0, 1.0], [1.0, 3.0]] + result = m.apply(UnaryOp.register_anonymous(inverse, is_udt=True)).new() + np.testing.assert_allclose(result[0].new().value, [[0.6, -0.2], [-0.2, 0.4]]) + p = Vector(p2, size=1) + p[0] = [1.0, 4.0] + result = p.apply(UnaryOp.register_anonymous(widths, is_udt=True)).new() + assert result[0].new().value.tolist() == [2.0, 8.0] + s = Vector(p6, size=1) + s[0] = np.arange(6.0) + by_rows = IndexUnaryOp.register_anonymous(rows, is_udt=True, ret_dtype=p3) + assert s.apply(by_rows, gb.Scalar.from_value(2.0)).new()[0].new().value.tolist() == [3, 5, 7] + checked = IndexUnaryOp.register_anonymous(need_t, is_udt=True, ret_dtype=p2) + for thunk in [gb.Scalar.from_value(3.0), 3.0]: + assert s.apply(checked, thunk).new()[0].new().value.tolist() == [0.0, 3.0] + # An op whose UDF raises for every element thunk types too, unchecked + iu = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=p2) + assert iu[p6, p6].return_type is p2 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_builtin_type_udf_wrappers_write_zero_before_the_udf(monkeypatch): + """Each wrapper a user's op gets for a built-in type writes 0 before it calls the UDF. + + The wrappers are made while the op registers, so they are caught from + ``numba.cfunc`` then, and each is called with an output of 0xAB bytes and a + UDF that always raises: the output must be 0 (False). Without the zero + writes it kept the 0xAB bytes, which a select reads as "keep". + """ + import ctypes + + import numba + + from graphblas.core.operator import IndexBinaryOp + + caught = [] + cfunc = numba.cfunc + + def catching_cfunc(*args, **kwargs): + decorate = cfunc(*args, **kwargs) + + def catch(func): + caught.append(decorate(func)) + return caught[-1] + + return catch + + # Numba types a return only if the raise is conditional + def raises_unary(x): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_binary(x, y): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_indexunary(x, i, j, t): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + def raises_indexbinary(x, ix, jx, y, iy, jy, t): # pragma: no cover (numba) + if x == x: + raise ValueError("raises") + return x + + monkeypatch.setattr(numba, "cfunc", catching_cfunc) + counts = {} + for cls, func in [ + (UnaryOp, raises_unary), + (BinaryOp, raises_binary), + (IndexUnaryOp, raises_indexunary), + (IndexBinaryOp, raises_indexbinary), + ]: + if cls is IndexBinaryOp and lib.__dict__.get("GxB_IndexBinaryOp_new") is None: + continue # pragma: no cover (SuiteSparse < 9.4) + start = len(caught) + cls.register_anonymous(func) + counts[cls.__name__] = len(caught) - start + monkeypatch.undo() + assert all(counts.values()), counts # one wrapper per built-in type the UDF compiles for + for wrapper in caught: + args = [] + out = None + for argtype, ctype in zip(wrapper._sig.args, wrapper.ctypes.argtypes, strict=True): + if not isinstance(argtype, numba.types.CPointer): + args.append(1) # an index + continue + np_type = numba.np.numpy_support.as_dtype(argtype.dtype) + buf = np.full(np_type.itemsize, 0xAB, np.uint8) if out is None else np.ones(1, np_type) + out = buf if out is None else out + args.append(ctypes.cast(buf.ctypes.data, ctype)) + wrapper.ctypes(*args) + assert not out.any(), wrapper._sig + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_op_whose_udf_raised_on_stand_ins_can_be_collected(): + """An op the probe records as raising on every stand-in is freed once dropped. + + The record kept the exception, whose traceback holds the frames that hold the + typed op, so the op, its parent and their GraphBLAS handles were never freed + (r612: 5 of 5 alive). It keeps the exception as text now. + """ + import gc + import weakref + + from graphblas.core.operator.udf import _RAISES_ON_SAMPLES + + r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_FreedStandIn6") + r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_FreedStandIn2") + v = Vector(r6, size=1) + v[0] = np.arange(6.0) + refs = [] + for _ in range(3): + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, t: x[:2] * t, is_udt=True, ret_dtype=r2 + ) + assert v.apply(iu, 3.0).new()[0].new().value.tolist() == [0.0, 3.0] + typed = iu[r6, r6] + assert typed in _RAISES_ON_SAMPLES + refs += [weakref.ref(iu), weakref.ref(typed)] + del iu, typed + gc.collect() + assert [ref() for ref in refs] == [None] * len(refs) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning") +def test_udt_select_takes_a_number_as_itself_when_it_cannot_take_an_element(): + """A literal thunk for a user's select that cannot take an element is typed as itself. + + ``x[0] > t`` with ``t`` an INT8[3] element returns three bools for the one + the select needs, so the wrapper raised for every entry and nothing was + kept (or, before the select's bool was written first, a stray byte decided). + The select's UDF is now run on stand-ins when it returns an array, and a + number beside it is typed as itself, as a Scalar of it is. + """ + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_SelOwnI3") + first_above = SelectOp.register_anonymous( + lambda x, i, j, t: x[0] > t, "_sel_own_first", is_udt=True + ) + v = Vector(i3, size=2) + v[0] = [1, 2, 3] + v[1] = [-4, 5, 6] + for thunk in [0, np.int8(0), 0.5, gb.Scalar.from_value(0)]: + kept = v.select(first_above, thunk).new() + assert kept.to_coo()[0].tolist() == [0] + # A select that takes an element still gets one + same = SelectOp.register_anonymous( + lambda x, i, j, t: (x == t).all(), "_sel_own_same", is_udt=True + ) + v[1] = [1, 1, 1] + assert v.select(same, 1).new().to_coo()[0].tolist() == [1] + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_user_op_takes_a_number_as_itself_when_it_cannot_take_an_element(): + """A number beside a user's op is typed as itself when the op cannot take it as an element. + + A literal beside a user's ``is_udt`` op becomes an element of the operand's + UDT, so ``x[:2] * t`` met a 6-element ``t``: its cfunc raised for every + element and the result held zeros, where a Scalar thunk worked. When the op + cannot be typed with an element, a Python or numpy number now takes its own + type, as a Scalar of it would. + """ + r6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_OwnNum6") + r2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_OwnNum2") + iu = IndexUnaryOp.register_anonymous( + lambda x, i, j, t: x[:2] * t, "_own_num_iu", is_udt=True, ret_dtype=r2 + ) + bop = BinaryOp.register_anonymous( + lambda x, y: x[:2] * y, "_own_num_b", is_udt=True, ret_dtype=r2 + ) + left = BinaryOp.register_anonymous( + lambda x, y: y[:2] * x, "_own_num_left", is_udt=True, ret_dtype=r2 + ) + v = Vector(r6, size=2) + v[0] = np.arange(6.0) + v[1] = np.arange(6.0) + 10 + expected = [[0.0, 3.0], [30.0, 33.0]] + for literal, own in [(3.0, FP64), (np.float64(3.0), FP64), (3, INT64), (np.float32(3), FP32)]: + for expr in [ + v.apply(iu, literal), + v.apply(bop, right=literal), + v.apply(left, left=literal), + ]: + assert own in {expr.op.type, expr.op._type2} + assert expr.new().to_coo()[1].tolist() == expected + # An op that can take an element still gets one, as udt.rst says + i3 = dtypes.register_anonymous(np.dtype((np.int8, (3,))), "_OwnNumI3") + add = IndexUnaryOp.register_anonymous(lambda x, i, j, t: x + t, "_own_num_add", is_udt=True) + w = Vector(i3, size=1) + w[0] = [1, 2, 3] + expr = w.apply(add, 1) + assert expr.op._type2 is i3 + assert expr.new()[0].new().value.tolist() == [2, 3, 4] + with pytest.raises(ValueError, match="does not fit _OwnNumI3"): + w.apply(add, 0.5) + # A record op that adds a number to each field takes it as itself + rec = dtypes.register_anonymous(np.dtype([("on_a", np.int32), ("on_b", np.float64)])) + each = BinaryOp.register_anonymous( + lambda x, y: (x.on_a + y, x.on_b + y), "_own_num_each", is_udt=True + ) + r = Vector(rec, size=1) + r[0] = (1, 2.0) + assert r.apply(each, right=2).new()[0].new().value.tolist() == (3, 4.0) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_record_array_field_roundtrip(): @@ -7273,3 +7629,508 @@ def test_deferred_commutes_to(): assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report assert lines.get("floordiv_ok") == "True", report assert lines.get("absfirst_ok") == "True", report + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_names_an_output_udt(): + """``ret_dtype`` names an output UDT that is not one of the operands. + + Without it the return type is inferred from what the UDF builds, and the + only names in scope are the input dtypes. That makes an op that shortens + its element, such as FP64[10] -> FP64[3], unreachable: the inferred type + is the input's, and the shape check then rejects the shorter array the + UDF returns. + """ + # The input shape is unique to this test; see the note in + # test_udt_array_udf_shape_errors. + ten = dtypes.register_anonymous(np.dtype((np.float64, (10,))), "_RetD10") + three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetD3") + + def _first_three(x): # pragma: no cover (numba) + return x[:3] + + without = UnaryOp.register_anonymous(_first_three, "_ret_dtype_without", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(3,\) when run on sample values"): + without[ten] + + op_ = UnaryOp.register_anonymous(_first_three, "_ret_dtype_with", is_udt=True, ret_dtype=three) + assert op_[ten].return_type is three + + # lazy=True must hand ret_dtype to the registration it defers, or the op + # that fires on first attribute access is built without it. + UnaryOp.register_new("_ret_dtype_lazy", _first_three, is_udt=True, lazy=True, ret_dtype=three) + try: + assert unary._ret_dtype_lazy[ten].return_type is three + finally: + # Same cleanup as test_udt_lazy_registration: the op lands in ``unary`` + # and ``op``, and both keep a ``_delayed`` entry until it fires. + for module in (unary, op): + vars(module).pop("_ret_dtype_lazy", None) + module._delayed.pop("_ret_dtype_lazy", None) + + v = Vector(ten, size=2) + v[0] = np.arange(10.0) + w = op_(v).new() + assert w.dtype is three + np.testing.assert_array_equal(w[0].new().value, np.arange(3.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_binary_and_record(): + """``ret_dtype`` on a binary op, including a type that is neither operand.""" + a4 = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_RetDBinA4") + b6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_RetDBinB6") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDBinOut2") + + def _head_two(x, y): # pragma: no cover (numba) + return x[:2] + y[:2] + + op_ = BinaryOp.register_anonymous(_head_two, "_ret_dtype_bin", is_udt=True, ret_dtype=out2) + assert op_[a4, b6].return_type is out2 + + # ret_dtype equal to an operand's dtype agrees with what inference picks. + def _plus(x, y): # pragma: no cover (numba) + return x + y + + same = BinaryOp.register_anonymous(_plus, "_ret_dtype_same", is_udt=True, ret_dtype=a4) + inferred = BinaryOp.register_anonymous(_plus, "_ret_dtype_inferred", is_udt=True) + assert same[a4, a4].return_type is inferred[a4, a4].return_type is a4 + + # A record UDT output that appears in no operand. + rec = dtypes.register_anonymous( + np.dtype([("lo", np.float64), ("hi", np.float64)], align=True), "_RetDRec" + ) + + def _bounds(x, y): # pragma: no cover (numba) + return (min(x, y), max(x, y)) + + recop = BinaryOp.register_anonymous(_bounds, "_ret_dtype_rec", is_udt=True, ret_dtype=rec) + assert recop[FP64, FP64].return_type is rec + + w = Vector(FP64, size=2) + w[0] = 3.0 + u = Vector(FP64, size=2) + u[0] = 1.0 + res = recop(w & u).new() + assert res.dtype is rec + assert tuple(res[0].new().value.tolist()) == (1.0, 3.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_errors(): + """``ret_dtype`` is rejected outside the UDT path and for unrecognized dtypes.""" + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk", is_udt=True, ret_dtype="NOPE") + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk2", is_udt=True, ret_dtype=object()) + + # The builtin path derives a return type per input dtype, so one fixed + # dtype cannot describe it. + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_builtin", ret_dtype=FP64) + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + BinaryOp.register_anonymous(lambda x, y: x + y, "_ret_dtype_builtin2", ret_dtype=FP64) + + # lazy=True must not defer the validation: a bad combination fails at the + # registration site, not at first attribute touch of the delayed op. + with pytest.raises(ValueError, match="parameterized=True"): + UnaryOp.register_new( + "_ret_dtype_lazy_param", + lambda x: x, + parameterized=True, + is_udt=True, + lazy=True, + ret_dtype=FP64, + ) + + # A parameterized operator builds its function when called; ret_dtype + # belongs to the register call for that function. + with pytest.raises(ValueError, match="does not work with parameterized=True"): + UnaryOp.register_anonymous( + lambda t=1: (lambda x: x + t), + "_ret_dtype_param", + parameterized=True, + is_udt=True, + ret_dtype=FP64, + ) + + # SelectOp takes no ret_dtype: GraphBLAS fixes its return type to BOOL. + with pytest.raises(TypeError, match="ret_dtype"): + SelectOp.register_anonymous( + lambda x, i, j, t: x > t, "_ret_dtype_select", is_udt=True, ret_dtype=BOOL + ) + + +def _ret_dtype_class_cases(): + """(class, namespace, UDF giving the first 3 of x, typed lookup) per class with ret_dtype.""" + cases = [ + (UnaryOp, "unary", lambda x: x[:3], lambda o, t: o[t]), + (BinaryOp, "binary", lambda x, y: x[:3] + y[:3], lambda o, t: o[t, t]), + (IndexUnaryOp, "indexunary", lambda x, i, j, th: x[:3] + i, lambda o, t: o[t, INT64]), + ] + if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: + from graphblas.core.operator import IndexBinaryOp + + cases.append( + ( + IndexBinaryOp, + "indexbinary", + lambda x, ix, jx, y, iy, jy, th: x[:3] + y[:3], + lambda o, t: o[t, t], + ) + ) + return cases + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_lazy_and_parameterized_for_each_op_class(): + """``lazy=True`` keeps ``ret_dtype``, and ``parameterized=True`` refuses it, for every class. + + Only UnaryOp was tested, so dropping ``ret_dtype`` from BinaryOp's or + IndexUnaryOp's deferred registration, or BinaryOp's parameterized check, + left every test passing. + """ + twelve = dtypes.register_anonymous(np.dtype((np.float64, (12,))), "_RetDCls12") + three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetDCls3") + for cls, modname, func, typed in _ret_dtype_class_cases(): + module = getattr(gb, modname) + name = f"_ret_dtype_lazy_{modname}" + cls.register_new(name, func, is_udt=True, lazy=True, ret_dtype=three) + try: + assert typed(getattr(module, name), twelve).return_type is three + finally: + for ns in (module, op): + vars(ns).pop(name, None) + ns._delayed.pop(name, None) + with pytest.raises(ValueError, match="does not work with parameterized=True"): + cls.register_anonymous( + lambda t=1, func=func: func, is_udt=True, parameterized=True, ret_dtype=three + ) + with pytest.raises(ValueError, match="parameterized=True"): + cls.register_new( + f"{name}_param", + lambda t=1, func=func: func, + is_udt=True, + parameterized=True, + lazy=True, + ret_dtype=three, + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_index_ops(): + """``ret_dtype`` reaches the IndexUnaryOp and IndexBinaryOp compile paths too.""" + # Shape (11,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in11 = dtypes.register_anonymous(np.dtype((np.float64, (11,))), "_RetDIdx11") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDIdx2") + + def _head_plus_row(x, i, j, t): # pragma: no cover (numba) + return x[:2] + i + + iu = IndexUnaryOp.register_anonymous( + _head_plus_row, "_ret_dtype_indexunary", is_udt=True, ret_dtype=out2 + ) + assert iu[in11, INT64].return_type is out2 + v = Vector(in11, size=3) + v[1] = np.arange(11.0) + v[2] = np.arange(11.0) + 10 + result = v.apply(iu, gb.Scalar.from_value(0)).new() + assert result.dtype is out2 + assert result.to_coo()[1].tolist() == [[1.0, 2.0], [12.0, 13.0]] + + if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: + from graphblas.core.operator import IndexBinaryOp + + def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x[:2] + y[:2] + + ib = IndexBinaryOp.register_anonymous( + _head_sum, "_ret_dtype_indexbinary", is_udt=True, ret_dtype=out2 + ) + assert ib[in11, in11].return_type is out2 + result = v.ewise_mult(v, ib(0.0)).new() + assert result.dtype is out2 + assert result.to_coo()[1].tolist() == [[0.0, 2.0], [20.0, 22.0]] + + # register_new also adds an IndexUnaryOp to ``select`` when it returns BOOL. + # A UDT op has no types until it is compiled, so a declared non-BOOL + # ret_dtype is what keeps this one out. + IndexUnaryOp.register_new("_ret_dtype_iu_named", _head_plus_row, is_udt=True, ret_dtype=out2) + try: + assert "_ret_dtype_iu_named" not in vars(select) + finally: + vars(indexunary).pop("_ret_dtype_iu_named", None) + vars(select).pop("_ret_dtype_iu_named", None) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_still_shape_checked(): + """The registration-time shape probe checks against the declared ``ret_dtype``. + + Short-circuiting the return-type inference must not take the probe with + it. Declaring the output type makes the check worth more, not less: the + inferred case can only ever compare a UDF against a type derived from + what it returned, while a declared type is an independent claim the UDF + can contradict. + """ + # Shape (14,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetDPrb14") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDPrb2") + + def _three(x): # pragma: no cover (numba) + return x[:3] + + # Three elements cannot fill a declared two-element output. + op_ = UnaryOp.register_anonymous(_three, "_ret_dtype_prb_bad", is_udt=True, ret_dtype=out2) + with pytest.raises(UdfParseError, match=r"shape \(3,\).*_RetDPrb2 elements are \(2,\)"): + op_[in14] + + # The check is fit-by-broadcast, not equality: a one-element return + # legitimately fills every slot of the declared element. + def _one(x): # pragma: no cover (numba) + return x[:1] + + ok = UnaryOp.register_anonymous(_one, "_ret_dtype_prb_bcast", is_udt=True, ret_dtype=out2) + assert ok[in14].return_type is out2 + + # An operand returned as-is keeps its extents in its Numba type, so the + # declared shape is checked without running the UDF (no "sample values"). + def _as_is(x): # pragma: no cover (numba) + return x + + as_is = UnaryOp.register_anonymous(_as_is, "_ret_dtype_prb_as_is", is_udt=True, ret_dtype=out2) + with pytest.raises(UdfParseError, match=r"shape \(14,\), but _RetDPrb2 elements are \(2,\)"): + as_is[in14] + + # The record half of the probe checks the declared type's array leaves. + rec = dtypes.register_anonymous( + np.dtype([("v", np.float64, (4,)), ("n", np.int64)], align=True), "_RetDPrbRec" + ) + + def _short_leaf(x, y): # pragma: no cover (numba) + return (x[:2], 1) + + recop = BinaryOp.register_anonymous( + _short_leaf, "_ret_dtype_prb_rec", is_udt=True, ret_dtype=rec + ) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field \['v'\] of _RetDPrbRec"): + recop[in14, in14] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +# _RetDMisNested's numpy repr is 141 chars; see test_udt_eq_nested_record_with_nan_leaf. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_ret_dtype_rejects_a_return_it_cannot_hold(): + """A declared ``ret_dtype`` is checked against what the UDF returns. + + Inference only names a type the return can be written as, and a declared + type skips it. Unchecked, a mismatch reached the wrapper, which either + failed to compile with a raw Numba error or, worse, ran and gave a wrong + answer with no error: a longer tuple lost its extra values, an array + written to a scalar kept its first element, and a tuple written to an + array element left it uninitialized. + """ + # Names and shapes unique to this test; see the note in + # test_udt_array_udf_shape_errors. + rec2 = dtypes.register_anonymous( + np.dtype([("rdm_lo", np.float64), ("rdm_hi", np.float64)], align=True), "_RetDMisRec2" + ) + rec3 = dtypes.register_anonymous( + np.dtype([("rdm_a", np.int64), ("rdm_b", np.int64), ("rdm_c", np.int64)], align=True), + "_RetDMisRec3", + ) + arr2 = dtypes.register_anonymous(np.dtype((np.float32, (2,))), "_RetDMisArr2") + arr3 = dtypes.register_anonymous(np.dtype((np.float32, (3,))), "_RetDMisArr3") + + def binary_rejects(func, ret_dtype, match): + op_ = BinaryOp.register_anonymous(func, is_udt=True, ret_dtype=ret_dtype) + with pytest.raises(UdfParseError, match=match): + op_[FP64, FP64] + + binary_rejects( + lambda x, y: (x, y, x + y), rec2, "tuple of length 3, but ret_dtype=_RetDMisRec2 has 2" + ) + binary_rejects(lambda x, y: (x,), rec2, "tuple of length 1, but ret_dtype=_RetDMisRec2 has 2") + binary_rejects( + lambda x, y: x + y, rec2, "returned float64, which cannot be written as ret_dtype" + ) + + # A nested record takes a flat tuple, one value per leaf, and the count in + # the message says so when it differs from the number of top-level fields. + nested = dtypes.register_anonymous( + np.dtype( + [("rdm_in", [("rdm_p", np.float64), ("rdm_q", np.float64)]), ("rdm_r", np.float64)], + align=True, + ), + "_RetDMisNested", + ) + binary_rejects( + lambda x, y: ((x, y), x + y), + nested, + "tuple of length 2, but ret_dtype=_RetDMisNested has 3 fields once nested records", + ) + flat = BinaryOp.register_anonymous(lambda x, y: (x, y, x + y), is_udt=True, ret_dtype=nested) + assert flat[FP64, FP64].return_type is nested + binary_rejects(lambda x, y: (x, y, x + y), arr2, "tuple of length 3, which cannot be written") + binary_rejects(lambda x, y: (x, y), FP64, "tuple of length 2, which cannot be written") + # A complex return into real values; Numba's own message lost why. Into BOOL it is + # written as bool(...), so it works + for ret_dtype in [FP64, arr2]: + binary_rejects(lambda x, y: x * 1j + y, ret_dtype, "complex128, .* Return a real value") + nonzero = BinaryOp.register_anonymous(lambda x, y: x * 1j + y, is_udt=True, ret_dtype=BOOL) + one, zero = Vector.from_coo([0], [1.5]), Vector.from_coo([0], [0.0]) + assert one.ewise_mult(zero, nonzero).new()[0].new().value is True + binary_rejects( + lambda x, y: x if x > y else None, FP64, "returned a value that can be None, which cannot" + ) + + def unary_rejects(func, ret_dtype, operand, match): + op_ = UnaryOp.register_anonymous(func, is_udt=True, ret_dtype=ret_dtype) + with pytest.raises(UdfParseError, match=match): + op_[operand] + + unary_rejects(lambda x: x * 1.0, FP64, arr3, "returned an array, which cannot be written") + unary_rejects(lambda x: x, rec3, rec2, r"record with fields \['rdm_lo', 'rdm_hi'\]") + unary_rejects(lambda x: x, arr3, rec2, "record with fields") + unary_rejects(lambda x: x, rec2, arr3, "returned an array, which cannot be written") + + # What a declared type can hold still works: one value per field, a value of + # the declared record type, a scalar that fills an array element, and a scalar + # cast to the declared type. + ok = BinaryOp.register_anonymous(lambda x, y: (x, y), is_udt=True, ret_dtype=rec2) + same = UnaryOp.register_anonymous(lambda x: x, is_udt=True, ret_dtype=rec2) + fill = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=arr2) + cast = BinaryOp.register_anonymous(lambda x, y: x + y, is_udt=True, ret_dtype=INT64) + w = Vector(FP64, size=1) + w[0] = 3.0 + u = Vector(FP64, size=1) + u[0] = 1.0 + assert tuple(w.ewise_mult(u, ok).new()[0].new().value.tolist()) == (3.0, 1.0) + r = Vector(rec2, size=1) + r[0] = (1.0, 2.0) + assert tuple(r.apply(same).new()[0].new().value.tolist()) == (1.0, 2.0) + np.testing.assert_array_equal(w.ewise_mult(u, fill).new()[0].new().value, [4.0, 4.0]) + assert w.ewise_mult(u, cast).new()[0].new().value == 4 + # So does an operand of a date, duration or text UDT returned as its own type, + # which said "Return a scalar" + for np_type, value in [ + ("M8[s]", np.datetime64(5, "s")), + ("m8[s]", np.timedelta64(5, "s")), + ("S5", b"abc"), + ]: + udt = dtypes.register_anonymous(np.dtype(np_type)) + own = UnaryOp.register_anonymous(lambda x: x, is_udt=True, ret_dtype=udt) + assert own[udt].return_type is udt + v = Vector(udt, size=1) + v[0] = value + assert v.apply(own).new()[0].new().value == value + + +def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) + return x[:5] + + +_RET_DTYPE_PICKLE_WRITER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb +from graphblas.core.operator.unary import UnaryOp +from graphblas.tests.test_op import _ret_dtype_pickle_udf + +print("package: " + gb.__file__) +five = gb.dtypes.register_anonymous(np.dtype((np.float32, (5,))), "_RetPickle5") +UnaryOp.register_new("_ret_dtype_pickled", _ret_dtype_pickle_udf, is_udt=True, ret_dtype=five) +anon = UnaryOp.register_anonymous( + _ret_dtype_pickle_udf, "_ret_dtype_pickled_anon", is_udt=True, ret_dtype=five +) +with open(sys.argv[1], "wb") as f: + pickle.dump((gb.unary._ret_dtype_pickled, anon), f) +print("wrote: ok") +""" + + +_RET_DTYPE_PICKLE_READER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb + +print("package: " + gb.__file__) +with open(sys.argv[1], "rb") as f: + named, anon = pickle.load(f) +ten = gb.dtypes.register_anonymous(np.dtype((np.float32, (10,))), "_RetPickle10") +print("named shape: " + str(named[ten].return_type.np_type.subdtype[1])) +print("anon shape: " + str(anon[ten].return_type.np_type.subdtype[1])) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_survives_pickle(tmp_path): + # In-process unpickling takes the _find shortcut and returns the very + # same object, so only a cross-process round trip exercises what pickle + # exists for: the reduce tuple must carry ret_dtype, or the reconstructed + # op re-registers without it and fails its own shape probe. Both sides + # run in subprocesses so the named registration never lands in this + # process's operator namespace (test_operator_types enumerates it). + payload = tmp_path / "ops.pkl" + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + for probe, checks in ( + (_RET_DTYPE_PICKLE_WRITER, {"wrote": "ok"}), + (_RET_DTYPE_PICKLE_READER, {"named shape": "(5,)", "anon shape": "(5,)"}), + ): + result = subprocess.run( + [sys.executable, "-c", probe, str(payload)], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + for key, want in checks.items(): + assert lines.get(key) == want, report + + # An op without ret_dtype pickles as the 3-tuple it always did, which + # versions without ret_dtype can still read. + plain = UnaryOp.register_anonymous(_ret_dtype_pickle_udf, "_ret_dtype_plain", is_udt=True) + assert len(plain.__reduce__()[1]) == 3 + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_named_op_reregisters_on_unpickle(): + # A named op unpickled where its name is not registered is registered again, + # with its ret_dtype. The name is removed when done, since test_operator_types + # enumerates the operator namespaces. + five = dtypes.register_anonymous(np.dtype((np.float32, (5,))), "_RetPickle5") + twelve = dtypes.register_anonymous(np.dtype((np.float32, (12,))), "_RetPickle12") + name = "_ret_dtype_reregistered" + try: + named = UnaryOp.register_new(name, _ret_dtype_pickle_udf, is_udt=True, ret_dtype=five) + payload = pickle.dumps(named) + assert pickle.loads(payload) is named + delattr(unary, name) + vars(op).pop(name) + new = pickle.loads(payload) + assert new is not named + assert getattr(unary, name) is new + assert new[twelve].return_type is five + finally: + vars(unary).pop(name, None) + vars(op).pop(name, None)