Repository: polarsignals/frostdb
Pull request #599 (2023-11-28 16:09:05)

pqarrow: change back to row-by-row conversion but write directly
Previously, we were building a scratchRows buffer column by column before writing to a writer. Profiles show that this buffer is not very efficiently reused, most likely since its scope is limited to a single compaction. These allocations end up being 7% of allocations on prod. This commit changes the arrow->parquet conversion back to row-by-row but avoids allocations. Benchmarks show that the runtime is not affected that much:
```
goos: darwin
goarch: arm64
pkg: github.com/polarsignals/frostdb/pqarrow
                 │ benchmainnoscratch │           benchnew           │
                 │       sec/op       │   sec/op     vs base         │
RecordsToFile-12          278.9µ ± 3%   278.5µ ± 1%  ~ (p=0.818 n=6)

                 │ benchmainnoscratch │              benchnew               │
                 │        B/op        │     B/op      vs base               │
RecordsToFile-12        350.53Ki ± 0%   62.81Ki ± 0%  -82.08% (p=0.002 n=6)

                 │ benchmainnoscratch │              benchnew              │
                 │     allocs/op      │  allocs/op   vs base               │
RecordsToFile-12          4.873k ± 0%   3.850k ± 0%  -20.99% (p=0.002 n=6)
```

asubiotto: Also added another commit to avoid calling `GetOneForMarshal` on int32/64 types since that allocates

Patch:
diff --git a/pqarrow/parquet.go b/pqarrow/parquet.go
index 7e20e2b..ac571e9 100644
--- a/pqarrow/parquet.go
+++ b/pqarrow/parquet.go
@@ -37,75 +37,81 @@ func ArrowScalarToParquetValue(sc scalar.Scalar) (parquet.Value, error) {
 	}
 }
 
