feat(table): roll parquet files based on actual compressed size (#759)

This change refactors datafile writing to use the actual written file
size as iceberg-java & iceberg-rust do instead of the in-memory size.

Author: twuebi

manifest.go

@@ -27,6 +27,7 @@ import (
 	"reflect"
 	"slices"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -1584,6 +1585,20 @@ const (
 	ParquetFile FileFormat = "PARQUET"
 )
 
+// FileFormatFromString parses a file format string (case-insensitive).
+func FileFormatFromString(s string) (FileFormat, error) {
+	switch strings.ToUpper(s) {
+	case string(ParquetFile):
+		return ParquetFile, nil
+	case string(OrcFile):
+		return OrcFile, nil
+	case string(AvroFile):
+		return AvroFile, nil
+	default:
+		return "", fmt.Errorf("unknown file format: %s", s)
+	}
+}
+
 type colMap[K, V any] struct {
 	Key   K `avro:"key"`
 	Value V `avro:"value"`

table/arrow_utils.go

@@ -1332,54 +1332,77 @@ func recordsToDataFiles(ctx context.Context, rootLocation string, meta *Metadata
 			yield(nil, err)
 		}
 	}
-	currentSpec, err := meta.CurrentSpec()
-	if err != nil {
-		return func(yield func(iceberg.DataFile, error) bool) {
-			yield(nil, err)
-		}
-	}
-	if currentSpec == nil {
-		return func(yield func(iceberg.DataFile, error) bool) {
-			yield(nil, fmt.Errorf("cannot write files without a current spec: %w", err))
-		}
-	}
 
 	cw := newConcurrentDataFileWriter(func(rootLocation string, fs iceio.WriteFileIO, meta *MetadataBuilder, props iceberg.Properties, opts ...dataFileWriterOption) (dataFileWriter, error) {
 		return newDataFileWriter(rootLocation, fs, meta, props, opts...)
 	})
-	nextCount, stopCount := iter.Pull(args.counter)
-	if currentSpec.IsUnpartitioned() {
-		tasks := func(yield func(WriteTask) bool) {
-			defer stopCount()
 
-			fileCount := 0
-			for batch := range binPackRecords(args.itr, defaultBinPackLookback, targetFileSize) {
-				cnt, _ := nextCount()
-				fileCount++
-				t := WriteTask{
-					Uuid:        *args.writeUUID,
-					ID:          cnt,
-					PartitionID: iceberg.UnpartitionedSpec.ID(),
-					FileCount:   fileCount,
-					Schema:      taskSchema,
-					Batches:     batch,
-				}
-				if !yield(t) {
-					return
-				}
-			}
-		}
+	factory, err := newWriterFactory(rootLocation, args, meta, taskSchema, targetFileSize)
+	if err != nil {
+		panic(err)
+	}
 
-		return cw.writeFiles(ctx, rootLocation, args.fs, meta, meta.props, nil, tasks)
+	if factory.currentSpec.IsUnpartitioned() {
+		return unpartitionedWrite(ctx, factory, args.itr)
 	}
 
-	factory := NewWriterFactory(rootLocation, args, meta, taskSchema, targetFileSize)
-	partitionWriter := newPartitionedFanoutWriter(*currentSpec, cw, meta.CurrentSchema(), args.itr, &factory)
+	partitionWriter := newPartitionedFanoutWriter(factory.currentSpec, cw, meta.CurrentSchema(), args.itr, factory)
 	workers := config.EnvConfig.MaxWorkers
 
 	return partitionWriter.Write(ctx, workers)
 }
 
+func unpartitionedWrite(ctx context.Context, factory *writerFactory, records iter.Seq2[arrow.RecordBatch, error]) iter.Seq2[iceberg.DataFile, error] {
+	outputCh := make(chan iceberg.DataFile, 1)
+	errCh := make(chan error, 1)
+
+	go func() {
+		defer close(outputCh)
+		defer factory.stopCount()
+
+		writer := factory.newRollingDataWriter(ctx, nil, "", nil, outputCh)
+		for rec, err := range records {
+			if err != nil {
+				errCh <- err
+				close(errCh)
+				writer.close()
+				writer.wg.Wait()
+
+				return
+			}
+			if err := writer.Add(rec); err != nil {
+				errCh <- err
+				close(errCh)
+				writer.close()
+				writer.wg.Wait()
+
+				return
+			}
+		}
+		close(writer.recordCh)
+		writer.wg.Wait()
+		if err := <-writer.errorCh; err != nil {
+			errCh <- err
+		}
+		close(errCh)
+	}()
+
+	return func(yield func(iceberg.DataFile, error) bool) {
+		defer func() {
+			for range outputCh {
+			}
+		}()
+		for df := range outputCh {
+			if !yield(df, nil) {
+				return
+			}
+		}
+		if err := <-errCh; err != nil {
+			yield(nil, err)
+		}
+	}
+}
+
 type partitionContext struct {
 	partitionData map[int]any
 	specID        int32
@@ -1448,8 +1471,14 @@ func positionDeleteRecordsToDataFiles(ctx context.Context, rootLocation string,
 
 		return cw.writeFiles(ctx, rootLocation, args.fs, meta, meta.props, nil, tasks)
 	}
-	writerFactory := NewWriterFactory(rootLocation, args, meta, iceberg.PositionalDeleteSchema, targetFileSize)
-	partitionWriter := newPositionDeletePartitionedFanoutWriter(latestMetadata, cw, partitionContextByFilePath, args.itr, &writerFactory)
+	factory, err := newWriterFactory(rootLocation, args, meta, iceberg.PositionalDeleteSchema, targetFileSize,
+		withContentType(iceberg.EntryContentPosDeletes),
+		withFactoryFileSchema(iceberg.PositionalDeleteSchema))
+	if err != nil {
+		panic(err)
+	}
+
+	partitionWriter := newPositionDeletePartitionedFanoutWriter(latestMetadata, cw, partitionContextByFilePath, args.itr, factory)
 	workers := config.EnvConfig.MaxWorkers
 
 	return partitionWriter.Write(ctx, workers)

table/internal/interfaces.go

@@ -73,12 +73,22 @@ type FileReader interface {
 	ReadTable(context.Context) (arrow.Table, error)
 }
 
+// FileWriter is an incremental single-file writer with open/write/close
+// lifecycle. It writes Arrow record batches and tracks bytes written for
+// rolling file decisions.
+type FileWriter interface {
+	Write(arrow.RecordBatch) error
+	BytesWritten() int64
+	Close() (iceberg.DataFile, error)
+}
+
 type FileFormat interface {
 	Open(context.Context, iceio.IO, string) (FileReader, error)
 	PathToIDMapping(*iceberg.Schema) (map[string]int, error)
 	DataFileStatsFromMeta(rdr Metadata, statsCols map[int]StatisticsCollector, colMapping map[string]int) *DataFileStatistics
 	GetWriteProperties(iceberg.Properties) any
 	WriteDataFile(ctx context.Context, fs iceio.WriteFileIO, partitionValues map[int]any, info WriteFileInfo, batches []arrow.RecordBatch) (iceberg.DataFile, error)
+	NewFileWriter(ctx context.Context, fs iceio.WriteFileIO, partitionValues map[int]any, info WriteFileInfo, arrowSchema *arrow.Schema) (FileWriter, error)
 }
 
 func GetFileFormat(format iceberg.FileFormat) FileFormat {

table/internal/parquet_files.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"maps"
 	"slices"
 	"strconv"
@@ -243,46 +244,103 @@ func (parquetFormat) GetWriteProperties(props iceberg.Properties) any {
 		parquet.WithCompressionLevel(compressionLevel))
 }
 
-func (p parquetFormat) WriteDataFile(ctx context.Context, fs iceio.WriteFileIO, partitionValues map[int]any, info WriteFileInfo, batches []arrow.RecordBatch) (_ iceberg.DataFile, err error) {
+func (p parquetFormat) WriteDataFile(ctx context.Context, fs iceio.WriteFileIO, partitionValues map[int]any, info WriteFileInfo, batches []arrow.RecordBatch) (iceberg.DataFile, error) {
+	w, err := p.NewFileWriter(ctx, fs, partitionValues, info, batches[0].Schema())
+	if err != nil {
+		return nil, err
+	}
+
+	for _, batch := range batches {
+		if err := w.Write(batch); err != nil {
+			w.Close()
+
+			return nil, err
+		}
+	}
+
+	return w.Close()
+}
+
+// ParquetFileWriter is an incremental single-file writer with open/write/close
+// lifecycle. It writes Arrow record batches to a Parquet file and tracks bytes
+// written for rolling file decisions.
+type ParquetFileWriter struct {
+	pqWriter   *pqarrow.FileWriter
+	counter    *internal.CountingWriter
+	fileCloser io.Closer
+	format     parquetFormat
+	info       WriteFileInfo
+	partition  map[int]any
+	colMapping map[string]int
+}
+
+// NewFileWriter creates a ParquetFileWriter that writes batches to a single
+// Parquet file. Call Write to append batches, BytesWritten to check actual
+// compressed file size, and Close to finalize and get the resulting DataFile.
+func (p parquetFormat) NewFileWriter(ctx context.Context, fs iceio.WriteFileIO,
+	partitionValues map[int]any, info WriteFileInfo, arrowSchema *arrow.Schema,
+) (FileWriter, error) {
 	fw, err := fs.Create(info.FileName)
 	if err != nil {
 		return nil, err
 	}
 
-	defer internal.CheckedClose(fw, &err)
+	colMapping, err := p.PathToIDMapping(info.FileSchema)
+	if err != nil {
+		fw.Close()
 
-	cntWriter := internal.CountingWriter{W: fw}
+		return nil, err
+	}
+
+	counter := &internal.CountingWriter{W: fw}
 	mem := compute.GetAllocator(ctx)
 	writerProps := parquet.NewWriterProperties(info.WriteProps.([]parquet.WriterProperty)...)
 	arrProps := pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem), pqarrow.WithStoreSchema())
 
-	writer, err := pqarrow.NewFileWriter(batches[0].Schema(), &cntWriter, writerProps, arrProps)
+	writer, err := pqarrow.NewFileWriter(arrowSchema, counter, writerProps, arrProps)
 	if err != nil {
+		fw.Close()
+
 		return nil, err
 	}
 
-	for _, batch := range batches {
-		if err := writer.WriteBuffered(batch); err != nil {
-			return nil, err
-		}
-	}
+	return &ParquetFileWriter{
+		pqWriter:   writer,
+		counter:    counter,
+		fileCloser: fw,
+		format:     p,
+		info:       info,
+		partition:  partitionValues,
+		colMapping: colMapping,
+	}, nil
+}
 
-	if err := writer.Close(); err != nil {
-		return nil, err
-	}
+// Write appends a record batch to the Parquet file.
+func (w *ParquetFileWriter) Write(batch arrow.RecordBatch) error {
+	return w.pqWriter.WriteBuffered(batch)
+}
 
-	filemeta, err := writer.FileMetadata()
-	if err != nil {
+// BytesWritten returns the number of bytes flushed to the output so far.
+func (w *ParquetFileWriter) BytesWritten() int64 {
+	return w.counter.Count
+}
+
+// Close finalizes the Parquet file and returns the resulting DataFile with
+// accurate file statistics and size.
+func (w *ParquetFileWriter) Close() (_ iceberg.DataFile, err error) {
+	defer internal.CheckedClose(w.fileCloser, &err)
+
+	if err = w.pqWriter.Close(); err != nil {
 		return nil, err
 	}
 
-	colMapping, err := p.PathToIDMapping(info.FileSchema)
+	filemeta, err := w.pqWriter.FileMetadata()
 	if err != nil {
 		return nil, err
 	}
 
-	return p.DataFileStatsFromMeta(filemeta, info.StatsCols, colMapping).
-		ToDataFile(info.FileSchema, info.Spec, info.FileName, iceberg.ParquetFile, info.Content, cntWriter.Count, partitionValues), nil
+	return w.format.DataFileStatsFromMeta(filemeta, w.info.StatsCols, w.colMapping).
+		ToDataFile(w.info.FileSchema, w.info.Spec, w.info.FileName, iceberg.ParquetFile, w.info.Content, w.counter.Count, w.partition), nil
 }
 
 type decAsIntAgg[T int32 | int64] struct {

table/partitioned_fanout_writer_test.go

@@ -130,15 +130,13 @@ func (s *FanoutWriterTestSuite) testTransformPartition(transform iceberg.Transfo
 		},
 	}
 
-	nameMapping := icebergSchema.NameMapping()
-	taskSchema, err := ArrowSchemaToIceberg(args.sc, false, nameMapping)
-	s.Require().NoError(err)
-
 	cw := newConcurrentDataFileWriter(func(rootLocation string, fs iceio.WriteFileIO, meta *MetadataBuilder, props iceberg.Properties, opts ...dataFileWriterOption) (dataFileWriter, error) {
 		return newDataFileWriter(rootLocation, fs, meta, props, opts...)
 	})
-	writerFactory := NewWriterFactory(loc, args, metaBuilder, icebergSchema, 1024*1024)
-	partitionWriter := newPartitionedFanoutWriter(spec, cw, taskSchema, args.itr, &writerFactory)
+	rollingDataWriters, err := newWriterFactory(loc, args, metaBuilder, icebergSchema, 1024*1024)
+	s.Require().NoError(err)
+
+	partitionWriter := newPartitionedFanoutWriter(spec, cw, icebergSchema, args.itr, rollingDataWriters)
 	workers := config.EnvConfig.MaxWorkers
 
 	dataFiles := partitionWriter.Write(s.ctx, workers)

table/pos_delete_partitioned_fanout_writer_test.go

@@ -129,14 +129,16 @@ func TestPositionDeletePartitionedFanoutWriterProcessBatch(t *testing.T) {
 			cw := newConcurrentDataFileWriter(func(rootLocation string, fs io.WriteFileIO, meta *MetadataBuilder, props iceberg.Properties, opts ...dataFileWriterOption) (dataFileWriter, error) {
 				return newPositionDeleteWriter(rootLocation, fs, meta, props, opts...)
 			})
-			writerFactory := NewWriterFactory(t.TempDir(), recordWritingArgs{
+			factory, err := newWriterFactory(t.TempDir(), recordWritingArgs{
 				fs:        &io.LocalFS{},
 				sc:        PositionalDeleteArrowSchema,
 				writeUUID: &writeUUID,
 				counter:   internal.Counter(0),
-			}, metadataBuilder, iceberg.PositionalDeleteSchema, 1024*1024)
-			writer := newPositionDeletePartitionedFanoutWriter(latestMeta, cw, tc.pathToPartitionContext, nil, &writerFactory)
+			}, metadataBuilder, iceberg.PositionalDeleteSchema, 1024*1024,
+				withContentType(iceberg.EntryContentPosDeletes),
+				withFactoryFileSchema(iceberg.PositionalDeleteSchema))
 			require.NoError(t, err)
+			writer := newPositionDeletePartitionedFanoutWriter(latestMeta, cw, tc.pathToPartitionContext, nil, factory)
 
 			dataFileCh := make(chan iceberg.DataFile, 10)
 			err = writer.processBatch(ctx, tc.input, dataFileCh)
@@ -147,7 +149,7 @@ func TestPositionDeletePartitionedFanoutWriterProcessBatch(t *testing.T) {
 			}
 			require.NoError(t, err)
 
-			err = writerFactory.closeAll()
+			err = factory.closeAll()
 			require.NoError(t, err)
 
 			close(dataFileCh)

table/properties.go

@@ -79,6 +79,9 @@ const (
 	MetadataCompressionKey     = "write.metadata.compression-codec"
 	MetadataCompressionDefault = "none"
 
+	WriteFormatDefaultKey     = "write.format.default"
+	WriteFormatDefaultDefault = "parquet"
+
 	WriteTargetFileSizeBytesKey     = "write.target-file-size-bytes"
 	WriteTargetFileSizeBytesDefault = 512 * 1024 * 1024 // 512 MB