From 89c8b6d2c33a9780d601c98555ac9e380641e903 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 28 Aug 2026 16:44:35 -0400 Subject: [PATCH] fix(parquet): stop BYTE_STREAM_SPLIT FLBA spaced decode aliasing output A BYTE_STREAM_SPLIT FIXED_LEN_BYTE_ARRAY column containing nulls and spanning more than one data page decoded values shifted by one position, silently and without an error. spacedExpand moves values into their spaced positions with copy and deliberately leaves the null slots alone. For the scalar column types that is fine, but ByteArray and FixedLenByteArray buffers hold slice headers, so copying leaves duplicate headers behind: a null slot and a valid slot end up referencing the same backing array. That is still harmless for a decoder that replaces the header, which is what PlainFixedLenByteArrayDecoder does. ByteStreamSplitFixedLenByteArrayDecoder instead writes through the caller's existing slice, so once the record reader reused its value buffer for the next page two output slots shared one backing array and clobbered each other. Add spacedExpandSwap, which swaps rather than copies so the buffer stays a permutation of its original elements, and use it from the byte-stream-split decoder. No slot aliases another and every slot keeps its reusable capacity, so the buffer reuse the decoder relies on is preserved. spacedExpand itself is untouched, leaving every other column type and the non-spaced path unaffected. Fixes #1255 --- .../encoding/byte_stream_split_decode_test.go | 95 +++++++++++++++ parquet/internal/encoding/decoder.go | 40 +++++++ .../encoding/fixed_len_byte_array_decoder.go | 4 +- .../byte_stream_split_multipage_test.go | 113 ++++++++++++++++++ 4 files changed, 251 insertions(+), 1 deletion(-) create mode 100644 parquet/pqarrow/byte_stream_split_multipage_test.go diff --git a/parquet/internal/encoding/byte_stream_split_decode_test.go b/parquet/internal/encoding/byte_stream_split_decode_test.go index e6c72bdc0..cd714fe93 100644 --- a/parquet/internal/encoding/byte_stream_split_decode_test.go +++ b/parquet/internal/encoding/byte_stream_split_decode_test.go @@ -19,11 +19,16 @@ package encoding import ( "bytes" "fmt" + "math/rand" "testing" "unsafe" + "github.com/apache/arrow-go/v18/arrow/bitutil" + "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/internal/utils" "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" ) func TestDecodeByteStreamSplitWidth4(t *testing.T) { @@ -393,3 +398,93 @@ func BenchmarkDecodeByteStreamSplitBatchFLBAWidth8(b *testing.B) { }) } } + +// TestByteStreamSplitFLBADecodeSpacedReusedBuffer guards the aliasing bug where decoding +// a second page into a buffer previously expanded by DecodeSpaced silently corrupted +// values: spacedExpand moves slice headers with copy, leaving duplicate headers behind in +// the null slots, and because this decoder writes through the caller's slices rather than +// replacing them, two output slots shared one backing array and clobbered each other. +func TestByteStreamSplitFLBADecodeSpacedReusedBuffer(t *testing.T) { + for _, width := range []int{2, 3, 4, 7, 8, 16} { + t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) { + // 5 slots, 2 nulls: slots 0, 2 and 4 are valid. + validBits := []byte{0b00010101} + const nullCount = 2 + + col := schema.NewColumn(schema.NewFixedLenByteArrayNode("v", parquet.Repetitions.Optional, int32(width), -1), 1, 0) + dec := NewDecoder(parquet.Types.FixedLenByteArray, parquet.Encodings.ByteStreamSplit, + col, memory.DefaultAllocator).(FixedLenByteArrayDecoder) + + // A single output buffer reused across both pages, as the record reader does. + out := make([]parquet.FixedLenByteArray, 5) + + for page, offset := range []byte{0, 100} { + values := make([]parquet.FixedLenByteArray, 3) + for i := range values { + values[i] = make(parquet.FixedLenByteArray, width) + for j := range values[i] { + values[i][j] = offset + byte(i*width+j) + } + } + + data := make([]byte, len(values)*width) + for vi, v := range values { + for bi, b := range v { + data[bi*len(values)+vi] = b + } + } + + require.NoError(t, dec.SetData(len(values), data)) + n, err := dec.DecodeSpaced(out, nullCount, validBits, 0) + require.NoError(t, err) + require.Equal(t, len(out), n) + + require.Equal(t, values[0], out[0], "page %d slot 0", page) + require.Equal(t, values[1], out[2], "page %d slot 2", page) + require.Equal(t, values[2], out[4], "page %d slot 4", page) + } + }) + } +} + +// TestSpacedExpandSwapMatchesSpacedExpand checks that swapping places values in exactly +// the same slots as copying, and additionally never leaves duplicate entries behind. +func TestSpacedExpandSwapMatchesSpacedExpand(t *testing.T) { + rng := rand.New(rand.NewSource(42)) + for iter := 0; iter < 5000; iter++ { + n := 1 + rng.Intn(200) + validBits := make([]byte, bitutil.BytesForBits(int64(n))) + nullCount, density := 0, rng.Float64() + for i := 0; i < n; i++ { + if rng.Float64() < density { + bitutil.ClearBit(validBits, i) + nullCount++ + } else { + bitutil.SetBit(validBits, i) + } + } + + // distinct sentinels so duplicates are detectable + copied, swapped := make([]int64, n), make([]int64, n) + for i := range copied { + copied[i], swapped[i] = int64(i+1), int64(i+1) + } + + spacedExpand(copied, nullCount, validBits, 0) + spacedExpandSwap(swapped, nullCount, validBits, 0) + + for i := 0; i < n; i++ { + if bitutil.BitIsSet(validBits, i) { + require.Equalf(t, copied[i], swapped[i], + "iter %d n=%d nulls=%d: valid slot %d differs", iter, n, nullCount, i) + } + } + + seen := make(map[int64]struct{}, n) + for _, v := range swapped { + seen[v] = struct{}{} + } + require.Lenf(t, seen, n, + "iter %d n=%d nulls=%d: swap left duplicate entries", iter, n, nullCount) + } +} diff --git a/parquet/internal/encoding/decoder.go b/parquet/internal/encoding/decoder.go index 3955bf312..7445a00e1 100644 --- a/parquet/internal/encoding/decoder.go +++ b/parquet/internal/encoding/decoder.go @@ -218,6 +218,46 @@ func (d *dictDecoder[T]) DecodeIndicesSpaced(numValues, nullCount int, validBits return n, nil } +// spacedExpandSwap is spacedExpand for decoders that write *through* caller-provided +// storage instead of replacing it. spacedExpand moves values with copy and leaves the +// null slots alone, which for the slice-header column types (ByteArray / +// FixedLenByteArray) leaves duplicate headers behind: a null slot and a valid slot end +// up referencing the same backing array. That is harmless for a decoder that replaces +// the header, but a decoder that writes through it would have two output slots share +// one array and clobber each other once the buffer is reused for a later page. +// +// Swapping instead of copying keeps the buffer a permutation of its original elements, +// so no slot aliases another and every slot keeps its reusable capacity. +func spacedExpandSwap[T parquet.ColumnTypes](buffer []T, nullCount int, validBits []byte, validBitsOffset int64) int { + numValues := len(buffer) + + idxDecode := int64(numValues - nullCount) + if idxDecode == 0 { + return numValues + } + + rdr := bitutils.NewReverseSetBitRunReader(validBits, validBitsOffset, int64(numValues)) + for { + run := rdr.NextRun() + if run.Length == 0 { + break + } + + idxDecode -= run.Length + // Once the decoded prefix is already aligned every remaining swap is a + // self-swap, so there is nothing left to do. Mirrors spacedExpand. + if idxDecode == run.Pos { + return numValues + } + for k := run.Length - 1; k >= 0; k-- { + dst, src := run.Pos+k, idxDecode+k + buffer[dst], buffer[src] = buffer[src], buffer[dst] + } + } + + return numValues +} + // spacedExpand is used to take a slice of data and utilize the bitmap provided to fill in nulls into the // correct slots according to the bitmap in order to produce a fully expanded result slice with nulls // in the correct slots. diff --git a/parquet/internal/encoding/fixed_len_byte_array_decoder.go b/parquet/internal/encoding/fixed_len_byte_array_decoder.go index 080e9a068..4ba5cd195 100644 --- a/parquet/internal/encoding/fixed_len_byte_array_decoder.go +++ b/parquet/internal/encoding/fixed_len_byte_array_decoder.go @@ -229,5 +229,7 @@ func (dec *ByteStreamSplitFixedLenByteArrayDecoder) DecodeSpaced(out []parquet.F return valuesRead, errors.New("parquet: number of values / definitions levels read did not match") } - return spacedExpand(out, nullCount, validBits, validBitsOffset), nil + // This decoder writes through the caller's slices, so it must not leave aliased + // headers behind for the next page; see spacedExpandSwap. + return spacedExpandSwap(out, nullCount, validBits, validBitsOffset), nil } diff --git a/parquet/pqarrow/byte_stream_split_multipage_test.go b/parquet/pqarrow/byte_stream_split_multipage_test.go new file mode 100644 index 000000000..b1886ea76 --- /dev/null +++ b/parquet/pqarrow/byte_stream_split_multipage_test.go @@ -0,0 +1,113 @@ +// 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_test + +import ( + "bytes" + "context" + "fmt" + "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/file" + "github.com/apache/arrow-go/v18/parquet/pqarrow" + "github.com/stretchr/testify/require" +) + +// TestByteStreamSplitFLBANullsMultiPage covers a BYTE_STREAM_SPLIT FIXED_LEN_BYTE_ARRAY +// column whose chunk spans several data pages and contains nulls. The record reader +// reuses one value buffer across pages, and DecodeSpaced previously left aliased slice +// headers in it, so values decoded from the second page onwards came back shifted. +func TestByteStreamSplitFLBANullsMultiPage(t *testing.T) { + for _, width := range []int{4, 17} { + t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) { + const nrows = 5000 + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + dt := &arrow.FixedSizeBinaryType{ByteWidth: width} + sc := arrow.NewSchema([]arrow.Field{{Name: "v", Type: dt, Nullable: true}}, nil) + + bldr := array.NewFixedSizeBinaryBuilder(mem, dt) + defer bldr.Release() + + expected := make([][]byte, nrows) + for i := range expected { + if i%7 == 3 { + bldr.AppendNull() + continue + } + v := make([]byte, width) + for j := range v { + v[j] = byte(i*width + j) + } + bldr.Append(v) + expected[i] = v + } + + arr := bldr.NewArray() + defer arr.Release() + rec := array.NewRecordBatch(sc, []arrow.Array{arr}, nrows) + defer rec.Release() + + var buf bytes.Buffer + props := parquet.NewWriterProperties( + parquet.WithAllocator(mem), + parquet.WithEncoding(parquet.Encodings.ByteStreamSplit), + parquet.WithDictionaryDefault(false), + // small pages so the column chunk spans more than one data page + parquet.WithDataPageSize(512), + parquet.WithBatchSize(128), + ) + w, err := pqarrow.NewFileWriter(sc, &buf, props, pqarrow.DefaultWriterProps()) + require.NoError(t, err) + require.NoError(t, w.Write(rec)) + require.NoError(t, w.Close()) + + rdr, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()), + file.WithReadProps(parquet.NewReaderProperties(mem))) + require.NoError(t, err) + defer rdr.Close() + + fr, err := pqarrow.NewFileReader(rdr, pqarrow.ArrowReadProperties{BatchSize: 137}, mem) + require.NoError(t, err) + tbl, err := fr.ReadTable(context.Background()) + require.NoError(t, err) + defer tbl.Release() + + require.EqualValues(t, nrows, tbl.NumRows()) + + row := 0 + for _, chunk := range tbl.Column(0).Data().Chunks() { + fsb := chunk.(*array.FixedSizeBinary) + for i := 0; i < fsb.Len(); i++ { + if expected[row] == nil { + require.Truef(t, fsb.IsNull(i), "row %d should be null", row) + } else { + require.Falsef(t, fsb.IsNull(i), "row %d should be valid", row) + require.Equalf(t, expected[row], fsb.Value(i), "row %d", row) + } + row++ + } + } + require.Equal(t, nrows, row) + }) + } +}