perf: Implement native string lowering for Numba layouts and fix silent data corruption

[ianna]

Fixes #2704

Summary

This PR replaces the Python C API path used by layout.py:string_numba_lower with a fully native implementation.

The lowering now handles both:

• Awkward "string" → Numba native unicode_type
• Awkward "bytestring" → Numba native Bytes

Previously, reading strings from an Awkward Array inside @numba.njit required acquiring the GIL, constructing temporary Python str/bytes objects, and boxing them back into Numba native values. This PR removes that round-trip entirely.

The new implementation:

• removes all use of get_python_api,
• eliminates GIL acquisition,
• avoids creating temporary Python str or bytes objects,
• avoids boxing and unboxing between Python and native representations,
• constructs Numba native unicode_type and Bytes values directly.

For Unicode strings, the implementation decodes UTF-8 directly into Numba's native Unicode representation using Numba's internal _empty_string and _set_code_point helpers. For bytestrings, it allocates an NRT-managed buffer and copies the raw bytes directly.

Motivation

This is the second step in removing Python C API dependencies from Awkward's Numba string support.

The previous PR #4249 implemented native lowering for appending Numba Unicode strings to an ArrayBuilder. This PR completes the reverse direction by implementing native lowering for reading Awkward strings and bytestrings into Numba values.

Besides removing unnecessary transitions between native and Python representations, this reduces GIL usage during JIT execution and makes the string lowering path fully native.

Tests

Added regression tests covering:

• ASCII strings
• empty strings
• multi-byte Unicode strings (Latin, CJK, emoji)
• mixed-width Unicode strings
• bytestrings
• common string operations performed inside @numba.njit

These tests verify that strings and bytestrings can be retrieved from Awkward Arrays and used natively within compiled Numba functions without creating intermediate Python objects.

[codecov]

Codecov Report

❌ Patch coverage is 25.28736% with 65 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.92%. Comparing base (908a790) to head (c7a57db).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/awkward/_connect/numba/layout.py	25.28%	65 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_connect/numba/layout.py	77.65% <25.28%> (-6.08%)	⬇️

[TaiSakuma]

Review of PR #4250 — perf: Implement native string lowering for Numba layouts

Overview

The PR rewrites string_numba_lower in src/awkward/_connect/numba/layout.py to construct Numba-native values directly instead of going through the Python C API (GIL acquisition, temporary PyObject*, unbox). Strings get a two-pass UTF-8 decoder built on Numba's internal unicode helpers; bytestrings get an NRT-allocated buffer plus copy. This is the read-direction counterpart of #4249 and implements exactly the approach sketched in issue #2704, including the _make_constant_bytes idiom the issue pointed to.

Main finding: this is a correctness fix, not only a perf change

The old code called string_from_kind_and_data with py_unicode_1byte_kind, which interprets the UTF-8 bytes as Latin-1 code points. Verified against the implementation on main: non-ASCII strings come back corrupted from @numba.njit — "café" returns 'café' and "🙂" returns 'ð\x9f\x99\x82'. The PR fixes this silent data corruption. Neither the title (perf:) nor the body mentions it, and the title type feeds the changelog — either retitling as fix: or adding the behavior fix prominently to the PR body would let it reach the release notes. The existing test_0124_strings_in_numba.py covers only ASCII, which is why the bug was never caught.

Verification performed

All checks below were run locally on macOS arm64 with Python 3.10 and Numba 0.64.0, loading this PR's layout.py on top of an editable install of main (rebinding only string_numba_lower/string_numba_typer):

• All string kinds round-trip: ASCII, Latin-1 range (café), 2-byte kind (κόσμος, 日本語), 4-byte kind (🙂), mixed widths (a🙂é日) — all correct, including len, comparison, .upper(), concatenation.
• Bytestrings with arbitrary non-UTF-8 bytes (b"\x00\x01\x7f\x80\xff") round-trip correctly.
• UTF-8 validation checked against the Unicode Table 3-7 well-formedness rules: overlong rejection (0xC0/0xC1 excluded; 0xE0 requires second ≥ 0xA0; 0xF0 requires second ≥ 0x90), surrogate rejection (0xED with second ≥ 0xA0), above-U+10FFFF rejection (0xF4 with second > 0x8F), continuation-byte and truncation checks — all correct and equivalent to CPython's strict decoder.
• No NRT leaks: with NUMBA_NRT_STATS=1, 2,000 iterations over string and bytestring arrays produced 10,002 allocations and 10,002 frees.
• Data model: the seven fields the bytes path populates match Numba 0.64.0's ArrayModel for Bytes exactly (meminfo, parent, nitems, itemsize, data, shape, strides); it even improves on Numba's own _make_constant_bytes by setting real shape/strides instead of nulls.
• The private helpers used (_empty_string(kind, length, is_ascii=0), _set_code_point, _codepoint_to_kind) exist with the expected signatures in Numba 0.64.0.

