Skip to content

Normalize strided inputs before calling Rust kernels - #42

Merged
endrebak merged 1 commit into
masterfrom
fix/accept-strided-arrays
Oct 26, 2026
Merged

endrebak merged 1 commit into
masterfrom
fix/accept-strided-arrays

Conversation

@endrebak

@endrebak endrebak commented Oct 25, 2026 •

Copy link
Copy Markdown
Collaborator

Normalize one-dimensional arrays to contiguous storage at the NumPy boundary, including coordinate/group arrays and strand/direction flags. Already-contiguous inputs are reused. This prevents slice-conversion errors from strided pandas 3.1 columns, including int64 coordinates that do not pass through an int32 cast.

Add 36 regression cases across 18 entry points at small and large coordinates. The ABI3 job now checks out and runs those tests; its smoke command uses python -I so the source checkout cannot shadow the installed wheel. Bumps both Python and Rust package versions to 0.2.8.

Validation: built a macOS ARM64 cp312-abi3 wheel and passed all 36 stride cases plus the isolated ABI3 smoke check. Python 3.12, NumPy 2.5.4 and pandas 3.1.0. The 13 PyRanges1 strided-column integration cases also passed with this wheel.

Benchmark (Apple M3 Pro, Python 3.12.14, NumPy 2.5.4, pandas 3.1.0; median of 9 rotating-order measurements): both pre/post wrappers use the same release-built Rust extension, isolating this normalization patch. For 1M shuffled intervals per side across 25 groups, with int64 coordinates fitting int32, exactly one overlap per query, and contiguous group IDs:

Operation Old contiguous New contiguous New strided pandas coordinates
count_overlaps 25.90 ms 25.78 ms 28.99 ms
overlaps 41.75 ms 41.31 ms 44.46 ms
nearest 56.68 ms 55.76 ms 58.97 ms
subtract 96.73 ms 96.15 ms 100.06 ms
merge 76.75 ms 75.57 ms 76.46 ms

No consistent material slowdown on contiguous inputs was observed. The strided path copies to contiguous Rust slices, adding a few milliseconds per million rows in these datasets. The previous wrapper rejected fully strided inputs, so this is a compatibility/copy-cost benchmark, not an old-strided speedup. Also checked 1,000-row calls, sorted input, strided groups, and coordinates beyond int32; every new-layout output matched the previous wrapper on contiguous copies. Input construction/factorization excluded; boundary normalization and output conversion included. Nearest uses k=1, ties="first". These single-host figures do not measure a Rust algorithm change.

Possible follow-up (not part of this patch): normalize coordinates after dtype selection to avoid an intermediate contiguous int64 copy immediately followed by an int32 conversion.

Follow-up benchmark of delaying coordinate normalization (scratch prototypes; no production changes): simply converting strided int64 views directly to the final dtype is slower because the existing repeated bounds checks then scan strided memory. Combining delayed normalization with reuse of the bounds already computed during dtype selection gives modest gains when coordinates fit int32.

On 1M shuffled intervals per side, stride-three coordinate columns, contiguous group IDs, same Apple M3 Pro/release Rust extension, median of 12 rotating-order measurements:

Stage / operation Current Deferred copy only Deferred + reused bounds
Coordinate preparation 6.216 ms 7.361 ms 4.695 ms
count_overlaps 29.338 ms 30.397 ms 28.336 ms
overlaps 46.997 ms 47.551 ms 45.057 ms
nearest 61.698 ms 61.815 ms 59.590 ms
subtract 97.772 ms 101.864 ms 97.543 ms
merge 75.896 ms 76.710 ms 75.554 ms

Coordinate-preparation peak allocations (tracemalloc, inputs preallocated, groups/kernel/output excluded) fall from 48 MB to 16 MB in the int64-to-int32 case. For coordinates beyond int32, the combined variant is mixed: coordinate preparation rises from 5.979 to 6.896 ms, count_overlaps from 36.979 to 37.886 ms, overlaps from 50.686 to 51.457 ms; other operations vary in either direction. This is not a clear unconditional replacement. All 36 strided-kernel cases and timed output comparisons match across the three wrappers. No validation was disabled; already computed bounds are explicitly reused.

Scaling follow-up (same M3 Pro host, rerun at 1M, 5M and 10M rows per side, 5 paired measurements after warmup): coordinate preparation scales almost exactly linearly in the int64-to-int32 case: current 6.645 / 33.377 / 65.841 ms; combined scratch prototype 5.113 / 24.517 / 50.986 ms. Coordinate-preparation peak allocation at 10M falls from 480 MB to 160 MB. Output comparisons pass and total system swap use is unchanged before/after the run.

Full-call combined-prototype times scale differently:

Operation 1M 5M 10M 10M / 1M
count_overlaps 25.655 ms 224.006 ms 607.079 ms 23.66x
overlaps 44.514 ms 268.242 ms 566.559 ms 12.73x
nearest 57.963 ms 350.231 ms 695.289 ms 12.00x
subtract 94.388 ms 648.089 ms 1306.844 ms 13.85x
merge 74.761 ms 412.777 ms 825.134 ms 11.04x

Additional direct Rust count-kernel experiment (no wrapper normalization, contiguous i32 coordinates/u32 groups, separate run):

Query order Target order 1M 10M Ratio
sorted sorted 6.445 ms 81.617 ms 12.66x
sorted shuffled 12.992 ms 200.717 ms 15.45x
shuffled sorted 15.459 ms 396.343 ms 25.64x
shuffled shuffled 22.005 ms 498.679 ms 22.66x

The release core 0.1.14 uses the same sorted-rank strategy at both sizes; exceptional and small-group paths do not apply. Shuffled-query processing is the strongest lead for the count-kernel scaling: indexed reads of original starts and scattered count writes plausibly lose locality at larger sizes. These experiments do not isolate individual phases, so this remains a profiling lead rather than a proven single cause. Neither scratch normalization prototype has been applied to production.

@endrebak
endrebak merged commit 10e1c2a into master Oct 26, 2026
16 checks passed
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