Repository navigation
Pass and receive numpy views in array-UDT UDFs - #589
Merged
Merged
Conversation
eriknw
marked this pull request as ready for review
August 4, 2026 16:07
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 4, 2026 16:12
660faec to
8164778
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 5, 2026 00:06
8164778 to
8960e9b
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 5, 2026 03:18
8960e9b to
3a7d0be
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
2 times, most recently
from
August 5, 2026 18:03
af88351 to
e3b42d6
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 5, 2026 18:05
e3b42d6 to
e16fb44
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 6, 2026 07:59
e16fb44 to
4ba7ef5
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 6, 2026 15:39
4ba7ef5 to
0a2ca8a
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
2 times, most recently
from
August 6, 2026 20:36
4bce522 to
9a39af4
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 6, 2026 20:41
9a39af4 to
9cfc310
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 7, 2026 02:48
9cfc310 to
6c21f4a
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 7, 2026 05:09
6c21f4a to
fa7b6bc
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
August 26, 2026 17:30
fa7b6bc to
1dc9a8d
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
September 18, 2026 03:09
1dc9a8d to
52b4882
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
September 18, 2026 03:46
52b4882 to
203ee60
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
September 18, 2026 17:34
203ee60 to
20f59fc
Compare
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
September 18, 2026 17:39
20f59fc to
b4afda6
Compare
A UDF over an array UDT wrote past the end of its output element. The cfunc
wrapper passed the raw element pointer and stored the return value with
`z_ptr[0] = ...`, so what landed in the buffer was Numba's NestedArray
descriptor (data pointer, shape, strides) rather than the element payload.
The descriptor is the wider of the two, so the store ran off the end of the
element into memory SuiteSparse owns. Measured against the code this
replaces, driving the wrapper's cfunc through ctypes over a guarded buffer:
(6,) np itemsize 48 | numba value type size 56 | overrun 8 bytes
(2, 3) np itemsize 48 | numba value type size 72 | overrun 24 bytes
z[:6] = [100.0, 101.0, 102.0, 103.0, 104.0, 105.0]
z[6:] = [5e-324, -1.0, -1.0, ...] # -1.0 is the guard byte pattern
indices written past the element: [6]
Both operands and the output are now bound with `numba.carray(ptr, shape)`
in the UDT's declared shape, so the UDF receives a numpy view it can index
(`x[i, j]`, `x.shape`) and the wrapper slice-assigns the result back with
`z[:] = ...`. Array-typed record leaves slice-assign for a second reason:
Numba's record-field setitem copies the destination's extent regardless of
the source's, so a short source was read past its end.
The return-type resolution has to move with the wrapper. Once the UDF
receives a view, a UDF that builds its result (`x + y`) instead of returning
an operand types as a plain Numba `Array`, which `lookup_dtype` does not
recognize; `_resolve_udt_return_type` gains an `Array` branch matching it
back to an input array UDT by base element type and rank. Splitting the two
apart would leave array-UDT UDFs that build a result broken outright.
An `Array` type carries `ndim` but not its extents, so a return whose shape
does not fit the UDT is not a type error. It is now memory-safe, but still
silent: the slice-assign's ValueError is raised inside a cfunc, which Numba
prints and swallows, leaving the element as SuiteSparse found it.
Binary ops outside _BUILTIN_UDT_BINARY_OPS (any, first, second) compile through the generic _numba_func branch of BinaryOp._compile_udt. The old wrapper there loaded and stored an array-UDT operand as a NestedArray value, which Numba models as its full array descriptor (meminfo, parent, nitems, itemsize, data, shape, strides): 56 bytes on 64-bit for a 1-D element, regardless of the 32-byte payload. SuiteSparse's generic reduce keeps a UDT accumulator in a stack array sized to the element, so each fold overflowed it by 24 bytes and clobbered a spilled GrB_Type pointer: segfault on some builds, SIGBUS or a silently wrong answer on others. The carray-based wrapper rework on this branch writes exactly itemsize bytes; this test keeps it that way. The test builds the wrapper for binary.any the same way _compile_udt does, compiles it with numba.cfunc, and calls it via ctypes on heap buffers with slack, so a regression trips an assert instead of corrupting a stack frame. The z guard sentinel (0xAB) must differ from the sentinel in y's trailing slack (0xCD): the descriptor load/store is a byte-preserving copy of the source element plus its trailing bytes, so with one shared sentinel the overflow would rewrite z's guard with identical values and go undetected. A public-path any-reduce smoke runs after the byte-level checks, so a regression fails the assert before reaching the code that can crash the process. Verified both directions: passes here (host numba 0.65.1 and a conda env with numba 0.66 plus python-suitesparse-graphblas 10.0.1.1), and on an export of main the assert trips with exactly 24 overflow bytes while the payload check still passes, demonstrating the blind spot a same-sentinel check would have.
A UDF that returns an array-UDT operand as-is types as that operand's NestedArray, which went through lookup_dtype. That registry keys on the Numba type, which a flat FP64[2, 3] and a layered FP64[3][2] share, so it named whichever UDT was registered last. binary.second, monoid.any, and assignment back into the operand's vector then failed with DomainMismatch. Match a NestedArray return to its operand by identity first, and raise if two distinct operands share its type. A NestedArray that matches no operand, such as a record's array field, is no longer matched by rank alone to an operand of a different length. The "matches no input array UDT" error now describes every array operand. Fold the two wrapper byte-guard tests into one, over a 1-D and a 2-D element. A user UDF and binary.any build the same wrapper, and only distinct sentinels see the old overrun reliably: it copied y's trailing bytes, so the dropped test caught it only because the memory after its buffer happened not to match the guard.
numpy keeps np.dtype((FP64[3], (2,))) distinct from FP64[2, 3], but nothing else does: same bytes, same expanded array, same Numba type. As separate DataTypes the nested one broke auto-lifted ops, monoids, and JIT, could not mix with its flat twin, and made lookup_dtype's Numba-type keys resolve to whichever was registered last. - register_anonymous and dtypes.ss.register_new flatten nested subarrays, including inside records - Drop UDT Numba types as registry keys; lookup_dtype converts a Record or NestedArray to its numpy dtype instead - Simplify _resolve_udt_return_type: known extents resolve directly, and only an array built in the UDF is matched to an operand by rank - View, rather than cast, values whose dtype differs only in nesting; numpy's cast between the two scrambles values - Fix stale field-type rules in the UDT guide and note that records with array fields skip the JIT
eriknw
force-pushed
the
09-array-udt-udf-views
branch
from
September 23, 2026 16:22
b4afda6 to
4026f2a
Compare
test_nested_subarray_record_values_load_intact failed on every CI job that
ran the suitesparse-vanilla backend without Numba:
KeyError: "any does not work with ({'li_k': INT64, 'li_pts': FP64[2, 3]},
{'li_k': INT64, 'li_pts': FP64[2, 3]}) types"
On that backend, Vector.build defaults dup_op to binary.any, and a UDT gets
binary.any only by compiling it with Numba. The values_to_numpy_buffer and
Scalar checks need neither, so they still run everywhere; only the Vector
half is gated.
SuiteSparse built without variable-length arrays (MSVC, so the Windows builds) rejects a UDT larger than 128 bytes with GrB_INVALID_VALUE before 9.0. The Windows job that drew psg 8.0.2 failed two tests on it: test_nested_subarrays_register_flat (FP32[3, 2, 7], 168 bytes) and test_udt_array_udf_known_extents_never_match_by_rank (FP64[21], 168 bytes). test_udt_array_udf_returning_operand_keeps_its_type registers 152- and 160-byte types as well; it is marked slow, and that job did not run slow tests. None of the three is about size, so keep their shapes and shrink the element type (INT16 or FP32) rather than skip them on SuiteSparse < 9.
eriknw
added a commit
that referenced
this pull request
Oct 3, 2026
* Reject UDT UDF returns that cannot fill the element A UDF over an array UDT whose return does not fit the element is not a type error. Numba's `Array` type carries `ndim` but not its extents, so returning `x[:2]` from a 9-element UDT types identically to returning `x`. The wrapper's `z[:] = ...` then raises inside the cfunc, where Numba prints the traceback and returns. GraphBLAS is handed no error, and the element keeps whatever was already in the buffer, so the caller gets a wrong answer and no exception. A record UDF that under-fills an array-typed leaf fails the same way and abandons the write part-way: every leaf after the short one, scalar leaves included, keeps what SuiteSparse had there. Extents can only be recovered by running the function, so `_get_udt_wrapper` and the IndexBinaryOp wrapper now call it once on stand-in operands and check that the result fits the UDT. Fitting is not the same as matching. Slice-assign broadcasts, so a `(1,)` return fills a `(6,)` element and a `(1, 3)` return fills both rows of a `(2, 3)` one; both work today and must keep working, and an equality check would reject them. `_fits_by_broadcast` applies numpy's assignment rule instead, and `test_udt_broadcast_matches_numba_slice_assign` pins it against Numba's own slice-assign in both directions, including the two ranks where broadcasting alone gives the wrong answer. The operands are ones rather than zeros for the same reason `dtypes._sample_values` avoids zeros: a UDF that divides by an operand would raise on a zero-filled probe, and a raising probe is treated as "cannot check" rather than as a failure. The check is best-effort, and this commit should not claim more. A UDF that raises on the probe values is not checked at all, so one that also returns a shape that does not fit still reaches the wrapper and still fails silently, exactly as it did before. The wrapper's slice-assign remains the only backstop that always runs. What the check buys is that the common case, a UDF that is simply wrong about the shape, becomes an error at registration instead of a wrong answer at apply time. Because the probe reports the shape for the values it was run on, the diagnostic says so. A UDF whose output shape depends on its operands' values rather than their types can otherwise be rejected on a shape it would never return for real data. This is a judgement call rather than a straight fix, and it is cheap to decline. Typing an op for a UDT now executes the user's function, where before it only compiled it. `OpBase.__contains__` is implemented as a typed lookup, so `udt in some_op` runs user code as a side effect of a membership test. Measured: registering an op runs no probe, since compilation stays lazy, but a single `udt in some_op` runs exactly one. That probe costs no extra compile, because it calls the function with an `ndarray` and the wrapper's `numba.carray` view needs that same specialization anyway. A/B with the check stubbed out: 93.4 ms median disabled against 92.1 ms enabled, two Numba specializations either way. Only array UDTs and records that actually have an array leaf are probed, so an ordinary record UDF is still typed without running user code. The array-UDT paragraph in `docs/user_guide/udt.rst`, added one commit earlier, described the rule as building "a new array of the same shape". It now states the rule the check enforces, and says the shape is learned by running the UDF. If this is declined, the array-UDT view change it sits on still stands. A return that does not fit stays memory-safe, because the slice-assign cannot overrun the element, and goes back to being silent. * Drop the flat-vs-layered ambiguity case from the shape-error test test_udt_array_udf_shape_errors expected a flat FP64[3, 4] and a layered FP64[4] x 3 to be separate UDTs, so that a UDF building an array from the pair was ambiguous. "Flatten nested subarray dtypes so each UDT layout has one DataType" (#589) made them the same DataType, so nothing is ambiguous and the test failed with "DID NOT RAISE UdfParseError" on every CI job with Numba. Keep the same-type check that block ended with, now on udt9. The flattening itself is tested where it is introduced. * Run a UDT UDF for its shape only when Numba cannot tell it The shape check ran the UDF on sample values for every array return and every record return with an array leaf. An operand, or operand field, returned as-is keeps its extents in Numba's `NestedArray` type, so check those from the type and run the UDF only for an array it builds (a plain `Array`). Returning an operand no longer runs user code when the op is typed, and fields returned into the wrong slots are caught even if the UDF would raise on the sample values. The two checkers become one, `_check_udf_fills_output`. Also reject a record leaf the UDF can return as `None` (typed `Optional`, or always `None`), scalar leaves included. On that path the wrapper's assignment raises inside the cfunc and the element comes back holding another element's values. The probe docstring's reason for using ones was wrong: UDFs compile with `error_model="numpy"`, so dividing by zero does not raise. Drop two test cases that repeated #589's resolver tests, and add a test pinning when the UDF is run. * Ignore the UDT repr size warning in the shape-check probe test Below SuiteSparse 9, register_anonymous stores a UDT's numpy repr in a type name capped at 128 characters, and warns when the repr is longer. _ProbeWhenRec's is 142, so the test failed on psg 8 draws that ran slow tests on the suitesparse backend. It is about the shape check, not serialization. * Reject record returns that cannot be written to their fields Two kinds of record UDF return got past the shape check. A tuple or list of the wrong length for an array field failed silently, like a wrong-shape array. Numba copies a sequence into a field element by element, only into a 1-D field and only at its exact length. It does not broadcast, so even a 1-tuple for a 3-element field raises inside the cfunc, and the element keeps whatever was in the buffer. A tuple's length is in its type; a list's is not, so a list is run on sample values like a built array. An array, tuple or list returned for a scalar field failed only at the wrapper compile, with Numba's long TypingError. It is now rejected from its type with a UdfParseError, as a field that can be None already is. Also test the probe's IndexUnaryOp path, which nothing covered. If the row and column were dropped from the probe call, the check silently did nothing. Pin that a scalar or tuple for an array field is checked without running the UDF, and that every leading 1 is dropped in the broadcast rule. udt.rst names the error and documents the record-field rules.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A UDF over an array UDT wrote past the end of its output element. The cfunc
wrapper passed the raw element pointer and stored the return value with
z_ptr[0] = ..., so what landed in the buffer was Numba's NestedArraydescriptor (data pointer, shape, strides) rather than the element payload.
The descriptor is the wider of the two, so the store ran off the end of the
element into memory SuiteSparse owns. Measured against the code this
replaces, driving the wrapper's cfunc through ctypes over a guarded buffer:
Both operands and the output are now bound with
numba.carray(ptr, shape)in the UDT's declared shape, so the UDF receives a numpy view it can index
(
x[i, j],x.shape) and the wrapper slice-assigns the result back withz[:] = .... Array-typed record leaves slice-assign for a second reason:Numba's record-field setitem copies the destination's extent regardless of
the source's, so a short source was read past its end.
The return-type resolution has to move with the wrapper. Once the UDF
receives a view, a UDF that builds its result (
x + y) instead of returningan operand types as a plain Numba
Array, whichlookup_dtypedoes notrecognize;
_resolve_udt_return_typegains anArraybranch matching itback to an input array UDT by base element type and rank. Splitting the two
apart would leave array-UDT UDFs that build a result broken outright.
An
Arraytype carriesndimbut not its extents, so a return whose shapedoes not fit the UDT is not a type error. It is now memory-safe, but still
silent: the slice-assign's ValueError is raised inside a cfunc, which Numba
prints and swallows, leaving the element as SuiteSparse found it.
Stack created with GitHub Stacks CLI • Give Feedback 💬