perf(parquet/pqarrow): write record ranges directly

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

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()
+}

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},