diff --git a/parquet/file/column_writer.go b/parquet/file/column_writer.go index 0cb4071e..bc8cd701 100644 --- a/parquet/file/column_writer.go +++ b/parquet/file/column_writer.go @@ -657,7 +657,12 @@ func (w *columnWriter) Close() (err error) { func (w *columnWriter) doBatches(total int64, repLevels []int16, action func(offset, batch int64)) { batchSize := w.props.WriteBatchSize() // if we're writing V1 data pages, have no replevels or the max replevel is 0 then just - // use the regular doBatches function + // use the regular doBatches function. + // + // The spec also requires row-aligned pages for V1 when an OffsetIndex is + // present (PageIndexEnabled). That gap is on the WriteBatch path and is + // out of scope for the DataPageV2 row alignment here; it is tracked in + // https://github.com/apache/arrow-go/issues/887. if w.props.DataPageVersion() == parquet.DataPageV1 || repLevels == nil || w.descr.MaxRepetitionLevel() == 0 { doBatches(total, batchSize, action) return @@ -675,6 +680,17 @@ func (w *columnWriter) doBatches(total int64, repLevels []int16, action func(off panic("columnwriter: not enough repetition levels for batch to write") } + // nothing to write: a non-nil but empty repLevels slice would index + // repLevels[0] below and panic (recovered as an opaque error). + if total == 0 { + return + } + + // drive the loop over exactly the number of levels we were asked to write. + // a caller passing a repLevels slice longer than total must not spill the + // extra levels into the column. + repLevels = repLevels[:total] + if repLevels[0] != 0 { panic("columnwriter: batch writing for V2 data pages must start at a row boundary") } @@ -684,16 +700,15 @@ func (w *columnWriter) doBatches(total int64, repLevels []int16, action func(off batchStart, batch int64 ) for batchStart = 0; batchStart+batchSize < int64(len(repLevels)); batchStart += batch { - // check one past the last value of the batch for if it's a new row - // if it's not, shrink the batch and feel back to the beginning of a - // previous row boundary to end on - batch = batchSize - for ; repLevels[batchStart+batch] != 0; batch-- { - } // batchStart <--> batch now begins and ends on a row boundary! + batch = alignBatchToRowBoundary(repLevels, batchStart, batchSize) action(batchStart, batch) } - action(batchStart, int64(len(repLevels))-batchStart) + // a grow-to-end alignment on an oversized final row can leave batchStart + // already at len(repLevels); skip the trailing zero-length action. + if batchStart < int64(len(repLevels)) { + action(batchStart, int64(len(repLevels))-batchStart) + } } func doBatches(total, batchSize int64, action func(offset, batch int64)) { @@ -706,6 +721,30 @@ func doBatches(total, batchSize int64, action func(offset, batch int64)) { } } +// alignBatchToRowBoundary adjusts batch so that repLevels[offset+batch] lands on +// a row boundary (repetition level 0) or the end of the level slice. A repeated +// row must never span a DataPageV2 page boundary, so it first shrinks toward the +// previous boundary. If there is no boundary at or before the requested split - +// the current row is wider than batch - it grows forward to the next one so the +// whole row stays in a single batch and the caller keeps making progress rather +// than looping on a zero-length batch. offset must already sit on a boundary. +func alignBatchToRowBoundary(repLevels []int16, offset, batch int64) int64 { + n := int64(len(repLevels)) + if offset+batch >= n { + return batch + } + b := batch + for b > 0 && repLevels[offset+b] != 0 { + b-- + } + if b > 0 { + return b + } + for b = batch; offset+b < n && repLevels[offset+b] != 0; b++ { + } + return b +} + func levelSliceOrNil(rep []int16, offset, batch int64) []int16 { if rep == nil { return nil diff --git a/parquet/file/column_writer_align_test.go b/parquet/file/column_writer_align_test.go new file mode 100644 index 00000000..dd772bbe --- /dev/null +++ b/parquet/file/column_writer_align_test.go @@ -0,0 +1,160 @@ +// 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 file + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestAlignBatchToRowBoundary exercises the row-boundary alignment used by the +// V2 writer batching. A repetition level of 0 marks the first level of a new +// row; every other level continues the current row. The helper must return a +// batch whose end (offset+batch) sits on such a boundary or on the end of the +// slice, shrinking toward the previous boundary and only growing forward when a +// single row is wider than the requested batch. +func TestAlignBatchToRowBoundary(t *testing.T) { + tests := []struct { + name string + repLevels []int16 + offset int64 + batch int64 + want int64 + }{ + { + // offset+batch already lands on a row start, so it is returned as-is. + name: "already on boundary", + repLevels: []int16{0, 1, 0, 1, 0, 1}, + offset: 0, + batch: 2, + want: 2, + }, + { + // rows are {0}, {1,2,3}, {4}; batch 3 ends mid-row at index 3, the + // nearest earlier boundary is index 1. + name: "shrink to previous boundary", + repLevels: []int16{0, 0, 1, 1, 0}, + offset: 0, + batch: 3, + want: 1, + }, + { + // the first row {0,1,2,3} is wider than batch 2 and has no earlier + // boundary to back off to, so the whole row is kept in one batch. + name: "grow forward when row wider than batch", + repLevels: []int16{0, 1, 1, 1, 0}, + offset: 0, + batch: 2, + want: 4, + }, + { + // growing forward runs to the end of the slice when the wide row is + // the last one. + name: "grow forward to end of slice", + repLevels: []int16{0, 1, 1, 1}, + offset: 0, + batch: 2, + want: 4, + }, + { + // offset+batch == len: at the tail there is nothing past the end to + // split, so the batch is returned untouched. + name: "batch reaches end untouched", + repLevels: []int16{0, 1, 0}, + offset: 0, + batch: 3, + want: 3, + }, + { + // offset+batch > len (a batch larger than the remaining data) is + // likewise returned untouched; callers only reach this at the tail. + name: "large batch past end untouched", + repLevels: []int16{0, 1, 0}, + offset: 0, + batch: 100, + want: 100, + }, + { + // a degenerate zero batch stays zero; callers clamp it up to 1. + name: "zero batch stays zero", + repLevels: []int16{0, 1, 0}, + offset: 0, + batch: 0, + want: 0, + }, + { + // batch 1 that ends on a boundary is unchanged. + name: "one batch on boundary", + repLevels: []int16{0, 0, 0}, + offset: 0, + batch: 1, + want: 1, + }, + { + // batch 1 that ends mid-row grows to keep the row {0,1} whole. + name: "one batch grows through a row", + repLevels: []int16{0, 1, 0}, + offset: 0, + batch: 1, + want: 2, + }, + { + // a non-zero offset that already sits on a boundary: the row at + // index 2 is {2,3,4}, wider than batch 2, so it grows to 3. + name: "non-zero offset grows current row", + repLevels: []int16{0, 1, 0, 1, 1, 0}, + offset: 2, + batch: 2, + want: 3, + }, + { + // a large batch spanning several whole rows that happens to land on + // a boundary is returned as-is. + name: "large batch on boundary spanning rows", + repLevels: []int16{0, 1, 0, 1, 0, 1}, + offset: 0, + batch: 4, + want: 4, + }, + { + // a large batch that lands mid-row shrinks back across the rows it + // oversteps to the last boundary at or before the split (index 2). + name: "large batch shrinks across rows", + repLevels: []int16{0, 1, 0, 1, 1, 1, 0}, + offset: 0, + batch: 5, + want: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := alignBatchToRowBoundary(tt.repLevels, tt.offset, tt.batch) + assert.Equal(t, tt.want, got) + + // Post-conditions that must hold for every non-tail result: the + // batch ends on a row boundary or the end of the slice, and it never + // splits inside the current row (offset+batch is never a rep>0 that + // has a rep==0 immediately reachable behind it within the batch). + if tt.offset+got < int64(len(tt.repLevels)) && got > 0 { + assert.Zero(t, tt.repLevels[tt.offset+got], + "aligned batch must end on a row boundary") + } + }) + } +} diff --git a/parquet/file/column_writer_types.gen.go b/parquet/file/column_writer_types.gen.go index 3d6b3a26..9a39615b 100644 --- a/parquet/file/column_writer_types.gen.go +++ b/parquet/file/column_writer_types.gen.go @@ -139,7 +139,7 @@ func (w *Int32ColumnChunkWriter) WriteBatchSpacedWithError(values []int32, defLe if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []int32 info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -178,7 +178,7 @@ func (w *Int32ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLevels dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -401,7 +401,7 @@ func (w *Int64ColumnChunkWriter) WriteBatchSpacedWithError(values []int64, defLe if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []int64 info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -440,7 +440,7 @@ func (w *Int64ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLevels dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -663,7 +663,7 @@ func (w *Int96ColumnChunkWriter) WriteBatchSpacedWithError(values []parquet.Int9 if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []parquet.Int96 info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -702,7 +702,7 @@ func (w *Int96ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLevels dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -925,7 +925,7 @@ func (w *Float32ColumnChunkWriter) WriteBatchSpacedWithError(values []float32, d if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []float32 info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -964,7 +964,7 @@ func (w *Float32ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLeve dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -1187,7 +1187,7 @@ func (w *Float64ColumnChunkWriter) WriteBatchSpacedWithError(values []float64, d if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []float64 info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -1226,7 +1226,7 @@ func (w *Float64ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLeve dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -1452,7 +1452,7 @@ func (w *BooleanColumnChunkWriter) WriteBatchSpacedWithError(values []bool, defL if defLevels == nil { length = len(values) } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []bool info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -1535,7 +1535,7 @@ func (w *BooleanColumnChunkWriter) WriteBitmapBatchSpacedWithError(bitmap []byte length = numValues } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -1571,7 +1571,7 @@ func (w *BooleanColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLeve dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -1833,11 +1833,11 @@ func (w *ByteArrayColumnChunkWriter) WriteBatch(values []parquet.ByteArray, defL } } - // V2 row-boundary alignment - if isV2WithRep && levelOffset+batch < n { - for batch > 1 && repLevels[levelOffset+batch] != 0 { - batch-- - } + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) } if batch < 1 { batch = 1 @@ -1903,6 +1903,8 @@ func (w *ByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values []parquet. const maxSafeBatchDataSize int64 = 1 << 30 // 1GB batchSize := w.props.WriteBatchSize() + isV2WithRep := w.props.DataPageVersion() != parquet.DataPageV1 && + repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) n := int64(length) @@ -1921,6 +1923,13 @@ func (w *ByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values []parquet. cumDataSize += valSize } } + + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) + } if batch < 1 { batch = 1 } @@ -1964,7 +1973,7 @@ func (w *ByteArrayColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLe dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -2163,11 +2172,11 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatch(values []parquet.FixedLe } } - // V2 row-boundary alignment - if isV2WithRep && levelOffset+batch < n { - for batch > 1 && repLevels[levelOffset+batch] != 0 { - batch-- - } + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) } if batch < 1 { batch = 1 @@ -2233,6 +2242,8 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values [] const maxSafeBatchDataSize int64 = 1 << 30 // 1GB batchSize := w.props.WriteBatchSize() + isV2WithRep := w.props.DataPageVersion() != parquet.DataPageV1 && + repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) n := int64(length) @@ -2251,6 +2262,13 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchSpacedWithError(values [] cumDataSize += valSize } } + + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) + } if batch < 1 { batch = 1 } @@ -2294,7 +2312,7 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteDictIndices(indices arrow.Arra dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) diff --git a/parquet/file/column_writer_types.gen.go.tmpl b/parquet/file/column_writer_types.gen.go.tmpl index 04799f2b..bf1f2c51 100644 --- a/parquet/file/column_writer_types.gen.go.tmpl +++ b/parquet/file/column_writer_types.gen.go.tmpl @@ -125,11 +125,11 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatch(values []{{.name}}, defLevels, r } } - // V2 row-boundary alignment - if isV2WithRep && levelOffset+batch < n { - for batch > 1 && repLevels[levelOffset+batch] != 0 { - batch-- - } + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) } if batch < 1 { batch = 1 @@ -214,6 +214,8 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatchSpacedWithError(values []{{.name} const maxSafeBatchDataSize int64 = 1 << 30 // 1GB batchSize := w.props.WriteBatchSize() + isV2WithRep := w.props.DataPageVersion() != parquet.DataPageV1 && + repLevels != nil && w.descr.MaxRepetitionLevel() > 0 levelOffset := int64(0) n := int64(length) @@ -236,6 +238,13 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatchSpacedWithError(values []{{.name} cumDataSize += valSize } } + + // V2 row-boundary alignment: a repeated row must not span a page boundary, + // so snap the batch onto the nearest row boundary (or keep a single wide row + // whole). See alignBatchToRowBoundary in column_writer.go. + if isV2WithRep { + batch = alignBatchToRowBoundary(repLevels, levelOffset, batch) + } if batch < 1 { batch = 1 } @@ -262,7 +271,7 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBatchSpacedWithError(values []{{.name} levelOffset += batch } {{- else}} - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { var vals []{{.name}} info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) @@ -347,7 +356,7 @@ func (w *{{.Name}}ColumnChunkWriter) WriteBitmapBatchSpacedWithError(bitmap []by length = numValues } - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) @@ -384,7 +393,7 @@ func (w *{{.Name}}ColumnChunkWriter) WriteDictIndices(indices arrow.Array, defLe dictEncoder := w.currentEncoder.(encoding.DictEncoder) - doBatches(int64(length), w.props.WriteBatchSize(), func(offset, batch int64) { + w.doBatches(int64(length), repLevels, func(offset, batch int64) { info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) diff --git a/parquet/file/column_writer_v2_row_boundary_test.go b/parquet/file/column_writer_v2_row_boundary_test.go new file mode 100644 index 00000000..89dcec7d --- /dev/null +++ b/parquet/file/column_writer_v2_row_boundary_test.go @@ -0,0 +1,522 @@ +// 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 file_test + +import ( + "bytes" + "context" + "fmt" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/bitutil" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/compress" + "github.com/apache/arrow-go/v18/parquet/file" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" + "github.com/apache/arrow-go/v18/parquet/pqarrow" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +// TestWriteBatchSpacedV2RowBoundary is a regression test for #882: when a +// repeated column is written through the spaced path under DataPageV2, batches +// were split by a fixed WriteBatchSize instead of on row boundaries. Because a +// page flush is only ever checked at a batch boundary, a fixed-size batch that +// ended mid-row caused a data page to end (and the next one to begin) in the +// middle of a repeated row. DataPageV2 requires every page to start at a row +// boundary (first repetition level == 0) for offset-index page pruning to be +// correct. +// +// The spaced write path now routes through the rep-level-aware +// columnWriter.doBatches (the same batching WriteBatch already used), so every +// batch — and therefore every page — lands on a row boundary. +func TestWriteBatchSpacedV2RowBoundary(t *testing.T) { + const ( + numRows = 1000 + valsPerRow = 4 + total = numRows * valsPerRow + // A batch size that is not a multiple of valsPerRow guarantees that a + // fixed-size split would land mid-row. + batchSize = 7 + dataPageSize = 1024 + ) + + root, err := schema.NewGroupNode("schema", parquet.Repetitions.Required, schema.FieldList{ + schema.NewInt32Node("v", parquet.Repetitions.Repeated, -1), + }, -1) + require.NoError(t, err) + + values := make([]int32, total) + defLevels := make([]int16, total) + repLevels := make([]int16, total) + for i := 0; i < total; i++ { + values[i] = int32(i) + defLevels[i] = 1 // present (maxDefLevel for a repeated leaf) + if i%valsPerRow == 0 { + repLevels[i] = 0 // start of a new row + } else { + repLevels[i] = 1 // continuation of the current repeated row + } + } + validBits := make([]byte, (total+7)/8) + for i := range validBits { + validBits[i] = 0xFF + } + + props := parquet.NewWriterProperties( + parquet.WithVersion(parquet.V2_LATEST), + parquet.WithDataPageVersion(parquet.DataPageV2), + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(batchSize), + parquet.WithDataPageSize(dataPageSize), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + + var buf bytes.Buffer + w := file.NewParquetWriter(&buf, root, file.WithWriterProps(props)) + rgw := w.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Int32ColumnChunkWriter).WriteBatchSpacedWithError(values, defLevels, repLevels, validBits, 0) + require.NoError(t, err) + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, w.Close()) + + assertPagesStartOnRowBoundary(t, buf.Bytes()) +} + +// assertPagesStartOnRowBoundary reads the first column chunk back and requires +// that every DataPageV2 begins at a row boundary (its first repetition level is +// 0). The test forces multiple pages, so a mid-row split will surface here. +func assertPagesStartOnRowBoundary(t *testing.T, raw []byte) { + t.Helper() + + r, err := file.NewParquetReader(bytes.NewReader(raw)) + require.NoError(t, err) + defer r.Close() + + maxRep := r.MetaData().Schema.Column(0).MaxRepetitionLevel() + require.Greater(t, maxRep, int16(0), "column must be repeated for this test to be meaningful") + + pr, err := r.RowGroup(0).GetColumnPageReader(0) + require.NoError(t, err) + + pageCount, v2Count := 0, 0 + for pr.Next() { + pageCount++ + v2, ok := pr.Page().(*file.DataPageV2) + if !ok { + continue + } + v2Count++ + + var dec encoding.LevelDecoder + require.NoError(t, dec.SetDataV2(v2.RepetitionLevelByteLen(), maxRep, int(v2.NumValues()), v2.Data())) + + levels := make([]int16, v2.NumValues()) + n, _ := dec.Decode(levels) + require.Greater(t, n, 0, "page %d decoded zero repetition levels", v2Count) + require.Zerof(t, levels[0], + "DataPageV2 #%d starts mid-row (first repetition level = %d, want 0)", v2Count, levels[0]) + } + require.NoError(t, pr.Err()) + require.Greater(t, v2Count, 1, "expected more than one DataPageV2 so a mid-row split could occur") +} + +// TestWriteSpacedV2NullableListRowBoundary is a regression test for the second +// half of #882 (surfaced by punkeel while reviewing #883): when a nullable list +// element is written through the spaced path under DataPageV2 and a single +// repeated row is wider than the write-batch size, the row-boundary alignment +// has no earlier boundary to shrink back to. For fixed-width elements (float64) +// the batch shrank to zero and the writer looped without making progress (a +// hang); for byte-array elements (string, binary, fixed_size_binary) the batch +// clamped at one and a DataPageV2 was still emitted starting mid-row. +// +// punkeel's key observation: the trigger requires an actual null element in the +// data, because only then does the Arrow writer route the leaf through +// WriteBatchSpaced instead of WriteBatch. Every row here ends with a null +// element, and each row (width 5) is wider than the batch size (4), so there is +// no row boundary to shrink back to within a batch. +func TestWriteSpacedV2NullableListRowBoundary(t *testing.T) { + cases := []struct { + name string + elemType arrow.DataType + }{ + {"string", arrow.BinaryTypes.String}, + {"binary", arrow.BinaryTypes.Binary}, + {"fixed_size_binary", &arrow.FixedSizeBinaryType{ByteWidth: 1}}, + {"float64", arrow.PrimitiveTypes.Float64}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + const ( + numRows = 20 + rowWidth = 300 // much wider than the batch size below + batch = 5 + // A page size well under one row's worth of data forces + // several page flushes inside every row, so any batch that + // ends mid-row surfaces as a mid-row page. + pageSize = 100 + ) + widths := make([]int, numRows) + for i := range widths { + widths[i] = rowWidth + } + raw := writeNullableListV2WithinTimeout(t, tc.elemType, false, widths, batch, pageSize, 30*time.Second) + assertPagesStartOnRowBoundary(t, raw) + }) + } +} + +// TestWriteSpacedV2NullableListWideRow is primarily a hang guard for punkeel's +// long-list case: a single nullable list whose length (20001) exceeds the +// write-batch size (20000). There is no row boundary before the requested split, +// so the fix must grow the batch forward to keep the whole row in one page rather +// than shrinking to zero and looping without progress. A fixed-width element type +// makes the pre-fix failure deterministic: the batch shrinks to zero and the +// writer hangs (a byte-array element would instead silently clamp to one, which +// the timeout would not catch), so the 30s timeout in +// writeNullableListV2WithinTimeout is the real regression guard here. With only +// one row the page-boundary assertions below cannot themselves detect a mid-row +// split (a single row can never be split) and pass even with the fix reverted; +// they exist only to confirm the wide row still produces a single valid +// DataPageV2 once the write completes. +func TestWriteSpacedV2NullableListWideRow(t *testing.T) { + const ( + listLen = 20001 + batch = 20000 + ) + raw := writeNullableListV2WithinTimeout(t, arrow.PrimitiveTypes.Float64, true, []int{listLen}, batch, 1024*1024, 30*time.Second) + + r, err := file.NewParquetReader(bytes.NewReader(raw)) + require.NoError(t, err) + defer r.Close() + require.EqualValues(t, 1, r.NumRows()) + + pr, err := r.RowGroup(0).GetColumnPageReader(0) + require.NoError(t, err) + maxRep := r.MetaData().Schema.Column(0).MaxRepetitionLevel() + v2Count := 0 + for pr.Next() { + v2, ok := pr.Page().(*file.DataPageV2) + if !ok { + continue + } + v2Count++ + var dec encoding.LevelDecoder + require.NoError(t, dec.SetDataV2(v2.RepetitionLevelByteLen(), maxRep, int(v2.NumValues()), v2.Data())) + levels := make([]int16, v2.NumValues()) + n, _ := dec.Decode(levels) + require.Greater(t, n, 0) + require.Zero(t, levels[0], "wide row page starts mid-row") + } + require.NoError(t, pr.Err()) + require.Equal(t, 1, v2Count, "a single wide row must occupy exactly one DataPageV2") +} + +// writeNullableListV2WithinTimeout builds a list Arrow column +// whose rows have the given widths (the last element of every row is null so +// the leaf writer takes the WriteBatchSpaced path), writes it as DataPageV2 with +// the given batch and page sizes, and returns the parquet bytes. The write runs +// under a timeout so that a regression to the no-progress batching loop (#882) +// fails the test instead of hanging the whole suite. On a hang the write +// goroutine is left spinning: the parquet writer exposes no cancellation hook +// into the batching loop, and once the timeout has failed the test the process +// exits, so the leak is bounded to the failing run. +func writeNullableListV2WithinTimeout(t *testing.T, elemType arrow.DataType, listNullable bool, rowWidths []int, batchSize, pageSize int64, timeout time.Duration) []byte { + t.Helper() + + type result struct { + b []byte + err error + } + done := make(chan result, 1) + go func() { + b, err := writeNullableListV2(elemType, listNullable, rowWidths, batchSize, pageSize) + done <- result{b, err} + }() + + select { + case r := <-done: + require.NoError(t, r.err) + return r.b + case <-time.After(timeout): + t.Fatalf("write did not finish within %s: batching made no progress on a row wider than the batch size (#882)", timeout) + return nil + } +} + +func writeNullableListV2(elemType arrow.DataType, listNullable bool, rowWidths []int, batchSize, pageSize int64) ([]byte, error) { + mem := memory.NewGoAllocator() + + listType := arrow.ListOf(elemType) // element field is nullable by default + field := arrow.Field{Name: "l", Type: listType, Nullable: listNullable} + arrowSchema := arrow.NewSchema([]arrow.Field{field}, nil) + + lb := array.NewBuilder(mem, listType).(*array.ListBuilder) + defer lb.Release() + appendElem := elemAppender(lb.ValueBuilder()) + + valueIdx := 0 + for _, w := range rowWidths { + lb.Append(true) + for k := 0; k < w; k++ { + if k == w-1 { + lb.ValueBuilder().AppendNull() // an actual null element (spaced trigger) + } else { + appendElem(valueIdx) + } + valueIdx++ + } + } + arr := lb.NewArray() + defer arr.Release() + + rec := array.NewRecordBatch(arrowSchema, []arrow.Array{arr}, int64(len(rowWidths))) + defer rec.Release() + + var buf bytes.Buffer + props := parquet.NewWriterProperties( + parquet.WithVersion(parquet.V2_LATEST), + parquet.WithDataPageVersion(parquet.DataPageV2), + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(batchSize), + parquet.WithDataPageSize(pageSize), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + pqw, err := pqarrow.NewFileWriter(arrowSchema, &buf, props, pqarrow.NewArrowWriterProperties()) + if err != nil { + return nil, err + } + if err := pqw.Write(rec); err != nil { + return nil, err + } + if err := pqw.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func elemAppender(vb array.Builder) func(int) { + switch b := vb.(type) { + case *array.StringBuilder: + return func(i int) { b.Append(fmt.Sprintf("v%d", i)) } + case *array.BinaryBuilder: + return func(i int) { b.Append([]byte(fmt.Sprintf("v%d", i))) } + case *array.FixedSizeBinaryBuilder: + return func(i int) { b.Append([]byte{byte(i)}) } + case *array.Float64Builder: + return func(i int) { b.Append(float64(i)) } + default: + panic(fmt.Sprintf("unsupported element builder %T", vb)) + } +} + +// TestWriteBitmapBatchSpacedV2RowBoundary covers the boolean bitmap spaced path, +// which WriteBitmapBatchSpacedWithError now routes through the rep-aware +// columnWriter.doBatches. Like the Int32 case, a repeated boolean column written +// with a batch size that is not a multiple of the row width would, before the +// fix, split a data page mid-row. Reverting alignBatchToRowBoundary to +// `return batch` makes this fail with "starts mid-row". +func TestWriteBitmapBatchSpacedV2RowBoundary(t *testing.T) { + const ( + numRows = 2000 + valsPerRow = 4 + total = numRows * valsPerRow + batchSize = 7 // not a multiple of valsPerRow: a fixed split lands mid-row + // A small page size relative to the batch stride keeps page flushes from + // coincidentally landing on row boundaries; with the fix reverted this + // produces 12 of 16 pages starting mid-row. + pageSize = 64 + ) + + root, err := schema.NewGroupNode("schema", parquet.Repetitions.Required, schema.FieldList{ + schema.NewBooleanNode("v", parquet.Repetitions.Repeated, -1), + }, -1) + require.NoError(t, err) + + bitmap := make([]byte, bitutil.BytesForBits(total)) + validBits := make([]byte, bitutil.BytesForBits(total)) + defLevels := make([]int16, total) + repLevels := make([]int16, total) + for i := 0; i < total; i++ { + defLevels[i] = 1 // present + bitutil.SetBit(validBits, i) + if i%3 == 0 { + bitutil.SetBit(bitmap, i) + } + if i%valsPerRow == 0 { + repLevels[i] = 0 + } else { + repLevels[i] = 1 + } + } + + props := parquet.NewWriterProperties( + parquet.WithVersion(parquet.V2_LATEST), + parquet.WithDataPageVersion(parquet.DataPageV2), + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(batchSize), + parquet.WithDataPageSize(pageSize), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + + var buf bytes.Buffer + w := file.NewParquetWriter(&buf, root, file.WithWriterProps(props)) + rgw := w.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.BooleanColumnChunkWriter).WriteBitmapBatchSpacedWithError(bitmap, 0, total, defLevels, repLevels, validBits, 0) + require.NoError(t, err) + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, w.Close()) + + assertPagesStartOnRowBoundary(t, buf.Bytes()) +} + +// TestWriteDictIndicesV2RowBoundary covers the dictionary path: WriteDictIndices +// is now routed through the rep-aware doBatches too. A repeated int32 column with +// low-cardinality values keeps the dictionary (no fallback to plain), while a +// small data page size and many rows force several data pages so a mid-row split +// is possible. The values also round-trip, guarding the variable batch sizes the +// fix introduces against a value-cursor slip. +func TestWriteDictIndicesV2RowBoundary(t *testing.T) { + const ( + numRows = 4000 + valsPerRow = 3 + // A high dictionary page-size limit keeps the column dictionary-encoded; + // a small data page size forces the indices across many data pages. + batchSize = 5 // not a multiple of valsPerRow + pageSize = 512 + ) + + mem := memory.NewGoAllocator() + listType := arrow.ListOf(arrow.PrimitiveTypes.Int32) + arrowSchema := arrow.NewSchema([]arrow.Field{{Name: "l", Type: listType, Nullable: false}}, nil) + + lb := array.NewBuilder(mem, listType).(*array.ListBuilder) + defer lb.Release() + vb := lb.ValueBuilder().(*array.Int32Builder) + for r := 0; r < numRows; r++ { + lb.Append(true) + for k := 0; k < valsPerRow; k++ { + vb.Append(int32((r + k) % 8)) // only 8 distinct values: dictionary stays + } + } + arr := lb.NewArray() + defer arr.Release() + + rec := array.NewRecordBatch(arrowSchema, []arrow.Array{arr}, int64(numRows)) + defer rec.Release() + + raw := writeRecordV2(t, rec, true, batchSize, pageSize) + assertPagesStartOnRowBoundary(t, raw) + requireRecordRoundTrips(t, arrowSchema, rec, raw) +} + +// TestWriteSpacedV2NestedListRoundTrip locks in a value round-trip for a +// maxRep=2 column (list>) written through the spaced path under +// DataPageV2 with small pages, the deepest-nesting case punkeel and zeroshade +// flagged. Because the fix makes batch sizes variable, this guards against a +// value cursor drifting out of step with the levels. +func TestWriteSpacedV2NestedListRoundTrip(t *testing.T) { + const ( + numRows = 500 + batchSize = 7 + pageSize = 512 + ) + + mem := memory.NewGoAllocator() + innerType := arrow.ListOf(arrow.PrimitiveTypes.Int32) + outerType := arrow.ListOf(innerType) + // Nullable inner element so the leaf takes the WriteBatchSpaced path. + arrowSchema := arrow.NewSchema([]arrow.Field{{Name: "l", Type: outerType, Nullable: true}}, nil) + + ob := array.NewBuilder(mem, outerType).(*array.ListBuilder) + defer ob.Release() + ib := ob.ValueBuilder().(*array.ListBuilder) + vb := ib.ValueBuilder().(*array.Int32Builder) + n := 0 + for r := 0; r < numRows; r++ { + ob.Append(true) + for i := 0; i < 3; i++ { + ib.Append(true) + for k := 0; k < 4; k++ { + if k == 3 { + vb.AppendNull() // an actual null element (spaced trigger) + } else { + vb.Append(int32(n)) + } + n++ + } + } + } + arr := ob.NewArray() + defer arr.Release() + + rec := array.NewRecordBatch(arrowSchema, []arrow.Array{arr}, int64(numRows)) + defer rec.Release() + + raw := writeRecordV2(t, rec, false, batchSize, pageSize) + assertPagesStartOnRowBoundary(t, raw) + requireRecordRoundTrips(t, arrowSchema, rec, raw) +} + +// writeRecordV2 writes a single record batch as DataPageV2 with the given batch +// and page sizes, optionally keeping dictionary encoding on, and returns the +// parquet bytes. +func writeRecordV2(t *testing.T, rec arrow.RecordBatch, dictionary bool, batchSize, pageSize int64) []byte { + t.Helper() + var buf bytes.Buffer + props := parquet.NewWriterProperties( + parquet.WithVersion(parquet.V2_LATEST), + parquet.WithDataPageVersion(parquet.DataPageV2), + parquet.WithDictionaryDefault(dictionary), + parquet.WithBatchSize(batchSize), + parquet.WithDataPageSize(pageSize), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + pqw, err := pqarrow.NewFileWriter(rec.Schema(), &buf, props, pqarrow.NewArrowWriterProperties()) + require.NoError(t, err) + require.NoError(t, pqw.Write(rec)) + require.NoError(t, pqw.Close()) + return buf.Bytes() +} + +// requireRecordRoundTrips reads the parquet bytes back through pqarrow and +// requires the single column to equal the source record's column. +func requireRecordRoundTrips(t *testing.T, arrowSchema *arrow.Schema, rec arrow.RecordBatch, raw []byte) { + t.Helper() + tbl, err := pqarrow.ReadTable(context.Background(), bytes.NewReader(raw), nil, pqarrow.ArrowReadProperties{}, memory.NewGoAllocator()) + require.NoError(t, err) + defer tbl.Release() + + require.EqualValues(t, rec.NumRows(), tbl.NumRows()) + want := rec.Column(0) + + // Concatenate the read-back chunks and equate with the source column. + merged, err := array.Concatenate(tbl.Column(0).Data().Chunks(), memory.NewGoAllocator()) + require.NoError(t, err) + defer merged.Release() + require.Truef(t, array.Equal(want, merged), "round-trip mismatch\n want: %v\n got: %v", want, merged) +}