diff --git a/data_file_codec_test.go b/data_file_codec_test.go index 7ca765849..3830af534 100644 --- a/data_file_codec_test.go +++ b/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{}, diff --git a/table/rewrite_data_files.go b/table/rewrite_data_files.go index a66e90da2..cfa9da606 100644 --- a/table/rewrite_data_files.go +++ b/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, diff --git a/table/rewrite_data_files_bench_test.go b/table/rewrite_data_files_bench_test.go new file mode 100644 index 000000000..99dd8ae88 --- /dev/null +++ b/table/rewrite_data_files_bench_test.go @@ -0,0 +1,250 @@ +// 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_test + +import ( + "context" + "fmt" + "io/fs" + "math/rand/v2" + "os" + "path/filepath" + "slices" + "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" + "github.com/apache/arrow-go/v18/parquet/pqarrow" + "github.com/apache/iceberg-go" + iceio "github.com/apache/iceberg-go/io" + "github.com/apache/iceberg-go/table" + "github.com/stretchr/testify/require" +) + +const ( + groupConcPartitions = 8 + groupConcFilesPerPart = 8 + groupConcRowsPerFile = 30000 + groupConcPayloadWords = 6 + groupConcPartitionField = 1000 +) + +func newGroupConcTable(tb testing.TB) *table.Table { + tb.Helper() + + location := filepath.ToSlash(tb.TempDir()) + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String, Required: false}, + iceberg.NestedField{ID: 3, Name: "payload", Type: iceberg.PrimitiveTypes.String, Required: false}, + iceberg.NestedField{ID: 4, Name: "score", Type: iceberg.PrimitiveTypes.Float64, Required: false}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{2}, FieldID: groupConcPartitionField, Transform: iceberg.IdentityTransform{}, Name: "data", + }) + meta, err := table.NewMetadata(schema, &spec, table.UnsortedSortOrder, location, + iceberg.Properties{table.PropertyFormatVersion: "2"}) + require.NoError(tb, err) + + cat := &partialProgressCatalog{metadata: meta} + + return table.New( + table.Identifier{"db", "group_concurrency_bench"}, + meta, location+"/metadata/v1.metadata.json", + func(context.Context) (iceio.IO, error) { return iceio.LocalFS{}, nil }, + cat, + ) +} + +func writeGroupConcFile(tb testing.TB, path string, sc *arrow.Schema, partition string, fileIdx int) int64 { + tb.Helper() + + mem := memory.DefaultAllocator + rng := rand.New(rand.NewPCG(uint64(fileIdx), 0x9e3779b97f4a7c15)) + + idB := array.NewInt64Builder(mem) + dataB := array.NewStringBuilder(mem) + payloadB := array.NewStringBuilder(mem) + scoreB := array.NewFloat64Builder(mem) + defer idB.Release() + defer dataB.Release() + defer payloadB.Release() + defer scoreB.Release() + + idB.Reserve(groupConcRowsPerFile) + dataB.Reserve(groupConcRowsPerFile) + payloadB.Reserve(groupConcRowsPerFile) + scoreB.Reserve(groupConcRowsPerFile) + + var scratch [groupConcPayloadWords]uint64 + for i := range groupConcRowsPerFile { + id := int64(fileIdx*groupConcRowsPerFile + i) + idB.Append(id) + dataB.Append(partition) + for w := range groupConcPayloadWords { + scratch[w] = rng.Uint64() + } + payloadB.Append(fmt.Sprintf("%016x%016x%016x%016x%016x%016x", + scratch[0], scratch[1], scratch[2], scratch[3], scratch[4], scratch[5])) + scoreB.Append(float64(id) * 0.5) + } + + rec := array.NewRecordBatch(sc, []arrow.Array{ + idB.NewArray(), dataB.NewArray(), payloadB.NewArray(), scoreB.NewArray(), + }, int64(groupConcRowsPerFile)) + defer rec.Release() + + fs := iceio.LocalFS{} + fw, err := fs.Create(path) + require.NoError(tb, err) + defer fw.Close() + + arrTable := array.NewTableFromRecords(sc, []arrow.RecordBatch{rec}) + defer arrTable.Release() + + props := parquet.NewWriterProperties(parquet.WithStats(true)) + require.NoError(tb, pqarrow.WriteTable(arrTable, fw, int64(groupConcRowsPerFile), props, pqarrow.DefaultWriterProps())) + + info, err := os.Stat(path) + require.NoError(tb, err) + + return info.Size() +} + +func planGroupConcGroups(tb testing.TB, tbl *table.Table) []table.CompactionTaskGroup { + tb.Helper() + + tasks, err := tbl.Scan().PlanFiles(context.Background()) + require.NoError(tb, err) + require.Len(tb, tasks, groupConcPartitions*groupConcFilesPerPart) + + byPart := make(map[string][]table.FileScanTask, groupConcPartitions) + for _, task := range tasks { + part, ok := task.File.Partition()[groupConcPartitionField].(string) + require.True(tb, ok) + byPart[part] = append(byPart[part], task) + } + + keys := make([]string, 0, len(byPart)) + for k := range byPart { + keys = append(keys, k) + } + slices.Sort(keys) + require.Len(tb, keys, groupConcPartitions) + + groups := make([]table.CompactionTaskGroup, 0, len(keys)) + for _, k := range keys { + require.Len(tb, byPart[k], groupConcFilesPerPart) + var total int64 + for _, task := range byPart[k] { + total += task.File.FileSizeBytes() + } + groups = append(groups, table.CompactionTaskGroup{ + PartitionKey: k, + Tasks: byPart[k], + TotalSizeBytes: total, + }) + } + + return groups +} + +func groupConcOutputPaths(tb testing.TB, location string, before map[string]struct{}) []string { + tb.Helper() + + var out []string + err := filepath.WalkDir(filepath.Join(location, "data"), func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() || filepath.Ext(path) != ".parquet" { + return nil + } + if _, ok := before[path]; !ok { + out = append(out, path) + } + + return nil + }) + require.NoError(tb, err) + + return out +} + +func BenchmarkRewriteDataFilesGroupConcurrency(b *testing.B) { + tbl := newGroupConcTable(b) + + arrowSc, err := table.SchemaToArrowSchema(tbl.Schema(), nil, false, false) + require.NoError(b, err) + + ctx := context.Background() + for p := range groupConcPartitions { + partition := fmt.Sprintf("p%d", p) + files := make([]iceberg.DataFile, 0, groupConcFilesPerPart) + for f := range groupConcFilesPerPart { + fileIdx := p*groupConcFilesPerPart + f + dataPath := tbl.Location() + "/data/" + fmt.Sprintf("p%d-file-%d.parquet", p, f) + size := writeGroupConcFile(b, dataPath, arrowSc, partition, fileIdx) + builder, err := iceberg.NewDataFileBuilder( + tbl.Spec(), iceberg.EntryContentData, dataPath, iceberg.ParquetFile, + map[int]any{groupConcPartitionField: partition}, nil, nil, groupConcRowsPerFile, size) + require.NoError(b, err) + files = append(files, builder.Build()) + } + txn := tbl.NewTransaction() + require.NoError(b, txn.AddDataFiles(ctx, files, nil)) + tbl, err = txn.Commit(ctx) + require.NoError(b, err) + } + + groups := planGroupConcGroups(b, tbl) + totalRows := int64(groupConcPartitions * groupConcFilesPerPart * groupConcRowsPerFile) + + for _, maxConcurrentGroups := range []int{1, 2, 4, 8} { + b.Run(fmt.Sprintf("MaxConcurrentGroups=%d", maxConcurrentGroups), func(b *testing.B) { + before, err := filepath.Glob(filepath.Join(tbl.Location(), "data", "*.parquet")) + require.NoError(b, err) + inputs := make(map[string]struct{}, len(before)) + for _, p := range before { + inputs[p] = struct{}{} + } + opts := table.RewriteDataFilesOptions{MaxConcurrentGroups: maxConcurrentGroups} + + b.ReportAllocs() + b.ResetTimer() + var completed int64 + for b.Loop() { + tx := tbl.NewTransaction() + result, err := tx.RewriteDataFiles(ctx, groups, opts) + require.NoError(b, err) + require.Equal(b, groupConcPartitions, result.RewrittenGroups) + completed++ + + b.StopTimer() + for _, p := range groupConcOutputPaths(b, tbl.Location(), inputs) { + require.NoError(b, os.Remove(p)) + } + b.StartTimer() + } + b.StopTimer() + b.ReportMetric(float64(totalRows*completed)/b.Elapsed().Seconds(), "rows/s") + }) + } +} diff --git a/table/rewrite_data_files_test.go b/table/rewrite_data_files_test.go index d9bc9dc6d..eae48b51a 100644 --- a/table/rewrite_data_files_test.go +++ b/table/rewrite_data_files_test.go @@ -18,13 +18,19 @@ package table_test import ( + "bytes" "context" "errors" "fmt" + "io/fs" + "log/slog" "os" "path/filepath" + "slices" "strings" + "sync" "testing" + "time" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" @@ -843,6 +849,25 @@ func parquetFiles(t *testing.T, location string) []string { return paths } +func allParquetFiles(t *testing.T, location string) []string { + t.Helper() + + var paths []string + err := filepath.WalkDir(filepath.Join(location, "data"), func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if !d.IsDir() && strings.HasSuffix(path, ".parquet") { + paths = append(paths, filepath.ToSlash(path)) + } + + return nil + }) + require.NoError(t, err) + + return paths +} + func newPartialProgressPartitionedTable(t *testing.T) *table.Table { t.Helper() @@ -1401,3 +1426,749 @@ func appendEqualityDelete(t *testing.T, tbl *table.Table, equalityFieldIDs []int return out } + +func newMaxConcPartitionedTable(t *testing.T, fs iceio.IO) *table.Table { + t.Helper() + + return newMaxConcPartitionedTableWithFSF(t, func(context.Context) (iceio.IO, error) { return fs, nil }) +} + +func newMaxConcPartitionedTableWithFSF(t *testing.T, fsF func(context.Context) (iceio.IO, error)) *table.Table { + t.Helper() + + location := filepath.ToSlash(t.TempDir()) + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String, Required: false}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{2}, FieldID: 1000, Transform: iceberg.IdentityTransform{}, Name: "data", + }) + meta, err := table.NewMetadata(schema, &spec, table.UnsortedSortOrder, location, + iceberg.Properties{table.PropertyFormatVersion: "2"}) + require.NoError(t, err) + + cat := &partialProgressCatalog{metadata: meta} + + return table.New( + table.Identifier{"db", "max_conc_test"}, + meta, location+"/metadata/v1.metadata.json", + fsF, + cat, + ) +} + +func addMaxConcPartitions(t *testing.T, tbl *table.Table, partitions, filesPerPartition, rowsPerFile int) *table.Table { + t.Helper() + + var nextID int64 = 1 + for p := range partitions { + partition := fmt.Sprintf("p%d", p) + for f := range filesPerPartition { + ids := make([]int64, rowsPerFile) + for r := range rowsPerFile { + ids[r] = nextID + nextID++ + } + tbl = addPartitionedRowsOnRef(t, tbl, table.MainBranch, fmt.Sprintf("p%d-%d", p, f), partition, ids...) + } + } + + return tbl +} + +func groupsByPartition(t *testing.T, tbl *table.Table) []table.CompactionTaskGroup { + t.Helper() + + tasks, err := tbl.Scan().PlanFiles(t.Context()) + require.NoError(t, err) + + byPart := make(map[string][]table.FileScanTask) + for _, task := range tasks { + part, ok := task.File.Partition()[1000].(string) + require.True(t, ok) + byPart[part] = append(byPart[part], task) + } + keys := make([]string, 0, len(byPart)) + for k := range byPart { + keys = append(keys, k) + } + slices.Sort(keys) + + groups := make([]table.CompactionTaskGroup, 0, len(keys)) + for _, k := range keys { + var total int64 + for _, task := range byPart[k] { + total += task.File.FileSizeBytes() + } + groups = append(groups, table.CompactionTaskGroup{ + PartitionKey: k, + Tasks: byPart[k], + TotalSizeBytes: total, + }) + } + + return groups +} + +func idsByPartitionValue(t *testing.T, tbl *table.Table) map[string][]int64 { + t.Helper() + + _, itr, err := tbl.Scan().ToArrowRecords(t.Context()) + require.NoError(t, err) + + out := make(map[string][]int64) + for rec, err := range itr { + require.NoError(t, err) + dataIdx := rec.Schema().FieldIndices("data") + require.NotEmpty(t, dataIdx) + idIdx := rec.Schema().FieldIndices("id") + require.NotEmpty(t, idIdx) + dataCol, ok := rec.Column(dataIdx[0]).(*array.String) + require.True(t, ok) + idCol, ok := rec.Column(idIdx[0]).(*array.Int64) + require.True(t, ok) + for i := range int(rec.NumRows()) { + out[dataCol.Value(i)] = append(out[dataCol.Value(i)], idCol.Value(i)) + } + rec.Release() + } + + return out +} + +func manifestDataPartitions(t *testing.T, tbl *table.Table) []string { + t.Helper() + + snap := tbl.CurrentSnapshot() + require.NotNil(t, snap) + fs, err := tbl.FS(t.Context()) + require.NoError(t, err) + manifests, err := snap.Manifests(fs) + require.NoError(t, err) + + var parts []string + for _, m := range manifests { + for e, err := range m.Entries(fs, false) { + require.NoError(t, err) + if e.Status() == iceberg.EntryStatusDELETED { + continue + } + df := e.DataFile() + if df.ContentType() != iceberg.EntryContentData { + continue + } + part, ok := df.Partition()[1000].(string) + require.True(t, ok) + parts = append(parts, part) + } + } + + return parts +} + +func manifestLiveDataPaths(t *testing.T, tbl *table.Table) []string { + t.Helper() + + snap := tbl.CurrentSnapshot() + require.NotNil(t, snap) + fs, err := tbl.FS(t.Context()) + require.NoError(t, err) + manifests, err := snap.Manifests(fs) + require.NoError(t, err) + + var paths []string + for _, m := range manifests { + for e, err := range m.Entries(fs, false) { + require.NoError(t, err) + if e.Status() == iceberg.EntryStatusDELETED { + continue + } + df := e.DataFile() + if df.ContentType() != iceberg.EntryContentData { + continue + } + paths = append(paths, df.FilePath()) + } + } + + return paths +} + +func TestRewriteDataFiles_MaxConcurrentGroupsMatchesSequential(t *testing.T) { + tblSeq := newMaxConcPartitionedTable(t, iceio.LocalFS{}) + tblSeq = addMaxConcPartitions(t, tblSeq, 8, 2, 5) + tblConc := newMaxConcPartitionedTable(t, iceio.LocalFS{}) + tblConc = addMaxConcPartitions(t, tblConc, 8, 2, 5) + + groupsSeq := groupsByPartition(t, tblSeq) + groupsConc := groupsByPartition(t, tblConc) + require.Len(t, groupsSeq, 8) + require.Len(t, groupsConc, 8) + + txSeq := tblSeq.NewTransaction() + resSeq, err := txSeq.RewriteDataFiles(t.Context(), groupsSeq, table.RewriteDataFilesOptions{}) + require.NoError(t, err) + committedSeq, err := txSeq.Commit(t.Context()) + require.NoError(t, err) + + txConc := tblConc.NewTransaction() + resConc, err := txConc.RewriteDataFiles(t.Context(), groupsConc, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + require.NoError(t, err) + committedConc, err := txConc.Commit(t.Context()) + require.NoError(t, err) + + assert.Equal(t, resSeq.RewrittenGroups, resConc.RewrittenGroups) + assert.Equal(t, resSeq.AddedDataFiles, resConc.AddedDataFiles) + assert.Equal(t, resSeq.RemovedDataFiles, resConc.RemovedDataFiles) + assert.Equal(t, resSeq.RemovedPositionDeleteFiles, resConc.RemovedPositionDeleteFiles) + assert.Equal(t, resSeq.RemovedEqualityDeleteFiles, resConc.RemovedEqualityDeleteFiles) + assert.Equal(t, resSeq.RemovedDeletionVectorFiles, resConc.RemovedDeletionVectorFiles) + assert.Equal(t, resSeq.BytesBefore, resConc.BytesBefore) + assert.Equal(t, 8, resConc.RewrittenGroups) + assert.Equal(t, 16, resConc.RemovedDataFiles) + assert.Equal(t, 8, resConc.AddedDataFiles) + + idsSeq := idsByPartitionValue(t, committedSeq) + idsConc := idsByPartitionValue(t, committedConc) + require.Len(t, idsConc, 8) + for p := range 8 { + key := fmt.Sprintf("p%d", p) + assert.ElementsMatch(t, idsSeq[key], idsConc[key]) + assert.Len(t, idsConc[key], 10) + } + + paths := manifestLiveDataPaths(t, committedConc) + require.Len(t, paths, 8) + assert.Len(t, map[string]struct{}{paths[0]: {}, paths[1]: {}, paths[2]: {}, paths[3]: {}, paths[4]: {}, paths[5]: {}, paths[6]: {}, paths[7]: {}}, 8) + onDisk := allParquetFiles(t, committedConc.Location()) + for _, p := range paths { + assert.Contains(t, onDisk, p) + } +} + +func TestRewriteDataFiles_MaxConcurrentGroupsNegativeRejected(t *testing.T) { + tbl := newRewriteTestTable(t) + + tx := tbl.NewTransaction() + _, err := tx.RewriteDataFiles(t.Context(), nil, table.RewriteDataFilesOptions{MaxConcurrentGroups: -1}) + require.ErrorIs(t, err, table.ErrInvalidOperation) + + txPartial := tbl.NewTransaction() + _, err = txPartial.RewriteDataFiles(t.Context(), nil, table.RewriteDataFilesOptions{PartialProgress: true, MaxConcurrentGroups: -1}) + require.ErrorIs(t, err, table.ErrInvalidOperation) +} + +func TestRewriteDataFiles_MaxConcurrentGroupsDeterministicOrder(t *testing.T) { + tblA := newMaxConcPartitionedTable(t, iceio.LocalFS{}) + tblA = addMaxConcPartitions(t, tblA, 8, 1, 5) + tblB := newMaxConcPartitionedTable(t, iceio.LocalFS{}) + tblB = addMaxConcPartitions(t, tblB, 8, 1, 5) + + groupsA := groupsByPartition(t, tblA) + groupsB := groupsByPartition(t, tblB) + + txA := tblA.NewTransaction() + _, err := txA.RewriteDataFiles(t.Context(), groupsA, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + require.NoError(t, err) + committedA, err := txA.Commit(t.Context()) + require.NoError(t, err) + + txB := tblB.NewTransaction() + _, err = txB.RewriteDataFiles(t.Context(), groupsB, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + require.NoError(t, err) + committedB, err := txB.Commit(t.Context()) + require.NoError(t, err) + + orderA := manifestDataPartitions(t, committedA) + orderB := manifestDataPartitions(t, committedB) + require.Len(t, orderA, 8) + require.Len(t, orderB, 8) + assert.Equal(t, orderA, orderB) + assert.Equal(t, []string{"p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"}, orderA) +} + +type failOpenIO struct { + iceio.LocalFS + mu sync.Mutex + failSubstr string + failErr error +} + +func (f *failOpenIO) setFail(substr string, err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.failSubstr = substr + f.failErr = err +} + +func (f *failOpenIO) Open(name string) (iceio.File, error) { + f.mu.Lock() + substr, failErr := f.failSubstr, f.failErr + f.mu.Unlock() + if substr != "" && strings.Contains(name, substr) { + return nil, failErr + } + + return f.LocalFS.Open(name) +} + +func TestRewriteDataFiles_MaxConcurrentGroupsGroupFailure(t *testing.T) { + injected := errors.New("injected compaction read failure") + + fsAtomic := &failOpenIO{} + tblAtomic := newMaxConcPartitionedTable(t, fsAtomic) + tblAtomic = addMaxConcPartitions(t, tblAtomic, 4, 1, 5) + groupsAtomic := groupsByPartition(t, tblAtomic) + require.Len(t, groupsAtomic, 4) + fsAtomic.setFail(groupsAtomic[2].Tasks[0].File.FilePath(), injected) + beforeAtomicFiles := allParquetFiles(t, tblAtomic.Location()) + + txAtomic := tblAtomic.NewTransaction() + _, err := txAtomic.RewriteDataFiles(t.Context(), groupsAtomic, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + require.Error(t, err) + assert.Contains(t, err.Error(), injected.Error()) + assert.ElementsMatch(t, beforeAtomicFiles, allParquetFiles(t, tblAtomic.Location())) + + fsPartial := &failOpenIO{} + tblPartial := newMaxConcPartitionedTable(t, fsPartial) + tblPartial = addMaxConcPartitions(t, tblPartial, 4, 1, 5) + groupsPartial := groupsByPartition(t, tblPartial) + require.Len(t, groupsPartial, 4) + fsPartial.setFail(groupsPartial[1].Tasks[0].File.FilePath(), injected) + beforeFiles := allParquetFiles(t, tblPartial.Location()) + + txPartial := tblPartial.NewTransaction() + result, err := txPartial.RewriteDataFiles(t.Context(), groupsPartial, table.RewriteDataFilesOptions{ + PartialProgress: true, + MaxCommits: 1, + MaxConcurrentGroups: 4, + }) + require.Error(t, err) + require.NotNil(t, result) + assert.Contains(t, err.Error(), injected.Error()) + assert.Empty(t, result.CompletedGroups) + assert.ElementsMatch(t, beforeFiles, allParquetFiles(t, tblPartial.Location())) +} + +type badContentFile struct { + iceberg.DataFile +} + +func (badContentFile) ContentType() iceberg.ManifestEntryContent { + return iceberg.ManifestEntryContent(99) +} + +func TestRewriteDataFiles_MaxConcurrentGroupsCommitFailureCleansOutputs(t *testing.T) { + for _, maxConc := range []int{4, 0} { + t.Run(fmt.Sprintf("maxConc=%d", maxConc), func(t *testing.T) { + tbl := newMaxConcPartitionedTable(t, iceio.LocalFS{}) + tbl = addMaxConcPartitions(t, tbl, 4, 1, 5) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 4) + before := allParquetFiles(t, tbl.Location()) + + tx := tbl.NewTransaction() + _, err := tx.RewriteDataFiles(t.Context(), groups, table.RewriteDataFilesOptions{ + MaxConcurrentGroups: maxConc, + ExtraDeleteFilesToRemove: []iceberg.DataFile{badContentFile{groups[0].Tasks[0].File}}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported content type") + assert.ElementsMatch(t, before, allParquetFiles(t, tbl.Location())) + }) + } +} + +type blockPathIO struct { + iceio.LocalFS + mu sync.Mutex + substr string + entered chan struct{} + release chan struct{} +} + +func (b *blockPathIO) Open(name string) (iceio.File, error) { + b.mu.Lock() + substr, entered, release := b.substr, b.entered, b.release + b.mu.Unlock() + if substr != "" && strings.Contains(name, substr) { + select { + case entered <- struct{}{}: + default: + } + <-release + } + + return b.LocalFS.Open(name) +} + +func cancelOnDoneFSF(t *testing.T, fs iceio.IO) func(context.Context) (iceio.IO, error) { + t.Helper() + + return func(ctx context.Context) (iceio.IO, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + return fs, nil + } +} + +func TestRewriteDataFiles_MaxConcurrentGroupsCleanupAfterCancelUsesOpenFS(t *testing.T) { + t.Run("sequential", func(t *testing.T) { + blocker := &blockPathIO{entered: make(chan struct{}, 1), release: make(chan struct{})} + tbl := newMaxConcPartitionedTableWithFSF(t, cancelOnDoneFSF(t, blocker)) + tbl = addMaxConcPartitions(t, tbl, 4, 1, 5) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 4) + blocker.mu.Lock() + blocker.substr = groups[1].Tasks[0].File.FilePath() + blocker.mu.Unlock() + before := allParquetFiles(t, tbl.Location()) + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + var rewriteErr error + go func() { + defer close(done) + tx := tbl.NewTransaction() + _, rewriteErr = tx.RewriteDataFiles(ctx, groups, table.RewriteDataFilesOptions{}) + }() + + select { + case <-blocker.entered: + case <-done: + t.Fatalf("rewrite finished before group 1 started, err=%v", rewriteErr) + case <-t.Context().Done(): + t.Fatal("test context done while waiting for group 1") + } + cancel() + close(blocker.release) + <-done + require.Error(t, rewriteErr) + assert.ErrorIs(t, rewriteErr, context.Canceled) + assert.ElementsMatch(t, before, allParquetFiles(t, tbl.Location())) + }) + + t.Run("concurrent", func(t *testing.T) { + blocker := &blockPathIO{entered: make(chan struct{}, 1), release: make(chan struct{})} + tbl := newMaxConcPartitionedTableWithFSF(t, cancelOnDoneFSF(t, blocker)) + tbl = addMaxConcPartitions(t, tbl, 8, 1, 10) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 8) + blocker.mu.Lock() + blocker.substr = groups[7].Tasks[0].File.FilePath() + blocker.mu.Unlock() + before := allParquetFiles(t, tbl.Location()) + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + var rewriteErr error + go func() { + defer close(done) + tx := tbl.NewTransaction() + _, rewriteErr = tx.RewriteDataFiles(ctx, groups, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + }() + + deadline := time.Now().Add(time.Minute) + for len(allParquetFiles(t, tbl.Location())) == len(before) { + select { + case <-done: + t.Fatalf("rewrite finished before any group wrote output, err=%v", rewriteErr) + default: + } + if time.Now().After(deadline) { + cancel() + close(blocker.release) + <-done + t.Fatal("timed out waiting for a group to write output") + } + time.Sleep(5 * time.Millisecond) + } + cancel() + close(blocker.release) + <-done + require.Error(t, rewriteErr) + assert.ErrorIs(t, rewriteErr, context.Canceled) + assert.ElementsMatch(t, before, allParquetFiles(t, tbl.Location())) + }) +} + +type orderFailIO struct { + iceio.LocalFS + mu sync.Mutex + slowSubstr string + slowErr error + fastSubstr string + fastErr error + slowEntered chan struct{} + fastFailed chan struct{} + slowSignaled bool + fastSignaled bool +} + +func (o *orderFailIO) Open(name string) (iceio.File, error) { + o.mu.Lock() + slow := o.slowSubstr != "" && strings.Contains(name, o.slowSubstr) + fast := o.fastSubstr != "" && strings.Contains(name, o.fastSubstr) + if slow && !o.slowSignaled { + o.slowSignaled = true + close(o.slowEntered) + } + slowErr, fastErr := o.slowErr, o.fastErr + slowEntered, fastFailed := o.slowEntered, o.fastFailed + o.mu.Unlock() + switch { + case slow: + <-fastFailed + + return nil, slowErr + case fast: + <-slowEntered + o.mu.Lock() + if !o.fastSignaled { + o.fastSignaled = true + close(fastFailed) + } + o.mu.Unlock() + + return nil, fastErr + default: + return o.LocalFS.Open(name) + } +} + +func TestRewriteDataFiles_MaxConcurrentGroupsFailureReturnsLowestIndexError(t *testing.T) { + slowErr := errors.New("injected slow group failure") + fastErr := errors.New("injected fast group failure") + + ordered := &orderFailIO{slowEntered: make(chan struct{}), fastFailed: make(chan struct{})} + tbl := newMaxConcPartitionedTable(t, ordered) + tbl = addMaxConcPartitions(t, tbl, 4, 1, 5) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 4) + ordered.mu.Lock() + ordered.slowSubstr = groups[1].Tasks[0].File.FilePath() + ordered.slowErr = slowErr + ordered.fastSubstr = groups[3].Tasks[0].File.FilePath() + ordered.fastErr = fastErr + ordered.mu.Unlock() + before := allParquetFiles(t, tbl.Location()) + + var buf bytes.Buffer + origLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}))) + defer slog.SetDefault(origLogger) + + tx := tbl.NewTransaction() + _, err := tx.RewriteDataFiles(t.Context(), groups, table.RewriteDataFilesOptions{MaxConcurrentGroups: 4}) + require.Error(t, err) + assert.Contains(t, err.Error(), slowErr.Error()) + assert.NotContains(t, err.Error(), fastErr.Error()) + assert.ElementsMatch(t, before, allParquetFiles(t, tbl.Location())) + + logged := buf.String() + assert.Contains(t, logged, "compaction group failed") + assert.Contains(t, logged, slowErr.Error()) + assert.Contains(t, logged, fastErr.Error()) +} + +type gateOpenIO struct { + iceio.LocalFS + mu sync.Mutex + enabled bool + entered chan struct{} + release chan struct{} +} + +func (g *gateOpenIO) enable() { + g.mu.Lock() + defer g.mu.Unlock() + g.enabled = true + g.entered = make(chan struct{}, 32) + g.release = make(chan struct{}) +} + +func (g *gateOpenIO) releaseAll() { + g.mu.Lock() + defer g.mu.Unlock() + select { + case <-g.release: + default: + close(g.release) + } +} + +func (g *gateOpenIO) Open(name string) (iceio.File, error) { + g.mu.Lock() + enabled, entered, release := g.enabled, g.entered, g.release + g.mu.Unlock() + if enabled && strings.Contains(name, "/data/") { + select { + case <-release: + default: + select { + case entered <- struct{}{}: + default: + } + <-release + } + } + + return g.LocalFS.Open(name) +} + +func TestRewriteDataFiles_MaxConcurrentGroupsContextCancel(t *testing.T) { + for _, partial := range []bool{false, true} { + t.Run(fmt.Sprintf("partial=%v", partial), func(t *testing.T) { + gate := &gateOpenIO{} + tbl := newMaxConcPartitionedTable(t, gate) + tbl = addMaxConcPartitions(t, tbl, 8, 1, 10) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 8) + gate.enable() + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + var rewriteErr error + go func() { + defer close(done) + tx := tbl.NewTransaction() + opts := table.RewriteDataFilesOptions{MaxConcurrentGroups: 4} + if partial { + opts.PartialProgress = true + opts.MaxCommits = 1 + } + _, rewriteErr = tx.RewriteDataFiles(ctx, groups, opts) + }() + + for range 4 { + select { + case <-gate.entered: + case <-done: + t.Fatalf("rewrite finished before 4 groups were in flight, err=%v", rewriteErr) + case <-t.Context().Done(): + t.Fatal("test context done while waiting for groups") + } + } + cancel() + gate.releaseAll() + <-done + require.Error(t, rewriteErr) + assert.ErrorIs(t, rewriteErr, context.Canceled) + assert.Equal(t, ctx.Err(), rewriteErr) + }) + } +} + +type countOpenFile struct { + iceio.File + owner *countOpenIO + once sync.Once +} + +func (f *countOpenFile) Close() error { + err := f.File.Close() + f.once.Do(func() { + f.owner.mu.Lock() + defer f.owner.mu.Unlock() + f.owner.cur-- + }) + + return err +} + +type countOpenIO struct { + iceio.LocalFS + mu sync.Mutex + cur int + peak int + barrier bool + overlap chan struct{} +} + +func (c *countOpenIO) Open(name string) (iceio.File, error) { + f, err := c.LocalFS.Open(name) + if err != nil { + return nil, err + } + if strings.Contains(name, "/data/") { + c.mu.Lock() + c.cur++ + if c.cur > c.peak { + c.peak = c.cur + } + if c.barrier && c.cur >= 2 { + select { + case <-c.overlap: + default: + close(c.overlap) + } + } + barrier, overlap := c.barrier, c.overlap + c.mu.Unlock() + if barrier { + <-overlap + } + + return &countOpenFile{File: f, owner: c}, nil + } + + return f, nil +} + +func (c *countOpenIO) reset() { + c.mu.Lock() + defer c.mu.Unlock() + c.cur = 0 + c.peak = 0 + c.overlap = make(chan struct{}) +} + +func (c *countOpenIO) setBarrier(enabled bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.barrier = enabled + if c.overlap == nil { + c.overlap = make(chan struct{}) + } +} + +func (c *countOpenIO) getPeak() int { + c.mu.Lock() + defer c.mu.Unlock() + + return c.peak +} + +func TestRewriteDataFiles_MaxConcurrentGroupsLimitsInFlight(t *testing.T) { + for _, maxConc := range []int{4, 0, 1} { + t.Run(fmt.Sprintf("maxConc=%d", maxConc), func(t *testing.T) { + counter := &countOpenIO{} + tbl := newMaxConcPartitionedTable(t, counter) + tbl = addMaxConcPartitions(t, tbl, 8, 1, 10) + groups := groupsByPartition(t, tbl) + require.Len(t, groups, 8) + counter.reset() + counter.setBarrier(maxConc > 1) + + tx := tbl.NewTransaction() + _, err := tx.RewriteDataFiles(t.Context(), groups, table.RewriteDataFilesOptions{ + MaxConcurrentGroups: maxConc, + GroupOptions: []table.CompactionGroupOption{table.WithCompactionScanConcurrency(1)}, + }) + require.NoError(t, err) + _, err = tx.Commit(t.Context()) + require.NoError(t, err) + + peak := counter.getPeak() + if maxConc > 1 { + assert.LessOrEqual(t, peak, maxConc) + assert.GreaterOrEqual(t, peak, 2) + } else { + assert.Equal(t, 1, peak) + } + }) + } +}