From 62acd46ecafd266dd6576d435b34a824f36a3c5e Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 9 Jul 2026 16:12:19 -0400 Subject: [PATCH] fix(parquet): clamp repLevels in byte-array DataPageV2 row-boundary alignment The ByteArray and FixedLenByteArray column writers run their own adaptive batching loop in WriteBatch and WriteBatchSpacedWithError. #883 added DataPageV2 row-boundary alignment there by calling alignBatchToRowBoundary with the caller's repLevels slice, but did not clamp it to the number of levels being written (n is derived from defLevels/values, not repLevels). columnWriter.doBatches, used by the numeric and boolean paths, already clamps repLevels to the write length for exactly this reason. When a caller passes repLevels longer than n, the byte-array paths could grow the final batch of a wide repeated row past n, slicing defLevels/values out of range (recovered as an opaque error) or spilling extra levels into the column. Clamp repLevels to n in both byte-array batching paths, mirroring the guards in columnWriter.doBatches (length check, empty-input short-circuit, clamp, row-boundary start check). The change is made in column_writer_types.gen.go.tmpl and regenerated. Adds TestWriteBatchByteArrayV2OversizedRepLevels covering WriteBatch and WriteBatchSpacedWithError for both ByteArray and FixedLenByteArray with an oversized repLevels slice; each case fails without the fix (slice bounds out of range) and passes with it. --- parquet/file/column_writer_types.gen.go | 62 ++++++++ parquet/file/column_writer_types.gen.go.tmpl | 31 ++++ .../column_writer_v2_row_boundary_test.go | 133 ++++++++++++++++++ 3 files changed, 226 insertions(+) diff --git a/parquet/file/column_writer_types.gen.go b/parquet/file/column_writer_types.gen.go index 9a39615b..4ae88bf6 100644 --- a/parquet/file/column_writer_types.gen.go +++ b/parquet/file/column_writer_types.gen.go @@ -1810,6 +1810,23 @@ func (w *ByteArrayColumnChunkWriter) WriteBatch(values []parquet.ByteArray, defL repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) + // Repeated DataPageV2 writes align batches on row boundaries using repLevels + // below. n comes from defLevels/values, not repLevels, so clamp repLevels to n + // to stop an oversized slice from growing a batch past n (spilling extra levels + // or reading out of range). Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) @@ -1908,6 +1925,20 @@ func (w *ByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values []parquet. levelOffset := int64(0) n := int64(length) + // Clamp repLevels to n; see WriteBatch. Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) @@ -2149,6 +2180,23 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatch(values []parquet.FixedLe repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) + // Repeated DataPageV2 writes align batches on row boundaries using repLevels + // below. n comes from defLevels/values, not repLevels, so clamp repLevels to n + // to stop an oversized slice from growing a batch past n (spilling extra levels + // or reading out of range). Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) @@ -2247,6 +2295,20 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values [] levelOffset := int64(0) n := int64(length) + // Clamp repLevels to n; see WriteBatch. Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) diff --git a/parquet/file/column_writer_types.gen.go.tmpl b/parquet/file/column_writer_types.gen.go.tmpl index bf1f2c51..d23bdfa3 100644 --- a/parquet/file/column_writer_types.gen.go.tmpl +++ b/parquet/file/column_writer_types.gen.go.tmpl @@ -98,6 +98,23 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatch(values []{{.name}}, defLevels, r repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) + // Repeated DataPageV2 writes align batches on row boundaries using repLevels + // below. n comes from defLevels/values, not repLevels, so clamp repLevels to n + // to stop an oversized slice from growing a batch past n (spilling extra levels + // or reading out of range). Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) @@ -219,6 +236,20 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatchSpacedWithError(values []{{.name} levelOffset := int64(0) n := int64(length) + // Clamp repLevels to n; see WriteBatch. Mirrors columnWriter.doBatches. + if isV2WithRep { + if int64(len(repLevels)) < n { + panic("columnwriter: not enough repetition levels for batch to write") + } + if n == 0 { + return + } + repLevels = repLevels[:n] + if repLevels[0] != 0 { + panic("columnwriter: batch writing for V2 data pages must start at a row boundary") + } + } + for levelOffset < n { remaining := n - levelOffset batch := min(remaining, batchSize) diff --git a/parquet/file/column_writer_v2_row_boundary_test.go b/parquet/file/column_writer_v2_row_boundary_test.go index 89dcec7d..aa8e4b39 100644 --- a/parquet/file/column_writer_v2_row_boundary_test.go +++ b/parquet/file/column_writer_v2_row_boundary_test.go @@ -520,3 +520,136 @@ func requireRecordRoundTrips(t *testing.T, arrowSchema *arrow.Schema, rec arrow. defer merged.Release() require.Truef(t, array.Equal(want, merged), "round-trip mismatch\n want: %v\n got: %v", want, merged) } + +// TestWriteBatchByteArrayV2OversizedRepLevels is a regression test for the +// byte-array and fixed-len-byte-array counterpart of #883's DataPageV2 +// row-boundary alignment. Those writers run their own adaptive batching loop and +// passed the caller's full repLevels slice to alignBatchToRowBoundary, while the +// write length n is derived from defLevels/values. When repLevels is longer than +// n - which the low-level typed-writer API permits, and which +// columnWriter.doBatches already clamps against for the numeric/boolean paths - +// the alignment of a wide final row could grow the last batch past n, slicing +// defLevels/values out of range (recovered as an opaque error) or spilling extra +// levels into the column. The byte-array paths must clamp repLevels to n too. +func TestWriteBatchByteArrayV2OversizedRepLevels(t *testing.T) { + const batchSize = 3 + + // n = 6 requested levels. repLevels is deliberately longer (8) and its final + // row is wide, so unclamped alignment would grow the last batch past n. + // Clamped to repLevels[:6] = {0,1,1,0,1,1} the data is two rows of three. + defLevels := []int16{1, 1, 1, 1, 1, 1} + repLevels := []int16{0, 1, 1, 0, 1, 1, 1, 0} + validBits := []byte{0xFF} + wantVals := []string{"a", "b", "c", "d", "e", "f"} + const wantRows = 2 + + byteArrayVals := []parquet.ByteArray{[]byte("a"), []byte("b"), []byte("c"), []byte("d"), []byte("e"), []byte("f")} + flbaVals := []parquet.FixedLenByteArray{[]byte("a"), []byte("b"), []byte("c"), []byte("d"), []byte("e"), []byte("f")} + + props := parquet.NewWriterProperties( + parquet.WithVersion(parquet.V2_LATEST), + parquet.WithDataPageVersion(parquet.DataPageV2), + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(batchSize), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + + cases := []struct { + name string + node *schema.PrimitiveNode + write func(cw file.ColumnChunkWriter) (int64, error) + read func(t *testing.T, ccr file.ColumnChunkReader) []string + }{ + { + name: "ByteArray/WriteBatch", + node: schema.NewByteArrayNode("v", parquet.Repetitions.Repeated, -1), + write: func(cw file.ColumnChunkWriter) (int64, error) { + return cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(byteArrayVals, defLevels, repLevels) + }, + read: readByteArrayValues, + }, + { + name: "ByteArray/WriteBatchSpaced", + node: schema.NewByteArrayNode("v", parquet.Repetitions.Repeated, -1), + write: func(cw file.ColumnChunkWriter) (int64, error) { + return cw.(*file.ByteArrayColumnChunkWriter).WriteBatchSpacedWithError(byteArrayVals, defLevels, repLevels, validBits, 0) + }, + read: readByteArrayValues, + }, + { + name: "FixedLenByteArray/WriteBatch", + node: schema.NewFixedLenByteArrayNode("v", parquet.Repetitions.Repeated, 1, -1), + write: func(cw file.ColumnChunkWriter) (int64, error) { + return cw.(*file.FixedLenByteArrayColumnChunkWriter).WriteBatch(flbaVals, defLevels, repLevels) + }, + read: readFixedLenByteArrayValues, + }, + { + name: "FixedLenByteArray/WriteBatchSpaced", + node: schema.NewFixedLenByteArrayNode("v", parquet.Repetitions.Repeated, 1, -1), + write: func(cw file.ColumnChunkWriter) (int64, error) { + return cw.(*file.FixedLenByteArrayColumnChunkWriter).WriteBatchSpacedWithError(flbaVals, defLevels, repLevels, validBits, 0) + }, + read: readFixedLenByteArrayValues, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + root, err := schema.NewGroupNode("schema", parquet.Repetitions.Required, schema.FieldList{tc.node}, -1) + require.NoError(t, err) + + var buf bytes.Buffer + w := file.NewParquetWriter(&buf, root, file.WithWriterProps(props)) + rgw := w.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + + valueOffset, err := tc.write(cw) + require.NoError(t, err, "oversized repLevels must not fail the write") + require.EqualValues(t, len(wantVals), valueOffset, "must write exactly the requested values, no spill past n") + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, w.Close()) + + r, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) + require.NoError(t, err) + defer r.Close() + + require.EqualValues(t, wantRows, r.NumRows()) + cc, err := r.MetaData().RowGroup(0).ColumnChunk(0) + require.NoError(t, err) + require.EqualValues(t, len(wantVals), cc.NumValues(), "column must hold exactly the requested values, no spill") + + ccr, err := r.RowGroup(0).Column(0) + require.NoError(t, err) + require.Equal(t, wantVals, tc.read(t, ccr)) + }) + } +} + +func readByteArrayValues(t *testing.T, ccr file.ColumnChunkReader) []string { + t.Helper() + r := ccr.(*file.ByteArrayColumnChunkReader) + vals := make([]parquet.ByteArray, 16) + _, read, err := r.ReadBatch(int64(len(vals)), vals, nil, nil) + require.NoError(t, err) + out := make([]string, read) + for i := 0; i < read; i++ { + out[i] = string(vals[i]) + } + return out +} + +func readFixedLenByteArrayValues(t *testing.T, ccr file.ColumnChunkReader) []string { + t.Helper() + r := ccr.(*file.FixedLenByteArrayColumnChunkReader) + vals := make([]parquet.FixedLenByteArray, 16) + _, read, err := r.ReadBatch(int64(len(vals)), vals, nil, nil) + require.NoError(t, err) + out := make([]string, read) + for i := 0; i < read; i++ { + out[i] = string(vals[i]) + } + return out +}