perf(table): run compaction groups concurrently in RewriteDataFiles

data_file_codec_test.go

@@ -194,6 +194,37 @@ func TestMarshalAvroEntryDecimalPartitionRoundTrip(t *testing.T) {
 	require.True(t, got.Equals(DecimalLiteral(want)))
 }
 
+func TestMarshalAvroEntryTimestampPartitionRoundTrip(t *testing.T) {
+	// Micro cases run first.
+	// A nano type sharing their schema-cache key would decode as Timestamp instead of TimestampNano.
+	for _, tc := range []struct {
+		typ  Type
+		want any
+	}{
+		{TimestampType{}, Timestamp(1_700_000_000_000_000)},
+		{TimestampTzType{}, Timestamp(1_700_000_000_000_000)},
+		{TimestampNsType{}, TimestampNano(1_700_000_000_000_000_123)},
+		{TimestampTzNsType{}, TimestampNano(1_700_000_000_000_000_123)},
+	} {
+		t.Run(tc.typ.String(), func(t *testing.T) {
+			schema := NewSchema(0, NestedField{ID: 1, Name: "ts", Type: tc.typ})
+			spec := NewPartitionSpecID(1, PartitionField{SourceIDs: []int{1}, FieldID: 1000, Name: "ts", Transform: IdentityTransform{}})
+
+			path := "s3://bucket/ns/tbl/data/ts.parquet"
+			builder, err := NewDataFileBuilder(spec, EntryContentData, path, ParquetFile, map[int]any{1000: tc.want}, nil, nil, 1, 1024)
+			require.NoError(t, err)
+			df, ok := builder.Build().(*dataFile)
+			require.True(t, ok)
+
+			encoded, err := df.MarshalAvroEntry(spec, schema, 3)
+			require.NoError(t, err)
+			decoded, err := unmarshalAvroDataFileEntry(encoded, spec, schema, 3)
+			require.NoError(t, err)
+			require.Equal(t, tc.want, decoded.Partition()[1000])
+		})
+	}
+}
+
 // snapshotAvroFields returns a deep copy of every avro-tagged field on
 // d, keyed by field name. Slices, maps, byte arrays, and pointer
 // targets are reconstructed so the snapshot is fully independent of d
