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 d51f6a33 perf(parquet): decode RLE booleans into bitmaps (#1246)
d51f6a33 is described below

commit d51f6a335da65de35f312ed8d55caf413368fce7
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 18:11:05 2026 +0200

    perf(parquet): decode RLE booleans into bitmaps (#1246)
    
    ## What does this PR do?
    
    - **Writes repeated RLE boolean runs directly** into the output bitmap
    - **Copies literal bit-packed runs directly** without a `[]uint64`
    staging buffer
    - Keeps partial reads and truncated-page errors aligned with the
    existing decoder
    
    ## Why?
    
    `RleBooleanDecoder.DecodeToBitmap` currently decodes through a
    `[1024]uint64` buffer and writes each value one at a time. This change
    uses bulk bitmap fills for repeated runs and packed bitmap copies for
    literal runs.
    
    ## Benchmarks
    
    Apple M1 Pro, `GOMAXPROCS=1`, `-benchmem`, 1M values:
    
    | pattern | before | after | change |
    | --- | ---: | ---: | ---: |
    | all true | 2.65 ms | 2.45 us | -99.9% |
    | alternating | 3.45 ms | 0.20 ms | -94.2% |
    | short runs | 3.50 ms | 0.20 ms | -94.3% |
    
    The benchmark also covers 64K values, unaligned output offsets, and
    nullable spaced decoding.
    
    ## Tests
    
    - `go test ./parquet/internal/encoding ./parquet/internal/utils
    ./parquet/pqarrow`
    - `go test -race ./parquet/internal/encoding ./parquet/internal/utils
    ./parquet/pqarrow`
    - `go test -tags noasm ./parquet/internal/encoding
    ./parquet/internal/utils`
---
 parquet/internal/encoding/boolean_decoder.go       |  40 ++---
 .../encoding/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(-)

diff --git a/parquet/internal/encoding/boolean_decoder.go 
b/parquet/internal/encoding/boolean_decoder.go
index 61acb3a7..a3b51665 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 00000000..e6064e6b
--- /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 00000000..44e2768a
--- /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 22f5dc49..85f012e2 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 44804c24..be5c01c7 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 623cd5f7..2d25ad95 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)

Reply via email to