perf(table): precompute data writer invariants

table/writer.go

@@ -64,8 +64,19 @@ type defaultDataFileWriter struct {
 	format           tblutils.FileFormat
 	props            iceberg.Properties
 	content          iceberg.ManifestEntryContent
-	meta             *MetadataBuilder
 	equalityFieldIDs []int
+	invariants       dataWriterInvariants
+}
+
+// dataWriterInvariants contains the write configuration shared by every task
+// handled by a defaultDataFileWriter.
+type dataWriterInvariants struct {
+	statsCols     map[int]tblutils.StatisticsCollector
+	writeProps    any
+	rowGroupBytes int64
+	spec          iceberg.PartitionSpec
+	extension     string
+	schemaOpts    SchemaOptions
 }
 
 type dataFileWriterOption func(writer *defaultDataFileWriter)
@@ -119,12 +130,43 @@ func newDataFileWriter(rootLocation string, fs io.WriteFileIO, meta *MetadataBui
 		format:     tblutils.GetFileFormat(fileFormat),
 		content:    iceberg.EntryContentData,
 		props:      props,
-		meta:       meta,
 	}
 	for _, apply := range opts {
 		apply(&w)
 	}
 
+	statsCols, err := computeStatsPlan(w.fileSchema, meta.props)
+	if err != nil {
+		return nil, err
+	}
+
+	currentSpec, err := meta.CurrentSpec()
+	if err != nil {
+		return nil, err
+	}
+
+	var rowGroupBytes int64
+	if w.fileFormat == iceberg.ParquetFile {
+		rowGroupBytes, err = tblutils.ParquetRowGroupTargetSizeBytes(w.props)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	w.invariants = dataWriterInvariants{
+		statsCols:     statsCols,
+		writeProps:    w.format.GetWriteProperties(w.props),
+		rowGroupBytes: rowGroupBytes,
+		spec:          *currentSpec,
+		extension:     strings.ToLower(string(w.fileFormat)),
+		schemaOpts: SchemaOptions{
+			DowncastTimestamp: true,
+			IncludeFieldIDs:   true,
+			UseWriteDefault:   true,
+			TableProperties:   meta.props,
+		},
+	}
+
 	return &w, nil
 }
 
@@ -144,48 +186,25 @@ func (w *defaultDataFileWriter) writeFile(ctx context.Context, partitionValues m
 	batches := make([]arrow.RecordBatch, len(task.Batches))
 	for i, b := range task.Batches {
 		rec, err := ToRequestedSchema(ctx, w.fileSchema,
-			task.Schema, b, SchemaOptions{
-				DowncastTimestamp: true,
-				IncludeFieldIDs:   true,
-				UseWriteDefault:   true,
-				TableProperties:   w.meta.props,
-			})
+			task.Schema, b, w.invariants.schemaOpts)
 		if err != nil {
 			return nil, err
 		}
 		batches[i] = rec
 		defer rec.Release()
 	}
 
-	statsCols, err := computeStatsPlan(w.fileSchema, w.meta.props)
-	if err != nil {
-		return nil, err
-	}
-
 	filePath := w.loc.NewDataLocation(
-		task.GenerateDataFileName(strings.ToLower(string(w.fileFormat))))
-
-	currentSpec, err := w.meta.CurrentSpec()
-	if err != nil {
-		return nil, err
-	}
-
-	var rowGroupTargetSizeBytes int64
-	if w.fileFormat == iceberg.ParquetFile {
-		rowGroupTargetSizeBytes, err = tblutils.ParquetRowGroupTargetSizeBytes(w.props)
-		if err != nil {
-			return nil, err
-		}
-	}
+		task.GenerateDataFileName(w.invariants.extension))
 
 	return w.format.WriteDataFile(ctx, w.fs, partitionValues, tblutils.WriteFileInfo{
 		FileSchema:       w.fileSchema,
 		Content:          w.content,
 		FileName:         filePath,
-		StatsCols:        statsCols,
-		WriteProps:       w.format.GetWriteProperties(w.props),
-		RowGroupBytes:    rowGroupTargetSizeBytes,
-		Spec:             *currentSpec,
+		StatsCols:        w.invariants.statsCols,
+		WriteProps:       w.invariants.writeProps,
+		RowGroupBytes:    w.invariants.rowGroupBytes,
+		Spec:             w.invariants.spec,
 		EqualityFieldIDs: w.equalityFieldIDs,
 		SortOrderID:      task.SortOrderID,
 	}, batches)

table/writer_invariants_bench_test.go

@@ -0,0 +1,91 @@
+// 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 table
+
+import (
+	"context"
+	"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/iceberg-go"
+	iceio "github.com/apache/iceberg-go/io"
+	tblutils "github.com/apache/iceberg-go/table/internal"
+	"github.com/google/uuid"
+)
+
+type benchmarkNoopDataFileFormat struct {
+	tblutils.FileFormat
+}
+
+func (benchmarkNoopDataFileFormat) WriteDataFile(context.Context, iceio.WriteFileIO, map[int]any, tblutils.WriteFileInfo, []arrow.RecordBatch) (iceberg.DataFile, error) {
+	return nil, nil
+}
+
+func BenchmarkDefaultDataFileWriter(b *testing.B) {
+	const (
+		fieldID = 1
+		rows    = 128
+	)
+
+	schema := iceberg.NewSchema(0, iceberg.NestedField{
+		ID: fieldID, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true,
+	})
+	spec := iceberg.NewPartitionSpec()
+	metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, b.TempDir(), iceberg.Properties{})
+	if err != nil {
+		b.Fatal(err)
+	}
+	metaBuilder, err := MetadataBuilderFromBase(metadata, "")
+	if err != nil {
+		b.Fatal(err)
+	}
+	arrowSchema, err := SchemaToArrowSchemaWithOptions(schema, ArrowSchemaOptions{IncludeFieldIDs: true})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	builder := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
+	for range rows {
+		builder.Field(0).(*array.Int64Builder).Append(1)
+	}
+	record := builder.NewRecordBatch()
+	builder.Release()
+	defer record.Release()
+
+	writer, err := newDataFileWriter(b.TempDir(), iceio.LocalFS{}, metaBuilder, iceberg.Properties{},
+		withFormat(benchmarkNoopDataFileFormat{FileFormat: tblutils.GetFileFormat(iceberg.ParquetFile)}))
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	writeUUID := uuid.MustParse("12345678-1234-1234-1234-123456789abc")
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; b.Loop(); i++ {
+		record.Retain()
+		_, err := writer.writeFile(b.Context(), nil, WriteTask{
+			Uuid: writeUUID, ID: i, FileCount: 1, Schema: schema,
+			Batches: []arrow.RecordBatch{record},
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}