@@ -357,6 +388,8 @@ func TestManifestEntrySchemaForMatchesPartitionAvroShape(t *testing.T) {
 		TimeType{},
 		TimestampType{},
 		TimestampTzType{},
+		TimestampNsType{},
+		TimestampTzNsType{},
 		UUIDType{},
 		BooleanType{},
 		BinaryType{},

table/rewrite_data_files.go

@@ -28,6 +28,7 @@ import (
 	"github.com/apache/iceberg-go"
 	iceberginternal "github.com/apache/iceberg-go/internal"
 	iceio "github.com/apache/iceberg-go/io"
+	"golang.org/x/sync/errgroup"
 )
 
 // RewriteResult summarizes a completed compaction.
@@ -217,6 +218,30 @@ type RewriteDataFilesOptions struct {
 	// size, scan concurrency). See the With* helpers returning
 	// [CompactionGroupOption].
 	GroupOptions []CompactionGroupOption
+
+	// MaxConcurrentGroups bounds how many compaction groups run at once.
+	// Zero and one both mean sequential execution, which is the default.
+	// Larger values run [ExecuteCompactionGroup] calls under a bounded
+	// errgroup and apply their results in the original group order, so
+	// manifests and [RewriteResult] are identical to a sequential run.
+	// A failure cancels the groups still running; the returned error
+	// is the lowest-index failure not caused by that cancellation, or
+	// the lowest-index failure when every group was canceled, so
+	// failures match a sequential run. Every group error is logged.
+	// Peak record-pipeline memory is MaxConcurrentGroups times the
+	// per-group bound stated on [WithCompactionArrowBatchSize]:
+	//
+	//	MaxConcurrentGroups x (workers x (rows in the largest task + n) + (recordBatchBufferSize + 2) x n)
+	//
+	// rows, where workers, n and recordBatchBufferSize are the per-group
+	// values. Multiply rows by the average row width in bytes for a byte
+	// estimate. Delete-side memory is outside this bound. File-open
+	// fan-out multiplies too: every group scans with up to
+	// [WithCompactionScanConcurrency] workers, so N groups open about N
+	// times the scan worker count in files at once. Size the two knobs
+	// together against connection and file descriptor limits. Negative
+	// values are rejected with [ErrInvalidOperation].
+	MaxConcurrentGroups int
 }
 
 // CompactionGroupOption configures a single [ExecuteCompactionGroup]
@@ -252,7 +277,10 @@ func WithCompactionTargetFileSize(size int64) CompactionGroupOption {
 // The scan runs min(n, number of tasks) workers, and each worker holds
 // the decoded batches of at most one task until the writer has taken
 // them, so the worker count multiplies the read-side term of the
-// memory bound stated on [WithCompactionArrowBatchSize].
+// memory bound stated on [WithCompactionArrowBatchSize]. It also
+// multiplies [RewriteDataFilesOptions.MaxConcurrentGroups] for file
+// opens: N groups scan with up to n workers each, so about N times n
+// files are open at once.
 func WithCompactionScanConcurrency(n int) CompactionGroupOption {
 	return func(c *compactionGroupConfig) {
 		c.scanConcurrency = n
@@ -337,6 +365,9 @@ func (t *Transaction) RewriteDataFiles(ctx context.Context, groups []CompactionT
 	if _, err := t.txnMeta(); err != nil {
 		return nil, err
 	}
+	if opts.MaxConcurrentGroups < 0 {
+		return nil, fmt.Errorf("%w: MaxConcurrentGroups must be non-negative", ErrInvalidOperation)
+	}
 	if opts.PartialProgress {
 		return t.rewriteDataFilesPartial(ctx, groups, opts)
 	}
@@ -348,31 +379,38 @@ func (t *Transaction) RewriteDataFiles(ctx context.Context, groups []CompactionT
 	rewrite := t.NewRewrite(opts.SnapshotProps)
 	stagedDeleteFiles := make(map[string]struct{})
 
-	for _, group := range groups {
-		if err := ctx.Err(); err != nil {
-			return result, err
-		}
-
-		if len(group.Tasks) == 0 {
-			continue
-		}
+	fs, err := t.tbl.fsF(ctx)
+	if err != nil {
+		return result, fmt.Errorf("open table IO for atomic rewrite: %w", err)
+	}
 
-		gr, err := ExecuteCompactionGroup(ctx, t.tbl, group, opts.GroupOptions...)
+	var applied []CompactionGroupResult
+	if opts.MaxConcurrentGroups > 1 {
+		results, err := executeCompactionGroups(ctx, t.tbl, groups, opts.GroupOptions, opts.MaxConcurrentGroups)
 		if err != nil {
-			return result, err
+			return result, cleanupAtomicRewriteOutputs(fs, results, err)
 		}
-
-		if len(gr.OldDataFiles) == 0 && len(gr.NewDataFiles) == 0 {
-			continue
+		applied = results
+		for _, gr := range results {
+			applyAtomicGroupResult(rewrite, result, stagedDeleteFiles, gr)
 		}
+	} else {
+		for _, group := range groups {
+			if err := ctx.Err(); err != nil {
+				return result, cleanupAtomicRewriteOutputs(fs, applied, err)
+			}
 
-		rewrite.ApplyResult(gr)
-		accumulateGroupMetrics(result, gr)
-		for _, df := range gr.SafePosDeletes {
-			stagedDeleteFiles[df.FilePath()] = struct{}{}
-		}
-		for _, df := range gr.SafeDeletionVectors {
-			stagedDeleteFiles[df.FilePath()] = struct{}{}
+			if len(group.Tasks) == 0 {
+				continue
+			}
+
+			gr, err := ExecuteCompactionGroup(ctx, t.tbl, group, opts.GroupOptions...)
+			if err != nil {
+				return result, cleanupAtomicRewriteOutputs(fs, append(applied, gr), err)
+			}
+			applied = append(applied, gr)
+
+			applyAtomicGroupResult(rewrite, result, stagedDeleteFiles, gr)
 		}
 	}
 
@@ -401,12 +439,96 @@ func (t *Transaction) RewriteDataFiles(ctx context.Context, groups []CompactionT
 	}
 
 	if err := rewrite.Commit(ctx); err != nil {
-		return result, fmt.Errorf("commit compaction: %w", err)
+		return result, cleanupAtomicRewriteOutputs(fs, applied, fmt.Errorf("commit compaction: %w", err))
 	}
 
 	return result, nil
 }
 
+func applyAtomicGroupResult(rewrite *RewriteFiles, result *RewriteResult, stagedDeleteFiles map[string]struct{}, gr CompactionGroupResult) {
+	if len(gr.OldDataFiles) == 0 && len(gr.NewDataFiles) == 0 {
+		return
+	}
+	rewrite.ApplyResult(gr)
+	accumulateGroupMetrics(result, gr)
+	for _, df := range gr.SafePosDeletes {
+		stagedDeleteFiles[df.FilePath()] = struct{}{}
+	}
+	for _, df := range gr.SafeDeletionVectors {
+		stagedDeleteFiles[df.FilePath()] = struct{}{}
+	}
+}
+
+func cleanupAtomicRewriteOutputs(fs iceio.IO, results []CompactionGroupResult, cause error) error {
+	if err := cleanupCompactionOutputs(fs, results); err != nil {
+		return errors.Join(cause, fmt.Errorf("clean up atomic rewrite outputs: %w", err))
+	}
+
+	return cause
+}
+
+func executeCompactionGroups(ctx context.Context, tbl *Table, groups []CompactionTaskGroup, groupOpts []CompactionGroupOption, maxConcurrentGroups int) ([]CompactionGroupResult, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	limit := min(maxConcurrentGroups, len(groups))
+	if limit < 1 {
+		limit = 1
+	}
+	var g errgroup.Group
+	g.SetLimit(limit)
+	runCtx, cancelRuns := context.WithCancel(ctx)
+	defer cancelRuns()
+	results := make([]CompactionGroupResult, len(groups))
+	groupErrs := make([]error, len(groups))
+	for i, group := range groups {
+		if len(group.Tasks) == 0 {
+			continue
+		}
+		g.Go(func() error {
+			gr, err := ExecuteCompactionGroup(runCtx, tbl, group, groupOpts...)
+			results[i] = gr
+			groupErrs[i] = err
+			if err != nil {
+				slog.Warn("compaction group failed", "index", i, "err", err)
+				cancelRuns()
+			}
+
+			return err
+		})
+	}
+	if err := g.Wait(); err != nil {
+		var firstErr, firstNonContextErr error
+		for _, groupErr := range groupErrs {
+			if groupErr == nil {
+				continue
+			}
+			if firstErr == nil {
+				firstErr = groupErr
+			}
+			if !errors.Is(groupErr, context.Canceled) && !errors.Is(groupErr, context.DeadlineExceeded) {
+				firstNonContextErr = groupErr
+
+				break
+			}
+		}
+		selected := firstNonContextErr
+		if selected == nil {
+			selected = firstErr
+		}
+		if selected == nil {
+			selected = err
+		}
+		if ctxErr := ctx.Err(); ctxErr != nil && (errors.Is(selected, context.Canceled) || errors.Is(selected, context.DeadlineExceeded)) {
+			return results, ctxErr
+		}
+
+		return results, selected
+	}
+
+	return results, nil
+}
+
 // ExecuteCompactionGroup reads a compaction group's tasks (with
 // deletes applied), writes consolidated output files via
 // [WriteRecords], and computes the position-delete files safe to
@@ -639,26 +761,26 @@ func (t *Transaction) rewriteDataFilesPartial(ctx context.Context, groups []Comp
 			return cause
 		}
 
-		for _, group := range batchGroups {
-			if err := ctx.Err(); err != nil {
-				return result, cleanupBatch(err)
-			}
-
-			gr, err := ExecuteCompactionGroup(ctx, current, group, opts.GroupOptions...)
+		if opts.MaxConcurrentGroups > 1 {
+			results, err := executeCompactionGroups(ctx, current, batchGroups, opts.GroupOptions, opts.MaxConcurrentGroups)
 			if err != nil {
-				return result, cleanupBatch(err, gr)
+				return result, cleanupBatch(err, results...)
 			}
-
-			if len(gr.OldDataFiles) == 0 && len(gr.NewDataFiles) == 0 {
-				continue
+			for _, gr := range results {
+				batchResults, rewrittenFiles = appendPartialGroupResult(batchResults, rewrittenPaths, rewrittenFiles, gr)
 			}
-			batchResults = append(batchResults, gr)
-			for _, df := range gr.OldDataFiles {
-				if _, ok := rewrittenPaths[df.FilePath()]; ok {
-					continue
+		} else {
+			for _, group := range batchGroups {
+				if err := ctx.Err(); err != nil {
+					return result, cleanupBatch(err)
 				}
-				rewrittenPaths[df.FilePath()] = struct{}{}
-				rewrittenFiles = append(rewrittenFiles, df)
+
+				gr, err := ExecuteCompactionGroup(ctx, current, group, opts.GroupOptions...)
+				if err != nil {
+					return result, cleanupBatch(err, gr)
+				}
+
+				batchResults, rewrittenFiles = appendPartialGroupResult(batchResults, rewrittenPaths, rewrittenFiles, gr)
 			}
 		}
 
@@ -752,6 +874,22 @@ func (t *Transaction) rewriteDataFilesPartial(ctx context.Context, groups []Comp
 	return result, nil
 }
 
+func appendPartialGroupResult(batchResults []CompactionGroupResult, rewrittenPaths map[string]struct{}, rewrittenFiles []iceberg.DataFile, gr CompactionGroupResult) ([]CompactionGroupResult, []iceberg.DataFile) {
+	if len(gr.OldDataFiles) == 0 && len(gr.NewDataFiles) == 0 {
+		return batchResults, rewrittenFiles
+	}
+	batchResults = append(batchResults, gr)
+	for _, df := range gr.OldDataFiles {
+		if _, ok := rewrittenPaths[df.FilePath()]; ok {
+			continue
+		}
+		rewrittenPaths[df.FilePath()] = struct{}{}
+		rewrittenFiles = append(rewrittenFiles, df)
+	}
+
+	return batchResults, rewrittenFiles
+}
+
 func recordCommittedRewriteBatch(
 	result *RewriteResult,
 	next *Table,