You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by at...@apache.org on 2026/09/23 21:17:57 UTC
(iceberg-go) branch main updated: perf(table): bound the scan reorder heap at scan concurrency above one (#2039)
This is an automated email from the ASF dual-hosted git repository.
laskoviymishka pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new dce9757b1 perf(table): bound the scan reorder heap at scan concurrency above one (#2039)
dce9757b1 is described below
commit dce9757b17bbb526d48c0594f1f3c84d8990b5d2
Author: KranzL <50...@users.noreply.github.com>
AuthorDate: Wed Sep 23 17:17:52 2026 -0400
perf(table): bound the scan reorder heap at scan concurrency above one (#2039)
* perf(table): add lagging-task reorder heap benchmark and test
BenchmarkArrowScanReorderHeapLaggingTask in
table/arrow_scanner_reorder_bench_test.go scans M single-batch data
files at WithMaxConcurrency(4) with the Open of task 0's file gated
until the other M - 1 files have been read, and reports the peak
HeapInuse delta while the gate is closed and the number of batches the
other files produced in that window. On Apple M3 Pro, 11 cores,
go1.25.9, -count=3, the delta is 4.7 MB at M = 8, 20 to 23 MB at
M = 32 and 96 MB at M = 128, with M - 1 batches held every time.
TestArrowScanReorderHeapBoundedWhileHeadTaskLags in
table/arrow_scanner_reorder_test.go runs the same scan inside a
testing/synctest bubble. synctest.Wait returns once every scan
goroutine is durably blocked, so the count of out-of-order batches the
pipeline holds is read at a fixed point with no sleeps, both before a
fix (all M - 1 other files are read) and after one (workers park on
credits). It asserts the count never exceeds numWorkers x K with K = 2
as reorderCreditsPerWorker. On main it fails with 31 held batches
against a bound of 8, then releases the gate, checks that every row
arrives and checks that the CheckedAllocator is empty.
Refs #2038.
Signed-off-by: KranzL <50...@users.noreply.github.com>
* perf(table): bound the scan reorder heap at scan concurrency above one
Each scan worker in recordBatchesFromTasksAndDeletes and in
makePositionDeleteRecordsForFilter now holds a task credit. The worker
takes the credit before it takes a task, the credit rides on the
task's Last record, and it comes back when the iterator's consumer
takes that record, when the deferred drain loop takes it, or when
MakeSequencedChanWithDiscard discards it at close. A worker waiting
for its credit also waits on the scan context, so cancellation still
stops it. Error records carry no credit.
A worker therefore cannot start its next task while any batch of its
previous task is still in the records channel, the reorder heap or
the sequenced channel. While the head task lags, each other worker
finishes one task and parks, so the heap holds at most numWorkers - 1
tasks instead of every task the free workers can reach. Decoding
inside a task is not throttled and the head task's worker never waits.
Per-batch credits, as #2038 proposed, were measured first on
BenchmarkArrowScanManyFilesAndBatches (32 files of 32 one-row batches,
concurrency 4): K = 2 took 141.6 ms/op against 69.4 ms/op on main,
K = 8 took 127.9 ms/op and K = 32 took 78.7 ms/op, because the batches
a free worker may run ahead are the scan's parallelism. With the
per-task credit that benchmark is within noise of main in interleaved
runs (benchstat, n = 10: +0.4%, -0.2% and +1.5% on the three
sub-benchmarks), and BenchmarkArrowScanReorderHeapLaggingTask holds 3
gated batches and 2.0 to 2.3 MB of peak heap delta for M = 8, 32 and
128, where main holds M - 1 batches and 4.7, 20 to 25 and 95 MB.
The doc on WithCompactionArrowBatchSize now states the bound that
holds, workers x (rows in the largest task + n) +
(recordBatchBufferSize + 2) x n rows, and WithCompactionScanConcurrency
says its worker count multiplies the read-side term.
TestArrowScanReorderHeapBoundedWhileHeadTaskLags asserts the bound of
(numWorkers - 1) x maxInFlightTasksPerWorker completed out-of-order
tasks. TestArrowScanReorderCancelWhileWorkersWaitForTaskCredits
cancels the scan while the free workers are parked and checks that the
iterator returns context.Canceled, every scan goroutine exits and every
batch is released.
Fixes #2038.
Signed-off-by: KranzL <50...@users.noreply.github.com>
* perf(table): hand the scan credit to exactly one record and test the compaction bound
recordSink now tracks whether its worker holds a credit and hands it to
exactly one record: the task's Last record in send, or the error record
in fail when no Last record was sent. processRecords can send Last and
then return an error from the reader or the deferred close, so fail
hands off nil in that case. No path abandons a credit any more, and a
worker that failed and went back to reserve would wait only until its
error record leaves the sequenced channel. The three delete-loader
error sites in recordBatchesFromTasksAndDeletes go through sink.fail,
so fail is the only path that puts an error record into records.
recordSink is a pointer built only by newRecordSink, so a nil sink
panics on first use instead of hanging on a nil channel send. The nil
guard in acquire is gone; taskCredits is declared before
enumeratedRecord.
Scan.ReadTasks now states the bound for the general scan: each worker
holds the decoded batches of at most one task until the iterator has
returned them, so the batches waiting behind a slow early task total at
most the worker count times the rows in the largest task, rounded up to
a multiple of the batch size, and one large file among small ones still
holds its whole decoded contents while it waits. ToArrowRecords and
WithMaxConcurrency point at that paragraph.
TestRecordSinkHandsOffCreditOnce covers fail without Last, Last followed
by fail, and that reserve blocks while the credit is out.
TestExecuteCompactionGroupRecordPipelineBounded drives
ExecuteCompactionGroup with four scan workers and a writer queue of two
over 32 single-batch files inside testing/synctest, gating task 0 in
Open and then the output file in Create, and checks 3 fully read tasks
while task 0 is gated and exactly 8 while the writer is gated, the two
terms of the bound on WithCompactionArrowBatchSize counted in tasks.
With the credit gate disabled the counts are 31 and 32.
Signed-off-by: KranzL <50...@users.noreply.github.com>
* test(table): write uncompressed parquet in the synctest compaction fixture
TestExecuteCompactionGroupRecordPipelineBounded wrote zstd parquet
inside its synctest bubble. arrow-go keeps zstd encoders in a global
sync.Pool and klauspost's Encoder creates its encoders channel on first
use, so the first zstd write inside the bubble left a pooled encoder
whose channel belongs to the bubble. When a later test reused that
encoder, go1.25.9 failed the whole package with "receive on synctest
channel from outside bubble" (seen in TestClusteredWriter on the
ubuntu and macos go1.25.9 jobs; go1.26.1 passes). Locally the pool was
usually cleared by GC before the next zstd write, which is why the
suite passed here; with GOGC=off it fails every time.
The reorder fixture now sets write.parquet.compression-codec to
uncompressed, so no codec pool is touched inside the bubble. With
GOGC=off the pair TestExecuteCompactionGroupRecordPipelineBounded and
TestClusteredWriter passes three times in a row.
Signed-off-by: KranzL <50...@users.noreply.github.com>
---------
Signed-off-by: KranzL <50...@users.noreply.github.com>
---
table/arrow_scanner.go | 131 ++++++---
table/arrow_scanner_pruning_internal_test.go | 6 +-
table/arrow_scanner_reorder_bench_test.go | 140 ++++++++++
table/arrow_scanner_reorder_test.go | 401 +++++++++++++++++++++++++++
table/rewrite_data_files.go | 30 +-
table/scanner.go | 10 +
table/table.go | 4 +-
table/transaction.go | 6 +-
8 files changed, 681 insertions(+), 47 deletions(-)
diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go
index 64778ab77..7b5469902 100644
--- a/table/arrow_scanner.go
+++ b/table/arrow_scanner.go
@@ -1241,10 +1241,73 @@ func (as *arrowScan) addTaskProjectedFieldIDs(invariants *arrowScanInvariants, t
return nil
}
+const maxInFlightTasksPerWorker = 1
+
+type taskCredits chan struct{}
+
+func newTaskCredits() taskCredits {
+ return make(taskCredits, maxInFlightTasksPerWorker)
+}
+
+func (c taskCredits) acquire(ctx context.Context) error {
+ select {
+ case c <- struct{}{}:
+ return nil
+ case <-ctx.Done():
+ return context.Cause(ctx)
+ }
+}
+
+func (c taskCredits) release() {
+ if c != nil {
+ <-c
+ }
+}
+
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
+}
+
+type recordSink struct {
+ out chan<- enumeratedRecord
+ credits taskCredits
+ held bool
+}
+
+func newRecordSink(out chan<- enumeratedRecord) *recordSink {
+ return &recordSink{out: out, credits: newTaskCredits()}
+}
+
+func (s *recordSink) reserve(ctx context.Context) error {
+ if err := s.credits.acquire(ctx); err != nil {
+ return err
+ }
+ s.held = true
+
+ return nil
+}
+
+func (s *recordSink) send(rec enumeratedRecord) {
+ if rec.Record.Last {
+ rec.credits = s.handOff()
+ }
+ s.out <- rec
+}
+
+func (s *recordSink) fail(task tblutils.Enumerated[FileScanTask], err error) {
+ s.out <- enumeratedRecord{Task: task, Err: err, credits: s.handOff()}
+}
+
+func (s *recordSink) handOff() taskCredits {
+ if !s.held {
+ return nil
+ }
+ s.held = false
+
+ return s.credits
}
func (as *arrowScan) prepareToRead(ctx context.Context, file iceberg.DataFile, invariants *arrowScanInvariants) (iceSchema *iceberg.Schema, colIndices []int, rdr tblutils.FileReader, err error) {
@@ -1682,9 +1745,9 @@ func (as *arrowScan) processRecords(
columns []int,
pipeline []recProcessFn,
posSource *rowPositionSource,
- out chan<- enumeratedRecord,
+ sink *recordSink,
) error {
- return as.processRecordsWithPlans(ctx, task, fileSchema, rowFilter, rdr, columns, pipeline, posSource, out, nil)
+ return as.processRecordsWithPlans(ctx, task, fileSchema, rowFilter, rdr, columns, pipeline, posSource, sink, nil)
}
// adjustParquetTaskRange keeps a planned range within the physical file when
@@ -1713,7 +1776,7 @@ func (as *arrowScan) processRecordsWithPlans(
columns []int,
pipeline []recProcessFn,
posSource *rowPositionSource,
- out chan<- enumeratedRecord,
+ sink *recordSink,
plans *compiledFileFilterPlans,
) (err error) {
var (
@@ -1828,9 +1891,9 @@ func (as *arrowScan) processRecordsWithPlans(
for recRdr.Next() {
if prev != nil {
- out <- enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
+ sink.send(enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
Value: prev, Index: idx, Last: false,
- }, Task: task}
+ }, Task: task})
idx++
}
@@ -1846,9 +1909,9 @@ func (as *arrowScan) processRecordsWithPlans(
}
if prev != nil {
- out <- enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
+ sink.send(enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
Value: prev, Index: idx, Last: true,
- }, Task: task}
+ }, Task: task})
} else {
// The reader produced no batches (e.g. every row group was pruned by
// stats). The sequenced channel in createIterator still needs this
@@ -1860,9 +1923,9 @@ func (as *arrowScan) processRecordsWithPlans(
if err != nil {
return err
}
- out <- enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
+ sink.send(enumeratedRecord{Record: tblutils.Enumerated[arrow.RecordBatch]{
Value: array.NewRecordBatch(emptySchema, nil, 0), Index: idx, Last: true,
- }, Task: task}
+ }, Task: task})
}
if recRdr.Err() != nil && recRdr.Err() != io.EOF {
@@ -1899,10 +1962,10 @@ func pruningFilterHasMissingInitialDefault(
return false, nil
}
-func (as *arrowScan) recordsFromTask(ctx context.Context, task tblutils.Enumerated[FileScanTask], out chan<- enumeratedRecord, positionalDeletes positionDeletes, dvBitmap *dv.RoaringPositionBitmap, eqDeleteSets []*equalityDeleteSet, invariants *arrowScanInvariants) (err error) {
+func (as *arrowScan) recordsFromTask(ctx context.Context, task tblutils.Enumerated[FileScanTask], sink *recordSink, positionalDeletes positionDeletes, dvBitmap *dv.RoaringPositionBitmap, eqDeleteSets []*equalityDeleteSet, invariants *arrowScanInvariants) (err error) {
defer func() {
if err != nil {
- out <- enumeratedRecord{Task: task, Err: err}
+ sink.fail(task, err)
}
}()
@@ -2020,9 +2083,9 @@ func (as *arrowScan) recordsFromTask(ctx context.Context, task tblutils.Enumerat
if err != nil {
return err
}
- out <- enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{
+ sink.send(enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{
Value: array.NewRecordBatch(emptySchema, nil, 0), Index: 0, Last: true,
- }}
+ }})
return err
}
@@ -2040,15 +2103,15 @@ func (as *arrowScan) recordsFromTask(ctx context.Context, task tblutils.Enumerat
})
})
- err = as.processRecordsWithPlans(ctx, task, iceSchema, rowFilter, rdr, colIndices, pipeline, posSource, out, filterPlans)
+ err = as.processRecordsWithPlans(ctx, task, iceSchema, rowFilter, rdr, colIndices, pipeline, posSource, sink, filterPlans)
return err
}
-func (as *arrowScan) producePosDeletesFromTask(ctx context.Context, task tblutils.Enumerated[FileScanTask], positionalDeletes positionDeletes, out chan<- enumeratedRecord, invariants *arrowScanInvariants) (err error) {
+func (as *arrowScan) producePosDeletesFromTask(ctx context.Context, task tblutils.Enumerated[FileScanTask], positionalDeletes positionDeletes, sink *recordSink, invariants *arrowScanInvariants) (err error) {
defer func() {
if err != nil {
- out <- enumeratedRecord{Task: task, Err: err}
+ sink.fail(task, err)
}
}()
@@ -2104,9 +2167,9 @@ func (as *arrowScan) producePosDeletesFromTask(ctx context.Context, task tblutil
if err != nil {
return err
}
- out <- enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{
+ sink.send(enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{
Value: array.NewRecordBatch(emptySchema, nil, 0), Index: 0, Last: true,
- }}
+ }})
return err
}
@@ -2120,7 +2183,7 @@ func (as *arrowScan) producePosDeletesFromTask(ctx context.Context, task tblutil
return ToRequestedSchema(ctx, iceberg.PositionalDeleteSchema, enrichedIcebergSchema, r, SchemaOptions{IncludeFieldIDs: true, UseLargeTypes: as.useLargeTypes})
})
- err = as.processRecordsWithPlans(ctx, task, iceSchema, as.boundRowFilter, rdr, colIndices, pipeline, posSource, out, filterPlans)
+ err = as.processRecordsWithPlans(ctx, task, iceSchema, as.boundRowFilter, rdr, colIndices, pipeline, posSource, sink, filterPlans)
return err
}
@@ -2161,6 +2224,7 @@ func createIteratorWithCleanup(ctx context.Context, numWorkers uint, records <-c
prev.Record.Last && next.Record.Index == 0
}
}, enumeratedRecord{Task: tblutils.Enumerated[FileScanTask]{Index: -1}}, func(rec enumeratedRecord) {
+ rec.credits.release()
if rec.Record.Value != nil {
rec.Record.Value.Release()
}
@@ -2171,6 +2235,7 @@ func createIteratorWithCleanup(ctx context.Context, numWorkers uint, records <-c
return func(yield func(arrow.RecordBatch, error) bool) {
defer func() {
for rec := range sequenced {
+ rec.credits.release()
if rec.Record.Value != nil {
rec.Record.Value.Release()
}
@@ -2201,6 +2266,7 @@ func createIteratorWithCleanup(ctx context.Context, numWorkers uint, records <-c
return
}
+ enum.credits.release()
if enum.Err != nil {
yield(nil, enum.Err)
@@ -2251,7 +2317,11 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
for range numWorkers {
go func() {
defer wg.Done()
+ sink := newRecordSink(records)
for {
+ if err := sink.reserve(scanCtx); err != nil {
+ return
+ }
select {
case <-scanCtx.Done():
return
@@ -2269,10 +2339,7 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
var err error
positionalDeletes, err = positionDeleteLoader.load(scanCtx, task.Value)
if err != nil {
- select {
- case records <- enumeratedRecord{Task: task, Err: err}:
- case <-scanCtx.Done():
- }
+ sink.fail(task, err)
cancel(err)
return
@@ -2281,26 +2348,20 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
dvBitmap, err := dvLoader.load(scanCtx, filePath)
if err != nil {
- select {
- case records <- enumeratedRecord{Task: task, Err: err}:
- case <-scanCtx.Done():
- }
+ sink.fail(task, err)
cancel(err)
return
}
eqDeleteSets, err := equalityDeleteLoader.load(scanCtx, task.Value)
if err != nil {
- select {
- case records <- enumeratedRecord{Task: task, Err: err}:
- case <-scanCtx.Done():
- }
+ sink.fail(task, err)
cancel(err)
return
}
- if err := as.recordsFromTask(scanCtx, task, records,
+ if err := as.recordsFromTask(scanCtx, task, sink,
positionalDeletes,
dvBitmap,
eqDeleteSets,
diff --git a/table/arrow_scanner_pruning_internal_test.go b/table/arrow_scanner_pruning_internal_test.go
index 668125afe..4af6adc80 100644
--- a/table/arrow_scanner_pruning_internal_test.go
+++ b/table/arrow_scanner_pruning_internal_test.go
@@ -114,7 +114,7 @@ func TestProcessRecordsUsesRowGroupFilterForPruning(t *testing.T) {
err = (&arrowScan{rowGroupFilter: rowGroupFilter, projectedSchema: fileSchema}).processRecords(
ctx,
tblutils.Enumerated[FileScanTask]{Value: FileScanTask{File: dataFileBuilder.Build()}},
- fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, out)
+ fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, newRecordSink(out))
require.NoError(t, err)
require.NotNil(t, reader.tester)
@@ -179,7 +179,7 @@ func TestProcessRecordsRebindsRowGroupFilterToPromotedFileSchema(t *testing.T) {
}).processRecords(
ctx,
tblutils.Enumerated[FileScanTask]{Value: FileScanTask{File: dataFileBuilder.Build()}},
- fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, out)
+ fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, newRecordSink(out))
require.NoError(t, err)
assert.True(t, reader.statsResult, "matching INT32 row-group stats should be retained")
@@ -227,7 +227,7 @@ func TestProcessRecordsDoesNotPruneMissingInitialDefault(t *testing.T) {
}).processRecords(
ctx,
tblutils.Enumerated[FileScanTask]{Value: FileScanTask{File: dataFileBuilder.Build()}},
- fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, out)
+ fileSchema, iceberg.AlwaysTrue{}, reader, []int{0}, nil, nil, newRecordSink(out))
require.NoError(t, err)
require.NotNil(t, reader.tester)
diff --git a/table/arrow_scanner_reorder_bench_test.go b/table/arrow_scanner_reorder_bench_test.go
new file mode 100644
index 000000000..d95968268
--- /dev/null
+++ b/table/arrow_scanner_reorder_bench_test.go
@@ -0,0 +1,140 @@
+// 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
+
+import (
+ "context"
+ "fmt"
+ "path/filepath"
+ "runtime"
+ "testing"
+ "time"
+)
+
+func sampleHeapInuse() uint64 {
+ runtime.GC()
+ var m runtime.MemStats
+ runtime.ReadMemStats(&m)
+
+ return m.HeapInuse
+}
+
+func BenchmarkArrowScanReorderHeapLaggingTask(b *testing.B) {
+ const (
+ rowsPerFile = 2048
+ numWorkers = 4
+ sampleInterval = 5 * time.Millisecond
+ idleBeforeOpen = 100 * time.Millisecond
+ )
+
+ for _, files := range []int{8, 32, 128} {
+ b.Run(fmt.Sprintf("files_%d", files), func(b *testing.B) {
+ dir := filepath.ToSlash(b.TempDir())
+ gateIO := &reorderGateIO{}
+ tbl, tasks := newReorderScanFixture(b, dir, files, rowsPerFile, gateIO)
+ headPath := tasks[0].File.FilePath()
+ ctx := context.Background()
+
+ baseline := sampleHeapInuse()
+ var (
+ peakDelta float64
+ gatedBatches int
+ )
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ closes := make(chan struct{}, files)
+ gateIO.arm(headPath, closes)
+
+ _, records, err := tbl.Scan(WithMaxConcurrency(numWorkers)).ReadTasks(ctx, tasks)
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ done := make(chan error, 1)
+ go func() {
+ var rows int64
+ for rec, err := range records {
+ if err != nil {
+ done <- err
+
+ return
+ }
+ rows += rec.NumRows()
+ rec.Release()
+ }
+ if rows != int64(files*rowsPerFile) {
+ done <- fmt.Errorf("read %d rows, want %d", rows, files*rowsPerFile)
+
+ return
+ }
+ done <- nil
+ }()
+
+ stopSampler := make(chan struct{})
+ peakCh := make(chan uint64, 1)
+ go func() {
+ ticker := time.NewTicker(sampleInterval)
+ defer ticker.Stop()
+ var peak uint64
+ for {
+ select {
+ case <-stopSampler:
+ peakCh <- max(peak, sampleHeapInuse())
+
+ return
+ case <-ticker.C:
+ peak = max(peak, sampleHeapInuse())
+ }
+ }
+ }()
+
+ closed := 0
+ idle := time.NewTimer(idleBeforeOpen)
+ gated:
+ for closed < files-1 {
+ select {
+ case <-closes:
+ closed++
+ idle.Reset(idleBeforeOpen)
+ case <-idle.C:
+ break gated
+ }
+ }
+ idle.Stop()
+
+ close(stopSampler)
+ peak := <-peakCh
+ gateIO.release()
+
+ if err := <-done; err != nil {
+ b.Fatal(err)
+ }
+
+ if peak > baseline {
+ peakDelta = max(peakDelta, float64(peak-baseline)/(1024*1024))
+ }
+ gatedBatches = max(gatedBatches, closed)
+ }
+
+ b.ReportMetric(peakDelta, "peak-heap-delta-MB")
+ b.ReportMetric(float64(gatedBatches), "gated-batches")
+ })
+ }
+}
diff --git a/table/arrow_scanner_reorder_test.go b/table/arrow_scanner_reorder_test.go
new file mode 100644
index 000000000..a754c759c
--- /dev/null
+++ b/table/arrow_scanner_reorder_test.go
@@ -0,0 +1,401 @@
+// 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
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "os"
+ "path/filepath"
+ "strings"
+ "sync"
+ "testing"
+ "testing/synctest"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "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"
+ tblutils "github.com/apache/iceberg-go/table/internal"
+ "github.com/stretchr/testify/require"
+)
+
+type reorderGateIO struct {
+ iceio.LocalFS
+
+ mu sync.Mutex
+ headPath string
+ gate chan struct{}
+ createGate chan struct{}
+ closes chan<- struct{}
+ released bool
+ closed int
+}
+
+func (g *reorderGateIO) arm(headPath string, closes chan<- struct{}) {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ g.headPath = headPath
+ g.gate = make(chan struct{})
+ g.closes = closes
+ g.released = false
+ g.closed = 0
+}
+
+func (g *reorderGateIO) Open(name string) (iceio.File, error) {
+ g.mu.Lock()
+ headPath, gate := g.headPath, g.gate
+ g.mu.Unlock()
+
+ if gate != nil && name == headPath {
+ <-gate
+ }
+
+ f, err := g.LocalFS.Open(name)
+ if err != nil {
+ return nil, err
+ }
+
+ return &reorderCountingFile{File: f, owner: g}, nil
+}
+
+func (g *reorderGateIO) armCreate() {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ g.createGate = make(chan struct{})
+}
+
+func (g *reorderGateIO) Create(name string) (iceio.FileWriter, error) {
+ g.mu.Lock()
+ gate := g.createGate
+ g.mu.Unlock()
+
+ if gate != nil {
+ <-gate
+ }
+
+ return g.LocalFS.Create(name)
+}
+
+func (g *reorderGateIO) releaseCreate() {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ if g.createGate != nil {
+ close(g.createGate)
+ g.createGate = nil
+ }
+}
+
+func (g *reorderGateIO) noteClose() {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ g.closed++
+ if g.closes != nil {
+ select {
+ case g.closes <- struct{}{}:
+ default:
+ }
+ }
+}
+
+func (g *reorderGateIO) gatedCloses() int {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ return g.closed
+}
+
+func (g *reorderGateIO) release() {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+
+ if g.released {
+ return
+ }
+ g.released = true
+ close(g.gate)
+}
+
+type reorderCountingFile struct {
+ iceio.File
+
+ owner *reorderGateIO
+ once sync.Once
+}
+
+func (f *reorderCountingFile) Close() error {
+ f.once.Do(f.owner.noteClose)
+
+ return f.File.Close()
+}
+
+func newReorderScanFixture(tb testing.TB, dir string, files, rowsPerFile int, fs iceio.IO) (*Table, []FileScanTask) {
+ tb.Helper()
+
+ 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: true},
+ )
+ arrowSchema, err := SchemaToArrowSchemaWithOptions(schema, ArrowSchemaOptions{IncludeFieldIDs: true})
+ require.NoError(tb, err)
+
+ meta, err := NewMetadata(schema, iceberg.UnpartitionedSpec, UnsortedSortOrder, dir,
+ iceberg.Properties{PropertyFormatVersion: "2", ParquetCompressionKey: "uncompressed"})
+ require.NoError(tb, err)
+
+ payload := strings.Repeat("x", 256)
+ tasks := make([]FileScanTask, files)
+ for i := range files {
+ path := fmt.Sprintf("%s/data/file-%03d.parquet", dir, i)
+ tasks[i].File = writeReorderDataFile(tb, arrowSchema, path, i*rowsPerFile, rowsPerFile, payload)
+ }
+
+ tbl := New(Identifier{"db", "reorder_heap"}, meta, dir+"/metadata/v1.metadata.json",
+ func(context.Context) (iceio.IO, error) { return fs, nil }, nil)
+
+ return tbl, tasks
+}
+
+func writeReorderDataFile(tb testing.TB, arrowSchema *arrow.Schema, path string, firstID, rows int, payload string) iceberg.DataFile {
+ tb.Helper()
+
+ mem := memory.DefaultAllocator
+ idBldr := array.NewInt64Builder(mem)
+ defer idBldr.Release()
+ dataBldr := array.NewStringBuilder(mem)
+ defer dataBldr.Release()
+ for i := range rows {
+ idBldr.Append(int64(firstID + i))
+ dataBldr.Append(payload)
+ }
+ idArr := idBldr.NewArray()
+ defer idArr.Release()
+ dataArr := dataBldr.NewArray()
+ defer dataArr.Release()
+
+ rec := array.NewRecordBatch(arrowSchema, []arrow.Array{idArr, dataArr}, int64(rows))
+ defer rec.Release()
+ arrowTable := array.NewTableFromRecords(arrowSchema, []arrow.RecordBatch{rec})
+ defer arrowTable.Release()
+
+ fw, err := iceio.LocalFS{}.Create(path)
+ require.NoError(tb, err)
+ defer fw.Close()
+ require.NoError(tb, pqarrow.WriteTable(arrowTable, fw, int64(rows),
+ parquet.NewWriterProperties(parquet.WithStats(true)), pqarrow.DefaultWriterProps()))
+
+ info, err := os.Stat(path)
+ require.NoError(tb, err)
+
+ bldr, err := iceberg.NewDataFileBuilder(*iceberg.UnpartitionedSpec, iceberg.EntryContentData,
+ path, iceberg.ParquetFile, nil, nil, nil, int64(rows), info.Size())
+ require.NoError(tb, err)
+
+ return bldr.Build()
+}
+
+func TestArrowScanReorderHeapBoundedWhileHeadTaskLags(t *testing.T) {
+ const (
+ files = 32
+ rowsPerFile = 256
+ numWorkers = 4
+ )
+
+ synctest.Test(t, func(t *testing.T) {
+ dir := filepath.ToSlash(t.TempDir())
+ gateIO := &reorderGateIO{}
+ tbl, tasks := newReorderScanFixture(t, dir, files, rowsPerFile, gateIO)
+ gateIO.arm(tasks[0].File.FilePath(), nil)
+
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ ctx := compute.WithAllocator(context.Background(), mem)
+ _, records, err := tbl.Scan(WithMaxConcurrency(numWorkers)).ReadTasks(ctx, tasks)
+ require.NoError(t, err)
+
+ type scanResult struct {
+ rows int64
+ err error
+ }
+ done := make(chan scanResult, 1)
+ go func() {
+ var rows int64
+ for rec, err := range records {
+ if err != nil {
+ done <- scanResult{rows: rows, err: err}
+
+ return
+ }
+ rows += rec.NumRows()
+ rec.Release()
+ }
+ done <- scanResult{rows: rows}
+ }()
+
+ synctest.Wait()
+
+ held := gateIO.gatedCloses()
+ bound := (numWorkers - 1) * maxInFlightTasksPerWorker
+ if held > bound {
+ t.Errorf("scan completed %d of the %d out-of-order tasks while task 0 was gated, %d bytes in Arrow buffers, bound is (numWorkers - 1) x maxInFlightTasksPerWorker = %d x %d = %d",
+ held, files-1, mem.CurrentAlloc(), numWorkers-1, maxInFlightTasksPerWorker, bound)
+ }
+
+ gateIO.release()
+ result := <-done
+ require.NoError(t, result.err)
+ require.Equal(t, int64(files*rowsPerFile), result.rows)
+ mem.AssertSize(t, 0)
+ })
+}
+
+func TestArrowScanReorderCancelWhileWorkersWaitForTaskCredits(t *testing.T) {
+ const (
+ files = 16
+ rowsPerFile = 256
+ numWorkers = 4
+ )
+
+ synctest.Test(t, func(t *testing.T) {
+ dir := filepath.ToSlash(t.TempDir())
+ gateIO := &reorderGateIO{}
+ tbl, tasks := newReorderScanFixture(t, dir, files, rowsPerFile, gateIO)
+ gateIO.arm(tasks[0].File.FilePath(), nil)
+
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ ctx, cancel := context.WithCancel(compute.WithAllocator(context.Background(), mem))
+ defer cancel()
+ _, records, err := tbl.Scan(WithMaxConcurrency(numWorkers)).ReadTasks(ctx, tasks)
+ require.NoError(t, err)
+
+ done := make(chan error, 1)
+ go func() {
+ for rec, err := range records {
+ if err != nil {
+ done <- err
+
+ return
+ }
+ rec.Release()
+ }
+ done <- nil
+ }()
+
+ synctest.Wait()
+ require.Equal(t, numWorkers-1, gateIO.gatedCloses())
+
+ cancel()
+ require.ErrorIs(t, <-done, context.Canceled)
+
+ gateIO.release()
+ synctest.Wait()
+ mem.AssertSize(t, 0)
+ })
+}
+
+func TestRecordSinkHandsOffCreditOnce(t *testing.T) {
+ ctx := context.Background()
+ task := tblutils.Enumerated[FileScanTask]{Index: 3}
+ out := make(chan enumeratedRecord, 4)
+ sink := newRecordSink(out)
+
+ require.NoError(t, sink.reserve(ctx))
+ sink.fail(task, errors.New("open failed"))
+ failed := <-out
+ require.Error(t, failed.Err)
+ require.NotNil(t, failed.credits)
+
+ blocked, cancel := context.WithCancel(ctx)
+ cancel()
+ require.ErrorIs(t, sink.reserve(blocked), context.Canceled)
+ failed.credits.release()
+ require.NoError(t, sink.reserve(ctx))
+
+ sink.send(enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{Index: 0}})
+ sink.send(enumeratedRecord{Task: task, Record: tblutils.Enumerated[arrow.RecordBatch]{Index: 1, Last: true}})
+ sink.fail(task, errors.New("close failed"))
+ first, last, afterLast := <-out, <-out, <-out
+ require.Nil(t, first.credits)
+ require.NotNil(t, last.credits)
+ require.Nil(t, afterLast.credits)
+ require.ErrorIs(t, sink.reserve(blocked), context.Canceled)
+ last.credits.release()
+ require.NoError(t, sink.reserve(ctx))
+}
+
+func TestExecuteCompactionGroupRecordPipelineBounded(t *testing.T) {
+ const (
+ files = 32
+ rowsPerFile = 256
+ numWorkers = 4
+ recordBatchBufferSize = 2
+ )
+
+ synctest.Test(t, func(t *testing.T) {
+ dir := filepath.ToSlash(t.TempDir())
+ gateIO := &reorderGateIO{}
+ tbl, tasks := newReorderScanFixture(t, dir, files, rowsPerFile, gateIO)
+ gateIO.arm(tasks[0].File.FilePath(), nil)
+ gateIO.armCreate()
+
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ type groupResult struct {
+ result CompactionGroupResult
+ err error
+ }
+ done := make(chan groupResult, 1)
+ go func() {
+ result, err := ExecuteCompactionGroup(ctx, tbl,
+ CompactionTaskGroup{PartitionKey: "unpartitioned", Tasks: tasks},
+ WithCompactionScanConcurrency(numWorkers),
+ WithCompactionRecordBatchBufferSize(recordBatchBufferSize))
+ done <- groupResult{result: result, err: err}
+ }()
+
+ synctest.Wait()
+ if got, want := gateIO.gatedCloses(), (numWorkers-1)*maxInFlightTasksPerWorker; got != want {
+ t.Errorf("%d tasks fully read while task 0 was gated in Open, want (workers - 1) x maxInFlightTasksPerWorker = %d", got, want)
+ }
+
+ gateIO.release()
+ synctest.Wait()
+ if got, want := gateIO.gatedCloses(), numWorkers*maxInFlightTasksPerWorker+recordBatchBufferSize+2; got != want {
+ t.Errorf("%d tasks fully read while the writer was gated in Create, want workers x maxInFlightTasksPerWorker + recordBatchBufferSize + 2 = %d", got, want)
+ }
+
+ gateIO.releaseCreate()
+ res := <-done
+ require.NoError(t, res.err)
+ var rows int64
+ for _, df := range res.result.NewDataFiles {
+ rows += df.Count()
+ }
+ require.Equal(t, int64(files*rowsPerFile), rows)
+ mem.AssertSize(t, 0)
+ })
+}
diff --git a/table/rewrite_data_files.go b/table/rewrite_data_files.go
index 81ac9fb83..a66e90da2 100644
--- a/table/rewrite_data_files.go
+++ b/table/rewrite_data_files.go
@@ -249,6 +249,10 @@ func WithCompactionTargetFileSize(size int64) CompactionGroupOption {
// WithCompactionScanConcurrency sets the scan concurrency used when
// reading the group's tasks. Forwarded to [Table.Scan] as
// [WithMaxConcurrency]. Zero (the default) means runtime.GOMAXPROCS.
+// 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].
func WithCompactionScanConcurrency(n int) CompactionGroupOption {
return func(c *compactionGroupConfig) {
c.scanConcurrency = n
@@ -257,13 +261,25 @@ func WithCompactionScanConcurrency(n int) CompactionGroupOption {
// WithCompactionArrowBatchSize caps the number of rows decoded per
// Arrow record batch while reading the group's tasks, forwarded to the
-// scan as [WithArrowBatchSize]. Together with
-// [WithCompactionRecordBatchBufferSize] it bounds the memory held by
-// the record pipeline specifically: buffered batches times rows per
-// batch. Delete-side memory is not covered — positional deletes and
-// deletion-vector bitmaps for the group's tasks are materialized up
-// front and sized by delete volume, not by these knobs. A non-positive
-// value keeps the table's read.parquet.batch-size property.
+// scan as [WithArrowBatchSize]. The record pipeline holds at most
+//
+// workers x (rows in the largest task + n) + (recordBatchBufferSize + 2) x n
+//
+// rows, where workers is the scan's worker count (see
+// [WithCompactionScanConcurrency]) and recordBatchBufferSize is the
+// [WithCompactionRecordBatchBufferSize] value. Each scan worker holds
+// the batches of at most one task at a time, in its hands, in the
+// records channel, in the reorder heap or in the sequenced channel,
+// and does not start its next task until the last batch of the
+// previous one has been consumed; a task's batches total at most its
+// row count rounded up to a multiple of n. The clustered writer holds
+// one batch while routing it, recordBatchBufferSize batches in its
+// input channel and one batch while encoding. Not covered: the Parquet
+// reader's own buffers, including the raw batch it last returned, and
+// delete-side memory, since positional deletes and deletion-vector
+// bitmaps for the group's tasks are materialized up front and sized by
+// delete volume, not by these knobs. A non-positive value keeps the
+// table's read.parquet.batch-size property.
func WithCompactionArrowBatchSize(n int) CompactionGroupOption {
return func(c *compactionGroupConfig) {
if n > 0 {
diff --git a/table/scanner.go b/table/scanner.go
index e86c811b4..7c1b1c7da 100644
--- a/table/scanner.go
+++ b/table/scanner.go
@@ -2139,6 +2139,7 @@ type FileScanTask struct {
//
// The purpose for returning the schema up front is to handle the case where there are no
// rows returned. The resulting Arrow Schema of the projection will still be known.
+// The memory held while tasks are decoded in parallel is described on [Scan.ReadTasks].
func (scan *Scan) ToArrowRecords(ctx context.Context) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
tasks, err := scan.planFiles(ctx, true)
if err != nil {
@@ -2159,6 +2160,15 @@ func (scan *Scan) ToArrowRecords(ctx context.Context) (*arrow.Schema, iter.Seq2[
// read errors are delivered through the iterator when a task referencing that file is
// reached; if no such task is processed, the file is not read and its error is not
// returned. The returned iterator is single-use.
+//
+// With [WithMaxConcurrency] above one, tasks are decoded in parallel and the
+// batches are returned in task order. Each worker holds the decoded batches of
+// at most one task until the iterator has returned them, so while an early task
+// is slow to open or decode, the batches waiting behind it total at most the
+// worker count times the rows in the largest task, rounded up to a multiple of
+// the batch size (see [WithArrowBatchSize]). One large file among small ones can
+// therefore still hold its whole decoded contents while it waits. The bound for
+// a compaction pipeline is stated on [WithCompactionArrowBatchSize].
func (scan *Scan) ReadTasks(ctx context.Context, tasks []FileScanTask) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
if atomic.LoadUint32(&scan.closed) != 0 {
return nil, nil, fmt.Errorf("%w: scan is closed", ErrInvalidOperation)
diff --git a/table/table.go b/table/table.go
index c54396652..39ffa8649 100644
--- a/table/table.go
+++ b/table/table.go
@@ -1264,7 +1264,9 @@ func WithLimit(n int64) ScanOption {
}
// WithMaxConcurrency sets the maximum concurrency for table scan and plan
-// operations. When unset it defaults to runtime.GOMAXPROCS.
+// operations. When unset it defaults to runtime.GOMAXPROCS. Above one, a
+// scan holds the decoded batches of up to n tasks at once while it returns
+// them in task order; [Scan.ReadTasks] states the bound.
func WithMaxConcurrency(n int) ScanOption {
if n <= 0 {
return noopOption
diff --git a/table/transaction.go b/table/transaction.go
index 6a96b7d5f..ab03f560a 100644
--- a/table/transaction.go
+++ b/table/transaction.go
@@ -3233,7 +3233,11 @@ func (t *Transaction) makePositionDeleteRecordsForFilter(ctx context.Context, fs
for range numWorkers {
go func() {
defer wg.Done()
+ sink := newRecordSink(records)
for {
+ if err := sink.reserve(ctx); err != nil {
+ return
+ }
select {
case <-ctx.Done():
return
@@ -3242,7 +3246,7 @@ func (t *Transaction) makePositionDeleteRecordsForFilter(ctx context.Context, fs
return
}
- if err := scanner.producePosDeletesFromTask(ctx, task, deletesPerFile[task.Value.File.FilePath()], records, invariants); err != nil {
+ if err := scanner.producePosDeletesFromTask(ctx, task, deletesPerFile[task.Value.File.FilePath()], sink, invariants); err != nil {
cancel(err)
return