laskoviymishka commented on code in PR #2041:
URL: https://github.com/apache/iceberg-go/pull/2041#discussion_r4075722925
##########
table/rewrite_data_files.go:
##########
@@ -332,31 +369,51 @@ 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 {
+ if opts.MaxConcurrency > 1 {
+ results, err := executeCompactionGroups(ctx, t.tbl, groups,
opts.GroupOptions, opts.MaxConcurrency)
+ if err != nil {
Review Comment:
The concurrent atomic path leaks output files on failure, and the
concurrency makes it strictly worse than the sequential loop.
When a group fails here we return without touching `results`, so every group
that already finished writing its compacted parquet is orphaned on disk
(nothing gets committed to a manifest). Sequentially that's bounded to the
groups before the failing one; with `MaxConcurrency`, up to N groups race ahead
and finish real writes before cancellation propagates, so the leak scales with
the knob.
The partial-progress branch already handles this via `cleanupBatch(err,
results...)`. I'd mirror it here: on error call `cleanupCompactionOutputs(fs,
results)` (fs from `t.tbl.fsF(ctx)`) before returning, and do the same in the
sequential `else` branch. wdyt?
##########
table/rewrite_data_files.go:
##########
@@ -391,6 +448,35 @@ func (t *Transaction) RewriteDataFiles(ctx
context.Context, groups []CompactionT
return result, nil
}
+func executeCompactionGroups(ctx context.Context, tbl *Table, groups
[]CompactionTaskGroup, groupOpts []CompactionGroupOption, maxConcurrency int)
([]CompactionGroupResult, error) {
+ if err := ctx.Err(); err != nil {
+ return nil, err
+ }
+ g, gctx := errgroup.WithContext(ctx)
+ g.SetLimit(min(maxConcurrency, len(groups)))
+ results := make([]CompactionGroupResult, len(groups))
+ for i, group := range groups {
+ if len(group.Tasks) == 0 {
+ continue
+ }
+ g.Go(func() error {
+ gr, err := ExecuteCompactionGroup(gctx, tbl, group,
groupOpts...)
+ results[i] = gr
+
+ return err
+ })
+ }
+ if err := g.Wait(); err != nil {
+ if ctx.Err() != nil {
Review Comment:
This masks the real failure whenever the caller's ctx happens to be done.
If a group fails for an unrelated reason (bad data file, disk full) while
the caller's ctx is also canceled, we discard the group's actual error and
return a bare context error, so an operator debugging a production failure sees
"context canceled" instead of the real cause.
I'd only substitute `ctx.Err()` when the group's own error is context-caused
(`errors.Is(err, context.Canceled) || errors.Is(err,
context.DeadlineExceeded)`), or just return `errors.Join(err, ctx.Err())` to
keep both. The cancel test can switch to `assert.ErrorIs` instead of asserting
equality. wdyt?
##########
table/rewrite_data_files.go:
##########
@@ -391,6 +448,35 @@ func (t *Transaction) RewriteDataFiles(ctx
context.Context, groups []CompactionT
return result, nil
}
+func executeCompactionGroups(ctx context.Context, tbl *Table, groups
[]CompactionTaskGroup, groupOpts []CompactionGroupOption, maxConcurrency int)
([]CompactionGroupResult, error) {
+ if err := ctx.Err(); err != nil {
+ return nil, err
+ }
+ g, gctx := errgroup.WithContext(ctx)
+ g.SetLimit(min(maxConcurrency, len(groups)))
Review Comment:
This relies entirely on both call sites gating with `MaxConcurrency > 1`.
`errgroup.SetLimit` treats 0 as "no goroutine may run" (every `g.Go` blocks
forever) and a negative as unbounded, so if a future caller or a refactor drops
that guard we get either a deadlock or a silently unbounded fan-out, and
neither is caught by the type system.
Since the function is unexported and cheap to harden, I'd clamp defensively:
`limit := min(maxConcurrency, len(groups)); if limit < 1 { limit = 1 }`.
Thoughts?
##########
table/rewrite_data_files.go:
##########
@@ -332,31 +369,51 @@ 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 {
+ if opts.MaxConcurrency > 1 {
+ results, err := executeCompactionGroups(ctx, t.tbl, groups,
opts.GroupOptions, opts.MaxConcurrency)
+ if err != nil {
return result, err
}
-
- if len(group.Tasks) == 0 {
- continue
+ for _, gr := range results {
Review Comment:
This per-group apply block (skip empty, `ApplyResult`,
`accumulateGroupMetrics`, stage `SafePosDeletes`/`SafeDeletionVectors`) is now
copy-pasted four times: atomic concurrent and sequential here, plus the two
partial-progress branches around :700. A future tweak to the accounting has to
land in all four or the concurrent and sequential paths silently diverge, which
is the exact thing `MaxConcurrency`'s doc promises won't happen.
Could we factor the per-group apply into a small helper shared by both
branches?
```go
func applyGroupResult(rewrite *Rewrite, result *RewriteResult, staged
map[string]struct{}, gr CompactionGroupResult) { ... }
```
Same shape for the partial accumulation. Thoughts?
##########
table/rewrite_data_files_test.go:
##########
@@ -1401,3 +1424,483 @@ func appendEqualityDelete(t *testing.T, tbl
*table.Table, equalityFieldIDs []int
return out
}
+
+func newMaxConcPartitionedTable(t *testing.T, fs iceio.IO) *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",
+ func(context.Context) (iceio.IO, error) { return fs, nil },
+ 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 rowsByPartitionValue(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)
+ idx := rec.Schema().FieldIndices("data")
+ require.NotEmpty(t, idx)
+ col, ok := rec.Column(idx[0]).(*array.String)
+ require.True(t, ok)
+ for i := range int(rec.NumRows()) {
+ out[col.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_MaxConcurrencyMatchesSequential(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{MaxConcurrency: 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)
+
+ assert.Equal(t, rowsByPartitionValue(t, committedSeq),
rowsByPartitionValue(t, committedConc))
+ for p := range 8 {
+ assert.Equal(t, int64(10), rowsByPartitionValue(t,
committedConc)[fmt.Sprintf("p%d", p)])
+ }
+
+ 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_MaxConcurrencyNegativeRejected(t *testing.T) {
+ tbl := newRewriteTestTable(t)
+
+ tx := tbl.NewTransaction()
+ _, err := tx.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+
+ txPartial := tbl.NewTransaction()
+ _, err = txPartial.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{PartialProgress: true, MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+}
+
+func TestRewriteDataFiles_MaxConcurrencyDeterministicOrder(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{MaxConcurrency: 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{MaxConcurrency: 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_MaxConcurrencyGroupFailure(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)
+
+ txAtomic := tblAtomic.NewTransaction()
+ _, err := txAtomic.RewriteDataFiles(t.Context(), groupsAtomic,
table.RewriteDataFilesOptions{MaxConcurrency: 4})
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), injected.Error())
+
+ 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,
+ MaxConcurrency: 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 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_MaxConcurrencyContextCancel(t *testing.T) {
+ for _, partial := range []bool{false, true} {
+ 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{MaxConcurrency: 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
+}
+
+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
+ }
+ c.mu.Unlock()
+ time.Sleep(20 * time.Millisecond)
Review Comment:
This peak-concurrency assertion leans on a fixed 20ms sleep to widen the
race window, which is a real wall-clock race: under `-race` or a loaded CI
runner it's plausible for fewer than 2 groups to overlap in that window and
`assert.GreaterOrEqual(peak, 2)` flakes.
The adjacent cancel test already uses a proper channel gate (`gateOpenIO`),
and the scanner reorder tests use `testing/synctest`. I'd reuse one of those
here for a deterministic sync point instead of the sleep. Thoughts?
##########
table/rewrite_data_files.go:
##########
@@ -217,6 +218,23 @@ type RewriteDataFilesOptions struct {
// size, scan concurrency). See the With* helpers returning
// [CompactionGroupOption].
GroupOptions []CompactionGroupOption
+
+ // MaxConcurrency 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.
+ // The first error cancels the groups still running and is returned.
+ // Peak record-pipeline memory is MaxConcurrency times the per-group
+ // bound stated on [WithCompactionArrowBatchSize]:
+ //
+ // MaxConcurrency 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. Negative values
+ // are rejected with [ErrInvalidOperation].
+ MaxConcurrency int
Review Comment:
One foot-gun worth heading off: this `MaxConcurrency` (groups in parallel)
sits right next to `WithCompactionScanConcurrency`, which forwards to
`Table.Scan` as `WithMaxConcurrency` (scan workers per group). Two different
knobs both called "max concurrency", and they compose multiplicatively, so
`MaxConcurrency=N` with the scan default gives roughly N x GOMAXPROCS
concurrent file opens, an easy connection/FD spike on object stores.
The doc here only quantifies memory. I'd at minimum cross-reference the two
in each other's doc and note the I/O fan-out; renaming to something like
`MaxConcurrentGroups` would remove the ambiguity entirely if we're open to it.
wdyt?
##########
table/rewrite_data_files_test.go:
##########
@@ -1401,3 +1424,483 @@ func appendEqualityDelete(t *testing.T, tbl
*table.Table, equalityFieldIDs []int
return out
}
+
+func newMaxConcPartitionedTable(t *testing.T, fs iceio.IO) *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",
+ func(context.Context) (iceio.IO, error) { return fs, nil },
+ 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 rowsByPartitionValue(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)
+ idx := rec.Schema().FieldIndices("data")
+ require.NotEmpty(t, idx)
+ col, ok := rec.Column(idx[0]).(*array.String)
+ require.True(t, ok)
+ for i := range int(rec.NumRows()) {
+ out[col.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_MaxConcurrencyMatchesSequential(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{MaxConcurrency: 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)
+
+ assert.Equal(t, rowsByPartitionValue(t, committedSeq),
rowsByPartitionValue(t, committedConc))
+ for p := range 8 {
+ assert.Equal(t, int64(10), rowsByPartitionValue(t,
committedConc)[fmt.Sprintf("p%d", p)])
+ }
+
+ 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_MaxConcurrencyNegativeRejected(t *testing.T) {
+ tbl := newRewriteTestTable(t)
+
+ tx := tbl.NewTransaction()
+ _, err := tx.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+
+ txPartial := tbl.NewTransaction()
+ _, err = txPartial.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{PartialProgress: true, MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+}
+
+func TestRewriteDataFiles_MaxConcurrencyDeterministicOrder(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{MaxConcurrency: 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{MaxConcurrency: 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_MaxConcurrencyGroupFailure(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)
+
+ txAtomic := tblAtomic.NewTransaction()
+ _, err := txAtomic.RewriteDataFiles(t.Context(), groupsAtomic,
table.RewriteDataFilesOptions{MaxConcurrency: 4})
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), injected.Error())
+
+ 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,
+ MaxConcurrency: 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 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_MaxConcurrencyContextCancel(t *testing.T) {
+ for _, partial := range []bool{false, true} {
+ 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{MaxConcurrency: 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
+}
+
+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
+ }
+ c.mu.Unlock()
+ time.Sleep(20 * time.Millisecond)
+
+ 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
+}
+
+func (c *countOpenIO) getPeak() int {
+ c.mu.Lock()
+ defer c.mu.Unlock()
+
+ return c.peak
+}
+
+func TestRewriteDataFiles_MaxConcurrencyLimitsInFlight(t *testing.T) {
+ for _, maxConc := range []int{4, 0, 1} {
Review Comment:
These table cases run in a bare `for` loop, so a `require.*` failure at
`maxConc=4` aborts the loop and the 0 and 1 cases never run, and a plain
failure won't say which case broke.
I'd wrap each iteration in `t.Run(fmt.Sprintf("maxConc=%d", maxConc), ...)`
(same for the `partial=%v` loop in the cancel test). Small thing, just makes
failures legible.
##########
table/arrow_scanner.go:
##########
@@ -1242,9 +1242,61 @@ func (as *arrowScan) addTaskProjectedFieldIDs(invariants
*arrowScanInvariants, t
}
type enumeratedRecord struct {
- Record tblutils.Enumerated[arrow.RecordBatch]
- Task tblutils.Enumerated[FileScanTask]
- Err error
+ Record tblutils.Enumerated[arrow.RecordBatch]
+ Task tblutils.Enumerated[FileScanTask]
+ Err error
+ credits taskCredits
+}
+
+const maxInFlightTasksPerWorker = 1
+
+type taskCredits chan struct{}
+
+func newTaskCredits() taskCredits {
+ return make(taskCredits, maxInFlightTasksPerWorker)
+}
+
+func (c taskCredits) acquire(ctx context.Context) error {
+ if c == nil {
+ return nil
+ }
+
+ select {
+ case c <- struct{}{}:
+ return nil
+ case <-ctx.Done():
+ return context.Cause(ctx)
+ }
+}
+
+func (c taskCredits) release() {
+ if c != nil {
+ <-c
+ }
+}
+
+type recordSink struct {
+ out chan<- enumeratedRecord
+ credits taskCredits
+}
+
+func newRecordSink(out chan<- enumeratedRecord) recordSink {
+ return recordSink{out: out, credits: newTaskCredits()}
+}
+
+func (s recordSink) reserve(ctx context.Context) error {
+ return s.credits.acquire(ctx)
+}
+
+func (s recordSink) send(rec enumeratedRecord) {
+ if rec.Record.Last {
Review Comment:
The credit is only stamped onto the record when `Record.Last` is true, so
it's only released downstream on the last batch. That's correct today, but it
leans on a subtle invariant: every worker error path (`recordsFromTask` / the
delete loaders) returns from the goroutine without ever emitting a
`Last`-marked record, so the acquired credit is just abandoned and the worker
never calls `reserve()` again on that slot.
The moment someone makes a worker resilient (skip a bad task and continue)
that turns into a real deadlock on the next `reserve()`, and nothing here
documents or guards the invariant. I'd make release explicit and independent of
`Last` (e.g. a `defer` around each task attempt, with send/discard as
idempotent no-ops), or at least add a loud comment. wdyt?
##########
table/rewrite_data_files_test.go:
##########
@@ -1401,3 +1424,483 @@ func appendEqualityDelete(t *testing.T, tbl
*table.Table, equalityFieldIDs []int
return out
}
+
+func newMaxConcPartitionedTable(t *testing.T, fs iceio.IO) *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",
+ func(context.Context) (iceio.IO, error) { return fs, nil },
+ 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 rowsByPartitionValue(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)
+ idx := rec.Schema().FieldIndices("data")
+ require.NotEmpty(t, idx)
+ col, ok := rec.Column(idx[0]).(*array.String)
+ require.True(t, ok)
+ for i := range int(rec.NumRows()) {
+ out[col.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_MaxConcurrencyMatchesSequential(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{MaxConcurrency: 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)
+
+ assert.Equal(t, rowsByPartitionValue(t, committedSeq),
rowsByPartitionValue(t, committedConc))
Review Comment:
This is the test carrying the core "concurrent == sequential" claim, but it
only compares per-partition row counts, so it'd still pass if a bug swapped
rows between files while keeping the counts right.
Since it's the primary evidence for the central claim, I'd make it compare
the actual per-partition ID sets (or a hash of the arrow data), not just
tallies. Not blocking, but cheap confidence for the one invariant that matters
most here.
##########
table/rewrite_data_files_test.go:
##########
@@ -1401,3 +1424,483 @@ func appendEqualityDelete(t *testing.T, tbl
*table.Table, equalityFieldIDs []int
return out
}
+
+func newMaxConcPartitionedTable(t *testing.T, fs iceio.IO) *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",
+ func(context.Context) (iceio.IO, error) { return fs, nil },
+ 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 rowsByPartitionValue(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)
+ idx := rec.Schema().FieldIndices("data")
+ require.NotEmpty(t, idx)
+ col, ok := rec.Column(idx[0]).(*array.String)
+ require.True(t, ok)
+ for i := range int(rec.NumRows()) {
+ out[col.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_MaxConcurrencyMatchesSequential(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{MaxConcurrency: 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)
+
+ assert.Equal(t, rowsByPartitionValue(t, committedSeq),
rowsByPartitionValue(t, committedConc))
+ for p := range 8 {
+ assert.Equal(t, int64(10), rowsByPartitionValue(t,
committedConc)[fmt.Sprintf("p%d", p)])
+ }
+
+ 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_MaxConcurrencyNegativeRejected(t *testing.T) {
+ tbl := newRewriteTestTable(t)
+
+ tx := tbl.NewTransaction()
+ _, err := tx.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+
+ txPartial := tbl.NewTransaction()
+ _, err = txPartial.RewriteDataFiles(t.Context(), nil,
table.RewriteDataFilesOptions{PartialProgress: true, MaxConcurrency: -1})
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+}
+
+func TestRewriteDataFiles_MaxConcurrencyDeterministicOrder(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{MaxConcurrency: 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{MaxConcurrency: 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_MaxConcurrencyGroupFailure(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)
+
+ txAtomic := tblAtomic.NewTransaction()
+ _, err := txAtomic.RewriteDataFiles(t.Context(), groupsAtomic,
table.RewriteDataFilesOptions{MaxConcurrency: 4})
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), injected.Error())
Review Comment:
This half only asserts the error, so it passes whether zero or three groups
leaked their output, which is exactly the regression the atomic path is exposed
to.
The partial sub-case below captures `beforeFiles := allParquetFiles(...)`
and does `assert.ElementsMatch(t, beforeFiles, allParquetFiles(...))` after.
I'd add the same before/after check here. Even if it just pins the current
(leaky) behavior for now, it makes the trade-off visible and stops a silent
regression.
##########
table/arrow_scanner.go:
##########
@@ -2190,6 +2244,7 @@ func createIteratorWithCleanup(ctx context.Context,
numWorkers uint, records <-c
return
}
+ enum.credits.release()
Review Comment:
Small mismatch with the doc claim on `WithCompactionArrowBatchSize`:
`release()` fires as soon as the record is popped off the sequenced channel,
before `yield` hands it to the caller, so the next worker can start reading its
file while the current batch is still being encoded by the consumer.
The real bound is "until the last batch leaves the reorder buffer", which is
a touch weaker than the "until consumed" wording. I'd either move `release()`
to after `yield` returns, or soften the doc. Not blocking, just want the doc to
match the behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
