diff --git a/parquet/internal/encoding/byte_stream_split_decode_test.go b/parquet/internal/encoding/byte_stream_split_decode_test.go index e6c72bdc..cd714fe9 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 3955bf31..7445a00e 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 080e9a06..4ba5cd19 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 00000000..b1886ea7 --- /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) + }) + } +}