fix(parquet): stop BYTE_STREAM_SPLIT FLBA spaced decode aliasing output

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)
+	}
+}

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.

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
 }

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)
+		})
+	}
+}