ARROW-17467: [Go] Aligned Bitmap Ops mess up the final byte when no t… (

apache#13915) …railing bits Authored-by: Matt Topol <[email protected]> Signed-off-by: Matt Topol <[email protected]>
zagto · Oct 7, 2022 · 84b3b84 · 84b3b84
1 parent b7da7b2
commit 84b3b84
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/go/arrow/bitutil/bitmaps.go b/go/arrow/bitutil/bitmaps.go
@@ -449,25 +449,35 @@ func alignedBitmapOp(op bitOp, left, right []byte, lOffset, rOffset int64, out [
 	left = left[lOffset/8:]
 	right = right[rOffset/8:]
 	out = out[outOffset/8:]
+	endMask := (lOffset + length%8)
 	switch nbytes {
 	case 0:
 		return
 	case 1: // everything within a single byte
 		// (length+lOffset%8) <= 8
-		mask := PrecedingBitmask[lOffset%8] | TrailingBitmask[(lOffset+length)%8]
+		mask := PrecedingBitmask[lOffset%8]
+		if endMask != 0 {
+			mask |= TrailingBitmask[(lOffset+length)%8]
+		}
 		out[0] = (out[0] & mask) | (op.opByte(left[0], right[0]) &^ mask)
 	case 2: // don't send zero length to opAligned
 		firstByteMask := PrecedingBitmask[lOffset%8]
 		out[0] = (out[0] & firstByteMask) | (op.opByte(left[0], right[0]) &^ firstByteMask)
-		lastByteMask := TrailingBitmask[(lOffset+length)%8]
+		lastByteMask := byte(0)
+		if endMask != 0 {
+			lastByteMask = TrailingBitmask[(lOffset+length)%8]
+		}
 		out[1] = (out[1] & lastByteMask) | (op.opByte(left[1], right[1]) &^ lastByteMask)
 	default:
 		firstByteMask := PrecedingBitmask[lOffset%8]
 		out[0] = (out[0] & firstByteMask) | (op.opByte(left[0], right[0]) &^ firstByteMask)
 
 		op.opAligned(left[1:nbytes-1], right[1:nbytes-1], out[1:nbytes-1])
 
-		lastByteMask := TrailingBitmask[(lOffset+length)%8]
+		lastByteMask := byte(0)
+		if endMask != 0 {
+			lastByteMask = TrailingBitmask[(lOffset+length)%8]
+		}
 		out[nbytes-1] = (out[nbytes-1] & lastByteMask) | (op.opByte(left[nbytes-1], right[nbytes-1]) &^ lastByteMask)
 	}
 }

diff --git a/go/arrow/bitutil/bitmaps_test.go b/go/arrow/bitutil/bitmaps_test.go
@@ -521,6 +521,22 @@ func TestBitmapOps(t *testing.T) {
 	suite.Run(t, new(BitmapOpSuite))
 }
 
+func TestSmallBitmapOp(t *testing.T) {
+	// 0b01111111 0b11001111
+	left := [2]byte{127, 207}
+	// 0b11111110 0b01111111
+	right := [2]byte{254, 127}
+	// 0b01111110 0b01001111
+	results := [2]byte{126, 79}
+
+	var out [2]byte
+	bitutil.BitmapAnd(left[:], right[:], 0, 0, out[:], 0, 8)
+	assert.Equal(t, results[:1], out[:1])
+
+	bitutil.BitmapAnd(left[:], right[:], 0, 0, out[:], 0, 16)
+	assert.Equal(t, results, out)
+}
+
 func createRandomBuffer(mem memory.Allocator, src *rand.Rand, nbytes int) []byte {
 	buf := mem.Allocate(nbytes)
 	src.Read(buf)