You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by ze...@apache.org on 2026/03/09 14:38:53 UTC
(iceberg-go) branch main updated: feat(table): roll parquet files based on actual compressed size (#759)
This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new cc8a6ffe feat(table): roll parquet files based on actual compressed size (#759)
cc8a6ffe is described below
commit cc8a6ffe7c4d1f83495539c25c6f756cd0b8dce0
Author: Tobias Pütz <to...@minio.io>
AuthorDate: Mon Mar 9 15:38:47 2026 +0100
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.
---
manifest.go | 15 +
table/arrow_utils.go | 103 ++++---
table/internal/interfaces.go | 10 +
table/internal/parquet_files.go | 92 ++++--
table/partitioned_fanout_writer_test.go | 10 +-
table/pos_delete_partitioned_fanout_writer_test.go | 10 +-
table/properties.go | 3 +
table/rolling_data_writer.go | 313 ++++++++++++++++-----
table/rolling_data_writer_test.go | 261 +++++++++++++++++
table/writer.go | 14 +-
10 files changed, 690 insertions(+), 141 deletions(-)
diff --git a/manifest.go b/manifest.go
index 3c135d1d..9ced13a1 100644
--- a/manifest.go
+++ b/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"`
diff --git a/table/arrow_utils.go b/table/arrow_utils.go
index 82d49e2e..dc548307 100644
--- a/table/arrow_utils.go
+++ b/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)
diff --git a/table/internal/interfaces.go b/table/internal/interfaces.go
index 48830154..4baf3bcc 100644
--- a/table/internal/interfaces.go
+++ b/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 {
diff --git a/table/internal/parquet_files.go b/table/internal/parquet_files.go
index 2d0334c5..1a80fd57 100644
--- a/table/internal/parquet_files.go
+++ b/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 {
diff --git a/table/partitioned_fanout_writer_test.go b/table/partitioned_fanout_writer_test.go
index b2082c84..77727bf1 100644
--- a/table/partitioned_fanout_writer_test.go
+++ b/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)
diff --git a/table/pos_delete_partitioned_fanout_writer_test.go b/table/pos_delete_partitioned_fanout_writer_test.go
index f0f1ff4c..c3a44210 100644
--- a/table/pos_delete_partitioned_fanout_writer_test.go
+++ b/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)
diff --git a/table/properties.go b/table/properties.go
index 560b0a1f..63c91761 100644
--- a/table/properties.go
+++ b/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
diff --git a/table/rolling_data_writer.go b/table/rolling_data_writer.go
index e8950e2f..2477705e 100644
--- a/table/rolling_data_writer.go
+++ b/table/rolling_data_writer.go
@@ -22,52 +22,209 @@ import (
"fmt"
"iter"
"net/url"
+ "strings"
"sync"
"sync/atomic"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ tblutils "github.com/apache/iceberg-go/table/internal"
+ "github.com/google/uuid"
)
-// WriterFactory manages the creation and lifecycle of RollingDataWriter instances
+// writerFactory manages the creation and lifecycle of RollingDataWriter instances
// for different partitions, providing shared configuration and coordination
// across all writers in a partitioned write operation.
type writerFactory struct {
- rootLocation string
- args recordWritingArgs
- meta *MetadataBuilder
- taskSchema *iceberg.Schema
- targetFileSize int64
- writers sync.Map
- nextCount func() (int, bool)
- stopCount func()
- partitionIDCounter atomic.Int64 // partitionIDCounter generates unique IDs for partitions
- mu sync.Mutex
+ rootLocation string
+ rootURL *url.URL
+ fs iceio.WriteFileIO
+ writeUUID *uuid.UUID
+ taskSchema *iceberg.Schema
+ targetFileSize int64
+
+ locProvider LocationProvider
+ fileSchema *iceberg.Schema
+ arrowSchema *arrow.Schema
+ writeProps any
+ statsCols map[int]tblutils.StatisticsCollector
+ currentSpec iceberg.PartitionSpec
+ fileFormat iceberg.FileFormat
+ format tblutils.FileFormat
+ content iceberg.ManifestEntryContent
+
+ writers sync.Map
+ partitionLocProviders sync.Map
+ nextCount func() (int, bool)
+ stopCount func()
+ countMu sync.Mutex
+ partitionIDCounter atomic.Int64
+ mu sync.Mutex
}
-// NewWriterFactory creates a new WriterFactory with the specified configuration
-// for managing rolling data writerFactory across partitions.
-func NewWriterFactory(rootLocation string, args recordWritingArgs, meta *MetadataBuilder, taskSchema *iceberg.Schema, targetFileSize int64) writerFactory {
+type writerFactoryOption func(*writerFactory)
+
+func withContentType(content iceberg.ManifestEntryContent) writerFactoryOption {
+ return func(w *writerFactory) {
+ w.content = content
+ }
+}
+
+func withFactoryFileSchema(schema *iceberg.Schema) writerFactoryOption {
+ return func(w *writerFactory) {
+ w.fileSchema = schema
+ arrowSc, err := SchemaToArrowSchema(schema, nil, true, false)
+ if err != nil {
+ panic(fmt.Sprintf("withFactoryFileSchema: failed to convert schema: %v", err))
+ }
+ w.arrowSchema = arrowSc
+ }
+}
+
+// newWriterFactory creates a writerFactory with precomputed, invariant write
+// configuration derived from the table metadata.
+func newWriterFactory(rootLocation string, args recordWritingArgs, meta *MetadataBuilder, taskSchema *iceberg.Schema, targetFileSize int64, opts ...writerFactoryOption) (*writerFactory, error) {
nextCount, stopCount := iter.Pull(args.counter)
- return writerFactory{
+ rootURL, err := url.Parse(rootLocation)
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+
+ locProvider, err := LoadLocationProvider(rootLocation, meta.props)
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+
+ fileSchema := meta.CurrentSchema()
+ sanitized, err := iceberg.SanitizeColumnNames(fileSchema)
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+ if !sanitized.Equals(fileSchema) {
+ fileSchema = sanitized
+ }
+
+ fileFormat, err := iceberg.FileFormatFromString(
+ iceberg.Properties(meta.props).Get(WriteFormatDefaultKey, WriteFormatDefaultDefault))
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+
+ format := tblutils.GetFileFormat(fileFormat)
+
+ arrowSchema, err := SchemaToArrowSchema(fileSchema, nil, true, false)
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+
+ currentSpec, err := meta.CurrentSpec()
+ if err != nil || currentSpec == nil {
+ stopCount()
+
+ return nil, fmt.Errorf("%w: cannot write files without a current spec", err)
+ }
+
+ f := &writerFactory{
rootLocation: rootLocation,
- args: args,
- meta: meta,
+ rootURL: rootURL,
+ fs: args.fs,
+ writeUUID: args.writeUUID,
taskSchema: taskSchema,
targetFileSize: targetFileSize,
+ locProvider: locProvider,
+ fileSchema: fileSchema,
+ arrowSchema: arrowSchema,
+ writeProps: format.GetWriteProperties(meta.props),
+ currentSpec: *currentSpec,
+ fileFormat: fileFormat,
+ format: format,
nextCount: nextCount,
stopCount: stopCount,
}
+ for _, apply := range opts {
+ apply(f)
+ }
+
+ f.statsCols, err = computeStatsPlan(f.fileSchema, meta.props)
+ if err != nil {
+ stopCount()
+
+ return nil, err
+ }
+
+ return f, nil
+}
+
+func (w *writerFactory) openFileWriter(ctx context.Context, partitionPath string,
+ partitionValues map[int]any, partitionID int, fileCount int,
+) (tblutils.FileWriter, error) {
+ w.countMu.Lock()
+ cnt, _ := w.nextCount()
+ w.countMu.Unlock()
+
+ fileName := WriteTask{
+ Uuid: *w.writeUUID,
+ ID: cnt,
+ PartitionID: partitionID,
+ FileCount: fileCount,
+ }.GenerateDataFileName(strings.ToLower(string(w.fileFormat)))
+
+ var filePath string
+ if partitionPath != "" {
+ partitionLoc, err := w.partitionLocProvider(partitionPath)
+ if err != nil {
+ return nil, err
+ }
+ filePath = partitionLoc.NewDataLocation(fileName)
+ } else {
+ filePath = w.locProvider.NewDataLocation(fileName)
+ }
+
+ return w.format.NewFileWriter(ctx, w.fs, partitionValues, tblutils.WriteFileInfo{
+ FileSchema: w.fileSchema,
+ FileName: filePath,
+ StatsCols: w.statsCols,
+ WriteProps: w.writeProps,
+ Spec: w.currentSpec,
+ Content: w.content,
+ }, w.arrowSchema)
}
-// RollingDataWriter accumulates Arrow records for a specific partition and flushes
-// them to data files when the target file size is reached, implementing a rolling
-// file strategy to manage file sizes.
+func (w *writerFactory) partitionLocProvider(partitionPath string) (LocationProvider, error) {
+ if cached, ok := w.partitionLocProviders.Load(partitionPath); ok {
+ return cached.(LocationProvider), nil
+ }
+
+ partitionDataPath := w.rootURL.JoinPath("data", partitionPath).String()
+ loc, err := LoadLocationProvider(w.rootLocation,
+ iceberg.Properties{WriteDataPathKey: partitionDataPath})
+ if err != nil {
+ return nil, err
+ }
+
+ w.partitionLocProviders.Store(partitionPath, loc)
+
+ return loc, nil
+}
+
+// RollingDataWriter writes Arrow records for a specific partition, rolling to
+// new data files when the actual compressed file size reaches the target.
type RollingDataWriter struct {
partitionKey string
- partitionID int // unique ID for this partition
- fileCount atomic.Int64 // counter for files in this partition
+ partitionID int
+ fileCount atomic.Int64
recordCh chan arrow.RecordBatch
errorCh chan error
factory *writerFactory
@@ -78,9 +235,7 @@ type RollingDataWriter struct {
concurrentWriter *concurrentDataFileWriter
}
-// NewRollingDataWriter creates a new RollingDataWriter for the specified partition
-// with the given partition values.
-func (w *writerFactory) NewRollingDataWriter(ctx context.Context, concurrentWriter *concurrentDataFileWriter, partition string, partitionValues map[int]any, outputDataFilesCh chan<- iceberg.DataFile) *RollingDataWriter {
+func (w *writerFactory) newRollingDataWriter(ctx context.Context, concurrentWriter *concurrentDataFileWriter, partition string, partitionValues map[int]any, outputDataFilesCh chan<- iceberg.DataFile) *RollingDataWriter {
ctx, cancel := context.WithCancel(ctx)
partitionID := int(w.partitionIDCounter.Add(1) - 1)
writer := &RollingDataWriter{
@@ -113,14 +268,13 @@ func (w *writerFactory) getOrCreateRollingDataWriter(ctx context.Context, concur
return nil, fmt.Errorf("invalid writer type for partition: %s", partition)
}
- writer := w.NewRollingDataWriter(ctx, concurrentWriter, partition, partitionValues, outputDataFilesCh)
+ writer := w.newRollingDataWriter(ctx, concurrentWriter, partition, partitionValues, outputDataFilesCh)
w.writers.Store(partition, writer)
return writer, nil
}
-// Add appends a record to the writer's buffer and flushes to a data file if the
-// target file size is reached.
+// Add appends a record to the writer's buffer.
func (r *RollingDataWriter) Add(record arrow.RecordBatch) error {
record.Retain()
select {
@@ -141,70 +295,79 @@ func (r *RollingDataWriter) stream(outputDataFilesCh chan<- iceberg.DataFile) {
defer r.wg.Done()
defer close(r.errorCh)
- recordIter := func(yield func(arrow.RecordBatch, error) bool) {
- for record := range r.recordCh {
- if !yield(record, nil) {
- return
- }
+ var currentWriter tblutils.FileWriter
+ defer func() {
+ if currentWriter != nil {
+ currentWriter.Close()
}
- }
-
- binPackedRecords := binPackRecords(recordIter, defaultBinPackLookback, r.factory.targetFileSize)
- for batch := range binPackedRecords {
- if err := r.flushToDataFile(batch, outputDataFilesCh); err != nil {
- select {
- case r.errorCh <- err:
- default:
- }
+ }()
- return
+ closeWriter := func() error {
+ if currentWriter == nil {
+ return nil
}
- }
-}
+ df, err := currentWriter.Close()
+ currentWriter = nil
+ if err != nil {
+ return err
+ }
+ outputDataFilesCh <- df
-func (r *RollingDataWriter) flushToDataFile(batch []arrow.RecordBatch, outputDataFilesCh chan<- iceberg.DataFile) error {
- if len(batch) == 0 {
return nil
}
- tasks := iter.Seq[WriteTask](func(yield func(WriteTask) bool) {
- cnt, _ := r.factory.nextCount()
- fileCount := int(r.fileCount.Add(1))
+ for record := range r.recordCh {
+ if currentWriter == nil {
+ fileCount := int(r.fileCount.Add(1))
+ var err error
+ currentWriter, err = r.factory.openFileWriter(
+ r.ctx, r.partitionKey, r.partitionValues,
+ r.partitionID, fileCount)
+ if err != nil {
+ record.Release()
+ r.sendError(err)
- yield(WriteTask{
- Uuid: *r.factory.args.writeUUID,
- ID: cnt,
- PartitionID: r.partitionID,
- FileCount: fileCount,
- Schema: r.factory.taskSchema,
- Batches: batch,
- })
- })
+ return
+ }
+ }
- parseDataLoc, err := url.Parse(r.factory.rootLocation)
- if err != nil {
- return fmt.Errorf("failed to parse rootLocation: %v", err)
- }
+ converted, err := ToRequestedSchema(r.ctx, r.factory.fileSchema,
+ r.factory.taskSchema, record, false, true, false)
+ if err != nil {
+ record.Release()
+ r.sendError(err)
- partitionMeta := *r.factory.meta
- if partitionMeta.props == nil {
- partitionMeta.props = make(map[string]string)
- }
- partitionMeta.props[WriteDataPathKey] = parseDataLoc.JoinPath("data").JoinPath(r.partitionKey).String()
+ return
+ }
- outputDataFiles := r.concurrentWriter.writeFiles(r.ctx, r.factory.rootLocation, r.factory.args.fs, &partitionMeta, partitionMeta.props, r.partitionValues, tasks)
- for dataFile, err := range outputDataFiles {
+ err = currentWriter.Write(converted)
+ converted.Release()
+ record.Release()
if err != nil {
- return err
+ r.sendError(err)
+
+ return
+ }
+
+ if currentWriter.BytesWritten() >= r.factory.targetFileSize {
+ if err := closeWriter(); err != nil {
+ r.sendError(err)
+
+ return
+ }
}
- outputDataFilesCh <- dataFile
}
- for _, rec := range batch {
- rec.Release()
+ if err := closeWriter(); err != nil {
+ r.sendError(err)
}
+}
- return nil
+func (r *RollingDataWriter) sendError(err error) {
+ select {
+ case r.errorCh <- err:
+ default:
+ }
}
func (r *RollingDataWriter) close() {
diff --git a/table/rolling_data_writer_test.go b/table/rolling_data_writer_test.go
new file mode 100644
index 00000000..40e4845a
--- /dev/null
+++ b/table/rolling_data_writer_test.go
@@ -0,0 +1,261 @@
+// 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"
+ "os"
+ "path/filepath"
+ "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"
+ "github.com/stretchr/testify/suite"
+)
+
+type RollingDataWriterTestSuite struct {
+ suite.Suite
+
+ mem memory.Allocator
+ ctx context.Context
+}
+
+func (s *RollingDataWriterTestSuite) SetupTest() {
+ s.ctx = context.Background()
+ s.mem = memory.NewCheckedAllocator(memory.NewGoAllocator())
+}
+
+func TestRollingDataWriter(t *testing.T) {
+ suite.Run(t, new(RollingDataWriterTestSuite))
+}
+
+func (s *RollingDataWriterTestSuite) createWriterFactory(loc string, arrSchema *arrow.Schema, targetFileSize int64) (*writerFactory, *iceberg.Schema) {
+ icebergSchema, err := ArrowSchemaToIcebergWithFreshIDs(arrSchema, false)
+ s.Require().NoError(err)
+
+ spec := iceberg.NewPartitionSpec()
+ meta, err := NewMetadata(icebergSchema, &spec, UnsortedSortOrder, loc, iceberg.Properties{})
+ s.Require().NoError(err)
+
+ metaBuilder, err := MetadataBuilderFromBase(meta, "")
+ s.Require().NoError(err)
+
+ writeUUID := uuid.New()
+ args := recordWritingArgs{
+ sc: arrSchema,
+ fs: iceio.LocalFS{},
+ writeUUID: &writeUUID,
+ counter: func(yield func(int) bool) {
+ for i := 0; ; i++ {
+ if !yield(i) {
+ break
+ }
+ }
+ },
+ }
+
+ factory, err := newWriterFactory(loc, args, metaBuilder, icebergSchema, targetFileSize)
+ s.Require().NoError(err)
+
+ return factory, icebergSchema
+}
+
+func (s *RollingDataWriterTestSuite) buildRecord(arrSchema *arrow.Schema, numRows int) arrow.RecordBatch {
+ bldr := array.NewRecordBuilder(s.mem, arrSchema)
+ defer bldr.Release()
+
+ for i := range numRows {
+ bldr.Field(0).(*array.Int32Builder).Append(int32(i))
+ bldr.Field(1).(*array.StringBuilder).Append("value")
+ }
+
+ return bldr.NewRecordBatch()
+}
+
+func (s *RollingDataWriterTestSuite) TestSingleFileUnderTarget() {
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ loc := filepath.ToSlash(s.T().TempDir())
+ factory, _ := s.createWriterFactory(loc, arrSchema, 1024*1024)
+ defer factory.closeAll()
+
+ outputCh := make(chan iceberg.DataFile, 10)
+ writer := factory.newRollingDataWriter(s.ctx, nil, "", nil, outputCh)
+
+ record := s.buildRecord(arrSchema, 5)
+ defer record.Release()
+
+ s.Require().NoError(writer.Add(record))
+ s.Require().NoError(writer.closeAndWait())
+ close(outputCh)
+
+ var dataFiles []iceberg.DataFile
+ for df := range outputCh {
+ dataFiles = append(dataFiles, df)
+ }
+
+ s.Require().Len(dataFiles, 1)
+ s.Equal(int64(5), dataFiles[0].Count())
+}
+
+func (s *RollingDataWriterTestSuite) TestRollsMultipleFiles() {
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ loc := filepath.ToSlash(s.T().TempDir())
+ factory, _ := s.createWriterFactory(loc, arrSchema, 1)
+ defer factory.closeAll()
+
+ outputCh := make(chan iceberg.DataFile, 100)
+ writer := factory.newRollingDataWriter(s.ctx, nil, "", nil, outputCh)
+
+ totalRows := int64(0)
+ for range 10 {
+ record := s.buildRecord(arrSchema, 50)
+ s.Require().NoError(writer.Add(record))
+ totalRows += record.NumRows()
+ record.Release()
+ }
+
+ s.Require().NoError(writer.closeAndWait())
+ close(outputCh)
+
+ var dataFiles []iceberg.DataFile
+ for df := range outputCh {
+ dataFiles = append(dataFiles, df)
+ }
+
+ s.Greater(len(dataFiles), 1, "expected multiple files when targetFileSize is 1 byte")
+
+ var actualRows int64
+ for _, df := range dataFiles {
+ actualRows += df.Count()
+ }
+ s.Equal(totalRows, actualRows)
+}
+
+func (s *RollingDataWriterTestSuite) TestBytesWrittenNoDoubleCountAcrossRowGroups() {
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ loc := filepath.ToSlash(s.T().TempDir())
+
+ icebergSchema, err := ArrowSchemaToIcebergWithFreshIDs(arrSchema, false)
+ s.Require().NoError(err)
+
+ spec := iceberg.NewPartitionSpec()
+ format := tblutils.GetFileFormat(iceberg.ParquetFile)
+ writeProps := format.GetWriteProperties(iceberg.Properties{
+ tblutils.ParquetRowGroupLimitKey: "10",
+ })
+
+ statsCols, err := computeStatsPlan(icebergSchema, iceberg.Properties{})
+ s.Require().NoError(err)
+
+ filePath := filepath.Join(loc, "test-no-double-count.parquet")
+ fw, err := format.NewFileWriter(s.ctx, iceio.LocalFS{}, nil, tblutils.WriteFileInfo{
+ FileSchema: icebergSchema,
+ FileName: filePath,
+ StatsCols: statsCols,
+ WriteProps: writeProps,
+ Spec: spec,
+ }, arrSchema)
+ s.Require().NoError(err)
+
+ var prevBytes int64
+ for batch := range 5 {
+ record := s.buildRecord(arrSchema, 100)
+ s.Require().NoError(fw.Write(record))
+ record.Release()
+
+ bytesWritten := fw.BytesWritten()
+ s.Greater(bytesWritten, prevBytes,
+ "BytesWritten must monotonically increase after batch %d", batch)
+ prevBytes = bytesWritten
+ }
+
+ bytesBeforeClose := fw.BytesWritten()
+
+ df, err := fw.Close()
+ s.Require().NoError(err)
+ s.Equal(int64(500), df.Count())
+
+ s.GreaterOrEqual(df.FileSizeBytes(), bytesBeforeClose,
+ "final file size should be >= BytesWritten before close (close flushes the last row group)")
+
+ stat, err := os.Stat(filePath)
+ s.Require().NoError(err)
+ s.Equal(df.FileSizeBytes(), stat.Size(),
+ "DataFile.FileSizeBytes should match actual file size on disk")
+}
+
+func (s *RollingDataWriterTestSuite) TestBytesWrittenReflectsCompressedSize() {
+ arrSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ loc := filepath.ToSlash(s.T().TempDir())
+
+ icebergSchema, err := ArrowSchemaToIcebergWithFreshIDs(arrSchema, false)
+ s.Require().NoError(err)
+
+ spec := iceberg.NewPartitionSpec()
+ format := tblutils.GetFileFormat(iceberg.ParquetFile)
+ writeProps := format.GetWriteProperties(iceberg.Properties{})
+
+ statsCols, err := computeStatsPlan(icebergSchema, iceberg.Properties{})
+ s.Require().NoError(err)
+
+ filePath := filepath.Join(loc, "test-bytes-written.parquet")
+ fw, err := format.NewFileWriter(s.ctx, iceio.LocalFS{}, nil, tblutils.WriteFileInfo{
+ FileSchema: icebergSchema,
+ FileName: filePath,
+ StatsCols: statsCols,
+ WriteProps: writeProps,
+ Spec: spec,
+ }, arrSchema)
+ s.Require().NoError(err)
+
+ record := s.buildRecord(arrSchema, 100)
+ defer record.Release()
+
+ s.Require().NoError(fw.Write(record))
+
+ bytesWritten := fw.BytesWritten()
+ s.Greater(bytesWritten, int64(0), "BytesWritten should be positive after writing data")
+
+ arrowBytes := int64(record.NumCols()) * record.NumRows() * 4
+ s.Less(bytesWritten, arrowBytes, "compressed parquet should be smaller than raw arrow estimate")
+
+ df, err := fw.Close()
+ s.Require().NoError(err)
+ s.Equal(int64(100), df.Count())
+}
diff --git a/table/writer.go b/table/writer.go
index cdda0357..32e39dbe 100644
--- a/table/writer.go
+++ b/table/writer.go
@@ -21,6 +21,7 @@ import (
"context"
"fmt"
"iter"
+ "strings"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/iceberg-go"
@@ -51,6 +52,7 @@ type defaultDataFileWriter struct {
loc LocationProvider
fs io.WriteFileIO
fileSchema *iceberg.Schema
+ fileFormat iceberg.FileFormat
format internal.FileFormat
props iceberg.Properties
content iceberg.ManifestEntryContent
@@ -82,11 +84,19 @@ func newDataFileWriter(rootLocation string, fs io.WriteFileIO, meta *MetadataBui
if err != nil {
return nil, err
}
+
+ fileFormat, err := iceberg.FileFormatFromString(
+ props.Get(WriteFormatDefaultKey, WriteFormatDefaultDefault))
+ if err != nil {
+ return nil, err
+ }
+
w := defaultDataFileWriter{
loc: locProvider,
fs: fs,
fileSchema: meta.CurrentSchema(),
- format: internal.GetFileFormat(iceberg.ParquetFile),
+ fileFormat: fileFormat,
+ format: internal.GetFileFormat(fileFormat),
content: iceberg.EntryContentData,
props: props,
meta: meta,
@@ -122,7 +132,7 @@ func (w *defaultDataFileWriter) writeFile(ctx context.Context, partitionValues m
}
filePath := w.loc.NewDataLocation(
- task.GenerateDataFileName("parquet"))
+ task.GenerateDataFileName(strings.ToLower(string(w.fileFormat))))
currentSpec, err := w.meta.CurrentSpec()
if err != nil {