[arrow] branch master updated: ARROW-14001: [Go] Fixing AppendBoolean function in BitmapWriter

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/master by this push:
     new 4131fde  ARROW-14001: [Go] Fixing AppendBoolean function in BitmapWriter
4131fde is described below

commit 4131fde463a73d937919f66424850ead32fe2ee1
Author: Matthew Topol <mt...@factset.com>
AuthorDate: Fri Sep 17 13:35:35 2021 -0400

    ARROW-14001: [Go] Fixing AppendBoolean function in BitmapWriter
    
    Also adding tests specifically for the AppendBoolean function in order to ensure it stays fixed.
    
    Closes #11158 from zeroshade/arrow-14001
    
    Authored-by: Matthew Topol <mt...@factset.com>
    Signed-off-by: Matthew Topol <mt...@factset.com>
---
 go/arrow/bitutil/bitmaps.go      | 15 +++++++++------
 go/arrow/bitutil/bitmaps_test.go | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/go/arrow/bitutil/bitmaps.go b/go/arrow/bitutil/bitmaps.go
index 206cde7..5ebe1ae 100644
--- a/go/arrow/bitutil/bitmaps.go
+++ b/go/arrow/bitutil/bitmaps.go
@@ -152,18 +152,21 @@ func (b *BitmapWriter) Next() {
 // AppendBools writes a series of booleans to the bitmapwriter and returns
 // the number of remaining bytes left in the buffer for writing.
 func (b *BitmapWriter) AppendBools(in []bool) int {
-	space := min(int(BytesForBits(int64(b.length-b.pos))), len(in))
+	space := min(b.length-b.pos, len(in))
+	if space == 0 {
+		return 0
+	}
 
+	bitOffset := bits.TrailingZeros32(uint32(b.bitMask))
 	// location that the first byte needs to be written to for appending
-	appslice := b.buf[int(b.byteOffset):]
+	appslice := b.buf[int(b.byteOffset) : b.byteOffset+int(BytesForBits(int64(bitOffset+space)))]
 	// update everything but curByte
-	bitOffset := bits.TrailingZeros32(uint32(b.bitMask))
 	appslice[0] = b.curByte
 	for i, b := range in[:space] {
 		if b {
-			SetBit(appslice, i)
+			SetBit(appslice, i+bitOffset)
 		} else {
-			ClearBit(appslice, i)
+			ClearBit(appslice, i+bitOffset)
 		}
 	}
 
@@ -172,7 +175,7 @@ func (b *BitmapWriter) AppendBools(in []bool) int {
 	b.byteOffset += (bitOffset + space) / 8
 	b.curByte = appslice[len(appslice)-1]
 
-	return int(space)
+	return space
 }
 
 // Finish flushes the final byte out to the byteslice in case it was not already
diff --git a/go/arrow/bitutil/bitmaps_test.go b/go/arrow/bitutil/bitmaps_test.go
index 211949c..25dbdfb 100644
--- a/go/arrow/bitutil/bitmaps_test.go
+++ b/go/arrow/bitutil/bitmaps_test.go
@@ -116,6 +116,12 @@ func TestBitmapWriter(t *testing.T) {
 		}
 		{
 			bitmap := []byte{fillByte, fillByte, fillByte, fillByte}
+			wr := bitutil.NewBitmapWriter(bitmap, 0, 12)
+			wr.AppendBools([]bool{false, true, true, false, true, true, false, false, false, true, false, true})
+			assert.Equal(t, []byte{0x36, (0x0A | (fillByte & 0xF0)), fillByte, fillByte}, bitmap)
+		}
+		{
+			bitmap := []byte{fillByte, fillByte, fillByte, fillByte}
 			wr := bitutil.NewBitmapWriter(bitmap, 3, 12)
 			writeToWriter([]int{0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 1}, wr)
 			// {0b10110..., 0b.1010001, ........, ........}
@@ -123,6 +129,14 @@ func TestBitmapWriter(t *testing.T) {
 		}
 		{
 			bitmap := []byte{fillByte, fillByte, fillByte, fillByte}
+			wr := bitutil.NewBitmapWriter(bitmap, 3, 12)
+			wr.AppendBools([]bool{false, true, true, false})
+			wr.AppendBools([]bool{true, true, false, false})
+			wr.AppendBools([]bool{false, true, false, true})
+			assert.Equal(t, []byte{0xb0 | (fillByte & 0x07), 0x51 | (fillByte & 0x80), fillByte, fillByte}, bitmap)
+		}
+		{
+			bitmap := []byte{fillByte, fillByte, fillByte, fillByte}
 			wr := bitutil.NewBitmapWriter(bitmap, 20, 12)
 			writeToWriter([]int{0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 1}, wr)
 			// {........, ........, 0b0110...., 0b10100011}