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 ef004bc9 fix(parquet): clamp repLevels in byte-array DataPageV2
row-boundary alignment (#919)
ef004bc9 is described below
commit ef004bc97dc61d788479be807a0ac0cd97951642
Author: Matt Topol <[email protected]>
AuthorDate: Fri Jul 10 12:30:00 2026 -0400
fix(parquet): clamp repLevels in byte-array DataPageV2 row-boundary
alignment (#919)
### Rationale for this change
The `ByteArray` and `FixedLenByteArray` column writers run their own
adaptive
batching loop in `WriteBatch` and `WriteBatchSpacedWithError`. #883
added
DataPageV2 row-boundary alignment to those loops by calling
`alignBatchToRowBoundary(repLevels, levelOffset, batch)`, but passed the
caller's
**full** `repLevels` slice. The write length `n` in those loops is
derived from
`defLevels`/`values`, **not** `repLevels`, and nothing requires
`len(repLevels) == n`.
`columnWriter.doBatches` (used by the numeric and boolean paths) already
clamps
`repLevels = repLevels[:total]` for exactly this reason — its comment
notes *"a
caller passing a repLevels slice longer than total must not spill the
extra
levels into the column."* The byte-array paths did not.
When a caller passes a `repLevels` slice longer than `n`,
`alignBatchToRowBoundary` (which grows/inspects using `len(repLevels)`)
can grow
the final batch of a wide repeated row past `n`. The writer then slices
`defLevels`/`values` out of range — recovered as an opaque error — or,
with
oversized backing buffers, spills extra levels/values into the column.
`len(repLevels) == n` is safe (the helper returns early at the tail);
the bug
requires `len(repLevels) > n`, which is reachable through the low-level
typed-writer API for repeated columns under DataPageV2.
### What changes are included in this PR?
- In `column_writer_types.gen.go.tmpl`, both `{{if .isByteArray}}`
blocks
(`WriteBatch` and `WriteBatchSpacedWithError`) now clamp `repLevels` to
`n`
before the batching loop, mirroring the guards already in
`columnWriter.doBatches`: length check, empty-input short-circuit,
`repLevels = repLevels[:n]`, and a `repLevels[0] == 0` row-boundary
start check.
- Regenerated `column_writer_types.gen.go` (4 sites: `ByteArray` and
`FixedLenByteArray` × `WriteBatch`/`WriteBatchSpacedWithError`).
### Are these changes tested?
Yes. `TestWriteBatchByteArrayV2OversizedRepLevels` writes a repeated
`ByteArray`
and `FixedLenByteArray` column under DataPageV2 (dictionary disabled,
small batch
size) through both `WriteBatch` and `WriteBatchSpacedWithError`, passing
a
`repLevels` slice longer than `len(defLevels)` with a wide final row. It
asserts
the write does not error, exactly the requested values are written (no
spill),
and the values round-trip. Each of the four cases fails without the fix
(`slice bounds out of range [:7] with capacity 6`) and passes with it.
`./parquet/file/...` and `./parquet/pqarrow/...` pass
(`PARQUET_TEST_DATA` set).
### Are there any user-facing changes?
No API changes. Repeated `BYTE_ARRAY`/`FIXED_LEN_BYTE_ARRAY` columns
written
under DataPageV2 with a `repLevels` slice longer than the
value/def-level count
no longer risk an out-of-range write error or spilled levels; the
byte-array
paths now match the numeric/boolean paths.
---
parquet/file/column_writer_types.gen.go | 62 ++++++++++
parquet/file/column_writer_types.gen.go.tmpl | 31 +++++
parquet/file/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
+}
