diff --git a/table/writer.go b/table/writer.go index 430ae860f..e97503e30 100644 --- a/table/writer.go +++ b/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,12 +186,7 @@ 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 } @@ -157,35 +194,17 @@ func (w *defaultDataFileWriter) writeFile(ctx context.Context, partitionValues m 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) diff --git a/table/writer_invariants_bench_test.go b/table/writer_invariants_bench_test.go new file mode 100644 index 000000000..1a3499875 --- /dev/null +++ b/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) + } + } +}