From 239bff0e88721fbabdc5128b7c6a6f15a97bfadb Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 28 Aug 2026 10:32:05 +0200 Subject: [PATCH] perf(parquet): decode RLE booleans into bitmaps --- parquet/internal/encoding/boolean_decoder.go | 40 ++--- .../boolean_decoder_benchmark_test.go | 123 +++++++++++++ .../encoding/boolean_decoder_bitmap_test.go | 170 ++++++++++++++++++ parquet/internal/utils/bit_reader.go | 69 +++++++ parquet/internal/utils/bit_reader_test.go | 71 ++++++++ parquet/internal/utils/rle.go | 36 ++++ 6 files changed, 479 insertions(+), 30 deletions(-) create mode 100644 parquet/internal/encoding/boolean_decoder_benchmark_test.go create mode 100644 parquet/internal/encoding/boolean_decoder_bitmap_test.go diff --git a/parquet/internal/encoding/boolean_decoder.go b/parquet/internal/encoding/boolean_decoder.go index 61acb3a75..a3b51665d 100644 --- a/parquet/internal/encoding/boolean_decoder.go +++ b/parquet/internal/encoding/boolean_decoder.go @@ -381,38 +381,18 @@ func (dec *RleBooleanDecoder) Decode(out []bool) (int, error) { 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 - } + if max == 0 { + return 0, nil } - writer.Finish() - dec.nvals -= max + decoded, err := dec.rleDec.GetBatchBitmap(out, int(outOffset), max) + dec.nvals -= decoded + if err != nil { + return decoded, err + } + if decoded != max { + return decoded, io.ErrUnexpectedEOF + } return max, nil } diff --git a/parquet/internal/encoding/boolean_decoder_benchmark_test.go b/parquet/internal/encoding/boolean_decoder_benchmark_test.go new file mode 100644 index 000000000..e6064e6b1 --- /dev/null +++ b/parquet/internal/encoding/boolean_decoder_benchmark_test.go @@ -0,0 +1,123 @@ +// 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 encoding_test + +import ( + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow/bitutil" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" +) + +func BenchmarkRleBooleanDecoderDecodeToBitmap(b *testing.B) { + patterns := []struct { + name string + value func(int) bool + }{ + {name: "all_true", value: func(int) bool { return true }}, + {name: "all_false", value: func(int) bool { return false }}, + {name: "alternating", value: func(i int) bool { return i%2 == 0 }}, + {name: "short_runs", value: func(i int) bool { return (i/7)%2 == 0 }}, + } + + for _, size := range []int{64 * 1024, 1024 * 1024} { + for _, pattern := range patterns { + b.Run(fmt.Sprintf("size_%d/%s", size, pattern.name), func(b *testing.B) { + values := makeBooleanValues(size, pattern.value) + data := encodeRleBooleanValues(b, values) + out := make([]byte, bitutil.BytesForBits(int64(size))) + dec := encoding.NewDecoder(parquet.Types.Boolean, parquet.Encodings.RLE, + nil, memory.DefaultAllocator).(encoding.BooleanBitmapDecoder) + + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for b.Loop() { + if err := dec.SetData(size, data); err != nil { + b.Fatal(err) + } + n, err := dec.DecodeToBitmap(out, 0, size) + if err != nil { + b.Fatal(err) + } + if n != size { + b.Fatalf("expected %d values, got %d", size, n) + } + } + }) + } + } +} + +func BenchmarkRleBooleanDecoderDecodeSpacedToBitmap(b *testing.B) { + patterns := []struct { + name string + nullEvery int + }{ + {name: "all_valid"}, + {name: "nullable_10pct", nullEvery: 10}, + {name: "nullable_50pct", nullEvery: 2}, + } + + for _, size := range []int{64 * 1024, 1024 * 1024} { + for _, pattern := range patterns { + b.Run(fmt.Sprintf("size_%d/%s", size, pattern.name), func(b *testing.B) { + logicalValues := makeBooleanValues(size, func(i int) bool { return i%2 == 0 }) + validity := make([]byte, bitutil.BytesForBits(int64(size))) + physicalValues := make([]bool, 0, size) + nullCount := 0 + for i, value := range logicalValues { + if pattern.nullEvery > 0 && i%pattern.nullEvery == 0 { + nullCount++ + continue + } + bitutil.SetBit(validity, i) + physicalValues = append(physicalValues, value) + } + if pattern.nullEvery == 0 { + for i := range logicalValues { + bitutil.SetBit(validity, i) + } + } + + data := encodeRleBooleanValues(b, physicalValues) + out := make([]byte, bitutil.BytesForBits(int64(size))) + dec := encoding.NewDecoder(parquet.Types.Boolean, parquet.Encodings.RLE, + nil, memory.DefaultAllocator).(encoding.BooleanBitmapDecoder) + + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for b.Loop() { + if err := dec.SetData(len(physicalValues), data); err != nil { + b.Fatal(err) + } + n, err := dec.DecodeSpacedToBitmap(out, 0, size, nullCount, validity, 0) + if err != nil { + b.Fatal(err) + } + if n != size { + b.Fatalf("expected %d values, got %d", size, n) + } + } + }) + } + } +} diff --git a/parquet/internal/encoding/boolean_decoder_bitmap_test.go b/parquet/internal/encoding/boolean_decoder_bitmap_test.go new file mode 100644 index 000000000..44e2768a6 --- /dev/null +++ b/parquet/internal/encoding/boolean_decoder_bitmap_test.go @@ -0,0 +1,170 @@ +// 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 encoding_test + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + "testing" + + "github.com/apache/arrow-go/v18/arrow/bitutil" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func encodeRleBooleanValues(t testing.TB, values []bool) []byte { + t.Helper() + enc := encoding.NewEncoder(parquet.Types.Boolean, parquet.Encodings.RLE, + false, nil, memory.DefaultAllocator).(encoding.BooleanEncoder) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + data := append([]byte(nil), buf.Bytes()...) + buf.Release() + return data +} + +func newRleBooleanBitmapDecoder(t testing.TB, nvalues int, data []byte) encoding.BooleanBitmapDecoder { + t.Helper() + dec := encoding.NewDecoder(parquet.Types.Boolean, parquet.Encodings.RLE, + nil, memory.DefaultAllocator).(encoding.BooleanBitmapDecoder) + require.NoError(t, dec.SetData(nvalues, data)) + return dec +} + +func makeBooleanValues(length int, value func(int) bool) []bool { + values := make([]bool, length) + for i := range values { + values[i] = value(i) + } + return values +} + +func assertBitmapMatches(t *testing.T, bitmap []byte, offset int64, expected []bool) { + t.Helper() + for i, want := range expected { + assert.Equal(t, want, bitutil.BitIsSet(bitmap, int(offset)+i), "bit %d", i) + } +} + +func TestRleBooleanDecoderDecodeToBitmap(t *testing.T) { + patterns := []struct { + name string + values []bool + }{ + {name: "repeated_true", values: makeBooleanValues(2048, func(int) bool { return true })}, + {name: "repeated_false", values: makeBooleanValues(2048, func(int) bool { return false })}, + {name: "alternating", values: makeBooleanValues(2048, func(i int) bool { return i%2 == 0 })}, + {name: "short_runs", values: makeBooleanValues(2048, func(i int) bool { return (i/7)%2 == 0 })}, + {name: "mixed", values: makeBooleanValues(2048, func(i int) bool { + switch { + case i < 9: + return true + case i < 23: + return false + default: + return i%5 == 0 + } + })}, + } + + for _, tc := range patterns { + t.Run(tc.name, func(t *testing.T) { + data := encodeRleBooleanValues(t, tc.values) + for _, outOffset := range []int64{0, 1, 7} { + t.Run(fmt.Sprintf("offset_%d", outOffset), func(t *testing.T) { + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(outOffset+int64(len(tc.values)+8)))) + before := append([]byte(nil), out...) + dec := newRleBooleanBitmapDecoder(t, len(tc.values), data) + + n, err := dec.DecodeToBitmap(out, outOffset, len(tc.values)) + require.NoError(t, err) + require.Equal(t, len(tc.values), n) + assertBitmapMatches(t, out, outOffset, tc.values) + + for i := int64(0); i < outOffset; i++ { + assert.Equal(t, bitutil.BitIsSet(before, int(i)), bitutil.BitIsSet(out, int(i)), "prefix bit %d", i) + } + for i := outOffset + int64(len(tc.values)); i < int64(len(out))*8; i++ { + assert.Equal(t, bitutil.BitIsSet(before, int(i)), bitutil.BitIsSet(out, int(i)), "suffix bit %d", i) + } + }) + } + }) + } +} + +func TestRleBooleanDecoderDecodeToBitmapConsecutiveCalls(t *testing.T) { + values := makeBooleanValues(2048, func(i int) bool { return i%2 == 0 }) + data := encodeRleBooleanValues(t, values) + dec := newRleBooleanBitmapDecoder(t, len(values), data) + + calls := []struct { + outOffset int64 + length int + }{ + {outOffset: 7, length: 5}, + {outOffset: 1, length: 1000}, + {outOffset: 0, length: len(values) - 1005}, + } + position := 0 + for i, call := range calls { + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(call.outOffset+int64(call.length+8)))) + n, err := dec.DecodeToBitmap(out, call.outOffset, call.length) + require.NoError(t, err, "call %d", i) + require.Equal(t, call.length, n, "call %d", i) + assertBitmapMatches(t, out, call.outOffset, values[position:position+call.length]) + position += call.length + } + assert.Equal(t, len(values), position) +} + +func TestRleBooleanDecoderDecodeToBitmapTruncatedLiteral(t *testing.T) { + for _, tc := range []struct { + name string + payloadSize int + err error + }{ + {name: "partial_group", payloadSize: 7, err: io.ErrUnexpectedEOF}, + {name: "group_boundary", payloadSize: 4, err: io.EOF}, + } { + t.Run(tc.name, func(t *testing.T) { + payload := append([]byte{17}, bytes.Repeat([]byte{0xff}, tc.payloadSize)...) + data := make([]byte, 4+len(payload)) + binary.LittleEndian.PutUint32(data[:4], uint32(len(payload))) + copy(data[4:], payload) + + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(72))) + before := append([]byte(nil), out...) + dec := newRleBooleanBitmapDecoder(t, 64, data) + n, err := dec.DecodeToBitmap(out, 7, 64) + require.ErrorIs(t, err, tc.err) + require.Equal(t, 32, n) + for i := 0; i < n; i++ { + assert.True(t, bitutil.BitIsSet(out, 7+i), "decoded bit %d", i) + } + for i := n; i < 64; i++ { + assert.Equal(t, bitutil.BitIsSet(before, 7+i), bitutil.BitIsSet(out, 7+i), "unreturned bit %d", i) + } + }) + } +} diff --git a/parquet/internal/utils/bit_reader.go b/parquet/internal/utils/bit_reader.go index 22f5dc493..85f012e22 100644 --- a/parquet/internal/utils/bit_reader.go +++ b/parquet/internal/utils/bit_reader.go @@ -658,6 +658,75 @@ func (b *BitReader) GetBatchLevels(bits uint, out []int16, maxLevel int16) (int, return i, maxCount, nil } +// GetBatchBitmap fills out with bit-packed boolean values without unpacking +// them into a slice of integers. The input must be packed one bit per value. +func (b *BitReader) GetBatchBitmap(out []byte, outOffset, length int) (int, error) { + if length == 0 { + return 0, nil + } + + i := 0 + // Match GetBatch's scalar handling when the current value is in the + // middle of the reader's buffered word. + for ; i < length && b.bitoffset != 0; i++ { + val, err := b.next(1) + if err != nil { + return i, err + } + bitutil.SetBitTo(out, outOffset+i, val != 0) + } + + if _, err := b.reader.Seek(b.byteoffset, io.SeekStart); err != nil { + return i, err + } + + // Read complete 32-value groups directly into the bitmap. Keep the + // temporary packed data so a short final read cannot modify bits that were + // not returned to the caller. + for i < length { + batch := min(buflen, length-i) + batch = batch / 32 * 32 + if batch == 0 { + break + } + + packedBytes := batch / 8 + packed := arrow.Uint32Traits.CastToBytes(b.unpackBuf[:packedBytes/4]) + nread, err := io.ReadFull(b.reader, packed[:packedBytes]) + if err == io.ErrUnexpectedEOF && nread%4 == 0 { + // Match unpack32's group-boundary behavior: an EOF between + // complete 32-value groups is reported as io.EOF. + err = io.EOF + } + completeBytes := nread / 4 * 4 + if completeBytes > 0 { + if (outOffset+i)%8 == 0 { + copy(out[(outOffset+i)/8:], packed[:completeBytes]) + } else { + bitutil.CopyBitmap(packed, 0, completeBytes*8, out, outOffset+i) + } + i += completeBytes * 8 + b.byteoffset += int64(completeBytes) + } + if err != nil { + return i, err + } + } + + if err := b.fillbuffer(); err != nil { + return i, err + } + for ; i < length; i++ { + val, err := b.next(1) + if err != nil { + return i, err + } + bitutil.SetBitTo(out, outOffset+i, val != 0) + } + + return length, nil +} + // GetValue returns a single value that is bit packed using width as the number of bits // and returns false if there weren't enough bits remaining. func (b *BitReader) GetValue(width int) (uint64, bool) { diff --git a/parquet/internal/utils/bit_reader_test.go b/parquet/internal/utils/bit_reader_test.go index 44804c243..be5c01c76 100644 --- a/parquet/internal/utils/bit_reader_test.go +++ b/parquet/internal/utils/bit_reader_test.go @@ -195,6 +195,77 @@ func TestBitReaderGetBatchBools(t *testing.T) { }) } +func TestBitReaderGetBatchBitmap(t *testing.T) { + data := bytes.Repeat([]byte{0xAA, 0xCC, 0xF0}, 128) + + for _, outOffset := range []int{0, 1, 7} { + t.Run(fmt.Sprintf("aligned/offset=%d", outOffset), func(t *testing.T) { + const length = 2048 + reader := utils.NewBitReader(bytes.NewReader(data)) + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(int64(outOffset+length+8)))) + before := append([]byte(nil), out...) + + n, err := reader.GetBatchBitmap(out, outOffset, length) + assert.NoError(t, err) + assert.Equal(t, length, n) + for i := 0; i < length; i++ { + assert.Equal(t, bitutil.BitIsSet(data, i), bitutil.BitIsSet(out, outOffset+i), "bit %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) + } + }) + } + + t.Run("unaligned_input", func(t *testing.T) { + reader := utils.NewBitReader(bytes.NewReader(data)) + _, ok := reader.GetValue(1) + assert.True(t, ok) + const length = 80 + const outOffset = 7 + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(outOffset+length+8))) + + n, err := reader.GetBatchBitmap(out, outOffset, length) + assert.NoError(t, err) + assert.Equal(t, length, n) + for i := 0; i < length; i++ { + assert.Equal(t, bitutil.BitIsSet(data, i+1), bitutil.BitIsSet(out, outOffset+i), "bit %d", i) + } + }) + + for _, tc := range []struct { + name string + inputSize int + length int + want int + err error + }{ + {name: "partial_group", inputSize: 7, length: 64, want: 32, err: io.ErrUnexpectedEOF}, + {name: "group_boundary", inputSize: 4, length: 64, want: 32, err: io.EOF}, + {name: "partial_scalar", inputSize: 1, length: 9, want: 8, err: io.ErrUnexpectedEOF}, + } { + t.Run("truncated/"+tc.name, func(t *testing.T) { + input := bytes.Repeat([]byte{0xff}, tc.inputSize) + reader := utils.NewBitReader(bytes.NewReader(input)) + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(int64(tc.length+8)))) + before := append([]byte(nil), out...) + + n, err := reader.GetBatchBitmap(out, 3, tc.length) + assert.ErrorIs(t, err, tc.err) + assert.Equal(t, tc.want, n) + for i := 0; i < n; i++ { + assert.True(t, bitutil.BitIsSet(out, 3+i), "decoded bit %d", i) + } + for i := n; i < tc.length; i++ { + assert.Equal(t, bitutil.BitIsSet(before, 3+i), bitutil.BitIsSet(out, 3+i), "unreturned bit %d", i) + } + }) + } +} + func TestBitReader(t *testing.T) { buf := []byte{0xAA, 0xCC} // 0b10101010 0b11001100 diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go index 623cd5f7b..2d25ad951 100644 --- a/parquet/internal/utils/rle.go +++ b/parquet/internal/utils/rle.go @@ -281,6 +281,42 @@ func (r *RleDecoder) GetBatchLevels(values []int16, maxLevel int16) (int, int64, return read, maxCount, nil } +// GetBatchBitmap decodes one-bit values directly into a bitmap. +func (r *RleDecoder) GetBatchBitmap(out []byte, outOffset, size int) (int, error) { + if r.bitWidth != 1 { + return 0, errors.New("bitmap decoding requires a bit width of 1") + } + + read := 0 + for read < size { + remain := size - read + + if r.repCount > 0 { + repbatch := min(remain, int(r.repCount)) + bitutil.SetBitsTo(out, int64(outOffset+read), int64(repbatch), r.curVal != 0) + + r.repCount -= int32(repbatch) + read += repbatch + } else if r.litCount > 0 { + litbatch := min(remain, int(r.litCount)) + n, err := r.r.GetBatchBitmap(out, outOffset+read, litbatch) + r.litCount -= int32(n) + read += n + if err != nil { + return read, err + } + if n != litbatch { + return read, nil + } + } else { + if !r.Next() { + return read, nil + } + } + } + return read, nil +} + func (r *RleDecoder) GetBatchSpaced(vals []uint64, nullcount int, validBits []byte, validBitsOffset int64) (int, error) { if nullcount == 0 { return r.GetBatch(vals)