diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go index cda8ad341..0203bec05 100644 --- a/table/arrow_scanner.go +++ b/table/arrow_scanner.go @@ -22,6 +22,7 @@ import ( "fmt" "io" "iter" + "maps" "slices" "strconv" "strings" @@ -942,6 +943,10 @@ type arrowScan struct { useLargeTypes bool concurrency int + + // arrowBatchSize, when positive, overrides the table's + // read.parquet.batch-size property for this scan's reads. + arrowBatchSize int } // preparedFileRead contains the physical schema projection shared by all @@ -2139,6 +2144,17 @@ func (as *arrowScan) GetRecords(ctx context.Context, tasks []FileScanTask) (*arr } tableProperties := as.metadata.Properties() + batchSize := as.options.Get(ParquetBatchSizeKey, "") + if as.arrowBatchSize > 0 { + batchSize = strconv.Itoa(as.arrowBatchSize) + } + if batchSize != "" { + tableProperties = maps.Clone(tableProperties) + if tableProperties == nil { + tableProperties = iceberg.Properties{} + } + tableProperties[ParquetBatchSizeKey] = batchSize + } ctx = tblutils.WithTableProperties(ctx, tableProperties) resultSchema, err := SchemaToArrowSchemaWithOptions(as.projectedSchema, ArrowSchemaOptions{ diff --git a/table/arrow_utils.go b/table/arrow_utils.go index 7fab00d6b..e9804e205 100644 --- a/table/arrow_utils.go +++ b/table/arrow_utils.go @@ -1933,6 +1933,11 @@ type recordWritingArgs struct { maxWriteWorkers int clustered bool factoryOpts []writerFactoryOption + + // recordBatchBufferSize overrides the rolling writers' record + // channel capacity; non-positive uses the default (see + // rollingDataWriterQueueCapacity). + recordBatchBufferSize int // existingDVs maps a data file path to the positions already recorded in // its current deletion vector. On the v3 DV write path these are folded // into the newly written DV so a data file that already had a DV ends up diff --git a/table/rewrite_data_files.go b/table/rewrite_data_files.go index 43b883167..81ac9fb83 100644 --- a/table/rewrite_data_files.go +++ b/table/rewrite_data_files.go @@ -224,8 +224,11 @@ type RewriteDataFilesOptions struct { type CompactionGroupOption func(*compactionGroupConfig) type compactionGroupConfig struct { - targetFileSize int64 - scanConcurrency int + targetFileSize int64 + scanConcurrency int + arrowBatchSize int + recordBatchBufferSize int + parquetRowGroupLimit int } // WithCompactionTargetFileSize sets the size target for output files @@ -252,6 +255,47 @@ func WithCompactionScanConcurrency(n int) CompactionGroupOption { } } +// WithCompactionArrowBatchSize caps the number of rows decoded per +// Arrow record batch while reading the group's tasks, forwarded to the +// scan as [WithArrowBatchSize]. Together with +// [WithCompactionRecordBatchBufferSize] it bounds the memory held by +// the record pipeline specifically: buffered batches times rows per +// batch. Delete-side memory is not covered — positional deletes and +// deletion-vector bitmaps for the group's tasks are materialized up +// front and sized by delete volume, not by these knobs. A non-positive +// value keeps the table's read.parquet.batch-size property. +func WithCompactionArrowBatchSize(n int) CompactionGroupOption { + return func(c *compactionGroupConfig) { + if n > 0 { + c.arrowBatchSize = n + } + } +} + +// WithCompactionRecordBatchBufferSize sets the capacity, in record +// batches, of the write pipeline's per-writer input buffer, forwarded +// to [WriteRecords] as [WithRecordBatchBufferSize]. The default is 64 +// batches. A non-positive value is ignored. +func WithCompactionRecordBatchBufferSize(n int) CompactionGroupOption { + return func(c *compactionGroupConfig) { + if n > 0 { + c.recordBatchBufferSize = n + } + } +} + +// WithCompactionParquetRowGroupLimit caps the rows per Parquet row +// group in the compacted output files, forwarded to [WriteRecords] as +// [WithParquetRowGroupLimit]. A non-positive value keeps the table's +// write.parquet.row-group-limit property. +func WithCompactionParquetRowGroupLimit(n int) CompactionGroupOption { + return func(c *compactionGroupConfig) { + if n > 0 { + c.parquetRowGroupLimit = n + } + } +} + // RewriteDataFiles compacts the given groups by reading data with // deletes applied, writing new consolidated files, and atomically // replacing the old files. Position delete files that are fully @@ -378,6 +422,9 @@ func ExecuteCompactionGroup(ctx context.Context, tbl *Table, group CompactionTas if cfg.scanConcurrency > 0 { scanOpts = append(scanOpts, WithMaxConcurrency(cfg.scanConcurrency)) } + if cfg.arrowBatchSize > 0 { + scanOpts = append(scanOpts, WithArrowBatchSize(cfg.arrowBatchSize)) + } // Preserve row lineage only when every source file in the group carries // it. A mixed group (some files with FirstRowID, some without — e.g. @@ -422,6 +469,12 @@ func ExecuteCompactionGroup(ctx context.Context, tbl *Table, group CompactionTas if cfg.targetFileSize > 0 { writeOpts = append(writeOpts, WithTargetFileSize(cfg.targetFileSize)) } + if cfg.recordBatchBufferSize > 0 { + writeOpts = append(writeOpts, WithRecordBatchBufferSize(cfg.recordBatchBufferSize)) + } + if cfg.parquetRowGroupLimit > 0 { + writeOpts = append(writeOpts, WithParquetRowGroupLimit(cfg.parquetRowGroupLimit)) + } if preserveLineage { // Rebuild the arrow schema from the projected iceberg schema so the // reserved row-lineage field IDs (_row_id, _last_updated_sequence_number) diff --git a/table/rewrite_data_files_test.go b/table/rewrite_data_files_test.go index ba1b1975c..d9bc9dc6d 100644 --- a/table/rewrite_data_files_test.go +++ b/table/rewrite_data_files_test.go @@ -29,6 +29,7 @@ import ( "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/arrow-go/v18/parquet/file" "github.com/apache/iceberg-go" iceio "github.com/apache/iceberg-go/io" "github.com/apache/iceberg-go/table" @@ -361,6 +362,71 @@ func TestExecuteCompactionGroup_TargetFileSizeForwarded(t *testing.T) { "without the option, the same group consolidates into a single file") } +// TestExecuteCompactionGroup_ParquetRowGroupLimitForwarded verifies +// that WithCompactionParquetRowGroupLimit reaches the underlying +// WriteRecords call: with the limit set, no output row group may +// exceed it, while the no-option baseline packs all rows into one. +// It doubles as the end-to-end pin for the option-forwarding block in +// ExecuteCompactionGroup (the row-group limit is the cheapest +// observable of the three forwarded tuning options). +func TestExecuteCompactionGroup_ParquetRowGroupLimitForwarded(t *testing.T) { + tbl := newRewriteTestTable(t) + + arrowSc, err := table.SchemaToArrowSchema(tbl.Schema(), nil, false, false) + require.NoError(t, err) + + var rows strings.Builder + rows.WriteString("[") + for i := range 100 { + if i > 0 { + rows.WriteString(",") + } + fmt.Fprintf(&rows, `{"id": %d, "data": "row-%d"}`, i+1, i+1) + } + rows.WriteString("]") + + dataPath := tbl.Location() + "/data/rg-limit.parquet" + writeParquetFile(t, dataPath, arrowSc, rows.String()) + tx := tbl.NewTransaction() + require.NoError(t, tx.AddFiles(t.Context(), []string{dataPath}, nil, false)) + tbl, err = tx.Commit(t.Context()) + require.NoError(t, err) + + tasks, err := tbl.Scan().PlanFiles(t.Context()) + require.NoError(t, err) + require.Len(t, tasks, 1) + + group := table.CompactionTaskGroup{ + PartitionKey: "single", + Tasks: []table.FileScanTask{tasks[0]}, + TotalSizeBytes: tasks[0].File.FileSizeBytes(), + } + + withLimit, err := table.ExecuteCompactionGroup(t.Context(), tbl, group, + table.WithCompactionParquetRowGroupLimit(10)) + require.NoError(t, err) + require.Len(t, withLimit.NewDataFiles, 1) + + limited, err := file.OpenParquetFile(strings.TrimPrefix(withLimit.NewDataFiles[0].FilePath(), "file://"), false) + require.NoError(t, err) + defer limited.Close() + assert.GreaterOrEqual(t, limited.NumRowGroups(), 10, + "100 rows with a 10-row limit need at least 10 row groups") + for rg := range limited.NumRowGroups() { + assert.LessOrEqual(t, limited.RowGroup(rg).NumRows(), int64(10)) + } + + baseline, err := table.ExecuteCompactionGroup(t.Context(), tbl, group) + require.NoError(t, err) + require.Len(t, baseline.NewDataFiles, 1) + + unlimited, err := file.OpenParquetFile(strings.TrimPrefix(baseline.NewDataFiles[0].FilePath(), "file://"), false) + require.NoError(t, err) + defer unlimited.Close() + assert.Equal(t, 1, unlimited.NumRowGroups(), + "without the option, 100 rows fit in one row group") +} + // TestExecuteCompactionGroup_ScanConcurrencyForwarded is a smoke test // confirming WithCompactionScanConcurrency is wired through without // breaking the read path. We can't easily observe scan parallelism diff --git a/table/rolling_data_writer.go b/table/rolling_data_writer.go index efef59213..e431d0f2b 100644 --- a/table/rolling_data_writer.go +++ b/table/rolling_data_writer.go @@ -43,6 +43,11 @@ import ( // streaming goroutine has already stopped. var ErrWriterClosed = errors.New("writer is closed") +// rollingDataWriterQueueCapacity is the default capacity of a rolling +// writer's record channel. Each buffered record batch is retained until +// the stream goroutine writes it, so this bound is what caps the memory +// a stalled writer can hold. Override per write with +// [WithRecordBatchBufferSize]. const rollingDataWriterQueueCapacity = 64 // writerFactory manages the creation and lifecycle of RollingDataWriter instances @@ -76,6 +81,10 @@ type writerFactory struct { shredEnabled bool shredBufferRows int + // recordBufferSize overrides the rolling writers' record channel + // capacity; non-positive uses rollingDataWriterQueueCapacity. + recordBufferSize int + writers sync.Map partitionLocProviders sync.Map nextCount func() (int, bool) @@ -194,23 +203,24 @@ func newWriterFactory(rootLocation string, args recordWritingArgs, meta *Metadat } f := &writerFactory{ - rootLocation: rootLocation, - rootURL: rootURL, - fs: args.fs, - writeUUID: args.writeUUID, - taskSchema: taskSchema, - targetFileSize: targetFileSize, - locProvider: locProvider, - tableProps: meta.props, - fileSchema: fileSchema, - arrowSchema: arrowSchema, - writeProps: format.GetWriteProperties(meta.props), - rowGroupBytes: rowGroupTargetSizeBytes, - currentSpec: *currentSpec, - fileFormat: fileFormat, - format: format, - nextCount: nextCount, - stopCount: stopCount, + rootLocation: rootLocation, + rootURL: rootURL, + fs: args.fs, + writeUUID: args.writeUUID, + taskSchema: taskSchema, + targetFileSize: targetFileSize, + locProvider: locProvider, + tableProps: meta.props, + fileSchema: fileSchema, + arrowSchema: arrowSchema, + writeProps: format.GetWriteProperties(meta.props), + rowGroupBytes: rowGroupTargetSizeBytes, + currentSpec: *currentSpec, + fileFormat: fileFormat, + format: format, + nextCount: nextCount, + stopCount: stopCount, + recordBufferSize: args.recordBatchBufferSize, } for _, apply := range opts { if err := apply(f); err != nil { @@ -353,13 +363,21 @@ type RollingDataWriter struct { noMoreSends bool } +func (w *writerFactory) recordQueueCapacity() int { + if w.recordBufferSize > 0 { + return w.recordBufferSize + } + + return rollingDataWriterQueueCapacity +} + func (w *writerFactory) newRollingDataWriter(ctx context.Context, 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{ partitionKey: partition, partitionID: partitionID, - recordCh: make(chan arrow.RecordBatch, rollingDataWriterQueueCapacity), + recordCh: make(chan arrow.RecordBatch, w.recordQueueCapacity()), errorCh: make(chan error, 1), factory: w, partitionValues: partitionValues, diff --git a/table/scanner.go b/table/scanner.go index c30c76ed8..3552bafc3 100644 --- a/table/scanner.go +++ b/table/scanner.go @@ -497,6 +497,13 @@ type Scan struct { concurrency int + // arrowBatchSize, when positive, caps the rows decoded per Arrow + // record batch, overriding the table's read.parquet.batch-size + // property. Set via WithArrowBatchSize; kept as a dedicated field + // (like limit and concurrency) so WithOptions replacing the options + // map cannot silently drop it. + arrowBatchSize int + reporter metrics.Reporter } @@ -1962,6 +1969,7 @@ func (scan *Scan) ReadTasks(ctx context.Context, tasks []FileScanTask) (*arrow.S rowLimit: scan.limit, options: scan.options, concurrency: scan.concurrency, + arrowBatchSize: scan.arrowBatchSize, }).GetRecords(ctx, readTasks) if err != nil { // No iterator to drive cleanup on a setup error, so release here. diff --git a/table/table.go b/table/table.go index 61df7c44f..62ee2d1e2 100644 --- a/table/table.go +++ b/table/table.go @@ -1304,6 +1304,24 @@ func WithRowLineage() ScanOption { } } +// WithArrowBatchSize caps the number of rows decoded per Arrow record +// batch when reading data files, overriding the table's +// read.parquet.batch-size property for this scan. Smaller batches bound +// the memory a scan holds per decoded batch, which matters when the +// consumer buffers batches (e.g. a compaction's read+write pipeline). +// The cap is stored on the scan itself rather than in the options map, +// so it applies regardless of ordering relative to [WithOptions]. A +// non-positive value is ignored. +func WithArrowBatchSize(n int) ScanOption { + if n <= 0 { + return noopOption + } + + return func(scan *Scan) { + scan.arrowBatchSize = n + } +} + func (t Table) Scan(opts ...ScanOption) *Scan { s := &Scan{ identifier: slices.Clone(t.identifier), diff --git a/table/write_read_tuning_test.go b/table/write_read_tuning_test.go new file mode 100644 index 000000000..4ab715e1f --- /dev/null +++ b/table/write_read_tuning_test.go @@ -0,0 +1,247 @@ +// 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 ( + "fmt" + "path/filepath" + "strings" + "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/arrow-go/v18/parquet/file" + "github.com/apache/iceberg-go" + iceio "github.com/apache/iceberg-go/io" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func manyRowsJSON(n int) string { + var sb strings.Builder + sb.WriteString("[") + for i := range n { + if i > 0 { + sb.WriteString(",") + } + fmt.Fprintf(&sb, `{"id":%d,"data":"row-%d"}`, i, i) + } + sb.WriteString("]") + + return sb.String() +} + +func TestWithArrowBatchSizeCapsDecodedBatches(t *testing.T) { + const numRows = 100 + const batchSize = 7 + tbl := buildV3TableWithRows(t, manyRowsJSON(numRows)) + + _, records, err := tbl.Scan(WithArrowBatchSize(batchSize)).ToArrowRecords(t.Context()) + require.NoError(t, err) + + var totalRows, batches int64 + for rec, err := range records { + require.NoError(t, err) + assert.LessOrEqual(t, rec.NumRows(), int64(batchSize)) + totalRows += rec.NumRows() + batches++ + rec.Release() + } + assert.Equal(t, int64(numRows), totalRows) + assert.Greater(t, batches, int64(numRows/batchSize)) + + // Control: the default batch size returns all rows in one batch. + _, records, err = tbl.Scan().ToArrowRecords(t.Context()) + require.NoError(t, err) + batches = 0 + for rec, err := range records { + require.NoError(t, err) + batches++ + rec.Release() + } + assert.Equal(t, int64(1), batches) +} + +func TestWithArrowBatchSizeIgnoresNonPositive(t *testing.T) { + tbl := buildV3TableWithRows(t, manyRowsJSON(3)) + + scan := tbl.Scan(WithArrowBatchSize(0), WithArrowBatchSize(-5)) + assert.Zero(t, scan.arrowBatchSize) +} + +func TestWithArrowBatchSizeSurvivesWithOptionsOrdering(t *testing.T) { + const numRows = 40 + const batchSize = 9 + tbl := buildV3TableWithRows(t, manyRowsJSON(numRows)) + + callerOpts := iceberg.Properties{"include_empty_files": "true"} + scan := tbl.Scan(WithArrowBatchSize(batchSize), WithOptions(callerOpts)) + assert.Equal(t, batchSize, scan.arrowBatchSize) + assert.Empty(t, callerOpts[ParquetBatchSizeKey]) + assert.Equal(t, "true", scan.options.Get("include_empty_files", "")) + + // The cap must apply even though WithOptions replaced the options + // map after WithArrowBatchSize ran. + _, records, err := scan.ToArrowRecords(t.Context()) + require.NoError(t, err) + var totalRows, batches int64 + for rec, err := range records { + require.NoError(t, err) + assert.LessOrEqual(t, rec.NumRows(), int64(batchSize)) + totalRows += rec.NumRows() + batches++ + rec.Release() + } + assert.Equal(t, int64(numRows), totalRows) + assert.Greater(t, batches, int64(1)) +} + +func TestRecordQueueCapacityDefaultAndOverride(t *testing.T) { + f := &writerFactory{} + assert.Equal(t, rollingDataWriterQueueCapacity, f.recordQueueCapacity()) + + f.recordBufferSize = 8 + assert.Equal(t, 8, f.recordQueueCapacity()) + + f.recordBufferSize = -1 + assert.Equal(t, rollingDataWriterQueueCapacity, f.recordQueueCapacity()) +} + +func newTuningTestWriterFactory(t *testing.T, bufferSize int) *writerFactory { + t.Helper() + + loc := filepath.ToSlash(t.TempDir()) + schema := simpleSchema() + meta, err := NewMetadata(schema, iceberg.UnpartitionedSpec, UnsortedSortOrder, loc, nil) + require.NoError(t, err) + builder, err := MetadataBuilderFromBase(meta, "") + require.NoError(t, err) + + writeUUID := uuid.New() + factory, err := newWriterFactory(loc, recordWritingArgs{ + fs: iceio.LocalFS{}, + writeUUID: &writeUUID, + counter: func(yield func(int) bool) { + for i := 0; ; i++ { + if !yield(i) { + return + } + } + }, + recordBatchBufferSize: bufferSize, + }, builder, schema, 1024*1024) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, factory.closeAll()) + }) + + return factory +} + +// TestRecordBatchBufferSizeReachesRollingWriter pins the full chain from +// recordWritingArgs through the writer factory to the rolling writer's +// channel capacity, which is the bound the option exists to control. +func TestRecordBatchBufferSizeReachesRollingWriter(t *testing.T) { + outputCh := make(chan iceberg.DataFile, 1) + + overridden := newTuningTestWriterFactory(t, 3) + w := overridden.newRollingDataWriter(t.Context(), "", nil, outputCh) + assert.Equal(t, 3, cap(w.recordCh)) + w.abortAndWait() + + defaulted := newTuningTestWriterFactory(t, 0) + dw := defaulted.newRollingDataWriter(t.Context(), "", nil, outputCh) + assert.Equal(t, rollingDataWriterQueueCapacity, cap(dw.recordCh)) + dw.abortAndWait() +} + +func TestWriteRecordTuningOptions(t *testing.T) { + var cfg writeRecordConfig + WithRecordBatchBufferSize(16)(&cfg) + WithParquetRowGroupLimit(1000)(&cfg) + assert.Equal(t, 16, cfg.recordBatchBufferSize) + assert.Equal(t, 1000, cfg.parquetRowGroupLimit) + + WithRecordBatchBufferSize(0)(&cfg) + WithParquetRowGroupLimit(-1)(&cfg) + assert.Equal(t, 16, cfg.recordBatchBufferSize, "non-positive values are ignored") + assert.Equal(t, 1000, cfg.parquetRowGroupLimit) +} + +func TestCompactionGroupTuningOptions(t *testing.T) { + var cfg compactionGroupConfig + for _, opt := range []CompactionGroupOption{ + WithCompactionArrowBatchSize(1024), + WithCompactionRecordBatchBufferSize(4), + WithCompactionParquetRowGroupLimit(500), + } { + opt(&cfg) + } + assert.Equal(t, 1024, cfg.arrowBatchSize) + assert.Equal(t, 4, cfg.recordBatchBufferSize) + assert.Equal(t, 500, cfg.parquetRowGroupLimit) + + for _, opt := range []CompactionGroupOption{ + WithCompactionArrowBatchSize(0), + WithCompactionRecordBatchBufferSize(-2), + WithCompactionParquetRowGroupLimit(0), + } { + opt(&cfg) + } + assert.Equal(t, 1024, cfg.arrowBatchSize, "non-positive values are ignored") + assert.Equal(t, 4, cfg.recordBatchBufferSize) + assert.Equal(t, 500, cfg.parquetRowGroupLimit) +} + +func TestWithParquetRowGroupLimitBoundsRowGroups(t *testing.T) { + tbl := buildV3TableWithRows(t, `[{"id":0,"data":"seed"}]`) + + arrowSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}, + {Name: "data", Type: arrow.BinaryTypes.String, Nullable: true}, + }, nil) + const numRows = 100 + data, err := array.TableFromJSON(memory.DefaultAllocator, arrowSchema, []string{manyRowsJSON(numRows)}) + require.NoError(t, err) + defer data.Release() + + rdr := array.NewTableReader(data, numRows) + defer rdr.Release() + + var paths []string + for df, err := range WriteRecords(t.Context(), tbl, arrowSchema, + array.IterFromReader(rdr), WithParquetRowGroupLimit(10)) { + require.NoError(t, err) + paths = append(paths, df.FilePath()) + } + require.NotEmpty(t, paths) + + var rowGroups int + for _, path := range paths { + pf, err := file.OpenParquetFile(strings.TrimPrefix(path, "file://"), false) + require.NoError(t, err) + for rg := range pf.NumRowGroups() { + assert.LessOrEqual(t, pf.RowGroup(rg).NumRows(), int64(10)) + } + rowGroups += pf.NumRowGroups() + require.NoError(t, pf.Close()) + } + assert.GreaterOrEqual(t, rowGroups, numRows/10) +} diff --git a/table/write_records.go b/table/write_records.go index 83c2936c7..c50360fee 100644 --- a/table/write_records.go +++ b/table/write_records.go @@ -34,12 +34,14 @@ import ( type WriteRecordOption func(*writeRecordConfig) type writeRecordConfig struct { - targetFileSize int64 - writeUUID *uuid.UUID - maxWriteWorkers int - clustered bool - fileSchema *iceberg.Schema - preserveRowLineage bool + targetFileSize int64 + writeUUID *uuid.UUID + maxWriteWorkers int + clustered bool + fileSchema *iceberg.Schema + preserveRowLineage bool + recordBatchBufferSize int + parquetRowGroupLimit int } // WithTargetFileSize overrides the table's default target file size. @@ -95,6 +97,32 @@ func WithClusteredWrite() WriteRecordOption { } } +// WithRecordBatchBufferSize sets the capacity, in record batches, of +// each rolling data writer's input channel. Every buffered batch is +// retained in memory until its writer consumes it, so this bound times +// the batch row count caps the memory a stalled writer can hold. The +// default is 64 batches. A non-positive value is ignored. +func WithRecordBatchBufferSize(n int) WriteRecordOption { + return func(c *writeRecordConfig) { + if n > 0 { + c.recordBatchBufferSize = n + } + } +} + +// WithParquetRowGroupLimit overrides the table's +// write.parquet.row-group-limit property for this write, capping the +// number of rows per Parquet row group in the output files. Smaller row +// groups bound the writer's buffered memory before each flush. A +// non-positive value is ignored. +func WithParquetRowGroupLimit(n int) WriteRecordOption { + return func(c *writeRecordConfig) { + if n > 0 { + c.parquetRowGroupLimit = n + } + } +} + // WithPreserveRowLineage sets the output file schema to include the v3 row- // lineage metadata columns (_row_id, _last_updated_sequence_number) so that // row identity is preserved through rewrites and compactions. The input @@ -194,6 +222,13 @@ func WriteRecords(ctx context.Context, tbl *Table, meta.props[WriteTargetFileSizeBytesKey] = strconv.FormatInt(cfg.targetFileSize, 10) } + if cfg.parquetRowGroupLimit > 0 { + if meta.props == nil { + meta.props = make(iceberg.Properties) + } + meta.props[ParquetRowGroupLimitKey] = strconv.Itoa(cfg.parquetRowGroupLimit) + } + releasing := func(yield func(arrow.RecordBatch, error) bool) { for rec, err := range records { if err != nil { @@ -211,12 +246,13 @@ func WriteRecords(ctx context.Context, tbl *Table, } args := recordWritingArgs{ - sc: schema, - itr: releasing, - fs: writeFS, - writeUUID: cfg.writeUUID, - maxWriteWorkers: cfg.maxWriteWorkers, - clustered: cfg.clustered, + sc: schema, + itr: releasing, + fs: writeFS, + writeUUID: cfg.writeUUID, + maxWriteWorkers: cfg.maxWriteWorkers, + clustered: cfg.clustered, + recordBatchBufferSize: cfg.recordBatchBufferSize, } if cfg.fileSchema != nil {