Library-agnostic Arrow C Data Interface for morton_index (arro3-core / PyCapsule)#94
Conversation
espg
left a comment
espg
left a comment
There was a problem hiding this comment.
🤖 from Claude (review)
Fresh-eyes adversarial review of the phase-1 Arrow C Data Interface work. The core FFI ownership dance is sound: the array is moved out of its capsule via std::ptr::replace(ptr, FFI_ArrowArray::empty()) and handed to from_ffi by value (which wraps it in an Arc and owns the release callback), while the schema is only borrowed via &*schema_ptr — matching from_ffi(array: FFI_ArrowArray, schema: &FFI_ArrowSchema) in arrow 53, so there's no double-free and no use-after-free on the happy path. Null↔sentinel is symmetric (export: word == 0 → validity-bitmap null; import: null → 0), and the != DataType::UInt64 guard runs before UInt64Array::from, so a compatible-but-wrong integer type is rejected cleanly rather than panicking. The arrow crate dep was approved on #93; no other new deps; module is well under the 1000-line limit; docstrings match the surrounding style.
Findings, ranked:
- High — the new test module is a no-op in CI. The module-level
pytest.importorskip("arro3.core")skips the entire file at collection time when arro3-core is absent, which (per the file's own docstring) is the CI configuration. So the numpy-only round-trip, pandas dunder, and pyarrow-interop acceptance tests never run in CI. Scope the skip toTestArro3Interoponly. - Medium — no capsule-name validation.
PyCapsule::pointer()uses the capsule's own name, so a mis-ordered/wrong-named capsule pair isn't caught and gets reinterpreted as the wrong FFI struct → UB reachable from safe Python. Validatename()against"arrow_array"/"arrow_schema"before dereferencing. - Low — sliced-array-with-offset import is untested. Logical indexing should handle a non-zero offset correctly, but there's no test pinning it; an offset regression would be silent corruption, so it warrants a fixture.
No blocking concerns on the null/extension-metadata semantics themselves. Comment-only review — not requesting changes.
Generated by Claude Code
|
|
||
| a3 = pytest.importorskip( | ||
| "arro3.core", reason="arro3-core not installed (local-only leg)" | ||
| ) |
There was a problem hiding this comment.
🤖 from Claude (review)
High — this module-level importorskip silently disables the whole test file in CI. pytest.importorskip("arro3.core") here runs at collection time, so when the import fails it raises Skip for the entire module, not just TestArro3Interop. Per the file's own docstring and #93, arro3-core is deliberately not in the CI test extra — which means in CI this line always fails and pytest skips every test above it too: TestNumpyOnlyRoundTrip, TestMortonIndexArrayDunders, and the whole TestPyarrowInterop class. So the export/import round-trip, null/sentinel, and pyarrow-interop acceptance checks that are supposed to run in CI never actually execute — the file collects to zero tests.
Reproduced locally:
def test_before(): assert True
z = pytest.importorskip("definitely_not_installed_xyz")
def test_after(): assert True
# -> "collected 0 items / 1 skipped" (test_before never ran)
Scope the skip to the arro3 class only — e.g. put pytest.importorskip("arro3.core", ...) inside a TestArro3Interop method or a class-scoped fixture / @pytest.mark.skipif, so the numpy and pyarrow legs still run when arro3 is absent.
Generated by Claude Code
| array_capsule: &Bound<'_, PyCapsule>, | ||
| ) -> PyResult<Py<PyAny>> { | ||
| let schema_ptr = schema_capsule.pointer() as *const FFI_ArrowSchema; | ||
| let array_ptr = array_capsule.pointer() as *mut FFI_ArrowArray; |
There was a problem hiding this comment.
🤖 from Claude (review)
Medium — no capsule-name validation; a wrong-named/mis-ordered capsule is UB, not an error. PyCapsule::pointer() in pyo3 0.22 retrieves the pointer using the capsule's own stored name (name_ptr_ignore_error(self) → PyCapsule_GetPointer(cap, PyCapsule_GetName(cap))), so it succeeds for any name. It does not check that array_capsule is named "arrow_array" or schema_capsule "arrow_schema". The Arrow PyCapsule spec requires the consumer to validate these names precisely because the structs are reinterpreted from raw pointers. As written, if a caller passes the pair swapped (schema where array is expected), line 113–114 reinterprets an FFI_ArrowSchema as FFI_ArrowArray, std::ptr::replace moves out garbage, and from_ffi dereferences it → memory corruption / use-after-free, all from safe Python. The non-null check on line 103 doesn't help. Suggest validating array_capsule.name()? == Some(c"arrow_array") and schema_capsule.name()? == Some(c"arrow_schema") before touching the pointers, raising PyValueError otherwise. (The ownership dance itself — moving the array out via replace(..empty()) and borrowing the schema by &* — is correct: from_ffi takes the array by value/Arc and the schema by &, so no double-free.)
Generated by Claude Code
| pa = pytest.importorskip("pyarrow") | ||
| words = _sample_words() | ||
| arr = pa.array(words, type=pa.uint64()) | ||
| np.testing.assert_array_equal(marrow.import_c_array(arr), words) |
There was a problem hiding this comment.
🤖 from Claude (review)
Low — sliced-array-with-offset import is untested (and is the exact case most likely to break). The import loop uses logical indexing (arr.value(i) / arr.is_null(i)), which does respect a non-zero Arrow offset, so this likely works — but there is no test that feeds an offset array through import_c_array. A one-liner would pin it, e.g. pa.array(words).slice(2) (or arro3 equivalent) → import_c_array(...) equals words[2:], including a sentinel landing inside the slice. Given the FFI reinterpretation, an offset bug here would be silent corruption rather than an exception, so it's worth a fixture. (Chunked arrays are correctly out of scope — those are __arrow_c_stream__, not __arrow_c_array__.)
Generated by Claude Code
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #94 +/- ##
==========================================
- Coverage 96.90% 96.46% -0.44%
==========================================
Files 27 28 +1
Lines 4680 4870 +190
==========================================
+ Hits 4535 4698 +163
- Misses 145 172 +27
Flags with carried forward coverage won't be shown. Click here to find out more.
Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
Merging this PR will improve performance by 18.63%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | fast_norm2mort_orders[6] |
327.8 ns | 269.4 ns | +21.65% |
| ⚡ | fast_norm2mort_orders[14] |
390.3 ns | 331.9 ns | +17.57% |
| ⚡ | fast_norm2mort_scalar |
406.9 ns | 348.6 ns | +16.73% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing claude/93-arrow-c-data (91d1869) with main (fee466b)2
Footnotes
-
1 benchmark was skipped, so the baseline result was used instead. If it was deleted from the codebase, click here and archive it to remove it from the performance reports. ↩
-
No successful run was found on
main(9342705) during the generation of this report, so fee466b was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩
add stream here too. Coverage bot is complaining-- is this the stuff that's deferred to #92 ? |
|
🤖 from Claude Stream added (phase 4, commit On the coverage bot — no, it's not the #92 stuff. #92 is purely the no-pyarrow CI workflow leg; it has nothing to do with this red check. codecov/patch is flagging 23 lines, all in The wrinkle: Everything else (the required test/lint/build checks) is green on the phase-4 commit. Generated by Claude Code |
espg
left a comment
espg
left a comment
There was a problem hiding this comment.
🤖 from Claude (review)
Fresh-eyes adversarial review of the phase-4 diff (0b2ae20..91d1869) — Arrow C stream import for morton_index. No blocking findings; no CLAUDE.md violations. Verified against arrow-array 53.4.1 source and by building + running the suite locally (maturin develop --release; test_arrow_cdata.py → 18 passed, 5 skipped [arro3/pandas legs, gated as designed]).
What I checked and confirmed sound:
- Unsafe C-stream FFI (
rust_mi_import_c_stream) — theptr::replace(..., empty())move-out, theDrop-based release of bothstreamandschemaon every exit path (including error/early returns), theis_released()end-of-stream check, the pre-initializedffi_arraymaking theget_next-error partial state a safe no-op-drop, andfrom_ffi_and_data_type(..., UInt64)honoring chunk offsets. No leak, no double-free. (inline) - Dispatch order in
import_c_array— for pyarrow,Arrayexposes only__arrow_c_array__andChunkedArrayonly__arrow_c_stream__(verified), so array-before-stream is correct and the documented "prefer contiguous" tie-break is sensible for any object exposing both. No wrong-path bug. - Null/sentinel + offset across chunks — preserved via
append_words's logical accessors; cross-boundary nulls and sliced-chunk offsets are both exercised by the tests (inline). - Extension-type schema gate —
DataType::try_from(&schema)yieldsUInt64for an extension-tagged uint64 stream, so the type check behaves as the comment claims. - CLAUDE.md conformance —
arrow_ffi.rsis 224 lines (well under ~1000); commitphase 4 of issue #93is terse/title-only per repo style; no new dependency (arrow/ffi_streampre-approved); docstrings match the surrounding freeform style; behavioral change has tests.
Two non-blocking notes left inline: (a) the new fn re-triggers the file's pre-existing clippy useless_conversion nit on PyResult<Py<PyAny>> (consistent with rust_mi_import_c_array, not newly introduced); (b) the pandas-facing chunked from_arrow/isna() test is skip-only without pandas (covered in CI extra; numpy-only chunked path carries coverage otherwise).
Generated by Claude Code
| Ok(numpy::PyArray1::from_vec_bound(py, words) | ||
| .into_any() | ||
| .unbind()) | ||
| } |
There was a problem hiding this comment.
🤖 from Claude (review)
Verified the unsafe C-stream FFI against arrow-array 53.4.1 and it holds up:
ptr::replace(ptr, FFI_ArrowArrayStream::empty())moves the stream out and leaves arelease: Nonestub, so the producer's capsule destructor is a no-op — no double-free. The owned localstreamdrops throughimpl Drop(which callsrelease) on every exit path, including therc != 0anddt != UInt64early returns, so the producer stream is always released.FFI_ArrowSchema(localschema) has its ownDrop-based release, so no schema leak either.- End-of-stream via
ffi_array.is_released()(==release.is_none()) matches the C Data Interface contract;ffi_arrayis pre-initialized toempty(), so on aget_nexterror (contract leavesoutuntouched) the local is a valid no-op-drop struct — safe partial state. from_ffi_and_data_type(ffi_array, DataType::UInt64)consumes by value and honorsarray.offset(), so a sliced chunk (non-zero offset) imports correctly;append_wordsthen uses logical accessors, preserving null→sentinel per chunk.DataType::try_from(&schema)returns the storage type (ignoringARROW:extension:*field metadata), so the extension-typed uint64 stream passes the type gate as the doc comment claims.
No blocking issue here. One nit: this fn re-triggers clippy useless_conversion on the PyResult<Py<PyAny>> return (same as the pre-existing rust_mi_import_c_array). It's a consistent pre-existing pattern across the file, so not introduced here — flagging only so the file's clippy noise is a known quantity.
Generated by Claude Code
| mia = mortie.MortonIndexArray.from_arrow(ca) | ||
| np.testing.assert_array_equal(mia._data, words) | ||
| np.testing.assert_array_equal(mia.isna(), words == SENTINEL) | ||
|
|
There was a problem hiding this comment.
🤖 from Claude (review)
Coverage is good and I confirmed it locally (18 passed / 5 skipped, pyarrow leg green):
test_import_chunked_streamchunks[:3],[3:5],[5:]over_words_with_gap()(sentinels at idx 2 and 5), so nulls land at a chunk-end (idx 2) and a chunk-start (idx 5) — the cross-boundary null case is exercised.- Because pyarrow slices are zero-copy,
arr[3:5]/arr[5:]carry non-zerooffset, so the per-chunk offset path is covered too (I verifiedslice().offset != 0). - Empty stream and non-uint64 rejection are both pinned.
One gap worth noting: test_from_arrow_chunked is the only thing exercising the pandas-facing MortonIndexArray.from_arrow(chunked) + isna() over stream data, and it importorskip("pandas")s away. That's fine for local runs (it runs under the CI pandas extra), but it means the chunked→isna() assertion has no coverage in a pandas-absent env — the numpy-only import_c_array chunked path (test_import_chunked_stream) does carry the load there, so this is a note, not a request.
Generated by Claude Code
1 participant
Closes #93.
What & why
Makes the
morton_indexArrow surface library-agnostic: any Arrow lib that speaks the PyCapsule C Data Interface — arro3-core (the carrier zagg runs on its Lambda worker, where pyarrow is intentionally absent), pyarrow, polars — can pull a typedmorton_indexcolumn zero-copy and hand one back, carrying a real Arrow extension type rather than shimming through plainuint64.Where #86's
mortie.arrowbuilds a pyarrowExtensionArrayand is gated on pyarrow (_require_pyarrow,fill_null,to_pandas) and has no export hook, this path builds the raw Arrow C structs in Rust (thearrowcrate's C Data Interface), so the runtime stays numpy-only and gets both directions.Approach follows the decisions on #93: approach (A) (Rust FFI), the umbrella
arrowcrate, pinned targetarro3-core==0.8.1, and the arro3-only test kept runnable-locally + documented (no.github/workflows/change here — bundled to #92).Key finding (the #93 item-2 gate)
arro3-core 0.8.1 round-trips the extension metadata (
ARROW:extension:name = mortie.morton_index) through the C-Data boundary — verified inTestArro3Interop. So we ship the real extension type; nouint64-storage fallback is needed on the Lambda layer.Design
src_rust/src/arrow_ffi.rs, via thearrowcrate):rust_mi_export_c_array(words)/rust_mi_export_c_schema(): build the(schema, array)capsules — aUInt64Arraywith a real validity bitmap from the empty sentinel (data == 0→ Arrow null), schema Field carrying the extension metadata.rust_mi_import_c_array(schema, array): single contiguous array → packed words. Validates each capsule's C-Data name before dereferencing; moves the FFI struct out leaving a released stub (no double-free); rejects non-uint64.rust_mi_import_c_stream(stream): a chunked column (__arrow_c_stream__). Its stream schema is the bare (non-struct)uint64, which the struct-orientedArrowArrayStreamReaderrejects, so this drives the C stream callbacks directly (schema once,get_nexteach chunk to the released end-marker, import each viafrom_ffi_and_data_type(_, UInt64)), concatenating.MortonIndexArray.__arrow_c_array__/__arrow_c_schema__/from_arrow(source)(mortie/morton_index.py).mortie.arrow.export_c_array/export_c_schema/import_c_array(mortie/arrow.py) — pyarrow-free;import_c_arraydispatches__arrow_c_array__(contiguous) vs__arrow_c_stream__(chunked) vs a raw capsule tuple.Phases
arro3install extra (local-only) +docs/arrow_interchange.md.__arrow_c_stream__) per review feedback; verified round-tripping pyarrow + arro3-coreChunkedArray, empty, and non-uint64 rejection.Self-review resolutions
Phase-1 review folded (commit
fold phase 1 self-review on #93): scoped thearro3skip toTestArro3Interoponly (a module-levelimportorskipwould have skipped the whole file in CI); added capsule-name validation (checked_pointer) against a memory-corruption path; added an offset-slice import test. (Phase-4 stream self-review in progress.)How it was tested
Local (fresh venv,
maturin develop --release,arro3-core==0.8.1):pytest mortie/tests/test_arrow_cdata.py— 23 passed locally; the arro3 leg skips when arro3-core is absent (the CI shape).pytest test_arrow.py test_morton_index.py— 57 passed, 1 skipped (no regression to the morton index surface followup: points= encode + __from_arrow__ #86 pyarrow skin).cargo test --lib— 188 passed.cargo fmtclean;flake8 mortie --select=E9,F63,F7,F82clean.Notes for review
uint64fallback), as done.TestArro3Interoplines (skipped in CI); source files are 100% patch-covered.codecov.yml's existingignore: mortie/tests/**/*isn't excluding files directly undermortie/tests/from patch — a one-line glob fix greens it; flagged on the thread for your go-ahead since it's CI config.mortie/arrow.pyalongside the pyarrow skin (import no pyarrow) — flagging the placement in case you'd prefer a separate module.Refs #86, #92.