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 bb2390b5 perf(parquet): reuse DataPageV2 assembly buffer for eager
writes (#1192)
bb2390b5 is described below
commit bb2390b5ad7eadc71236f39579d80514689132af
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 18:59:46 2026 +0200
perf(parquet): reuse DataPageV2 assembly buffer for eager writes (#1192)
### Rationale for this change
`buildDataPageV2` assembles the definition/repetition levels and encoded
values into a new `bytes.Buffer` for every page. When dictionary
encoding is disabled or has fallen back, the page is written
synchronously and the assembly storage can be reused. The current
allocation pattern adds page-sized churn.
### What changes are included in this PR?
- Reuse `columnWriter.uncompressedData` for eagerly written DataPageV2
pages.
- Keep an owned buffer for pages retained while dictionary encoding is
active.
- Add a benchmark for 64 eager 64 KiB pages with uncompressed and Snappy
codecs.
- Exercise the eager DataPageV2 path in the existing multi-page
round-trip test.
Medians from 6 runs on an Apple M1 Pro were:
| codec | metric | main | this PR | change |
|---|---|---:|---:|---:|
| uncompressed | time/op | 1.08 ms | 650 us | -39.6% |
| uncompressed | B/op | 4,547,602 | 383,763 | -91.6% |
| uncompressed | allocs/op | 461 | 321 | -30.4% |
| Snappy | time/op | 9.05 ms | 8.88 ms | -1.9% |
| Snappy | B/op | 3,027,460 | 489,403 | -83.8% |
| Snappy | allocs/op | 459 | 325 | -29.2% |
### Are these changes tested?
- `go test ./parquet/... -count=1`
- `go test ./parquet/file -run
"^TestBufferedMultiPageDisabledDictionary$" -count=1`
- `go test ./parquet/file -run "^$" -bench
"^BenchmarkWriteDataPageV2Eager$" -benchmem -count=6`
The multi-page DataPageV2 test covers the eager path, while the existing
dictionary coverage keeps the retained-page path covered.
### Are there any user-facing changes?
No. The Parquet output and public API are unchanged.
---
parquet/file/column_writer.go | 25 ++++++--
parquet/file/file_writer_test.go | 103 +++++++++++++++++++-------------
parquet/file/writer_performance_test.go | 54 +++++++++++++++++
3 files changed, 133 insertions(+), 49 deletions(-)
diff --git a/parquet/file/column_writer.go b/parquet/file/column_writer.go
index 5da1663b..0baaa582 100644
--- a/parquet/file/column_writer.go
+++ b/parquet/file/column_writer.go
@@ -149,6 +149,7 @@ type columnWriter struct {
defLevelSink *encoding.PooledBufferWriter
repLevelSink *encoding.PooledBufferWriter
+ // Scratch for eagerly assembled pages. It retains the largest page
until Close.
uncompressedData bytes.Buffer
compressedTemp *bytes.Buffer
@@ -398,9 +399,20 @@ func (w *columnWriter) buildDataPageV2(defLevelsRLESize,
repLevelsRLESize, uncom
}
// concatenate uncompressed levels and the possibly compressed values
- var combined bytes.Buffer
- combined.Grow(int(int64(defLevelsRLESize) + int64(repLevelsRLESize) +
int64(len(data))))
- w.concatBuffers(defLevelsRLESize, repLevelsRLESize, data, &combined)
+ bufferedPage := w.hasDict && !w.fallbackToNonDict
+ combinedSize := int(int64(defLevelsRLESize) + int64(repLevelsRLESize) +
int64(len(data)))
+ var combined []byte
+ if bufferedPage {
+ var owned bytes.Buffer
+ owned.Grow(combinedSize)
+ w.concatBuffers(defLevelsRLESize, repLevelsRLESize, data,
&owned)
+ combined = owned.Bytes()
+ } else {
+ w.uncompressedData.Reset()
+ w.uncompressedData.Grow(combinedSize)
+ w.concatBuffers(defLevelsRLESize, repLevelsRLESize, data,
&w.uncompressedData)
+ combined = w.uncompressedData.Bytes()
+ }
pageStats, err := w.getPageStatistics()
if err != nil {
@@ -417,7 +429,7 @@ func (w *columnWriter) buildDataPageV2(defLevelsRLESize,
repLevelsRLESize, uncom
repLevelsByteLen := int32(repLevelsRLESize)
firstRowIndex := int64(w.rowsWritten)
- page :=
NewDataPageV2WithConfig(memory.NewBufferBytes(combined.Bytes()), nullCount,
numRows, defLevelsByteLen, repLevelsByteLen,
+ page := NewDataPageV2WithConfig(memory.NewBufferBytes(combined),
nullCount, numRows, defLevelsByteLen, repLevelsByteLen,
w.pager.HasCompressor(), DataPageConfig{
Num: numValues,
Encoding: w.encoding,
@@ -425,11 +437,11 @@ func (w *columnWriter) buildDataPageV2(defLevelsRLESize,
repLevelsRLESize, uncom
Stats: pageStats,
FirstRowIndex: firstRowIndex,
})
- if w.hasDict && !w.fallbackToNonDict {
+ if bufferedPage {
w.totalCompressedBytes += int64(page.buf.Len()) // + sizeof
pageheader
w.pages = append(w.pages, page)
} else {
- w.totalCompressedBytes += int64(combined.Len())
+ w.totalCompressedBytes += int64(len(combined))
defer page.Release()
return w.WriteDataPage(page)
}
@@ -686,6 +698,7 @@ func (w *columnWriter) resetPageStatistics() {
func (w *columnWriter) Close() (err error) {
if !w.closed {
w.closed = true
+ defer func() { w.uncompressedData = bytes.Buffer{} }()
if w.hasDict && !w.fallbackToNonDict {
if err = w.WriteDictionaryPage(); err != nil {
return err
diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go
index 2a28abde..87080779 100644
--- a/parquet/file/file_writer_test.go
+++ b/parquet/file/file_writer_test.go
@@ -310,49 +310,66 @@ func TestBufferedMultiPageDisabledDictionary(t
*testing.T) {
valueCount = 10000
pageSize = 16384
)
- var (
- sink = encoding.NewBufferWriter(0, memory.DefaultAllocator)
- props =
parquet.NewWriterProperties(parquet.WithDictionaryDefault(false),
parquet.WithDataPageSize(pageSize))
- sc, _ = schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
- schema.NewInt32Node("col",
parquet.Repetitions.Required, -1),
- }, -1)
- )
-
- writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props))
- rgWriter := writer.AppendBufferedRowGroup()
- cwr, _ := rgWriter.Column(0)
- cw := cwr.(*file.Int32ColumnChunkWriter)
- valuesIn := make([]int32, 0, valueCount)
- for i := int32(0); i < valueCount; i++ {
- valuesIn = append(valuesIn, (i%100)+1)
- }
- cw.WriteBatch(valuesIn, nil, nil)
- rgWriter.Close()
- writer.Close()
- buffer := sink.Finish()
- defer buffer.Release()
-
- reader, err := file.NewParquetReader(bytes.NewReader(buffer.Bytes()))
- assert.NoError(t, err)
-
- assert.EqualValues(t, 1, reader.NumRowGroups())
- valuesOut := make([]int32, valueCount)
-
- for r := 0; r < reader.NumRowGroups(); r++ {
- rgr := reader.RowGroup(r)
- assert.EqualValues(t, 1, rgr.NumColumns())
- assert.EqualValues(t, valueCount, rgr.NumRows())
-
- var totalRead int64
- col, err := rgr.Column(0)
- assert.NoError(t, err)
- colReader := col.(*file.Int32ColumnChunkReader)
- for colReader.HasNext() {
- total, _, _ :=
colReader.ReadBatch(valueCount-totalRead, valuesOut[totalRead:], nil, nil)
- totalRead += total
- }
- assert.EqualValues(t, valueCount, totalRead)
- assert.Equal(t, valuesIn, valuesOut)
+ for _, pageVersion := range []struct {
+ name string
+ version parquet.DataPageVersion
+ }{
+ {name: "v1", version: parquet.DataPageV1},
+ {name: "v2", version: parquet.DataPageV2},
+ } {
+ t.Run(pageVersion.name, func(t *testing.T) {
+ sink := encoding.NewBufferWriter(0,
memory.DefaultAllocator)
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(false),
+
parquet.WithDataPageVersion(pageVersion.version),
+ parquet.WithDataPageSize(pageSize),
+ )
+ sc, _ := schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.NewInt32Node("col",
parquet.Repetitions.Optional, -1),
+ }, -1)
+
+ writer := file.NewParquetWriter(sink, sc,
file.WithWriterProps(props))
+ rgWriter := writer.AppendBufferedRowGroup()
+ cwr, _ := rgWriter.Column(0)
+ cw := cwr.(*file.Int32ColumnChunkWriter)
+ valuesIn := make([]int32, 0, valueCount)
+ defLevels := make([]int16, 0, valueCount)
+ for i := int32(0); i < valueCount; i++ {
+ valuesIn = append(valuesIn, (i%100)+1)
+ defLevels = append(defLevels, 1)
+ }
+ _, err := cw.WriteBatch(valuesIn, defLevels, nil)
+ assert.NoError(t, err)
+ assert.NoError(t, rgWriter.Close())
+ assert.NoError(t, writer.Close())
+ buffer := sink.Finish()
+ defer buffer.Release()
+
+ reader, err :=
file.NewParquetReader(bytes.NewReader(buffer.Bytes()))
+ assert.NoError(t, err)
+ defer reader.Close()
+
+ assert.EqualValues(t, 1, reader.NumRowGroups())
+ valuesOut := make([]int32, valueCount)
+
+ for r := 0; r < reader.NumRowGroups(); r++ {
+ rgr := reader.RowGroup(r)
+ assert.EqualValues(t, 1, rgr.NumColumns())
+ assert.EqualValues(t, valueCount, rgr.NumRows())
+
+ var totalRead int64
+ col, err := rgr.Column(0)
+ assert.NoError(t, err)
+ colReader := col.(*file.Int32ColumnChunkReader)
+ for colReader.HasNext() {
+ total, _, err :=
colReader.ReadBatch(valueCount-totalRead, valuesOut[totalRead:], nil, nil)
+ assert.NoError(t, err)
+ totalRead += total
+ }
+ assert.EqualValues(t, valueCount, totalRead)
+ assert.Equal(t, valuesIn, valuesOut)
+ }
+ })
}
}
diff --git a/parquet/file/writer_performance_test.go
b/parquet/file/writer_performance_test.go
index 0d823769..3f379688 100644
--- a/parquet/file/writer_performance_test.go
+++ b/parquet/file/writer_performance_test.go
@@ -19,10 +19,12 @@ package file_test
import (
"bytes"
"fmt"
+ "io"
"testing"
"github.com/apache/arrow-go/v18/arrow"
"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/schema"
)
@@ -105,6 +107,58 @@ func BenchmarkWriteDictionaryBloomFilter(b *testing.B) {
}
}
+func BenchmarkWriteDataPageV2Eager(b *testing.B) {
+ sc := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.Must(schema.NewPrimitiveNode("data",
parquet.Repetitions.Required, parquet.Types.Int64, -1, -1)),
+ }, -1)))
+
+ const (
+ pageSize = 64 * 1024
+ numPages = 64
+ )
+ values := make([]int64, pageSize/arrow.Int64SizeBytes*numPages)
+ for i := range values {
+ values[i] = int64(i)
+ }
+
+ for _, codec := range
[]compress.Compression{compress.Codecs.Uncompressed, compress.Codecs.Snappy} {
+ b.Run(codec.String(), func(b *testing.B) {
+ props := parquet.NewWriterProperties(
+ parquet.WithStats(false),
+ parquet.WithDataPageVersion(parquet.DataPageV2),
+ parquet.WithDictionaryDefault(false),
+ parquet.WithCompression(codec),
+
parquet.WithBatchSize(int64(pageSize/arrow.Int64SizeBytes)),
+ parquet.WithDataPageSize(pageSize),
+ )
+
+ b.SetBytes(int64(len(values) * arrow.Int64SizeBytes))
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ writer := file.NewParquetWriter(io.Discard,
sc.Root(), file.WithWriterProps(props))
+ rgw := writer.AppendRowGroup()
+ colWriter, err := rgw.NextColumn()
+ if err != nil {
+ b.Fatal(err)
+ }
+ if _, err :=
colWriter.(*file.Int64ColumnChunkWriter).WriteBatch(values, nil, nil); err !=
nil {
+ b.Fatal(err)
+ }
+ if err := colWriter.Close(); err != nil {
+ b.Fatal(err)
+ }
+ if err := rgw.Close(); err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.Close(); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+}
+
// Benchmark writing small ByteArray values (typical case)
// This tests the common scenario where values are small (< 1KB)
func BenchmarkWriteSmallByteArrayValues(b *testing.B) {
