From 414537ebe67fd0de122ed2a8eb1bfbac296cfb6d Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Wed, 15 Jul 2026 05:38:34 +0200
Subject: [PATCH] fix(parquet): handle short boolean batches

---
 parquet/internal/utils/bit_reader.go      |  63 ++++++++++++--
 parquet/internal/utils/bit_reader_test.go | 101 ++++++++++++++++++++++
 2 files changed, 158 insertions(+), 6 deletions(-)

diff --git a/parquet/internal/utils/bit_reader.go b/parquet/internal/utils/bit_reader.go
index ced0f265a..5feaa8fa0 100644
--- a/parquet/internal/utils/bit_reader.go
+++ b/parquet/internal/utils/bit_reader.go
@@ -313,6 +313,11 @@ func (b *BitReader) GetBatchBools(out []bool) (int, error) {
 	i := 0
 	// read until we are byte-aligned
 	for ; i < length && b.bitoffset != 0; i++ {
+		if i == 0 || b.bitoffset%8 == 0 {
+			if err := b.ensureBoolBitAvailable(); err != nil {
+				return i, err
+			}
+		}
 		val, err := b.next(bits)
 		out[i] = val != 0
 		if err != nil {
@@ -323,22 +328,53 @@ func (b *BitReader) GetBatchBools(out []bool) (int, error) {
 	b.reader.Seek(b.byteoffset, io.SeekStart)
 	buf := arrow.Uint32Traits.CastToBytes(b.unpackBuf[:])
 	blen := buflen * 8
-	for i < length {
+	for length-i >= 8 {
 		// grab byte-aligned bits in a loop since it's more efficient than going
 		// bit by bit when you can grab 8 bools at a time.
 		unpackSize := utils.Min(blen, length-i) / 8 * 8
-		n, err := b.reader.Read(buf[:bitutil.BytesForBits(int64(unpackSize))])
-		if err != nil {
-			return i, err
+		bytesToRead := int(bitutil.BytesForBits(int64(unpackSize)))
+		n := 0
+		for n < bytesToRead {
+			nread, err := b.reader.Read(buf[n:bytesToRead])
+			n += nread
+			if n == bytesToRead {
+				break
+			}
+			if err != nil {
+				if n > 0 {
+					BytesToBools(buf[:n], out[i:])
+				}
+				i += n * 8
+				b.byteoffset += int64(n)
+				if err == io.EOF && n > 0 {
+					err = io.ErrUnexpectedEOF
+				}
+				return i, err
+			}
+			if nread == 0 {
+				if n > 0 {
+					BytesToBools(buf[:n], out[i:])
+				}
+				i += n * 8
+				b.byteoffset += int64(n)
+				return i, io.ErrNoProgress
+			}
 		}
 		BytesToBools(buf[:n], out[i:])
-		i += unpackSize
+		i += n * 8
 		b.byteoffset += int64(n)
 	}
 
-	b.fillbuffer()
+	if err := b.fillbuffer(); err != nil {
+		return i, err
+	}
 	// grab the trailing bits
 	for ; i < length; i++ {
+		if b.bitoffset%8 == 0 {
+			if err := b.ensureBoolBitAvailable(); err != nil {
+				return i, err
+			}
+		}
 		val, err := b.next(bits)
 		out[i] = val != 0
 		if err != nil {
@@ -349,6 +385,21 @@ func (b *BitReader) GetBatchBools(out []bool) (int, error) {
 	return i, nil
 }
 
+func (b *BitReader) ensureBoolBitAvailable() error {
+	var probe [1]byte
+	n, err := b.reader.ReadAt(probe[:], b.byteoffset+int64(b.bitoffset/8))
+	if n == 1 {
+		return nil
+	}
+	if err == nil {
+		return io.ErrNoProgress
+	}
+	if errors.Is(err, io.EOF) {
+		return io.ErrUnexpectedEOF
+	}
+	return err
+}
+
 func (b *BitReader) Discard(bits uint, n int) (int, error) {
 	if bits > 64 {
 		return 0, errors.New("must be 64 bits or less per read")
diff --git a/parquet/internal/utils/bit_reader_test.go b/parquet/internal/utils/bit_reader_test.go
index 366ae4085..b03849710 100644
--- a/parquet/internal/utils/bit_reader_test.go
+++ b/parquet/internal/utils/bit_reader_test.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+	"io"
 	"math"
 	"math/bits"
 	"math/rand/v2"
@@ -78,6 +79,106 @@ func BenchmarkBitWriter(b *testing.B) {
 	}
 }
 
+type stalledReader struct {
+	*bytes.Reader
+}
+
+func (*stalledReader) Read([]byte) (int, error) { return 0, nil }
+
+func TestBitReaderGetBatchBools(t *testing.T) {
+	data := []byte{0xAA, 0xCC, 0xF0}
+	for length := 0; length <= 17; length++ {
+		t.Run(fmt.Sprintf("aligned/%d", length), func(t *testing.T) {
+			reader := utils.NewBitReader(bytes.NewReader(data))
+			out := make([]bool, length)
+			n, err := reader.GetBatchBools(out)
+			assert.NoError(t, err)
+			assert.Equal(t, length, n)
+			for i, got := range out {
+				assert.Equal(t, data[i/8]&(1<<uint(i%8)) != 0, got)
+			}
+		})
+
+		t.Run(fmt.Sprintf("unaligned/%d", length), func(t *testing.T) {
+			reader := utils.NewBitReader(bytes.NewReader(data))
+			_, ok := reader.GetValue(1)
+			assert.True(t, ok)
+
+			out := make([]bool, length)
+			n, err := reader.GetBatchBools(out)
+			assert.NoError(t, err)
+			assert.Equal(t, length, n)
+			for i, got := range out {
+				bit := i + 1
+				assert.Equal(t, data[bit/8]&(1<<uint(bit%8)) != 0, got)
+			}
+		})
+	}
+
+	t.Run("truncated", func(t *testing.T) {
+		reader := utils.NewBitReader(bytes.NewReader(data[:1]))
+		out := make([]bool, 16)
+		n, err := reader.GetBatchBools(out)
+		assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+		assert.Equal(t, 8, n)
+		for i, got := range out[:n] {
+			assert.Equal(t, data[0]&(1<<uint(i)) != 0, got)
+		}
+	})
+
+	t.Run("truncated remainder", func(t *testing.T) {
+		for _, tc := range []struct {
+			length, inputBytes, expected int
+		}{
+			{length: 1, inputBytes: 0, expected: 0},
+			{length: 7, inputBytes: 0, expected: 0},
+			{length: 9, inputBytes: 1, expected: 8},
+			{length: 15, inputBytes: 1, expected: 8},
+			{length: 17, inputBytes: 2, expected: 16},
+		} {
+			t.Run(fmt.Sprint(tc.length), func(t *testing.T) {
+				reader := utils.NewBitReader(bytes.NewReader(data[:tc.inputBytes]))
+				out := make([]bool, tc.length)
+				for i := range out {
+					out[i] = true
+				}
+
+				n, err := reader.GetBatchBools(out)
+				assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+				assert.Equal(t, tc.expected, n)
+				for i, got := range out[:n] {
+					assert.Equal(t, data[i/8]&(1<<uint(i%8)) != 0, got)
+				}
+				for _, got := range out[n:] {
+					assert.True(t, got)
+				}
+			})
+		}
+	})
+
+	t.Run("truncated unaligned", func(t *testing.T) {
+		reader := utils.NewBitReader(bytes.NewReader(data[:1]))
+		_, ok := reader.GetValue(1)
+		assert.True(t, ok)
+
+		out := make([]bool, 8)
+		n, err := reader.GetBatchBools(out)
+		assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+		assert.Equal(t, 7, n)
+		for i, got := range out[:n] {
+			bit := i + 1
+			assert.Equal(t, data[bit/8]&(1<<uint(bit%8)) != 0, got)
+		}
+	})
+
+	t.Run("no progress", func(t *testing.T) {
+		reader := utils.NewBitReader(&stalledReader{bytes.NewReader(nil)})
+		n, err := reader.GetBatchBools(make([]bool, 8))
+		assert.ErrorIs(t, err, io.ErrNoProgress)
+		assert.Zero(t, n)
+	})
+}
+
 func TestBitReader(t *testing.T) {
 	buf := []byte{0xAA, 0xCC} // 0b10101010 0b11001100