KranzL commented on code in PR #2041:
URL: https://github.com/apache/iceberg-go/pull/2041#discussion_r4087862500
##########
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:
Done in 30ffe36. The concurrent branch passes its results to
cleanupAtomicRewriteOutputs at table/rewrite_data_files.go:382, and the
sequential branch tracks applied results and cleans them at :391 and :400. The
helper at :453 opens the IO through t.tbl.fsF and joins a cleanup failure with
the cause instead of replacing it.
##########
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:
Done in ad900a5. The atomic half captures beforeAtomicFiles and asserts
ElementsMatch after the error at table/rewrite_data_files_test.go:1722, the
same check the partial half uses. go test ./table/ -race -count=10 -run
'TestRewriteDataFiles_MaxConcurrentGroups(GroupFailure|LimitsInFlight|ContextCancel)'
passes in 11.2s.
##########
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:
Done in 30ffe36, first option. The caller context error replaces the group
error only when the group error is itself context-caused, at
table/rewrite_data_files.go:488. A data failure that races a cancellation keeps
its real error now. The cancel test still asserts exact equality at
table/rewrite_data_files_test.go:1828.
##########
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:
Done in 30ffe36 at table/rewrite_data_files.go:469-472, the clamp as
suggested.
##########
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:
This is #2039's code. The same point was raised there and reworked in
149c395: reserve marks the credit held and the sink hands it to exactly one
record, the Last record in send or the error record in fail when no Last record
was sent, and TestRecordSinkHandsOffCreditOnce covers fail without Last, Last
followed by fail, and reserve blocking while the credit is out. The thread on
#2039 has the details. This hunk drops out of the diff when #2039 merges and
this rebases.
##########
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:
Done in 30ffe36. Both atomic branches call applyAtomicGroupResult at
table/rewrite_data_files.go:439, and both partial branches call
appendPartialGroupResult at :843.
##########
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:
Renamed in 50480ae. The field is MaxConcurrentGroups at
table/rewrite_data_files.go:241, both docs cross-reference each other, and both
state the N times scan-workers file-open fan-out (:235-239 and :277-280). Issue
#2040 and the PR body use the new name.
--
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]
