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 16:57:20 UTC

(arrow-go) branch main updated: perf(parquet): batch delta binary packed decoding (#1169)

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 63833f78 perf(parquet): batch delta binary packed decoding (#1169)
63833f78 is described below

commit 63833f783b1e52cb0bb0dc0b4e72ad6a2077abff
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Tue Aug 18 18:57:12 2026 +0200

    perf(parquet): batch delta binary packed decoding (#1169)
    
    ### Rationale for this change
    
    DELTA_BINARY_PACKED decoding called `BitReader.GetBatch` once for every
    value in a miniblock. The miniblock is already the unit decoded by this
    code, so this missed the existing batch unpackers.
    
    ### What changes are included in this PR?
    
    - decode each miniblock into reusable scratch with one `GetBatch` call
    - keep 33 to 64 bit values on the generic uint64 path instead of the 32
    bit SIMD unpackers
    - read SIMD input directly into pooled scratch so batched calls do not
    add allocations
    - add int32 and int64 benchmarks with nonzero packed widths
    - add coverage for batched values from 33 through 64 bits
    
    ### Are these changes tested?
    
    Yes. I ran the full Parquet test tree with the parquet-testing fixtures,
    the internal utility and encoding tests with `noasm`, and a linux/amd64
    cross-build.
    
    Apple M1 Pro results for 65,536 values with `GOMAXPROCS=1` and 10
    samples:
    
    | Workload | Time |
    | --- | ---: |
    | int32 small deltas | -58.86% |
    | int32 alternating wide deltas | -66.04% |
    | int64 timestamp-like deltas | -59.72% |
    | int64 deltas wider than 32 bits | -33.89% |
    
    Allocations stay at 2 allocs/op in all four benchmarks.
---
 parquet/internal/encoding/delta_bit_packing.go     | 31 ++++---
 .../encoding/delta_bit_packing_validation_test.go  | 43 ++++++++++
 .../internal/encoding/encoding_benchmarks_test.go  | 96 +++++++++++++++++++++-
 parquet/internal/encoding/encoding_test.go         | 36 ++++++++
 parquet/internal/utils/bit_packing_avx2_amd64.go   | 11 +--
 parquet/internal/utils/bit_packing_neon_arm64.go   | 11 +--
 parquet/internal/utils/bit_reader.go               | 34 ++++----
 parquet/internal/utils/bit_reader_test.go          | 64 +++++++++++++++
 8 files changed, 286 insertions(+), 40 deletions(-)

diff --git a/parquet/internal/encoding/delta_bit_packing.go b/parquet/internal/encoding/delta_bit_packing.go
index 262a2d4e..e8b3a30f 100644
--- a/parquet/internal/encoding/delta_bit_packing.go
+++ b/parquet/internal/encoding/delta_bit_packing.go
@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"math"
 	"math/bits"
 
@@ -30,6 +31,8 @@ import (
 	"github.com/apache/arrow-go/v18/parquet/internal/utils"
 )
 
+const deltaBitPackScratchSize = 1024
+
 // see the deltaBitPack encoder for a description of the encoding format that is
 // used for delta-bitpacking.
 type deltaBitPackDecoder[T int32 | int64] struct {
@@ -181,23 +184,31 @@ func (d *deltaBitPackDecoder[T]) unpackNextMini() error {
 			d.miniBlockValues = append(d.miniBlockValues, T(d.lastVal))
 		}
 	} else {
-		if cap(d.deltaBuf) < 1 {
-			d.deltaBuf = make([]uint64, 1)
+		scratchSize := min(n, deltaBitPackScratchSize)
+		if cap(d.deltaBuf) < scratchSize {
+			d.deltaBuf = make([]uint64, scratchSize)
 		}
-		d.deltaBuf = d.deltaBuf[:1]
+		d.deltaBuf = d.deltaBuf[:scratchSize]
 		minDelta := d.minDelta
 
-		for j := 0; j < n; j++ {
-			nread, err := d.bitdecoder.GetBatch(width, d.deltaBuf)
-			if err != nil {
+		for remaining := n; remaining > 0; {
+			batchSize := min(remaining, len(d.deltaBuf))
+			nread, err := d.bitdecoder.GetBatch(width, d.deltaBuf[:batchSize])
+			if nread != batchSize {
+				if err == nil || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
+					return io.ErrUnexpectedEOF
+				}
 				return err
 			}
-			if nread != 1 {
-				return errors.New("parquet: eof exception")
+			if err != nil {
+				return err
 			}
 
-			d.lastVal += int64(d.deltaBuf[0]) + minDelta
-			d.miniBlockValues = append(d.miniBlockValues, T(d.lastVal))
+			for _, delta := range d.deltaBuf[:nread] {
+				d.lastVal += int64(delta) + minDelta
+				d.miniBlockValues = append(d.miniBlockValues, T(d.lastVal))
+			}
+			remaining -= batchSize
 		}
 	}
 	d.miniBlockIdx++
diff --git a/parquet/internal/encoding/delta_bit_packing_validation_test.go b/parquet/internal/encoding/delta_bit_packing_validation_test.go
index 90b4733c..030aa96a 100644
--- a/parquet/internal/encoding/delta_bit_packing_validation_test.go
+++ b/parquet/internal/encoding/delta_bit_packing_validation_test.go
@@ -17,6 +17,7 @@
 package encoding
 
 import (
+	"io"
 	"testing"
 
 	"github.com/apache/arrow-go/v18/arrow/memory"
@@ -74,3 +75,45 @@ func TestDeltaBitPackDecoderValidatesUsedBitWidth(t *testing.T) {
 		})
 	}
 }
+
+func TestDeltaBitPackDecoderRejectsTruncatedPackedMiniblock(t *testing.T) {
+	// Header: 1024 values/block, 1 miniblock, 1025 values, first value 0.
+	// The miniblock has a one-bit width, but only its first 32 packed values are
+	// present. The missing values must not be reported as a clean EOF.
+	data := []byte{
+		128, 8, // block size
+		1,      // miniblocks per block
+		129, 8, // total values
+		0, // first value
+		0, // minimum delta
+		1, // first miniblock bit width
+		0, 0, 0, 0,
+	}
+
+	dec := NewDecoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked, nil, memory.DefaultAllocator)
+	require.NoError(t, dec.SetData(1025, data))
+
+	_, err := dec.Discard(1025)
+	require.ErrorIs(t, err, io.ErrUnexpectedEOF)
+}
+
+func TestDeltaBitPackDecoderBoundsPackedScratch(t *testing.T) {
+	// A malformed header can declare a large miniblock without providing any
+	// packed values. The reusable packed-value scratch must stay bounded instead
+	// of growing to the attacker-controlled miniblock size.
+	data := []byte{
+		128, 32, // block size: 4096 values
+		1,       // miniblocks per block
+		129, 32, // total values: first value plus one block
+		0, // first value
+		0, // minimum delta
+		1, // first miniblock bit width
+	}
+
+	dec := NewDecoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked, nil, memory.DefaultAllocator).(*deltaBitPackDecoder[int64])
+	require.NoError(t, dec.SetData(4097, data))
+
+	_, err := dec.Discard(2)
+	require.Error(t, err)
+	require.LessOrEqual(t, cap(dec.deltaBuf), deltaBitPackScratchSize)
+}
diff --git a/parquet/internal/encoding/encoding_benchmarks_test.go b/parquet/internal/encoding/encoding_benchmarks_test.go
index 09c176d9..781967a9 100644
--- a/parquet/internal/encoding/encoding_benchmarks_test.go
+++ b/parquet/internal/encoding/encoding_benchmarks_test.go
@@ -735,13 +735,49 @@ func BenchmarkDeltaBinaryPackedEncodingInt32(b *testing.B) {
 }
 
 func BenchmarkDeltaBinaryPackedDecodingInt32(b *testing.B) {
-	for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 {
-		b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) {
-			output := make([]int32, sz)
-			values := make([]int32, sz)
+	patterns := []struct {
+		name   string
+		values func(int) []int32
+	}{
+		{"constant", func(size int) []int32 {
+			values := make([]int32, size)
 			for idx := range values {
 				values[idx] = 64
 			}
+			return values
+		}},
+		{"varying-small-deltas", func(size int) []int32 {
+			values := make([]int32, size)
+			for idx := 1; idx < size; idx++ {
+				values[idx] = values[idx-1] + int32(idx%8)
+			}
+			return values
+		}},
+		{"alternating-wide-deltas", func(size int) []int32 {
+			values := make([]int32, size)
+			for idx := 1; idx < size; idx++ {
+				if idx%2 == 0 {
+					values[idx] = values[idx-1] + 1_000_000
+				} else {
+					values[idx] = values[idx-1] - 1_000_000
+				}
+			}
+			return values
+		}},
+	}
+
+	for _, pattern := range patterns {
+		b.Run(pattern.name, func(b *testing.B) {
+			benchmarkDeltaBinaryPackedDecodingInt32(b, pattern.values)
+		})
+	}
+}
+
+func benchmarkDeltaBinaryPackedDecodingInt32(b *testing.B, makeValues func(int) []int32) {
+	for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 {
+		b.Run(fmt.Sprintf("len-%d", sz), func(b *testing.B) {
+			output := make([]int32, sz)
+			values := makeValues(sz)
 			encoder := encoding.NewEncoder(parquet.Types.Int32, parquet.Encodings.DeltaBinaryPacked,
 				false, nil, memory.DefaultAllocator).(encoding.Int32Encoder)
 			encoder.Put(values)
@@ -751,6 +787,7 @@ func BenchmarkDeltaBinaryPackedDecodingInt32(b *testing.B) {
 			decoder := encoding.NewDecoder(parquet.Types.Int32, parquet.Encodings.DeltaBinaryPacked, nil, memory.DefaultAllocator)
 			b.ResetTimer()
 			b.SetBytes(int64(len(values) * arrow.Int32SizeBytes))
+			b.ReportAllocs()
 			for n := 0; n < b.N; n++ {
 				decoder.SetData(sz, buf.Bytes())
 				decoder.(encoding.Int32Decoder).Decode(output)
@@ -759,6 +796,57 @@ func BenchmarkDeltaBinaryPackedDecodingInt32(b *testing.B) {
 	}
 }
 
+func BenchmarkDeltaBinaryPackedDecodingInt64(b *testing.B) {
+	patterns := []struct {
+		name   string
+		values func(int) []int64
+	}{
+		{"varying-timestamp-deltas", func(size int) []int64 {
+			values := make([]int64, size)
+			for idx := 1; idx < size; idx++ {
+				values[idx] = values[idx-1] + 1_000_000 + int64(idx%1024)
+			}
+			return values
+		}},
+		{"alternating-wide-deltas", func(size int) []int64 {
+			values := make([]int64, size)
+			for idx := 1; idx < size; idx++ {
+				if idx%2 == 0 {
+					values[idx] = values[idx-1] + 1_000_000_000_000
+				} else {
+					values[idx] = values[idx-1] - 1_000_000_000_000
+				}
+			}
+			return values
+		}},
+	}
+
+	for _, pattern := range patterns {
+		b.Run(pattern.name, func(b *testing.B) {
+			for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 {
+				b.Run(fmt.Sprintf("len-%d", sz), func(b *testing.B) {
+					output := make([]int64, sz)
+					values := pattern.values(sz)
+					encoder := encoding.NewEncoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked,
+						false, nil, memory.DefaultAllocator).(encoding.Int64Encoder)
+					encoder.Put(values)
+					buf, _ := encoder.FlushValues()
+					defer buf.Release()
+
+					decoder := encoding.NewDecoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked, nil, memory.DefaultAllocator)
+					b.ResetTimer()
+					b.SetBytes(int64(len(values) * arrow.Int64SizeBytes))
+					b.ReportAllocs()
+					for n := 0; n < b.N; n++ {
+						decoder.SetData(sz, buf.Bytes())
+						decoder.(encoding.Int64Decoder).Decode(output)
+					}
+				})
+			}
+		})
+	}
+}
+
 // Extended MemoTable benchmarks for int64
 func BenchmarkMemoTableInt64(b *testing.B) {
 	tests := []struct {
diff --git a/parquet/internal/encoding/encoding_test.go b/parquet/internal/encoding/encoding_test.go
index 7a73df08..3afac422 100644
--- a/parquet/internal/encoding/encoding_test.go
+++ b/parquet/internal/encoding/encoding_test.go
@@ -956,6 +956,42 @@ func TestWriteDeltaBitPackedInt64(t *testing.T) {
 	})
 }
 
+func TestDeltaBitPackedInt64DecodeAfterDiscard(t *testing.T) {
+	column := schema.NewColumn(schema.NewInt64Node("int64", parquet.Repetitions.Required, -1), 0, 0)
+	values := make([]int64, 257)
+	for idx := 1; idx < len(values); idx++ {
+		delta := int64(1<<40) + int64(idx%17)
+		if idx%2 != 0 {
+			delta = -delta
+		}
+		values[idx] = values[idx-1] + delta
+	}
+
+	enc := encoding.NewEncoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked, false, column, memory.DefaultAllocator)
+	enc.(encoding.Int64Encoder).Put(values)
+	buf, err := enc.FlushValues()
+	require.NoError(t, err)
+	defer buf.Release()
+
+	for _, discard := range []int{0, 1, 31, 32, 33, 127, 128, 129, 256} {
+		t.Run(fmt.Sprintf("discard=%d", discard), func(t *testing.T) {
+			dec := encoding.NewDecoder(parquet.Types.Int64, parquet.Encodings.DeltaBinaryPacked, column, memory.DefaultAllocator)
+			int64Dec := dec.(encoding.Int64Decoder)
+			require.NoError(t, int64Dec.SetData(len(values), buf.Bytes()))
+
+			n, err := int64Dec.Discard(discard)
+			require.NoError(t, err)
+			require.Equal(t, discard, n)
+
+			out := make([]int64, len(values)-discard)
+			n, err = int64Dec.Decode(out)
+			require.NoError(t, err)
+			require.Equal(t, len(out), n)
+			assert.Equal(t, values[discard:], out)
+		})
+	}
+}
+
 func TestDeltaLengthByteArrayEncoding(t *testing.T) {
 	column := schema.NewColumn(schema.NewByteArrayNode("bytearray", parquet.Repetitions.Required, -1), 0, 0)
 
diff --git a/parquet/internal/utils/bit_packing_avx2_amd64.go b/parquet/internal/utils/bit_packing_avx2_amd64.go
index 64fa86d9..74eb936b 100644
--- a/parquet/internal/utils/bit_packing_avx2_amd64.go
+++ b/parquet/internal/utils/bit_packing_avx2_amd64.go
@@ -42,17 +42,18 @@ func unpack32Avx2(in io.Reader, out []uint32, nbits int) (int, error) {
 	defer bufferPool.Put(buffer)
 	buffer.Reset()
 	buffer.Grow(n)
-	nread, err := io.CopyN(buffer, in, int64(n))
-	if err == io.EOF && int(nread)%(nbits*4) != 0 {
-		err = io.ErrUnexpectedEOF
+	packed := buffer.AvailableBuffer()[:n]
+	nread, err := io.ReadFull(in, packed)
+	if err == io.ErrUnexpectedEOF && nread%(nbits*4) == 0 {
+		err = io.EOF
 	}
-	completeBatch := int(nread) * 8 / nbits / 32 * 32
+	completeBatch := nread * 8 / nbits / 32 * 32
 	if completeBatch == 0 {
 		return 0, err
 	}
 
 	var (
-		input  = unsafe.Pointer(&buffer.Bytes()[0])
+		input  = unsafe.Pointer(&packed[0])
 		output = unsafe.Pointer(&out[0])
 	)
 
diff --git a/parquet/internal/utils/bit_packing_neon_arm64.go b/parquet/internal/utils/bit_packing_neon_arm64.go
index e474f64c..fa204657 100644
--- a/parquet/internal/utils/bit_packing_neon_arm64.go
+++ b/parquet/internal/utils/bit_packing_neon_arm64.go
@@ -42,17 +42,18 @@ func unpack32NEON(in io.Reader, out []uint32, nbits int) (int, error) {
 	defer bufferPool.Put(buffer)
 	buffer.Reset()
 	buffer.Grow(n)
-	nread, err := io.CopyN(buffer, in, int64(n))
-	if err == io.EOF && int(nread)%(nbits*4) != 0 {
-		err = io.ErrUnexpectedEOF
+	packed := buffer.AvailableBuffer()[:n]
+	nread, err := io.ReadFull(in, packed)
+	if err == io.ErrUnexpectedEOF && nread%(nbits*4) == 0 {
+		err = io.EOF
 	}
-	completeBatch := int(nread) * 8 / nbits / 32 * 32
+	completeBatch := nread * 8 / nbits / 32 * 32
 	if completeBatch == 0 {
 		return 0, err
 	}
 
 	var (
-		input  = unsafe.Pointer(&buffer.Bytes()[0])
+		input  = unsafe.Pointer(&packed[0])
 		output = unsafe.Pointer(&out[0])
 	)
 
diff --git a/parquet/internal/utils/bit_reader.go b/parquet/internal/utils/bit_reader.go
index 862394cc..39575d0c 100644
--- a/parquet/internal/utils/bit_reader.go
+++ b/parquet/internal/utils/bit_reader.go
@@ -555,24 +555,26 @@ func (b *BitReader) GetBatch(bits uint, out []uint64) (int, error) {
 		}
 	}
 
-	if _, err := b.reader.Seek(b.byteoffset, io.SeekStart); err != nil {
-		return i, err
-	}
-	for i < length {
-		// unpack groups of 32 bytes at a time into a buffer since it's more efficient
-		unpackSize := utils.Min(buflen, length-i)
-		numUnpacked, err := unpack32(b.reader, b.unpackBuf[:unpackSize], int(bits))
-
-		for k := 0; k < numUnpacked; k++ {
-			out[i+k] = uint64(b.unpackBuf[k])
-		}
-		i += numUnpacked
-		b.byteoffset += int64(numUnpacked * int(bits) / 8)
-		if err != nil {
+	if bits <= 32 {
+		if _, err := b.reader.Seek(b.byteoffset, io.SeekStart); err != nil {
 			return i, err
 		}
-		if numUnpacked == 0 {
-			break
+		for i < length {
+			// unpack groups of 32 bytes at a time into a buffer since it's more efficient
+			unpackSize := utils.Min(buflen, length-i)
+			numUnpacked, err := unpack32(b.reader, b.unpackBuf[:unpackSize], int(bits))
+
+			for k := 0; k < numUnpacked; k++ {
+				out[i+k] = uint64(b.unpackBuf[k])
+			}
+			i += numUnpacked
+			b.byteoffset += int64(numUnpacked * int(bits) / 8)
+			if err != nil {
+				return i, err
+			}
+			if numUnpacked == 0 {
+				break
+			}
 		}
 	}
 
diff --git a/parquet/internal/utils/bit_reader_test.go b/parquet/internal/utils/bit_reader_test.go
index 502bbbfb..0468ff13 100644
--- a/parquet/internal/utils/bit_reader_test.go
+++ b/parquet/internal/utils/bit_reader_test.go
@@ -360,6 +360,70 @@ func TestBitArrayVals(t *testing.T) {
 	}
 }
 
+func TestBitReaderGetBatchWideValues(t *testing.T) {
+	for width := uint(33); width <= 64; width++ {
+		t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) {
+			const nvalues = 32
+			values := make([]uint64, nvalues)
+			mask := uint64(math.MaxUint64)
+			if width < 64 {
+				mask = 1<<width - 1
+			}
+			for idx := range values {
+				values[idx] = uint64(idx) * 0x9e3779b97f4a7c15 & mask
+			}
+
+			buf := make([]byte, bitutil.BytesForBits(int64(width*nvalues)))
+			writer := utils.NewBitWriter(utils.NewWriterAtBuffer(buf))
+			for _, value := range values {
+				assert.NoError(t, writer.WriteValue(value, width))
+			}
+			writer.Flush(false)
+
+			reader := utils.NewBitReader(bytes.NewReader(buf))
+			actual := make([]uint64, nvalues)
+			n, err := reader.GetBatch(width, actual)
+			assert.NoError(t, err)
+			assert.Equal(t, nvalues, n)
+			assert.Equal(t, values, actual)
+		})
+	}
+}
+
+func TestBitReaderGetBatchKeepsFollowingValues(t *testing.T) {
+	for width := uint(1); width <= 64; width++ {
+		t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) {
+			const nvalues = 65
+			values := make([]uint64, nvalues)
+			mask := uint64(math.MaxUint64)
+			if width < 64 {
+				mask = 1<<width - 1
+			}
+			for idx := range values {
+				values[idx] = (uint64(idx)*0x9e3779b97f4a7c15 + uint64(idx/3)) & mask
+			}
+
+			buf := make([]byte, bitutil.BytesForBits(int64(width*nvalues)))
+			writer := utils.NewBitWriter(utils.NewWriterAtBuffer(buf))
+			for _, value := range values {
+				assert.NoError(t, writer.WriteValue(value, width))
+			}
+			writer.Flush(false)
+
+			reader := utils.NewBitReader(bytes.NewReader(buf))
+			actual := make([]uint64, nvalues-1)
+			n, err := reader.GetBatch(width, actual)
+			assert.NoError(t, err)
+			assert.Equal(t, nvalues-1, n)
+			assert.Equal(t, values[:nvalues-1], actual)
+
+			value, ok := reader.GetValue(int(width))
+			assert.True(t, ok)
+			assert.Equal(t, values[nvalues-1], value)
+		})
+	}
+}
+
 func TestBitReaderRejectsTruncatedPackedBatches(t *testing.T) {
 	for width := 1; width <= 32; width++ {
 		for _, batchSize := range []int{32, 64} {