fix(compute): handle comparison output offsets#956
Conversation
5c80111 to
632c218
Compare
zeroshade
left a comment
There was a problem hiding this comment.
The pure-Go fix here is correct and nicely tested — handling the leading prefix := offset % 8 sub-batch before the aligned batch loop is exactly right, and the temp-buffer handling is clean.
The concern is that this only fixes the pure-Go path. On the default amd64 build the comparison kernels dispatch to the generated SIMD assembly (scalar_comparison_avx2_amd64.s / scalar_comparison_sse4_amd64.s, generated from _lib/scalar_comparison.cc), which still carries the original output-offset bug:
num_batches = length / kBatchSizeis computed from the originallengthbefore any prefix handling, and the tail loop iterates over the originallength.- With a non-zero output offset (e.g.
length=1, offset%8==7) the SIMD path can read past the input and/or corrupt bits beyond the intended output span.
Because the new tests call the pure-Go helpers directly, they pass green even though the default amd64 dispatch is still broken.
Requested changes before merge:
- Apply the same prefix/offset handling in
_lib/scalar_comparison.ccand regeneratescalar_comparison_avx2_amd64.s/scalar_comparison_sse4_amd64.s. - Add a kernel-level test that exercises the real dispatch (not the pure-Go helper directly) with a non-zero output offset, so the amd64 SIMD path is actually covered.
| ) | ||
|
|
||
| tmpOutSlice := tmpOutput[:] | ||
| if prefix := offset % 8; prefix != 0 { |
There was a problem hiding this comment.
This prefix handling is correct for the pure-Go path — but the same fix needs to land in _lib/scalar_comparison.cc and the regenerated *_amd64.s, which is what the default amd64 build actually dispatches to. On amd64, num_batches/tail are still derived from the original length, so a non-zero output offset still over-reads / corrupts bits. See the main review comment for the full ask.
632c218 to
9e738ea
Compare
zeroshade
left a comment
There was a problem hiding this comment.
Verified both requested changes are addressed:
-
amd64 SIMD path fixed.
_lib/scalar_comparison.ccnow clamps the prefix to the available value count, decrementslength, early-returns when the prefix consumes everything, and recomputesnum_batches(and the tail) after the prefix — applied to all three shapes (arr_arr, arr_scalar, scalar_arr). All AVX2/SSE4.sfiles were regenerated via c2goasm. -
Kernel-level test on the real dispatch.
TestPrimitiveComparisonsWithOutputOffsetexercises both the pure-Go helper and the runtime dispatcher (genCompareKernel, which routes to SIMD on amd64) across all ops, operand shapes, offsets 0–7, and lengths through 65, asserting bits outside the output span are untouched.
Confirmed locally: passes on -tags noasm, default, -race, and GOARCH=amd64 (SSE4 via Rosetta). As a regression check I ran the new test against the pre-fix assembly under GOARCH=amd64 and it fails exactly on the out-of-span corruption — confirming the test genuinely covers the SIMD path and the regenerated assembly is what fixes it.
LGTM.
Rationale for this change
Primitive comparison kernels handle a leading partial output byte before processing aligned 32-value batches. The existing code calculated batch counts from the original length and assumed there were always enough input values to fill that prefix. With a non-byte-aligned output offset, short inputs could read beyond the input or modify bits outside the intended output span.
The issue affected both the pure-Go path and the generated amd64 SIMD implementations.
What changes are included in this PR?
Are these changes tested?
Yes. The regression verifies that bits outside the output range remain unchanged. The following checks pass:
go test ./arrow/compute/... -count=1GOARCH=amd64 go test ./arrow/compute/... -count=1go test -tags noasm ./arrow/compute/... -count=1go test -race ./arrow/compute/internal/kernels -run TestPrimitiveComparisonsWithOutputOffset -count=1Are there any user-facing changes?
Primitive comparisons now handle non-byte-aligned result offsets consistently across pure-Go and amd64 SIMD builds. There is no API change.