Skip to content

fix(compute): handle comparison output offsets#956

Merged
zeroshade merged 1 commit into
apache:mainfrom
fallintoplace:fix/comparison-offset-bounds
Jul 18, 2026
Merged

fix(compute): handle comparison output offsets#956
zeroshade merged 1 commit into
apache:mainfrom
fallintoplace:fix/comparison-offset-bounds

Conversation

@fallintoplace

@fallintoplace fallintoplace commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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?

  • Clamp prefix processing to the available value count.
  • Recalculate batches and the tail after consuming the prefix.
  • Pass exactly one batch to the pure-Go array-array operation.
  • Apply the same fix to the C++ source and regenerate the AVX2 and SSE4 assembly.
  • Cover offsets 0 through 7 and lengths through 65 for all primitive comparison operations and operand shapes, using both the pure-Go implementation and the real runtime dispatcher.

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=1
  • GOARCH=amd64 go test ./arrow/compute/... -count=1
  • go test -tags noasm ./arrow/compute/... -count=1
  • go test -race ./arrow/compute/internal/kernels -run TestPrimitiveComparisonsWithOutputOffset -count=1
  • 20 repeated amd64 dispatcher regression runs

Are 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.

@fallintoplace
fallintoplace requested a review from zeroshade as a code owner July 16, 2026 12:59
@fallintoplace
fallintoplace force-pushed the fix/comparison-offset-bounds branch from 5c80111 to 632c218 Compare July 17, 2026 00:15

@zeroshade zeroshade left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 / kBatchSize is computed from the original length before any prefix handling, and the tail loop iterates over the original length.
  • 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:

  1. Apply the same prefix/offset handling in _lib/scalar_comparison.cc and regenerate scalar_comparison_avx2_amd64.s / scalar_comparison_sse4_amd64.s.
  2. 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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@fallintoplace
fallintoplace force-pushed the fix/comparison-offset-bounds branch from 632c218 to 9e738ea Compare July 18, 2026 03:09
@fallintoplace fallintoplace changed the title fix(compute): handle pure-Go comparison output offsets fix(compute): handle comparison output offsets Jul 18, 2026

@zeroshade zeroshade left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified both requested changes are addressed:

  1. amd64 SIMD path fixed. _lib/scalar_comparison.cc now clamps the prefix to the available value count, decrements length, early-returns when the prefix consumes everything, and recomputes num_batches (and the tail) after the prefix — applied to all three shapes (arr_arr, arr_scalar, scalar_arr). All AVX2/SSE4 .s files were regenerated via c2goasm.

  2. Kernel-level test on the real dispatch. TestPrimitiveComparisonsWithOutputOffset exercises 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.

@zeroshade
zeroshade merged commit 76109be into apache:main Jul 18, 2026
22 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.

2 participants