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 05fa1adb perf(parquet/pqarrow): write record ranges directly (#1248)
05fa1adb is described below
commit 05fa1adb6c37a6d48b4c296900fe84a319a3812e
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 17:42:07 2026 +0200
perf(parquet/pqarrow): write record ranges directly (#1248)
## Summary
- **Write row-group ranges directly from a record batch.**
- Reuse one `*arrow.Chunked` wrapper per column when a batch spans row
groups.
- Keep the existing path when a batch fits in one row group.
- Preserve buffered row-group boundaries, including full and zero-row
cases.
- Add nested/null round-trip coverage and benchmarks.
## Benchmark
Local Apple M1 Pro. `-benchtime=500ms -count=3`. 8 int64 columns, 8,192
rows, row-group size 256.
| Method | main | change | allocs/op |
| --- | ---: | ---: | ---: |
| `Write` | ~2.31 ms, 4.95 MB | ~2.11 ms, 4.87 MB | 21,905 -> 20,761 |
| `WriteBuffered` | ~2.54 ms, 6.69 MB | ~2.45 ms, 6.61 MB | 25,407 ->
24,263 |
The wall-time result varies a bit between runs. The allocation reduction
is consistent.
## Tests
- `go test ./parquet/pqarrow -run
'^(TestFileWriter|TestWriteOnClosedFileWriter|TestBufferedRecWrite|TestByteArrayStatisticsStreamingReleaseBetweenBatches)$'`
- `go test -race ./parquet/pqarrow -run
'^(TestFileWriterRangeWritesPreserveData|TestFileWriterZeroRowRecord)$'`
- `go vet -composites=false ./parquet/pqarrow`
- `go test ./... -run '^$'`
- `GOOS=linux GOARCH=386 go test -c -o /dev/null ./parquet/pqarrow`
---
parquet/pqarrow/file_writer.go | 105 +++++++++++++-------
parquet/pqarrow/file_writer_range_bench_test.go | 102 +++++++++++++++++++
parquet/pqarrow/file_writer_test.go | 125 ++++++++++++++++++++++++
3 files changed, 297 insertions(+), 35 deletions(-)
diff --git a/parquet/pqarrow/file_writer.go b/parquet/pqarrow/file_writer.go
index 31967b6f..08bf351d 100644
--- a/parquet/pqarrow/file_writer.go
+++ b/parquet/pqarrow/file_writer.go
@@ -315,7 +315,6 @@ func (fw *FileWriter) WriteBuffered(rec arrow.RecordBatch)
error {
}
var (
- recList []arrow.RecordBatch
maxRows = fw.wr.Properties().MaxRowGroupLength()
curRows int
err error
@@ -329,31 +328,33 @@ func (fw *FileWriter) WriteBuffered(rec
arrow.RecordBatch) error {
return err
}
}
-
if int64(curRows)+rec.NumRows() <= maxRows {
- recList = []arrow.RecordBatch{rec}
- } else {
- recList = []arrow.RecordBatch{rec.NewSlice(0,
maxRows-int64(curRows))}
- defer recList[0].Release()
- for offset := maxRows - int64(curRows); offset < rec.NumRows();
offset += maxRows {
- s := rec.NewSlice(offset, offset+utils.Min(maxRows,
rec.NumRows()-offset))
- defer s.Release()
- recList = append(recList, s)
+ if err := fw.writeRecordBatchColumns(rec); err != nil {
+ fw.Close()
+ return err
}
+ fw.colIdx = 0
+ return nil
}
- for idx, r := range recList {
- if idx > 0 {
- if err := fw.NewBufferedRowGroupChecked(); err != nil {
- return err
- }
+ columns := newChunkedRecordColumns(rec)
+ defer releaseChunkedRecordColumns(columns)
+
+ firstSize := maxRows - int64(curRows)
+
+ for offset, size := int64(0), firstSize; ; {
+ if err := fw.writeRecordBatchRange(columns, offset, size); err
!= nil {
+ fw.Close()
+ return err
}
- for i := 0; i < int(r.NumCols()); i++ {
- if err := fw.WriteColumnData(r.Column(i)); err != nil {
- fw.Close()
- return err
- }
+ if offset+size >= rec.NumRows() {
+ break
+ }
+ if err := fw.NewBufferedRowGroupChecked(); err != nil {
+ return err
}
+ offset += size
+ size = utils.Min(maxRows, rec.NumRows()-offset)
}
fw.colIdx = 0
return nil
@@ -375,34 +376,68 @@ func (fw *FileWriter) Write(rec arrow.RecordBatch) error {
return fmt.Errorf("record schema does not match writer's.
\nrecord: %s\nwriter: %s", rec.Schema(), fw.schema)
}
- var recList []arrow.RecordBatch
rowgroupLen := fw.wr.Properties().MaxRowGroupLength()
- if rec.NumRows() > rowgroupLen {
- recList = make([]arrow.RecordBatch, 0)
- for offset := int64(0); offset < rec.NumRows(); offset +=
rowgroupLen {
- s := rec.NewSlice(offset, offset+utils.Min(rowgroupLen,
rec.NumRows()-offset))
- defer s.Release()
- recList = append(recList, s)
+ if rec.NumRows() <= rowgroupLen {
+ if err := fw.NewRowGroupChecked(); err != nil {
+ return err
}
- } else {
- recList = []arrow.RecordBatch{rec}
+ if err := fw.writeRecordBatchColumns(rec); err != nil {
+ fw.Close()
+ return err
+ }
+ fw.colIdx = 0
+ return fw.rgw.Close()
}
- for _, r := range recList {
+ columns := newChunkedRecordColumns(rec)
+ defer releaseChunkedRecordColumns(columns)
+
+ for offset := int64(0); offset < rec.NumRows(); offset += rowgroupLen {
if err := fw.NewRowGroupChecked(); err != nil {
return err
}
- for i := 0; i < int(r.NumCols()); i++ {
- if err := fw.WriteColumnData(r.Column(i)); err != nil {
- fw.Close()
- return err
- }
+ size := utils.Min(rowgroupLen, rec.NumRows()-offset)
+ if err := fw.writeRecordBatchRange(columns, offset, size); err
!= nil {
+ fw.Close()
+ return err
}
}
fw.colIdx = 0
return fw.rgw.Close()
}
+func newChunkedRecordColumns(rec arrow.RecordBatch) []*arrow.Chunked {
+ columns := make([]*arrow.Chunked, int(rec.NumCols()))
+ for i, column := range rec.Columns() {
+ columns[i] = arrow.NewChunked(column.DataType(),
[]arrow.Array{column})
+ }
+ return columns
+}
+
+func releaseChunkedRecordColumns(columns []*arrow.Chunked) {
+ for _, column := range columns {
+ column.Release()
+ }
+}
+
+func (fw *FileWriter) writeRecordBatchColumns(rec arrow.RecordBatch) error {
+ for i := 0; i < int(rec.NumCols()); i++ {
+ if err := fw.WriteColumnData(rec.Column(i)); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+func (fw *FileWriter) writeRecordBatchRange(columns []*arrow.Chunked, offset,
size int64) error {
+ for _, column := range columns {
+ if err := fw.WriteColumnChunked(column, offset, size); err !=
nil {
+ return err
+ }
+ }
+ return nil
+}
+
// WriteTable writes an arrow table to the underlying file using chunkSize to
determine
// the size to break at for making row groups. Writing a table will always
create a new
// row group for each chunk of chunkSize rows in the table. Calling this with
0 rows will
diff --git a/parquet/pqarrow/file_writer_range_bench_test.go
b/parquet/pqarrow/file_writer_range_bench_test.go
new file mode 100644
index 00000000..e20d375a
--- /dev/null
+++ b/parquet/pqarrow/file_writer_range_bench_test.go
@@ -0,0 +1,102 @@
+// 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"
+ "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/pqarrow"
+)
+
+func BenchmarkFileWriterRecordBatchRanges(b *testing.B) {
+ tests := []struct {
+ name string
+ numCols int
+ numRows int
+ rowGroup int64
+ }{
+ {name: "1col_16rows_rg16", numCols: 1, numRows: 16, rowGroup:
16},
+ {name: "8cols_8192rows_rg256", numCols: 8, numRows: 8192,
rowGroup: 256},
+ {name: "32cols_4096rows_rg64", numCols: 32, numRows: 4096,
rowGroup: 64},
+ }
+
+ for _, test := range tests {
+ schema, record := makeRangeWriteRecord(test.numCols,
test.numRows)
+ b.Run(test.name, func(b *testing.B) {
+ for _, method := range []struct {
+ name string
+ write func(*pqarrow.FileWriter,
arrow.RecordBatch) error
+ }{
+ {name: "Write", write:
(*pqarrow.FileWriter).Write},
+ {name: "WriteBuffered", write:
(*pqarrow.FileWriter).WriteBuffered},
+ } {
+ b.Run(method.name, func(b *testing.B) {
+ props := parquet.NewWriterProperties(
+
parquet.WithDictionaryDefault(false),
+
parquet.WithMaxRowGroupLength(test.rowGroup),
+ )
+ arrProps := pqarrow.DefaultWriterProps()
+ var output bytes.Buffer
+ output.Grow(test.numCols * test.numRows
* 8)
+
+ b.ReportAllocs()
+ b.SetBytes(int64(test.numCols *
test.numRows * 8))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ output.Reset()
+ writer, err :=
pqarrow.NewFileWriter(schema, &output, props, arrProps)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if err := method.write(writer,
record); err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.Close(); err
!= nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+ })
+ record.Release()
+ }
+}
+
+func makeRangeWriteRecord(numCols, numRows int) (*arrow.Schema,
arrow.RecordBatch) {
+ fields := make([]arrow.Field, numCols)
+ for i := range fields {
+ fields[i] = arrow.Field{Name: fmt.Sprintf("column%d", i), Type:
arrow.PrimitiveTypes.Int64}
+ }
+ schema := arrow.NewSchema(fields, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+
+ values := make([]int64, numRows)
+ for col := 0; col < numCols; col++ {
+ for row := range values {
+ values[row] = int64(row*numCols + col)
+ }
+ builder.Field(col).(*array.Int64Builder).AppendValues(values,
nil)
+ }
+ return schema, builder.NewRecordBatch()
+}
diff --git a/parquet/pqarrow/file_writer_test.go
b/parquet/pqarrow/file_writer_test.go
index b3672b0d..96b10f7f 100644
--- a/parquet/pqarrow/file_writer_test.go
+++ b/parquet/pqarrow/file_writer_test.go
@@ -18,6 +18,7 @@ package pqarrow_test
import (
"bytes"
+ "context"
"math"
"strings"
"testing"
@@ -26,6 +27,7 @@ import (
"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/assert"
"github.com/stretchr/testify/require"
@@ -209,6 +211,129 @@ func TestFileWriterTotalBytesBuffered(t *testing.T) {
assert.Equal(t, int64(1120), writer.TotalBytesWritten())
}
+func TestFileWriterRangeWritesPreserveData(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "number", Type: arrow.PrimitiveTypes.Int64, Nullable:
true},
+ {Name: "text", Type: arrow.BinaryTypes.String, Nullable: true},
+ {Name: "values", Type:
arrow.ListOf(arrow.PrimitiveTypes.Int32), Nullable: true},
+ }, nil)
+ record, _, err := array.RecordFromJSON(memory.DefaultAllocator, schema,
strings.NewReader(`[
+ {"number": 1, "text": "one", "values": [1, 2]},
+ {"number": null, "text": "two", "values": []},
+ {"number": 3, "text": null, "values": null},
+ {"number": 4, "text": "four", "values": [4]},
+ {"number": 5, "text": "five", "values": [5, 6, 7]}
+ ]`))
+ require.NoError(t, err)
+ defer record.Release()
+
+ writeAndRead := func(t *testing.T, write func(*pqarrow.FileWriter)
error) {
+ t.Helper()
+
+ var output bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(
+ schema,
+ &output,
+
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(2)),
+ pqarrow.DefaultWriterProps(),
+ )
+ require.NoError(t, err)
+ require.NoError(t, write(writer))
+ require.NoError(t, writer.Close())
+
+ reader, err :=
file.NewParquetReader(bytes.NewReader(output.Bytes()))
+ require.NoError(t, err)
+ require.Equal(t, 3, reader.NumRowGroups())
+ require.Equal(t, int64(5), reader.NumRows())
+ require.NoError(t, reader.Close())
+
+ got, err := pqarrow.ReadTable(context.Background(),
bytes.NewReader(output.Bytes()), nil, pqarrow.ArrowReadProperties{},
memory.DefaultAllocator)
+ require.NoError(t, err)
+ defer got.Release()
+ require.Equal(t, int64(5), got.NumRows())
+ for i := 0; i < int(record.NumCols()); i++ {
+ expected :=
arrow.NewChunked(record.Column(i).DataType(), []arrow.Array{record.Column(i)})
+ require.Truef(t, array.ChunkedEqual(expected,
got.Column(i).Data()), "column %d differs", i)
+ expected.Release()
+ }
+ }
+
+ t.Run("Write", func(t *testing.T) {
+ writeAndRead(t, func(writer *pqarrow.FileWriter) error {
+ return writer.Write(record)
+ })
+ })
+
+ t.Run("WriteBuffered", func(t *testing.T) {
+ writeAndRead(t, func(writer *pqarrow.FileWriter) error {
+ return writer.WriteBuffered(record)
+ })
+ })
+
+ t.Run("WriteBufferedAcrossCalls", func(t *testing.T) {
+ first := record.NewSlice(0, 1)
+ defer first.Release()
+ second := record.NewSlice(1, record.NumRows())
+ defer second.Release()
+
+ writeAndRead(t, func(writer *pqarrow.FileWriter) error {
+ if err := writer.WriteBuffered(first); err != nil {
+ return err
+ }
+ return writer.WriteBuffered(second)
+ })
+ })
+
+ t.Run("WriteBufferedAtFullBoundary", func(t *testing.T) {
+ first := record.NewSlice(0, 2)
+ defer first.Release()
+ second := record.NewSlice(2, record.NumRows())
+ defer second.Release()
+
+ writeAndRead(t, func(writer *pqarrow.FileWriter) error {
+ if err := writer.WriteBuffered(first); err != nil {
+ return err
+ }
+ return writer.WriteBuffered(second)
+ })
+ })
+}
+
+func TestFileWriterZeroRowRecord(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int64}}, nil)
+ builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ record := builder.NewRecordBatch()
+ builder.Release()
+ defer record.Release()
+
+ for _, test := range []struct {
+ name string
+ write func(*pqarrow.FileWriter, arrow.RecordBatch) error
+ }{
+ {name: "Write", write: (*pqarrow.FileWriter).Write},
+ {name: "WriteBuffered", write:
(*pqarrow.FileWriter).WriteBuffered},
+ } {
+ t.Run(test.name, func(t *testing.T) {
+ var output bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(
+ schema,
+ &output,
+
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(2)),
+ pqarrow.DefaultWriterProps(),
+ )
+ require.NoError(t, err)
+ require.NoError(t, test.write(writer, record))
+ require.NoError(t, writer.Close())
+
+ reader, err :=
file.NewParquetReader(bytes.NewReader(output.Bytes()))
+ require.NoError(t, err)
+ require.Equal(t, 1, reader.NumRowGroups())
+ require.Equal(t, int64(0), reader.NumRows())
+ require.NoError(t, reader.Close())
+ })
+ }
+}
+
func TestWriteOnClosedFileWriter(t *testing.T) {
schema := arrow.NewSchema([]arrow.Field{
{Name: "one", Nullable: true, Type:
arrow.PrimitiveTypes.Float64},