Performance: mixed result

Best of 5 runs in the same environment, 140,000 getitems inside @numba.njit (mostly-ASCII words, one 40-char string):

Path	Old	PR	Change
strings	8.39 ms	11.42 ms	~1.4× slower
bytestrings	7.54 ms	3.93 ms	~1.9× faster

The ASCII string path regresses single-threaded: the decoder walks the bytes twice with a per-codepoint function call, while the old path was a memcpy (under the GIL). The genuine wins are correctness for non-ASCII, no GIL acquisition (which serialized prange loops and matters on free-threaded builds — the 3.14t nogil CI jobs pass), and no temporary Python objects.

Suggested (non-blocking) improvement: the first pass already computes is_ascii; when it is 1, replace the second _set_code_point pass with a direct byte copy. That should recover the ASCII regression, which is the common case for HEP data. Adding inline="always" to register_jitable on _decode_utf8_codepoint may also help.

Behavior changes to document

• Non-ASCII strings now decode correctly (previously corrupted).
• Invalid or truncated UTF-8 now raises ValueError inside the JIT function (previously returned garbled text silently). This is defensible — Awkward "string" content is defined as UTF-8 — but it is a new exception path for arrays built from raw buffers, and CPython would raise UnicodeDecodeError rather than ValueError (the latter is a reasonable substitute given nopython-mode constraints). Worth one line in the PR body.

Public API change

No. Both changed functions live in awkward._connect.numba.layout, private by convention, and string_numba_* is referenced nowhere else in src/ or docs/. The typer is untouched, so the types seen inside @numba.njit (numba.types.string, numba.cpython.charseq.bytes_type) and all njit signatures are unchanged. Only runtime behavior changes, as described above.

Existing tests modified

No. The diff contains exactly two files: layout.py (modified) and tests/test_4250_numba_native_strings.py (new). No existing assertion or expected value changes anywhere, and the ASCII-only test_0124_strings_in_numba.py passes unmodified because ASCII behavior is identical.

Test coverage

Good: all code-point widths, mixed-width strings, empty string/bytes, arbitrary bytes, and string operations. Missing: a test asserting ValueError on invalid and truncated UTF-8 — those branches are a large part of the new code, and their absence is likely why codecov/patch fails. The file name follows the test_XXXX_description.py convention.

Risks

• Numba private API dependence: _empty_string, _set_code_point, _codepoint_to_kind are underscore-private and could move in a future Numba release. They have been stable for years, issue Stop making Numba strings through PyObject* #2704 explicitly recommended them, and register_and_check already gates the Numba version, so this is acceptable — a short comment noting the dependence would help future maintainers.
• CI: the only failures are the s390x job (Docker Hub connection timeout during image pull — infrastructure, not the PR; worth a re-run to confirm on big-endian) and codecov/patch (uncovered error branches, addressed by the test suggestion above). All other jobs pass, including both nogil builds and GPU tests.

Style nits (minor)

• The final return output._getvalue() is dedented outside the else: even though only that branch reaches it (the string branch returns earlier); moving it inside the else would read more clearly.
• The one-argument-per-line formatting for two-argument calls (e.g. cgutils.create_struct_proxy(rettype)(context, builder,)) is more vertical than the surrounding file, but it passes the formatter, so this is cosmetic.

Verdict

The implementation is correct and leak-free under everything tested here, and it fixes real silent string corruption on top of the stated goal. Three changes are worth requesting before merge, none structural:

• Surface the correctness fix in the title or body (changelog visibility).
• Add invalid-UTF-8 tests.
• Either add the ASCII fast path or note the measured single-threaded ASCII regression in the PR so it is a known trade-off.

🤖 Generated with Claude Code

[TaiSakuma]

Thank you.

I request these two of the three changes describied in #4250 (comment).

• Surface the correctness fix in the title or body (changelog visibility).
• Add invalid-UTF-8 tests.

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR4250

[ianna]

Thank you.

I request these two of the three changes describied in #4250 (comment).

* Surface the correctness fix in the title or body (changelog visibility).

* Add invalid-UTF-8 tests.

Thanks @TaiSakuma for a detailed review! Both requests are addressed: the title is fixed and the tests are added.

[TaiSakuma]

I thought the title should be "fix:" rather than "perf:". I feel it generally makes more sense to use "fix:" if a PR fixes a bug and improves performance.

But it is fine as is. I had AI look through the repo history and the Conventional Commits doc. There is no such precedence or ranking between types.

[ianna]

I thought the title should be "fix:" rather than "perf:". I feel it generally makes more sense to use "fix:" if a PR fixes a bug and improves performance.

But it is fine as is. I had AI look through the repo history and the Conventional Commits doc. There is no such precedence or ranking between types.

I don’t mind one way or another. I marked it as performance based on the original issue, it turns out there was also a bug to fix 😀

[TaiSakuma]

Thank you for the updates.