You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ze...@apache.org on 2026/08/18 21:56:56 UTC
(arrow-go) branch main updated: perf(parquet/pqarrow): decode booleans directly to bitmaps (#1182)
This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new ef86032e perf(parquet/pqarrow): decode booleans directly to bitmaps (#1182)
ef86032e is described below
commit ef86032eeaf8758438c8cc22bf1e3908060452ca
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Tue Aug 18 23:56:51 2026 +0200
perf(parquet/pqarrow): decode booleans directly to bitmaps (#1182)
## What does this PR do?
- decodes Parquet BOOLEAN values directly into a packed bitmap
- supports both PLAIN and RLE boolean pages
- expands nullable values inside the bitmap without creating a `[]bool`
- passes the record reader value bitmap directly to the Arrow Boolean
array
## Why?
The read path currently stores decoded booleans as one byte per value.
It then allocates an Arrow bitmap and packs the values again in
`transferBool`.
This finishes wiring the direct bitmap decoder API from #707 into the
pqarrow read path. The record reader now keeps boolean values packed
from decoding through Arrow array construction.
The change also fixes decoder state after a partial PLAIN bitmap decode,
so the next call resumes at the correct source bit.
## Benchmarks
Apple M1 Pro, 1M values, uncompressed Parquet, `GOMAXPROCS=1`:
| case | before | after | change | before B/op | after B/op | change |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| dense | 1.48 ms | 0.15 ms | -90% | 1.61 MB | 0.56 MB | -65% |
| 50% null | 6.72 ms | 4.63 ms | -31% | 4.11 MB | 3.08 MB | -25% |
```text
GOMAXPROCS=1 go test ./parquet/pqarrow -run '^$' -bench '^BenchmarkBooleanBitmapRead$' -benchmem -benchtime=500ms -count=5
```
## Tests
- added direct bitmap tests for consecutive PLAIN decode calls
- added nullable bitmap tests for PLAIN and RLE with unaligned offsets
- ran `go test ./parquet/...`
- ran race tests for the touched encoding, file, and pqarrow packages
- ran `go vet` for the touched packages
---
parquet/file/record_reader.go | 92 ++++++++++++++--
parquet/internal/encoding/boolean_decoder.go | 147 ++++++++++++++++++++++++-
parquet/internal/encoding/encoding_test.go | 126 ++++++++++++++++++++++
parquet/internal/encoding/typed_encoder.go | 9 ++
parquet/pqarrow/boolean_bitmap_bench_test.go | 80 ++++++++++++++
parquet/pqarrow/boolean_bitmap_test.go | 155 +++++++++++++++++++++++++++
parquet/pqarrow/column_readers.go | 31 +++---
7 files changed, 616 insertions(+), 24 deletions(-)
diff --git a/parquet/file/record_reader.go b/parquet/file/record_reader.go
index 62090443..213883f7 100644
--- a/parquet/file/record_reader.go
+++ b/parquet/file/record_reader.go
@@ -21,7 +21,6 @@ import (
"fmt"
"sync"
"sync/atomic"
- "unsafe"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
@@ -85,6 +84,16 @@ type RecordReader interface {
SeekToRow(int64) error
}
+// BooleanRecordReader provides access to the packed value bitmap used for
+// Boolean columns. The existing Values and ReleaseValues methods continue to
+// return one byte per value for compatibility.
+type BooleanRecordReader interface {
+ RecordReader
+ // ReleaseValueBitmap transfers the packed Boolean value bitmap to the
+ // caller. A new buffer will be allocated on subsequent reads.
+ ReleaseValueBitmap() *memory.Buffer
+}
+
// RecordReaderWithError is implemented by record readers that expose failures
// while replacing their page reader. It supplements RecordReader without
// changing the existing no-error SetPageReader method.
@@ -140,6 +149,7 @@ type primitiveRecordReader struct {
valuesCap int64
nullCount int64
values *memory.Buffer
+ legacyValues *memory.Buffer
validBits *memory.Buffer
mem memory.Allocator
@@ -174,6 +184,10 @@ func (pr *primitiveRecordReader) Release() {
pr.values.Release()
pr.values = nil
}
+ if pr.legacyValues != nil {
+ pr.legacyValues.Release()
+ pr.legacyValues = nil
+ }
if pr.validBits != nil {
pr.validBits.Release()
pr.validBits = nil
@@ -196,8 +210,50 @@ func (pr *primitiveRecordReader) ReleaseValidBits() *memory.Buffer {
return res
}
+func (pr *primitiveRecordReader) ReleaseValueBitmap() *memory.Buffer {
+ res := pr.values
+ res.Resize(int(bitutil.BytesForBits(pr.valuesWritten)))
+ pr.values = memory.NewResizableBuffer(pr.mem)
+ pr.valuesCap = 0
+ if pr.legacyValues != nil {
+ pr.legacyValues.ResizeNoShrink(0)
+ }
+ return res
+}
+
+func (pr *primitiveRecordReader) booleanValues() *memory.Buffer {
+ if pr.legacyValues == nil {
+ pr.legacyValues = memory.NewResizableBuffer(pr.mem)
+ }
+
+ pr.legacyValues.Resize(int(pr.valuesWritten))
+ values := pr.legacyValues.Bytes()
+ bitmap := pr.values.Bytes()
+ if len(bitmap) < int(bitutil.BytesForBits(pr.valuesWritten)) {
+ pr.legacyValues.Resize(0)
+ return pr.legacyValues
+ }
+ for i := range values {
+ if bitutil.BitIsSet(bitmap, i) {
+ values[i] = 1
+ } else {
+ values[i] = 0
+ }
+ }
+ return pr.legacyValues
+}
+
func (pr *primitiveRecordReader) ReleaseValues() (res *memory.Buffer) {
res = pr.values
+ if pr.Descriptor().PhysicalType() == parquet.Types.Boolean {
+ res = pr.booleanValues()
+ pr.legacyValues = memory.NewResizableBuffer(pr.mem)
+ pr.values.Release()
+ pr.values = memory.NewResizableBuffer(pr.mem)
+ pr.valuesCap = 0
+ return
+ }
+
nbytes, err := pr.numBytesForValues(pr.valuesWritten)
if err != nil {
panic(err)
@@ -217,10 +273,18 @@ func (pr *primitiveRecordReader) IncrementWritten(w, n int64) {
}
func (pr *primitiveRecordReader) GetValidBits() []byte { return pr.validBits.Bytes() }
func (pr *primitiveRecordReader) ValuesWritten() int64 { return pr.valuesWritten }
-func (pr *primitiveRecordReader) Values() []byte { return pr.values.Bytes() }
+func (pr *primitiveRecordReader) Values() []byte {
+ if pr.Descriptor().PhysicalType() == parquet.Types.Boolean {
+ return pr.booleanValues().Bytes()
+ }
+ return pr.values.Bytes()
+}
func (pr *primitiveRecordReader) ResetValues() {
if pr.valuesWritten > 0 {
pr.values.ResizeNoShrink(0)
+ if pr.legacyValues != nil {
+ pr.legacyValues.ResizeNoShrink(0)
+ }
pr.validBits.ResizeNoShrink(0)
pr.valuesWritten = 0
pr.valuesCap = 0
@@ -229,6 +293,9 @@ func (pr *primitiveRecordReader) ResetValues() {
}
func (pr *primitiveRecordReader) numBytesForValues(nitems int64) (num int64, err error) {
+ if pr.Descriptor().PhysicalType() == parquet.Types.Boolean {
+ return bitutil.BytesForBits(nitems), nil
+ }
typeSize := int64(pr.Descriptor().PhysicalType().ByteSize())
var ok bool
if num, ok = utils.Mul64(nitems, typeSize); !ok {
@@ -264,9 +331,8 @@ func (pr *primitiveRecordReader) ReserveValues(extra int64, hasNullable bool) er
func (pr *primitiveRecordReader) ReadValuesDense(toRead int64) (err error) {
switch cr := pr.ColumnChunkReader.(type) {
case *BooleanColumnChunkReader:
- data := pr.values.Bytes()[int(pr.valuesWritten):]
- values := *(*[]bool)(unsafe.Pointer(&data))
- _, err = cr.curDecoder.(encoding.BooleanDecoder).Decode(values[:toRead])
+ _, err = cr.curDecoder.(encoding.BooleanBitmapDecoder).DecodeToBitmap(
+ pr.values.Bytes(), pr.valuesWritten, int(toRead))
case *Int32ColumnChunkReader:
values := arrow.Int32Traits.CastFromBytes(pr.values.Bytes())[int(pr.valuesWritten):]
_, err = cr.curDecoder.(encoding.Int32Decoder).Decode(values[:toRead])
@@ -300,9 +366,8 @@ func (pr *primitiveRecordReader) ReadValuesSpaced(valuesWithNulls, nullCount int
switch cr := pr.ColumnChunkReader.(type) {
case *BooleanColumnChunkReader:
- data := pr.values.Bytes()[int(pr.valuesWritten):]
- values := *(*[]bool)(unsafe.Pointer(&data))
- _, err = cr.curDecoder.(encoding.BooleanDecoder).DecodeSpaced(values[:int(valuesWithNulls)], int(nullCount), validBits, offset)
+ _, err = cr.curDecoder.(encoding.BooleanBitmapDecoder).DecodeSpacedToBitmap(
+ pr.values.Bytes(), offset, int(valuesWithNulls), int(nullCount), validBits, offset)
case *Int32ColumnChunkReader:
values := arrow.Int32Traits.CastFromBytes(pr.values.Bytes())[int(pr.valuesWritten):]
_, err = cr.curDecoder.(encoding.Int32Decoder).DecodeSpaced(values[:int(valuesWithNulls)], int(nullCount), validBits, offset)
@@ -364,6 +429,14 @@ func (b *binaryRecordReader) ReserveData(nbytes int64) {
b.recordReaderImpl.(binaryRecordReaderImpl).ReserveData(nbytes)
}
+type booleanRecordReader struct {
+ *recordReader
+}
+
+func (b *booleanRecordReader) ReleaseValueBitmap() *memory.Buffer {
+ return b.recordReaderImpl.(*primitiveRecordReader).ReleaseValueBitmap()
+}
+
func newRecordReader(descr *schema.Column, info LevelInfo, mem memory.Allocator, bufferPool *sync.Pool) RecordReader {
if mem == nil {
mem = memory.DefaultAllocator
@@ -376,6 +449,9 @@ func newRecordReader(descr *schema.Column, info LevelInfo, mem memory.Allocator,
repLevels: memory.NewResizableBuffer(mem),
}
rr.refCount.Add(1)
+ if descr.PhysicalType() == parquet.Types.Boolean {
+ return &booleanRecordReader{recordReader: rr}
+ }
return rr
}
diff --git a/parquet/internal/encoding/boolean_decoder.go b/parquet/internal/encoding/boolean_decoder.go
index 84924742..61acb3a7 100644
--- a/parquet/internal/encoding/boolean_decoder.go
+++ b/parquet/internal/encoding/boolean_decoder.go
@@ -24,6 +24,7 @@ import (
"io"
"github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/internal/bitutils"
shared_utils "github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/internal/utils"
@@ -169,7 +170,8 @@ func (dec *PlainBooleanDecoder) DecodeToBitmap(out []byte, outOffset int64, leng
dstSlice[fullBytes] = (dstSlice[fullBytes] &^ mask) | (lastByte & mask)
}
- dec.data = dec.data[bytesToCopy:]
+ dec.data = dec.data[fullBytes:]
+ dec.bitOffset = trailingBits
dec.nvals -= max
return max, nil
}
@@ -205,6 +207,107 @@ func (dec *PlainBooleanDecoder) DecodeSpaced(out []bool, nullCount int, validBit
return dec.Decode(out)
}
+func copyBitmapChunk(buf []byte, srcOffset, dstOffset, length int) {
+ srcByte, srcBit := srcOffset/8, srcOffset%8
+ value := uint16(buf[srcByte]) >> srcBit
+ if srcBit+length > 8 {
+ value |= uint16(buf[srcByte+1]) << (8 - srcBit)
+ }
+ value &= uint16(1<<length) - 1
+
+ dstByte, dstBit := dstOffset/8, dstOffset%8
+ first := min(length, 8-dstBit)
+ firstMask := byte((1<<first)-1) << dstBit
+ buf[dstByte] = (buf[dstByte] &^ firstMask) | (byte(value) << dstBit & firstMask)
+ if first < length {
+ remaining := length - first
+ secondMask := byte((1 << remaining) - 1)
+ buf[dstByte+1] = (buf[dstByte+1] &^ secondMask) | (byte(value>>first) & secondMask)
+ }
+}
+
+func copyBitmapWithinBuffer(buf []byte, srcOffset, dstOffset, length int) {
+ if length == 0 || srcOffset == dstOffset {
+ return
+ }
+
+ if dstOffset < srcOffset {
+ for copied := 0; copied < length; {
+ n := min(length-copied, 8)
+ copyBitmapChunk(buf, srcOffset+copied, dstOffset+copied, n)
+ copied += n
+ }
+ return
+ }
+
+ for copied := length; copied > 0; {
+ n := min(copied, 8)
+ copied -= n
+ copyBitmapChunk(buf, srcOffset+copied, dstOffset+copied, n)
+ }
+}
+
+func expandSpacedBitmapPerBit(out []byte, outOffset int64, length, nullCount int,
+ validBits []byte, validBitsOffset int64) {
+ physicalIndex := 0
+ physicalOffset := int(outOffset) + nullCount
+ for logicalIndex := 0; logicalIndex < length; logicalIndex++ {
+ destination := int(outOffset) + logicalIndex
+ if bitutil.BitIsSet(validBits, int(validBitsOffset)+logicalIndex) {
+ value := bitutil.BitIsSet(out, physicalOffset+physicalIndex)
+ bitutil.SetBitTo(out, destination, value)
+ physicalIndex++
+ }
+ }
+}
+
+func decodeSpacedToBitmap(dec BooleanBitmapDecoder, out []byte, outOffset int64,
+ length, nullCount int, validBits []byte, validBitsOffset int64) (int, error) {
+ if nullCount == 0 {
+ return dec.DecodeToBitmap(out, outOffset, length)
+ }
+
+ valuesToRead := length - nullCount
+ valuesRead, err := dec.DecodeToBitmap(out, outOffset+int64(nullCount), valuesToRead)
+ if err != nil {
+ return valuesRead, err
+ }
+ if valuesRead != valuesToRead {
+ return valuesRead, errors.New("parquet: boolean decoder: number of values / definition levels read did not match")
+ }
+ perBitThreshold := length / 8
+ if length%8 != 0 {
+ perBitThreshold++
+ }
+ if nullCount >= perBitThreshold {
+ expandSpacedBitmapPerBit(out, outOffset, length, nullCount, validBits, validBitsOffset)
+ return length, nil
+ }
+
+ // Expand the packed physical values into their logical positions.
+ // Decoding after the null slots means each destination run is at or before
+ // its source run, so bitmap copies remain safe while the buffers overlap.
+ physicalIndex := int64(0)
+ runs := bitutils.NewSetBitRunReader(validBits, validBitsOffset, int64(length))
+ for {
+ run := runs.NextRun()
+ if run.Length == 0 {
+ break
+ }
+
+ copyBitmapWithinBuffer(out,
+ int(outOffset+int64(nullCount)+physicalIndex),
+ int(outOffset+run.Pos), int(run.Length))
+ physicalIndex += run.Length
+ }
+ return length, nil
+}
+
+func (dec *PlainBooleanDecoder) DecodeSpacedToBitmap(out []byte, outOffset int64,
+ length, nullCount int, validBits []byte, validBitsOffset int64) (int, error) {
+ return decodeSpacedToBitmap(dec, out, outOffset, length, nullCount, validBits, validBitsOffset)
+}
+
type RleBooleanDecoder struct {
decoder
@@ -276,6 +379,43 @@ func (dec *RleBooleanDecoder) Decode(out []bool) (int, error) {
return max, nil
}
+func (dec *RleBooleanDecoder) DecodeToBitmap(out []byte, outOffset int64, length int) (int, error) {
+ max := shared_utils.Min(length, dec.nvals)
+ writer := bitutil.NewBitmapWriter(out, int(outOffset), max)
+
+ var (
+ buf [1024]uint64
+ n = max
+ )
+ for n > 0 {
+ batch := shared_utils.Min(len(buf), n)
+ decoded, err := dec.rleDec.GetBatch(buf[:batch])
+ for _, value := range buf[:decoded] {
+ if value != 0 {
+ writer.Set()
+ } else {
+ writer.Clear()
+ }
+ writer.Next()
+ }
+ n -= decoded
+ if err != nil {
+ writer.Finish()
+ dec.nvals -= max - n
+ return max - n, err
+ }
+ if decoded != batch {
+ writer.Finish()
+ dec.nvals -= max - n
+ return max - n, io.ErrUnexpectedEOF
+ }
+ }
+
+ writer.Finish()
+ dec.nvals -= max
+ return max, nil
+}
+
func (dec *RleBooleanDecoder) DecodeSpaced(out []bool, nullCount int, validBits []byte, validBitsOffset int64) (int, error) {
if nullCount > 0 {
toRead := len(out) - nullCount
@@ -290,3 +430,8 @@ func (dec *RleBooleanDecoder) DecodeSpaced(out []bool, nullCount int, validBits
}
return dec.Decode(out)
}
+
+func (dec *RleBooleanDecoder) DecodeSpacedToBitmap(out []byte, outOffset int64,
+ length, nullCount int, validBits []byte, validBitsOffset int64) (int, error) {
+ return decodeSpacedToBitmap(dec, out, outOffset, length, nullCount, validBits, validBitsOffset)
+}
diff --git a/parquet/internal/encoding/encoding_test.go b/parquet/internal/encoding/encoding_test.go
index 3afac422..4b527332 100644
--- a/parquet/internal/encoding/encoding_test.go
+++ b/parquet/internal/encoding/encoding_test.go
@@ -1322,6 +1322,132 @@ func TestBooleanPlainDecoderDecodeToBitmapUnaligned(t *testing.T) {
}
}
+func TestBooleanPlainDecoderDecodeToBitmapConsecutiveCalls(t *testing.T) {
+ descr := schema.NewColumn(schema.NewBooleanNode("bool", parquet.Repetitions.Optional, -1), 0, 0)
+ enc := encoding.NewEncoder(parquet.Types.Boolean, parquet.Encodings.Plain, false, descr, memory.DefaultAllocator)
+ expected := []bool{true, false, true, true, false, false, true, false, true, true}
+ enc.(encoding.BooleanEncoder).Put(expected)
+ buf, err := enc.FlushValues()
+ require.NoError(t, err)
+
+ dec := encoding.NewDecoder(parquet.Types.Boolean, parquet.Encodings.Plain, descr, memory.DefaultAllocator)
+ require.NoError(t, dec.SetData(len(expected), buf.Buf()))
+ bdec := dec.(encoding.BooleanBitmapDecoder)
+
+ out := make([]byte, bitutil.BytesForBits(16))
+ n, err := bdec.DecodeToBitmap(out, 0, 3)
+ require.NoError(t, err)
+ require.Equal(t, 3, n)
+ n, err = bdec.DecodeToBitmap(out, 8, 7)
+ require.NoError(t, err)
+ require.Equal(t, 7, n)
+
+ for i, value := range expected[:3] {
+ assert.Equal(t, value, bitutil.BitIsSet(out, i))
+ }
+ for i, value := range expected[3:] {
+ assert.Equal(t, value, bitutil.BitIsSet(out, i+8))
+ }
+}
+
+func TestBooleanDecoderDecodeSpacedToBitmap(t *testing.T) {
+ descr := schema.NewColumn(schema.NewBooleanNode("bool", parquet.Repetitions.Optional, -1), 0, 0)
+ const validBitsOffset = 3
+ valid := []bool{false, true, true, false, true, true, false, true, true, false, true}
+ validBits := make([]byte, bitutil.BytesForBits(validBitsOffset+int64(len(valid))))
+ for i, isValid := range valid {
+ if isValid {
+ bitutil.SetBit(validBits, validBitsOffset+i)
+ }
+ }
+ logicalValues := []bool{false, true, false, false, true, true, false, false, true, false, true}
+ physicalValues := make([]bool, 0, len(logicalValues))
+ for i, value := range logicalValues {
+ if valid[i] {
+ physicalValues = append(physicalValues, value)
+ }
+ }
+
+ for _, encodingType := range []parquet.Encoding{parquet.Encodings.Plain, parquet.Encodings.RLE} {
+ t.Run(encodingType.String(), func(t *testing.T) {
+ enc := encoding.NewEncoder(parquet.Types.Boolean, encodingType, false, descr, memory.DefaultAllocator)
+ enc.(encoding.BooleanEncoder).Put(physicalValues)
+ buf, err := enc.FlushValues()
+ require.NoError(t, err)
+
+ dec := encoding.NewDecoder(parquet.Types.Boolean, encodingType, descr, memory.DefaultAllocator)
+ require.NoError(t, dec.SetData(len(physicalValues), buf.Buf()))
+ bdec := dec.(encoding.BooleanBitmapDecoder)
+
+ out := []byte{0xff, 0xff}
+ n, err := bdec.DecodeSpacedToBitmap(out, 2, len(logicalValues),
+ len(logicalValues)-len(physicalValues), validBits, validBitsOffset)
+ require.NoError(t, err)
+ require.Equal(t, len(logicalValues), n)
+ for i, value := range logicalValues {
+ if valid[i] {
+ assert.Equal(t, value, bitutil.BitIsSet(out, i+2), "value %d", i)
+ }
+ }
+ assert.True(t, bitutil.BitIsSet(out, 0))
+ assert.True(t, bitutil.BitIsSet(out, 1))
+ })
+ }
+}
+
+func TestBooleanDecoderDecodeSpacedToBitmapSparseNulls(t *testing.T) {
+ descr := schema.NewColumn(schema.NewBooleanNode("bool", parquet.Repetitions.Optional, -1), 0, 0)
+ const (
+ validBitsOffset = 3
+ length = 32
+ nullIndex = 11
+ outOffset = 2
+ )
+
+ logicalValues := make([]bool, length)
+ validBits := make([]byte, bitutil.BytesForBits(validBitsOffset+length))
+ physicalValues := make([]bool, 0, length-1)
+ for i := range logicalValues {
+ logicalValues[i] = i%3 == 0
+ if i != nullIndex {
+ bitutil.SetBit(validBits, validBitsOffset+i)
+ physicalValues = append(physicalValues, logicalValues[i])
+ }
+ }
+
+ for _, encodingType := range []parquet.Encoding{parquet.Encodings.Plain, parquet.Encodings.RLE} {
+ t.Run(encodingType.String(), func(t *testing.T) {
+ enc := encoding.NewEncoder(parquet.Types.Boolean, encodingType, false, descr, memory.DefaultAllocator)
+ enc.(encoding.BooleanEncoder).Put(physicalValues)
+ buf, err := enc.FlushValues()
+ require.NoError(t, err)
+
+ dec := encoding.NewDecoder(parquet.Types.Boolean, encodingType, descr, memory.DefaultAllocator)
+ require.NoError(t, dec.SetData(len(physicalValues), buf.Buf()))
+ bdec := dec.(encoding.BooleanBitmapDecoder)
+
+ out := []byte{0xa5, 0xa5, 0xa5, 0xa5, 0xa5}
+ before := append([]byte(nil), out...)
+ n, err := bdec.DecodeSpacedToBitmap(out, outOffset, length, 1, validBits, validBitsOffset)
+ require.NoError(t, err)
+ require.Equal(t, length, n)
+
+ for i, value := range logicalValues {
+ got := bitutil.BitIsSet(out, outOffset+i)
+ if i != nullIndex {
+ assert.Equal(t, value, got, "value %d", i)
+ }
+ }
+ for i := 0; i < outOffset; i++ {
+ assert.Equal(t, bitutil.BitIsSet(before, i), bitutil.BitIsSet(out, i), "prefix bit %d", i)
+ }
+ for i := outOffset + length; i < len(out)*8; i++ {
+ assert.Equal(t, bitutil.BitIsSet(before, i), bitutil.BitIsSet(out, i), "suffix bit %d", i)
+ }
+ })
+ }
+}
+
func TestPlainBooleanEncoderPutSpacedBitmapBasic(t *testing.T) {
descr := schema.NewColumn(schema.NewBooleanNode("bool", parquet.Repetitions.Optional, -1), 0, 0)
enc := encoding.NewEncoder(parquet.Types.Boolean, parquet.Encodings.Plain, false, descr, memory.DefaultAllocator)
diff --git a/parquet/internal/encoding/typed_encoder.go b/parquet/internal/encoding/typed_encoder.go
index 3d72ea60..79b887a3 100644
--- a/parquet/internal/encoding/typed_encoder.go
+++ b/parquet/internal/encoding/typed_encoder.go
@@ -56,6 +56,15 @@ type Decoder[T parquet.ColumnTypes] interface {
}
type BooleanDecoder = Decoder[bool]
+
+// BooleanBitmapDecoder decodes boolean values directly into a packed bitmap.
+type BooleanBitmapDecoder interface {
+ BooleanDecoder
+ DecodeToBitmap(out []byte, outOffset int64, length int) (int, error)
+ DecodeSpacedToBitmap(out []byte, outOffset int64, length, nullCount int,
+ validBits []byte, validBitsOffset int64) (int, error)
+}
+
type Int32Decoder = Decoder[int32]
type Int64Decoder = Decoder[int64]
type Int96Decoder = Decoder[parquet.Int96]
diff --git a/parquet/pqarrow/boolean_bitmap_bench_test.go b/parquet/pqarrow/boolean_bitmap_bench_test.go
index 1531e9c3..8309bc68 100644
--- a/parquet/pqarrow/boolean_bitmap_bench_test.go
+++ b/parquet/pqarrow/boolean_bitmap_bench_test.go
@@ -18,6 +18,7 @@ package pqarrow
import (
"bytes"
+ "context"
"testing"
"github.com/apache/arrow-go/v18/arrow"
@@ -25,6 +26,7 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/compress"
+ "github.com/apache/arrow-go/v18/parquet/file"
)
// Benchmark writing boolean columns with direct bitmap path
@@ -106,12 +108,90 @@ func benchmarkBooleanWrite(b *testing.B, size int, nullable bool) {
b.SetBytes(int64(size / 8)) // Report bits as bytes for throughput
}
+func BenchmarkBooleanBitmapRead(b *testing.B) {
+ patterns := []struct {
+ name string
+ nullEvery int
+ }{
+ {name: "dense"},
+ {name: "nullable-1pct", nullEvery: 100},
+ {name: "nullable-5pct", nullEvery: 20},
+ {name: "nullable-10pct", nullEvery: 10},
+ {name: "nullable-50pct", nullEvery: 2},
+ }
+
+ for _, size := range []int{65536, 1000000} {
+ for _, pattern := range patterns {
+ b.Run(formatSize(size)+"/"+pattern.name, func(b *testing.B) {
+ benchmarkBooleanRead(b, size, pattern.nullEvery)
+ })
+ }
+ }
+}
+
+func benchmarkBooleanRead(b *testing.B, size, nullEvery int) {
+ mem := memory.NewGoAllocator()
+ arrowSchema := arrow.NewSchema([]arrow.Field{{
+ Name: "bools", Type: arrow.FixedWidthTypes.Boolean, Nullable: nullEvery > 0,
+ }}, nil)
+
+ bldr := array.NewBooleanBuilder(mem)
+ defer bldr.Release()
+ for i := 0; i < size; i++ {
+ if nullEvery > 0 && i%nullEvery == 0 {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(i%2 == 0)
+ }
+ }
+ arr := bldr.NewBooleanArray()
+ defer arr.Release()
+ rec := array.NewRecordBatch(arrowSchema, []arrow.Array{arr}, int64(size))
+ defer rec.Release()
+
+ var buf bytes.Buffer
+ writer, err := NewFileWriter(arrowSchema, &buf,
+ parquet.NewWriterProperties(parquet.WithCompression(compress.Codecs.Uncompressed)),
+ NewArrowWriterProperties(WithAllocator(mem)))
+ if err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.WriteBuffered(rec); err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.Close(); err != nil {
+ b.Fatal(err)
+ }
+
+ b.ResetTimer()
+ b.ReportAllocs()
+ b.SetBytes(int64(size / 8))
+ ctx := context.Background()
+ for i := 0; i < b.N; i++ {
+ pf, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()))
+ if err != nil {
+ b.Fatal(err)
+ }
+ reader, err := NewFileReader(pf, ArrowReadProperties{}, mem)
+ if err != nil {
+ b.Fatal(err)
+ }
+ tbl, err := reader.ReadTable(ctx)
+ if err != nil {
+ b.Fatal(err)
+ }
+ tbl.Release()
+ }
+}
+
func formatSize(size int) string {
switch {
case size >= 1000000:
return "1M"
case size >= 100000:
return "100K"
+ case size >= 65536:
+ return "64K"
case size >= 10000:
return "10K"
case size >= 1000:
diff --git a/parquet/pqarrow/boolean_bitmap_test.go b/parquet/pqarrow/boolean_bitmap_test.go
new file mode 100644
index 00000000..db346b73
--- /dev/null
+++ b/parquet/pqarrow/boolean_bitmap_test.go
@@ -0,0 +1,155 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pqarrow
+
+import (
+ "bytes"
+ "context"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/compress"
+ "github.com/apache/arrow-go/v18/parquet/file"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestBooleanBitmapReadAcrossPages(t *testing.T) {
+ const numValues = 128
+ for _, tc := range []struct {
+ name string
+ encoding parquet.Encoding
+ }{
+ {name: "plain", encoding: parquet.Encodings.Plain},
+ {name: "rle", encoding: parquet.Encodings.RLE},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{{
+ Name: "bools",
+ Type: arrow.FixedWidthTypes.Boolean,
+ Nullable: true,
+ }}, nil)
+
+ bldr := array.NewBooleanBuilder(mem)
+ defer bldr.Release()
+ for i := 0; i < numValues; i++ {
+ if i%5 == 0 {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(i%2 == 0)
+ }
+ }
+ expected := bldr.NewBooleanArray()
+ defer expected.Release()
+
+ var buf bytes.Buffer
+ writer, err := NewFileWriter(schema, &buf,
+ parquet.NewWriterProperties(
+ parquet.WithCompression(compress.Codecs.Uncompressed),
+ parquet.WithDataPageSize(1),
+ parquet.WithEncoding(tc.encoding),
+ ),
+ NewArrowWriterProperties(WithAllocator(mem)),
+ )
+ require.NoError(t, err)
+
+ record := array.NewRecordBatch(schema, []arrow.Array{expected}, numValues)
+ require.NoError(t, writer.WriteBuffered(record))
+ record.Release()
+ require.NoError(t, writer.Close())
+
+ parquetReader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()),
+ file.WithReadProps(parquet.NewReaderProperties(mem)))
+ require.NoError(t, err)
+ defer parquetReader.Close()
+
+ reader, err := NewFileReader(parquetReader, ArrowReadProperties{}, mem)
+ require.NoError(t, err)
+
+ column, err := reader.GetColumn(context.Background(), 0)
+ require.NoError(t, err)
+ defer column.Release()
+
+ got, err := column.NextBatch(numValues)
+ require.NoError(t, err)
+ defer got.Release()
+
+ require.Len(t, got.Chunks(), 1)
+ require.True(t, array.Equal(expected, got.Chunk(0)))
+ })
+ }
+}
+
+func TestBooleanRecordReaderValuesRemainByteEncoded(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const numValues = 10
+ schema := arrow.NewSchema([]arrow.Field{{
+ Name: "bools",
+ Type: arrow.FixedWidthTypes.Boolean,
+ Nullable: false,
+ }}, nil)
+ want := []byte{1, 0, 1, 1, 0, 0, 1, 0, 1, 1}
+
+ bldr := array.NewBooleanBuilder(mem)
+ defer bldr.Release()
+ for _, value := range want {
+ bldr.Append(value != 0)
+ }
+ expected := bldr.NewBooleanArray()
+ defer expected.Release()
+
+ var buf bytes.Buffer
+ writer, err := NewFileWriter(schema, &buf,
+ parquet.NewWriterProperties(parquet.WithCompression(compress.Codecs.Uncompressed)),
+ NewArrowWriterProperties(WithAllocator(mem)),
+ )
+ require.NoError(t, err)
+ record := array.NewRecordBatch(schema, []arrow.Array{expected}, numValues)
+ require.NoError(t, writer.WriteBuffered(record))
+ record.Release()
+ require.NoError(t, writer.Close())
+
+ parquetReader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()),
+ file.WithReadProps(parquet.NewReaderProperties(mem)))
+ require.NoError(t, err)
+ defer parquetReader.Close()
+
+ pageReader, err := parquetReader.RowGroup(0).GetColumnPageReader(0)
+ require.NoError(t, err)
+ rr := file.NewRecordReader(
+ parquetReader.MetaData().Schema.Column(0),
+ file.LevelInfo{}, arrow.FixedWidthTypes.Boolean, mem, parquetReader.BufferPool())
+ defer rr.Release()
+ rr.SetPageReader(pageReader)
+ require.NoError(t, rr.Reserve(numValues))
+ read, err := rr.ReadRecords(numValues)
+ require.NoError(t, err)
+ require.EqualValues(t, numValues, read)
+
+ require.Equal(t, want, rr.Values())
+ values := rr.ReleaseValues()
+ defer values.Release()
+ require.Equal(t, want, values.Bytes())
+}
diff --git a/parquet/pqarrow/column_readers.go b/parquet/pqarrow/column_readers.go
index 1ab16bd7..1c175391 100644
--- a/parquet/pqarrow/column_readers.go
+++ b/parquet/pqarrow/column_readers.go
@@ -24,7 +24,6 @@ import (
"sync"
"sync/atomic"
"time"
- "unsafe"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
@@ -846,26 +845,28 @@ func transferInt(rdr file.RecordReader, dt arrow.DataType) arrow.ArrayData {
}
func transferBool(rdr file.RecordReader) arrow.ArrayData {
- // TODO(mtopol): optimize this so we don't convert bitmap to []bool back to bitmap
length := rdr.ValuesWritten()
- data := make([]byte, int(bitutil.BytesForBits(int64(length))))
- bytedata := rdr.Values()
- values := *(*[]bool)(unsafe.Pointer(&bytedata))
-
- for idx, v := range values[:length] {
- if v {
- bitutil.SetBit(data, idx)
- }
- }
-
bitmap := rdr.ReleaseValidBits()
if bitmap != nil {
defer bitmap.Release()
}
- bb := memory.NewBufferBytes(data)
- defer bb.Release()
+ var values *memory.Buffer
+ if boolReader, ok := rdr.(file.BooleanRecordReader); ok {
+ values = boolReader.ReleaseValueBitmap()
+ } else {
+ // Keep the bridge compatible with RecordReader implementations that do
+ // not expose the packed Boolean fast path.
+ data := make([]byte, int(bitutil.BytesForBits(int64(length))))
+ for idx, value := range rdr.Values()[:length] {
+ if value != 0 {
+ bitutil.SetBit(data, idx)
+ }
+ }
+ values = memory.NewBufferBytes(data)
+ }
+ defer values.Release()
return array.NewData(&arrow.BooleanType{}, length, []*memory.Buffer{
- bitmap, bb,
+ bitmap, values,
}, nil, int(rdr.NullCount()), 0)
}