Skip to content

Reject UDT UDF returns that cannot fill the element - #590

Open
eriknw wants to merge 3 commits into
41-ci-dep-pools-refreshfrom
10-udt-udf-shape-checks
Open

eriknw wants to merge 3 commits into
41-ci-dep-pools-refreshfrom
10-udt-udf-shape-checks

Conversation

@eriknw

@eriknw eriknw commented Aug 4, 2026

Copy link
Copy Markdown
Member

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.


Stack created with GitHub Stacks CLI • Give Feedback 💬

@eriknw
eriknw marked this pull request as ready for review August 4, 2026 16:07
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 2 times, most recently from 3676d9c to 9ab0b14 Compare August 5, 2026 00:06
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from 9ab0b14 to f5c4f3b Compare August 5, 2026 03:18
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from f5c4f3b to 5a10a1e Compare August 5, 2026 17:44
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from 5a10a1e to 9fc6fa2 Compare August 5, 2026 18:03
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 2 times, most recently from 58c5cc1 to bc96df1 Compare August 6, 2026 07:59
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from bc96df1 to b719ac2 Compare August 6, 2026 15:39
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 2 times, most recently from cb1874e to db8e82d Compare August 6, 2026 20:36
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from db8e82d to 1e2baf3 Compare August 6, 2026 20:42
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 2 times, most recently from 0ab8b52 to 4d91df6 Compare August 7, 2026 05:09
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from 4d91df6 to f209a6f Compare August 26, 2026 17:30
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from f209a6f to 91c7fca Compare September 18, 2026 03:09
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from 91c7fca to e01f5f5 Compare September 18, 2026 03:46
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from e01f5f5 to 4c49bfd Compare September 18, 2026 17:34
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 2 times, most recently from 6f0e968 to ddcf648 Compare September 23, 2026 16:22
@eriknw
eriknw removed this pull request from stack #627 September 23, 2026 18:18
@eriknw
eriknw changed the base branch from 09-array-udt-udf-views to 41-ci-dep-pools-refresh September 23, 2026 18:18
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from ddcf648 to 6f0f862 Compare September 23, 2026 18:19
@eriknw
eriknw added this pull request to stack #634 September 23, 2026 18:19
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from 6f0f862 to 519441d Compare September 24, 2026 20:06
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch 3 times, most recently from 9d92d84 to a689494 Compare September 26, 2026 05:23
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.
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.
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.
@eriknw
eriknw force-pushed the 10-udt-udf-shape-checks branch from a689494 to 0793d8b Compare September 26, 2026 05:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant