Add tunables bounding compaction pipeline memory

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{

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

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)

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

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,

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.

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),