+// singlePassThroughWriter is used to keep a reference to the rows written to
+// a parquet writer when only converting a single row. Calling WriteRows more
+// than once is unsupported.
+type singlePassThroughWriter struct {
+	rows []parquet.Row
+}
+
+var _ dynparquet.ParquetWriter = (*singlePassThroughWriter)(nil)
+
+func (w *singlePassThroughWriter) Schema() *parquet.Schema { return nil }
+
+func (w *singlePassThroughWriter) Write(_ []any) (int, error) { panic("use WriteRows instead") }
+
+func (w *singlePassThroughWriter) WriteRows(rows []parquet.Row) (int, error) {
+	if w.rows != nil {
+		panic("cannot call WriteRows more than once")
+	}
+	w.rows = rows
+	return len(rows), nil
+}
+
+func (w *singlePassThroughWriter) Flush() error { return nil }
+
+func (w *singlePassThroughWriter) Close() error { return nil }
+
+func (w *singlePassThroughWriter) Reset(_ io.Writer) {}
+
 // RecordToRow converts an arrow record with dynamic columns into a row using a dynamic parquet schema.
 func RecordToRow(schema *dynparquet.Schema, final *parquet.Schema, record arrow.Record, index int) (parquet.Row, error) {
-	rows, err := recordToRows(schema, record, index, index+1, final.Fields(), nil)
-	if err != nil {
+	w := &singlePassThroughWriter{}
+	if err := recordToRows(w, schema.IsDynamicColumn, record, index, index+1, final.Fields()); err != nil {
 		return nil, err
 	}
-	if len(rows) != 1 {
-		return nil, fmt.Errorf("converted wrong number of rows, expected 1, got %d", len(rows))
-	}
-	return rows[0], err
+	return w.rows[0], nil
 }
 
-// recordToRows converts a full arrow record to a slice of parquet rows.
-// scratchRows is used to append the conversion results to.
+type dynamicColumnVerifier func(string) bool
+
+// recordToRows converts a full arrow record to parquet rows which are written
+// to the parquet writer.
 // The caller should use recordStart=0 and recordEnd=record.NumRows() to convert
 // the entire record. Alternatively, the caller may only convert a subset of
 // rows by specifying a range of [recordStart, recordEnd).
-func recordToRows(
-	schema *dynparquet.Schema,
-	record arrow.Record,
-	recordStart, recordEnd int,
-	finalFields []parquet.Field,
-	scratchRows []parquet.Row,
-) ([]parquet.Row, error) {
+func recordToRows(w dynparquet.ParquetWriter, dcv dynamicColumnVerifier, record arrow.Record, recordStart, recordEnd int, finalFields []parquet.Field) error {
 	numRows := recordEnd - recordStart
-	if cap(scratchRows) < numRows {
-		scratchRows = make([]parquet.Row, 0, numRows)
-	}
-	scratchRows = scratchRows[:numRows]
-	for i := range scratchRows {
-		if cap(scratchRows[i]) < len(finalFields) {
-			scratchRows[i] = make([]parquet.Value, 0, len(finalFields))
-		}
-		scratchRows[i] = scratchRows[i][:0]
-	}
-
-	// This code converts the arrow record column-by-column for better cache
-	// locality.
-	for j, f := range finalFields {
-		columnIndexes := record.Schema().FieldIndices(f.Name())
-		if len(columnIndexes) == 0 {
-			// Column not found in record, append null to all rows.
-			for i := 0; i < numRows; i++ {
-				scratchRows[i] = append(
-					scratchRows[i], parquet.ValueOf(nil).Level(0, 0, j),
-				)
+	row := make([]parquet.Value, 0, len(finalFields))
+	for i := 0; i < numRows; i++ {
+		row = row[:0]
+		for j, f := range finalFields {
+			columnIndexes := record.Schema().FieldIndices(f.Name())
+			if len(columnIndexes) == 0 {
+				// Column not found in record, append null to row.
+				row = append(row, parquet.ValueOf(nil).Level(0, 0, j))
+				continue
 			}
-			continue
-		}
 
-		def := 0
-		if schema.IsDynamicColumn(f.Name()) {
-			def = 1
-		}
+			def := 0
+			if dcv(f.Name()) {
+				def = 1
+			}
 
-		col := record.Column(columnIndexes[0])
-		switch arr := col.(type) {
-		case *array.List:
-			dictionaryList, binaryDictionaryList, err := arrowutils.ToConcreteList(arr)
-			if err != nil {
-				return nil, err
+			col := record.Column(columnIndexes[0])
+			indexIntoRecord := recordStart + i
+			if col.IsNull(indexIntoRecord) {
+				row = append(
+					row, parquet.ValueOf(nil).Level(0, 0, j),
+				)
+				continue
 			}
-			for i := 0; i < numRows; i++ {
-				indexIntoRecord := recordStart + i
-				if col.IsNull(indexIntoRecord) {
-					scratchRows[i] = append(
-						scratchRows[i], parquet.ValueOf(nil).Level(0, 0, j),
-					)
-					continue
+
+			switch arr := col.(type) {
+			case *array.List:
+				dictionaryList, binaryDictionaryList, err := arrowutils.ToConcreteList(arr)
+				if err != nil {
+					return err
 				}
 
 				start, end := arr.ValueOffsets(indexIntoRecord)
@@ -115,38 +121,37 @@ func recordToRows(
 					if k != start {
 						rep = 1
 					}
-					scratchRows[i] = append(scratchRows[i], parquet.ByteArrayValue(v).Level(rep, def+1, j))
+					row = append(row, parquet.ByteArrayValue(v).Level(rep, def+1, j))
 				}
-			}
-		default:
-			for i := 0; i < numRows; i++ {
-				indexIntoRecord := recordStart + i
-				if col.IsNull(indexIntoRecord) {
-					scratchRows[i] = append(
-						scratchRows[i], parquet.ValueOf(nil).Level(0, 0, j),
-					)
-					continue
-				}
-				switch arr := col.(type) {
-				case *array.Dictionary:
-					// GetOneForMashal is expensive for binary types because it causes an allocation when converting interface{} to []byte.
-					// So instead we direct cast the array types to get a []byte.
-					vidx := arr.GetValueIndex(indexIntoRecord)
-					switch dict := arr.Dictionary().(type) {
-					case *array.Binary:
-						scratchRows[i] = append(scratchRows[i], parquet.ByteArrayValue(dict.Value(vidx)).Level(0, def, j))
-					default:
-						scratchRows[i] = append(scratchRows[i], parquet.ValueOf(dict.GetOneForMarshal(vidx)).Level(0, def, j))
-					}
+			case *array.Dictionary:
+				vidx := arr.GetValueIndex(indexIntoRecord)
+				switch dict := arr.Dictionary().(type) {
+				case *array.Binary:
+					row = append(row, parquet.ByteArrayValue(dict.Value(vidx)).Level(0, def, j))
 				default:
-					scratchRows[i] = append(
-						scratchRows[i], parquet.ValueOf(col.GetOneForMarshal(indexIntoRecord)).Level(0, def, j),
-					)
+					row = append(row, parquet.ValueOf(dict.GetOneForMarshal(vidx)).Level(0, def, j))
 				}
+			case *array.Int32:
+				row = append(
+					row, parquet.Int32Value(arr.Value(indexIntoRecord)).Level(0, def, j),
+				)
+			case *array.Int64:
+				row = append(
+					row, parquet.Int64Value(arr.Value(indexIntoRecord)).Level(0, def, j),
+				)
+			default:
+				row = append(
+					row, parquet.ValueOf(col.GetOneForMarshal(indexIntoRecord)).Level(0, def, j),
+				)
 			}
 		}
+
+		if _, err := w.WriteRows([]parquet.Row{row}); err != nil {
+			return err
+		}
 	}
-	return scratchRows, nil
+
+	return nil
 }
 
 func RecordDynamicCols(record arrow.Record) map[string][]string {
@@ -195,13 +200,8 @@ func RecordsToFile(schema *dynparquet.Schema, w dynparquet.ParquetWriter, recs [
 
 	finalFields := ps.Schema.Fields()
 
-	var rows []parquet.Row
 	for _, r := range recs {
-		rows, err = recordToRows(schema, r, 0, int(r.NumRows()), finalFields, rows)
-		if err != nil {
-			return err
-		}
-		if _, err := w.WriteRows(rows); err != nil {
+		if err := recordToRows(w, schema.IsDynamicColumn, r, 0, int(r.NumRows()), finalFields); err != nil {
 			return err
 		}
 	}


Test patch:
diff --git a/pqarrow/parquet_test.go b/pqarrow/parquet_test.go
new file mode 100644
index 0000000..8c7ba5d
--- /dev/null
+++ b/pqarrow/parquet_test.go
@@ -0,0 +1,88 @@
+package pqarrow
+
+import (
+	"fmt"
+	"io"
+	"testing"
+
+	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/apache/arrow/go/v14/arrow/array"
+	"github.com/apache/arrow/go/v14/arrow/memory"
+	"github.com/parquet-go/parquet-go"
+	"github.com/stretchr/testify/require"
+)
+
+type noopWriter struct{}
+
+func (w noopWriter) Schema() *parquet.Schema { return nil }
+
+func (w noopWriter) Write(r []any) (int, error) { return len(r), nil }
+
+func (w noopWriter) WriteRows(r []parquet.Row) (int, error) { return len(r), nil }
+
+func (w noopWriter) Flush() error { return nil }
+
+func (w noopWriter) Close() error { return nil }
+
+func (w noopWriter) Reset(_ io.Writer) {}
+
+func BenchmarkRecordsToFile(b *testing.B) {
+	b.ReportAllocs()
+	schema := arrow.NewSchema([]arrow.Field{
+		{Name: "int_column", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "string_column", Type: arrow.BinaryTypes.String},
+		{
+			Name: "dictionary_column",
+			Type: &arrow.DictionaryType{
+				ValueType: arrow.BinaryTypes.String,
+				IndexType: arrow.PrimitiveTypes.Int32,
+			},
+		},
+		{
+			Name: "list_column",
+			Type: arrow.ListOf(&arrow.DictionaryType{
+				ValueType: arrow.BinaryTypes.Binary,
+				IndexType: arrow.PrimitiveTypes.Int32,
+			}),
+		},
+	}, nil)
+
+	parquetFields := parquet.Group{}
+	for _, f := range schema.Fields() {
+		// No need to create actual nodes since the code only looks for field
+		// names.
+		parquetFields[f.Name] = parquet.Node(nil)
+	}
+
+	checked := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer checked.AssertSize(b, 0)
+	builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+	defer builder.Release()
+
+	const numRows = 1024
+	b.ResetTimer()
+	for i := 0; i < numRows; i++ {
+		builder.Field(0).(*array.Int64Builder).Append(int64(i))
+
+		builder.Field(1).(*array.StringBuilder).Append(fmt.Sprintf("%d", i))
+
+		dictBuilder := builder.Field(2).(*array.BinaryDictionaryBuilder)
+		require.NoError(b, dictBuilder.Append([]byte(fmt.Sprintf("%d-key", i))))
+
+		listBuilder := builder.Field(3).(*array.ListBuilder)
+		listBuilder.Append(true)
+		valueBuilder := listBuilder.ValueBuilder().(*array.BinaryDictionaryBuilder)
+		require.NoError(b, valueBuilder.Append([]byte(fmt.Sprintf("string-1-%d", i))))
+		require.NoError(b, valueBuilder.Append([]byte(fmt.Sprintf("string-2-%d", i))))
+		require.NoError(b, valueBuilder.Append([]byte(fmt.Sprintf("string-3-%d", i))))
+	}
+
+	record := builder.NewRecord()
+	for i := 0; i < b.N; i++ {
+		if err := recordToRows(
+			noopWriter{}, func(string) bool { return false }, record, 0, numRows, parquetFields.Fields(),
+		); err != nil {
+			b.Fatal(err)
+		}
+	}
+}