This is an automated email from the ASF dual-hosted git repository.
zeroshade 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 00fbb576e feat(table): sort-order rewrite strategy for data file
compaction (#2345)
00fbb576e is described below
commit 00fbb576ea4e8e9d87757189cc690f3c137dcbca
Author: Saurabh Kumar Ojha <[email protected]>
AuthorDate: Fri Oct 30 01:36:37 2026 +0530
feat(table): sort-order rewrite strategy for data file compaction (#2345)
* feat(table): add sort-order data file compaction
Compaction only bin-packed: ExecuteCompactionGroup streamed each
group's rows straight into WriteRecords, which sorts every record batch
on its own by the table's default sort order. Output files were
therefore sorted per batch, not per file, overlapped each other on the
sort key, and recorded no sort_order_id.
WithCompactionSortOrder switches a group rewrite to sort. The group's
rows are materialized, split by the table's current partition spec,
and each partition is sorted as a whole with sort_indices. Sorted rows
are gathered from the source batches one output batch at a time, so no
column is concatenated beyond one output batch (Arrow's chunked Take
would concatenate whole struct, list and decimal columns per call).
The writer is told the stream is presorted: it skips the per-batch
sort, and since every rolled file is a contiguous run of the stream,
each file records the order's id as sort_order_id when the order is
one of the table's sort orders and every field is a single-source,
order-preserving transform.
Splitting by partition before sorting keeps the stream clustered for
the clustered writer even when a group spans files written under an
older partition spec.
compaction.Config.RewriteAll mirrors Java's rewrite-all so files that
size-based selection skips as optimal can be sorted too. The CLI gains
`compact run --sort` (table default sort order) and `--rewrite-all`.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* feat(table): k-way merge presorted files in sort compaction
When every file of a sort group records the order's sort_order_id and
was written under the current partition spec, ExecuteCompactionGroup
now reads each file as its own stream and k-way merges the streams per
partition instead of materializing and sorting the whole group. Memory
drops from the decoded group to about two batches per file. Each merge
step copies the longest run of the minimum stream that precedes the
runner-up as one zero-copy slice, so files with disjoint ranges are
merged a batch at a time. A file whose rows are not actually sorted
fails the group rather than producing output that claims an order it
does not have. Scan.readTasksPerTask backs the merge: one stream per
task sharing one set of delete loaders, so a shared position or
equality delete file is read once.
Sorting now uses one Iceberg-ordered row comparator for both the
in-memory sort and the merge. Arrow's sort kernel places NaN next to
nulls, while the spec puts NaN after +Inf ascending and first
descending. Bound to a partition's batches the comparator matches the
kernel's speed (1M rows: ~347ms vs ~345ms int64, ~548ms vs ~570ms
string).
sort_order_id is now claimed only when every field but the last is
identity. Ordering by the source column of a coarser transform such as
day(ts) also orders the rows the transform ties on, so for
[day(ts), id] the rows were not sorted by id within a day.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* refactor(table): use only container/heap's public API in the sort merge
Pop the minimum cursor, peek the new root as the runner-up and push the
cursor back, instead of reading the root's children by index.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* feat(compaction): cap groups by MaxGroupSizeBytes
PlanCompaction bin-packed each partition's candidates to
TargetFileSizeBytes, so a group held about one output file's worth of
input. That suits bin-pack, but a sort rewrite orders rows only within
a group: with one-file groups every sorted output file of a partition
overlapped the others on the sort key, and files from an earlier sort
rewrite (already target-sized) each landed in a group of their own, so
the k-way merge path never saw two of them together.
compaction.Config.MaxGroupSizeBytes, mirroring Java's
max-file-group-size-bytes, caps a group's input size independently of
the output file size; ExecuteCompactionGroup already rolls output at
the target size, so one group yields many sorted, non-overlapping
files. Zero keeps the previous packing, so existing plans are
unchanged. The CLI gains --max-group-size on compact analyze and run.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* refactor(table): select compaction rewrite strategy behind an interface
ExecuteCompactionGroup branched on the sort order in three places. Move
the read-and-arrange step and the writer options into an internal
rewriteStrategy with bin-pack (default) and sort implementations, chosen
once from the CompactionGroupOptions. The public API is unchanged.
Also correct the planner comment: PlanCompaction only groups candidates
by size; the rewrite strategy is chosen per group.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* docs(compaction): explain why MaxGroupSizeBytes is separate from the
target file size
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* docs(compaction): trim MaxGroupSizeBytes comment
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
* fix(table): address sort compaction merge review
- Merge only orders whose fields are all identity. The merge compares
source values, and a file sorted by [day(ts)] may validly be unsorted
by ts within a day, so such files failed the group on every run. They
now take the in-memory path; output still records the id when it is
claimable.
- Cap the merge fan-in at the scan concurrency. Each file is its own
stream, so a partition with more files than the scan concurrency now
takes the in-memory path instead of opening every file at once. Docs
state the real cost: up to about five batches per merged file.
- Gather interleaved runs instead of slicing each. Overlapping inputs
give mostly one-row runs, and concatenating a slice per run made the
merge ~20x slower than sorting in memory. Runs are now recorded as
(batch, range) and each output batch is built with the existing Take
gather; a batch that is one run stays a zero-copy slice. On 2 x 131k
rows x 32 int64 columns: interleaved 1.32s -> 64ms, disjoint 3ms
unchanged (BenchmarkMergeInterleaved/Disjoint).
- api.md: sort_order_id is recorded only when the order is claimable;
the in-memory path can reach about twice the group's decoded size and
is always taken by the first sort rewrite.
Co-authored-by: Cursor <[email protected]>
Signed-off-by: saurabhojha <[email protected]>
---------
Signed-off-by: saurabhojha <[email protected]>
Co-authored-by: Cursor <[email protected]>
---
cmd/iceberg/args_test.go | 8 +-
cmd/iceberg/compact.go | 13 +
cmd/iceberg/main.go | 14 +-
table/arrow_scanner.go | 84 ++-
table/compaction/compaction.go | 53 +-
table/compaction/compaction_test.go | 65 ++
table/rewrite_data_files.go | 86 ++-
table/rolling_data_writer.go | 23 +-
table/scanner.go | 109 ++--
table/sort_compaction.go | 1009 ++++++++++++++++++++++++++++++++
table/sort_compaction_bench_test.go | 124 ++++
table/sort_compaction_internal_test.go | 460 +++++++++++++++
table/sort_compaction_test.go | 308 ++++++++++
table/write_records.go | 19 +
website/src/api.md | 16 +
15 files changed, 2334 insertions(+), 57 deletions(-)
diff --git a/cmd/iceberg/args_test.go b/cmd/iceberg/args_test.go
index 42dda0fef..1283511df 100644
--- a/cmd/iceberg/args_test.go
+++ b/cmd/iceberg/args_test.go
@@ -205,23 +205,27 @@ func TestArgsParsing(t *testing.T) {
},
{
name: "compact analyze",
- args: []string{"compact", "analyze", "prod.db.events",
"--target-file-size", "134217728"},
+ args: []string{"compact", "analyze", "prod.db.events",
"--target-file-size", "134217728", "--rewrite-all"},
check: func(t *testing.T, a Args) {
require.NotNil(t, a.Compact)
require.NotNil(t, a.Compact.Analyze)
assert.Equal(t, "prod.db.events",
a.Compact.Analyze.TableID)
assert.Equal(t, int64(134217728),
a.Compact.Analyze.TargetFileSize)
+ assert.True(t, a.Compact.Analyze.RewriteAll)
},
},
{
name: "compact run with all flags",
- args: []string{"compact", "run", "prod.db.events",
"--partial-progress", "--preserve-dead-equality-deletes"},
+ args: []string{"compact", "run", "prod.db.events",
"--partial-progress", "--preserve-dead-equality-deletes", "--sort",
"--rewrite-all", "--max-group-size", "10737418240"},
check: func(t *testing.T, a Args) {
require.NotNil(t, a.Compact)
require.NotNil(t, a.Compact.Run)
assert.Equal(t, "prod.db.events",
a.Compact.Run.TableID)
assert.True(t, a.Compact.Run.PartialProgress)
assert.True(t,
a.Compact.Run.PreserveDeadEqualityDeletes)
+ assert.True(t, a.Compact.Run.Sort)
+ assert.True(t, a.Compact.Run.RewriteAll)
+ assert.Equal(t, int64(10737418240),
a.Compact.Run.MaxGroupSize)
},
},
{
diff --git a/cmd/iceberg/compact.go b/cmd/iceberg/compact.go
index 3b689acbe..aa02ae156 100644
--- a/cmd/iceberg/compact.go
+++ b/cmd/iceberg/compact.go
@@ -34,12 +34,22 @@ type compactConfig struct {
analyzeOnly bool
partialProgress bool
preserveDeadEqualityDeletes bool
+ sort bool
+ rewriteAll bool
+ maxGroupSize int64
}
func compact(ctx context.Context, output Output, cat catalog.Catalog, cfg
compactConfig) {
tbl := loadTable(ctx, output, cat, cfg.tableID)
+ if cfg.sort && tbl.SortOrder().IsUnsorted() {
+ output.Error(fmt.Errorf("--sort requires a default sort order
on table %s", cfg.tableID))
+ osExit(1)
+ }
+
compCfg := compaction.DefaultConfig()
+ compCfg.RewriteAll = cfg.rewriteAll
+ compCfg.MaxGroupSizeBytes = cfg.maxGroupSize
if cfg.targetFileSize > 0 {
compCfg.TargetFileSizeBytes = cfg.targetFileSize
// Derive min/max using the same ratios as DefaultConfig
(75%/180%).
@@ -126,6 +136,9 @@ func compactRun(ctx context.Context, output Output, tbl
*table.Table, plan compa
rewriteOpts := table.RewriteDataFilesOptions{
PartialProgress: cfg.partialProgress,
}
+ if cfg.sort {
+ rewriteOpts.GroupOptions = append(rewriteOpts.GroupOptions,
table.WithCompactionSortOrder(tbl.SortOrder()))
+ }
// Cleanup of dead equality-delete files. The executor only
// orchestrates the commit; the policy + walk live in
diff --git a/cmd/iceberg/main.go b/cmd/iceberg/main.go
index 608bf8fa6..ad1b5406f 100644
--- a/cmd/iceberg/main.go
+++ b/cmd/iceberg/main.go
@@ -163,6 +163,8 @@ type PropertiesCmd struct {
type CompactAnalyzeCmd struct {
TableID string `arg:"positional,required" help:"full path to a
table"`
TargetFileSize int64 `arg:"--target-file-size" default:"0"
help:"target output file size in bytes"`
+ RewriteAll bool `arg:"--rewrite-all" help:"select every file
regardless of size, deletes or minimum group size"`
+ MaxGroupSize int64 `arg:"--max-group-size" default:"0" help:"maximum
input bytes per group (default: the target file size)"`
}
type CompactRunCmd struct {
@@ -170,11 +172,14 @@ type CompactRunCmd struct {
TargetFileSize int64 `arg:"--target-file-size"
default:"0" help:"target output file size in bytes"`
PartialProgress bool `arg:"--partial-progress"
help:"stage each group as a separate snapshot"`
PreserveDeadEqualityDeletes bool
`arg:"--preserve-dead-equality-deletes" help:"keep equality-delete files that
are provably dead after the rewrite (default: drop them — recommended for
sustained CDC workloads)"`
+ Sort bool `arg:"--sort" help:"sort each group
by the table's default sort order instead of bin-packing it"`
+ RewriteAll bool `arg:"--rewrite-all" help:"select
every file regardless of size, deletes or minimum group size"`
+ MaxGroupSize int64 `arg:"--max-group-size" default:"0"
help:"maximum input bytes per group (default: the target file size); with
--sort, raise it so a group covers a partition, within memory"`
}
type CompactCmd struct {
Analyze *CompactAnalyzeCmd `arg:"subcommand:analyze" help:"analyze
compaction plan"`
- Run *CompactRunCmd `arg:"subcommand:run" help:"run bin-pack
compaction"`
+ Run *CompactRunCmd `arg:"subcommand:run" help:"run bin-pack or
sort compaction"`
}
type RewriteManifestsCmd struct {
@@ -197,7 +202,7 @@ type Args struct {
Files *FilesCmd `arg:"subcommand:files"
help:"list all files of a table"`
Rename *RenameCmd `arg:"subcommand:rename"
help:"rename a table"`
Properties *PropertiesCmd `arg:"subcommand:properties"
help:"manage properties on tables/namespaces"`
- Compact *CompactCmd `arg:"subcommand:compact"
help:"analyze or run bin-pack compaction"`
+ Compact *CompactCmd `arg:"subcommand:compact"
help:"analyze or run bin-pack or sort compaction"`
RewriteManifests *RewriteManifestsCmd
`arg:"subcommand:rewrite-manifests" help:"rewrite (compact) table manifests"`
Info *InfoCmd `arg:"subcommand:info" help:"show
single-screen table summary"`
Snapshots *SnapshotsCmd `arg:"subcommand:snapshots"
help:"list table snapshots"`
@@ -752,6 +757,8 @@ func runCompact(ctx context.Context, output Output, cat
catalog.Catalog, cmd *Co
tableID: cmd.Analyze.TableID,
targetFileSize: cmd.Analyze.TargetFileSize,
analyzeOnly: true,
+ rewriteAll: cmd.Analyze.RewriteAll,
+ maxGroupSize: cmd.Analyze.MaxGroupSize,
})
case cmd.Run != nil:
compact(ctx, output, cat, compactConfig{
@@ -759,6 +766,9 @@ func runCompact(ctx context.Context, output Output, cat
catalog.Catalog, cmd *Co
targetFileSize: cmd.Run.TargetFileSize,
partialProgress: cmd.Run.PartialProgress,
preserveDeadEqualityDeletes:
cmd.Run.PreserveDeadEqualityDeletes,
+ sort: cmd.Run.Sort,
+ rewriteAll: cmd.Run.RewriteAll,
+ maxGroupSize: cmd.Run.MaxGroupSize,
})
}
}
diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go
index 8c884041b..ad7e68105 100644
--- a/table/arrow_scanner.go
+++ b/table/arrow_scanner.go
@@ -3046,6 +3046,18 @@ func newScanTaskChannel(taskCount, numWorkers int) chan
tblutils.Enumerated[File
}
func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context,
tasks []FileScanTask, positionDeleteLoader *lazyPositionDeleteLoader, dvLoader
*lazyDeletionVectorLoader, equalityDeleteLoader *lazyEqualityDeleteLoader,
invariants *arrowScanInvariants) iter.Seq2[arrow.RecordBatch, error] {
+ var cleanup func()
+ if positionDeleteLoader != nil {
+ cleanup = positionDeleteLoader.release
+ }
+
+ return as.recordBatchesWithCleanup(ctx, tasks, 0, positionDeleteLoader,
dvLoader, equalityDeleteLoader, invariants, cleanup)
+}
+
+// recordBatchesWithCleanup is recordBatchesFromTasksAndDeletes with the
+// cleanup run when iteration ends chosen by the caller. tasks start at
+// loaderOffset in the task slice positionDeleteLoader was built from.
+func (as *arrowScan) recordBatchesWithCleanup(ctx context.Context, tasks
[]FileScanTask, loaderOffset int, positionDeleteLoader
*lazyPositionDeleteLoader, dvLoader *lazyDeletionVectorLoader,
equalityDeleteLoader *lazyEqualityDeleteLoader, invariants
*arrowScanInvariants, cleanup func()) iter.Seq2[arrow.RecordBatch, error] {
return func(yield func(arrow.RecordBatch, error) bool) {
extSet := substrait.NewExtensionSet()
scanCtx, cancel :=
context.WithCancelCause(exprs.WithExtensionIDSet(ctx, extSet))
@@ -3078,7 +3090,7 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx
context.Context, tasks
var positionalDeleteBitmap
*dv.RoaringPositionBitmap
if positionDeleteLoader != nil {
var err error
- positionalDeleteBitmap,
err = positionDeleteLoader.load(scanCtx, task.Index, filePath)
+ positionalDeleteBitmap,
err = positionDeleteLoader.load(scanCtx, loaderOffset+task.Index, filePath)
if err != nil {
sink.fail(task,
err)
cancel(err)
@@ -3134,10 +3146,6 @@ func (as *arrowScan)
recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
}
}()
- var cleanup func()
- if positionDeleteLoader != nil {
- cleanup = positionDeleteLoader.release
- }
createIteratorWithCleanup(scanCtx, uint(numWorkers), records,
nil,
cancel, as.rowLimit, cleanup)(yield)
}
@@ -3150,6 +3158,59 @@ func (as *arrowScan)
recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
// termination may finish the scan before the error is observed.
Deletion-vector
// errors are returned before the iterator is created.
func (as *arrowScan) GetRecords(ctx context.Context, tasks []FileScanTask)
(*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
+ resultSchema, reads, err := as.prepareTaskReads(ctx, tasks)
+ if err != nil {
+ return nil, nil, err
+ }
+ if reads == nil {
+ return resultSchema, func(yield func(arrow.RecordBatch, error)
bool) {}, nil
+ }
+
+ return resultSchema, as.recordBatchesFromTasksAndDeletes(reads.ctx,
tasks,
+ reads.positionDeletes, reads.deletionVectors,
reads.equalityDeletes, reads.invariants), nil
+}
+
+// GetRecordsPerTask is GetRecords returning one single-use stream per task,
+// in task order. The streams share one set of delete-file loaders, so a
+// delete file that applies to several tasks is read once, and they may be
+// read concurrently and in any order. release frees the shared position
+// deletes and must be called once, after the last stream is done.
+func (as *arrowScan) GetRecordsPerTask(ctx context.Context, tasks
[]FileScanTask) (*arrow.Schema, []iter.Seq2[arrow.RecordBatch, error], func(),
error) {
+ resultSchema, reads, err := as.prepareTaskReads(ctx, tasks)
+ if err != nil {
+ return nil, nil, nil, err
+ }
+
+ streams := make([]iter.Seq2[arrow.RecordBatch, error], len(tasks))
+ if reads == nil {
+ for i := range streams {
+ streams[i] = func(yield func(arrow.RecordBatch, error)
bool) {}
+ }
+
+ return resultSchema, streams, func() {}, nil
+ }
+
+ for i := range tasks {
+ streams[i] = as.recordBatchesWithCleanup(reads.ctx,
tasks[i:i+1], i,
+ reads.positionDeletes, reads.deletionVectors,
reads.equalityDeletes, reads.invariants, nil)
+ }
+
+ return resultSchema, streams, reads.positionDeletes.release, nil
+}
+
+// taskReads is the state shared by every task read of one scan: the
+// prepared context and the lazily loaded delete files.
+type taskReads struct {
+ ctx context.Context
+ invariants *arrowScanInvariants
+ positionDeletes *lazyPositionDeleteLoader
+ deletionVectors *lazyDeletionVectorLoader
+ equalityDeletes *lazyEqualityDeleteLoader
+}
+
+// prepareTaskReads resolves the projected Arrow schema and the state to
+// read tasks with. It returns nil reads when the row limit is zero.
+func (as *arrowScan) prepareTaskReads(ctx context.Context, tasks
[]FileScanTask) (*arrow.Schema, *taskReads, error) {
var err error
as.useLargeTypes, err =
strconv.ParseBool(as.options.Get(ScanOptionArrowUseLargeTypes, "false"))
if err != nil {
@@ -3179,7 +3240,7 @@ func (as *arrowScan) GetRecords(ctx context.Context,
tasks []FileScanTask) (*arr
}
if as.rowLimit == 0 {
- return resultSchema, func(yield func(arrow.RecordBatch, error)
bool) {}, nil
+ return resultSchema, nil, nil
}
if len(tasks) > 1 {
@@ -3216,8 +3277,11 @@ func (as *arrowScan) GetRecords(ctx context.Context,
tasks []FileScanTask) (*arr
}
equalityDeleteLoader.addFieldIDs(invariants.projectedIDs)
- positionDeleteLoader := newLazyPositionDeleteLoader(as.fs, tasks)
-
- return resultSchema, as.recordBatchesFromTasksAndDeletes(ctx, tasks,
- positionDeleteLoader, dvLoader, equalityDeleteLoader,
invariants), nil
+ return resultSchema, &taskReads{
+ ctx: ctx,
+ invariants: invariants,
+ positionDeletes: newLazyPositionDeleteLoader(as.fs, tasks),
+ deletionVectors: dvLoader,
+ equalityDeletes: equalityDeleteLoader,
+ }, nil
}
diff --git a/table/compaction/compaction.go b/table/compaction/compaction.go
index c1c5d026f..1af817f3c 100644
--- a/table/compaction/compaction.go
+++ b/table/compaction/compaction.go
@@ -87,6 +87,31 @@ type Config struct {
// into the snapshot walk +
RewriteDataFilesOptions.ExtraDeleteFilesToRemove
// — see CollectDeadEqualityDeletesWithSpecs.
PreserveDeadEqualityDeletes bool
+
+ // RewriteAll selects every file for rewriting regardless of size or
+ // deletes and keeps every packed group regardless of MinInputFiles.
+ // Mirrors Java's rewrite-all option. Use it with
+ // [table.WithCompactionSortOrder] to sort files that the size-based
+ // selection would skip as already optimal.
+ RewriteAll bool
+
+ // MaxGroupSizeBytes caps the total input size of one group. Candidates
+ // of a partition are packed into groups of up to this size, and each
+ // group is rewritten into files of about TargetFileSizeBytes, so group
+ // size and output file size are set independently. When zero, groups
+ // are packed to TargetFileSizeBytes (always positive; see Validate) and
+ // each group yields about one file. Mirrors Java's
+ // max-file-group-size-bytes.
+ //
+ // Bin-pack rewrites are unaffected by the group size beyond how work is
+ // split. Sort rewrites ([table.WithCompactionSortOrder]) depend on it:
+ // rows are ordered only within a group, so a partition split across
+ // several groups yields files that each span the partition's whole key
+ // range, and readers cannot skip any of them by min/max statistics. For
+ // example, 10 GB of small files with a 512 MB target is about 20
+ // separately sorted groups whose files all overlap; one 10 GB group is
+ // sorted once and written as about 20 files with disjoint key ranges.
+ MaxGroupSizeBytes int64
}
const (
@@ -142,10 +167,23 @@ func (cfg Config) Validate() error {
if cfg.DeleteRatioThreshold < 0 || cfg.DeleteRatioThreshold > 1 {
return fmt.Errorf("delete ratio threshold must be in [0,1], got
%g", cfg.DeleteRatioThreshold)
}
+ if cfg.MaxGroupSizeBytes != 0 && cfg.MaxGroupSizeBytes <
cfg.TargetFileSizeBytes {
+ return fmt.Errorf("max group size (%d) must be zero or at least
the target file size (%d)",
+ cfg.MaxGroupSizeBytes, cfg.TargetFileSizeBytes)
+ }
return nil
}
+// groupSizeBytes is the weight candidates are bin-packed to.
+func (cfg Config) groupSizeBytes() int64 {
+ if cfg.MaxGroupSizeBytes > 0 {
+ return cfg.MaxGroupSizeBytes
+ }
+
+ return cfg.TargetFileSizeBytes
+}
+
// Group represents a set of files in the same partition
// that should be compacted together.
type Group struct {
@@ -189,7 +227,8 @@ type Plan struct {
// PlanCompaction analyzes the given scan tasks and produces a Plan.
// Tasks are grouped by partition, classified as candidates or skipped,
-// and candidates are bin-packed into groups targeting TargetFileSizeBytes.
+// and candidates are bin-packed into groups of up to MaxGroupSizeBytes
+// (TargetFileSizeBytes when zero).
// Tasks should come from one coherent scan plan; split tasks for the same file
// are collapsed, retaining the first task's per-file delete and lineage data.
// Residual filters select files to rewrite, but are not applied to file
contents.
@@ -242,7 +281,7 @@ func (cfg Config) PlanCompaction(tasks
[]table.FileScanTask) (Plan, error) {
var partitionBuckets []partitionBucket
for _, t := range wholeFileTasks {
- if !cfg.isCandidate(t) {
+ if !cfg.RewriteAll && !cfg.isCandidate(t) {
plan.SkippedFiles++
continue
@@ -258,15 +297,17 @@ func (cfg Config) PlanCompaction(tasks
[]table.FileScanTask) (Plan, error) {
partitionBuckets[index].candidates =
append(partitionBuckets[index].candidates, t)
}
- // Bin-pack candidates per partition.
+ // Group candidates per partition by size. This only splits the work;
+ // how each group's rows are rewritten (bin-pack or sort) is chosen per
+ // group with [table.CompactionGroupOption]s.
packer := internal.SlicePacker[table.FileScanTask]{
- TargetWeight: cfg.TargetFileSizeBytes,
+ TargetWeight: cfg.groupSizeBytes(),
Lookback: int(cfg.PackingLookback),
LargestBinFirst: false,
}
for _, bucket := range partitionBuckets {
- if len(bucket.candidates) < int(cfg.MinInputFiles) {
+ if !cfg.RewriteAll && len(bucket.candidates) <
int(cfg.MinInputFiles) {
plan.SkippedFiles += len(bucket.candidates)
continue
@@ -277,7 +318,7 @@ func (cfg Config) PlanCompaction(tasks
[]table.FileScanTask) (Plan, error) {
})
for _, bin := range bins {
- if len(bin) < int(cfg.MinInputFiles) {
+ if !cfg.RewriteAll && len(bin) < int(cfg.MinInputFiles)
{
plan.SkippedFiles += len(bin)
continue
diff --git a/table/compaction/compaction_test.go
b/table/compaction/compaction_test.go
index dbba3cf9a..ce4f6f952 100644
--- a/table/compaction/compaction_test.go
+++ b/table/compaction/compaction_test.go
@@ -278,6 +278,71 @@ func TestPlanCompaction_AllFilesOptimal(t *testing.T) {
assert.Equal(t, 5, plan.TotalInputFiles)
}
+func TestPlanCompaction_RewriteAll(t *testing.T) {
+ cfg := compaction.DefaultConfig()
+ cfg.RewriteAll = true
+
+ // Optimal and oversized files, too few per partition for MinInputFiles:
+ // size-based selection would skip every one of them.
+ tasks := []table.FileScanTask{
+ makeTask(newPartitionedDataFile("a-optimal.parquet", 500,
map[int]any{1000: int32(1)}), 0, 0),
+ makeTask(newPartitionedDataFile("a-oversized.parquet", 2000,
map[int]any{1000: int32(1)}), 0, 0),
+ makeTask(newPartitionedDataFile("b-optimal.parquet", 500,
map[int]any{1000: int32(2)}), 0, 0),
+ }
+
+ plan, err := cfg.PlanCompaction(tasks)
+ require.NoError(t, err)
+
+ assert.Equal(t, 0, plan.SkippedFiles)
+ var grouped int
+ for _, g := range plan.Groups {
+ grouped += len(g.Tasks)
+ }
+ assert.Equal(t, len(tasks), grouped)
+
+ cfg.RewriteAll = false
+ plan, err = cfg.PlanCompaction(tasks)
+ require.NoError(t, err)
+ assert.Empty(t, plan.Groups)
+ assert.Equal(t, len(tasks), plan.SkippedFiles)
+}
+
+func TestPlanCompaction_MaxGroupSize(t *testing.T) {
+ cfg := compaction.DefaultConfig()
+ cfg.RewriteAll = true
+
+ // Four right-sized files of one partition and one of another.
+ var tasks []table.FileScanTask
+ for i := range 4 {
+ tasks = append(tasks,
makeTask(newPartitionedDataFile(fmt.Sprintf("a-%d.parquet", i), 500,
map[int]any{1000: int32(1)}), 0, 0))
+ }
+ tasks = append(tasks, makeTask(newPartitionedDataFile("b-0.parquet",
500, map[int]any{1000: int32(2)}), 0, 0))
+
+ // Packed to one target file, every file lands in a group of its own.
+ plan, err := cfg.PlanCompaction(tasks)
+ require.NoError(t, err)
+ assert.Len(t, plan.Groups, 5)
+
+ // A group cap above the partition size keeps each partition whole, and
+ // each group still splits into target-sized output files.
+ cfg.MaxGroupSizeBytes = 10 * 1024 * 1024 * 1024
+ plan, err = cfg.PlanCompaction(tasks)
+ require.NoError(t, err)
+ require.Len(t, plan.Groups, 2)
+ sizes := []int{len(plan.Groups[0].Tasks), len(plan.Groups[1].Tasks)}
+ assert.ElementsMatch(t, []int{4, 1}, sizes)
+ assert.Equal(t, 5, plan.EstOutputFiles)
+}
+
+func TestValidate_MaxGroupSizeBelowTarget(t *testing.T) {
+ cfg := compaction.DefaultConfig()
+ cfg.MaxGroupSizeBytes = cfg.TargetFileSizeBytes - 1
+ assert.ErrorContains(t, cfg.Validate(), "max group size")
+
+ cfg.MaxGroupSizeBytes = cfg.TargetFileSizeBytes
+ assert.NoError(t, cfg.Validate())
+}
+
func TestPlanCompaction_SmallFiles(t *testing.T) {
cfg := compaction.Config{
TargetFileSizeBytes: 100 * 1024 * 1024,
diff --git a/table/rewrite_data_files.go b/table/rewrite_data_files.go
index cda61f6c7..e27933545 100644
--- a/table/rewrite_data_files.go
+++ b/table/rewrite_data_files.go
@@ -22,12 +22,15 @@ import (
"errors"
"fmt"
iofs "io/fs"
+ "iter"
"log/slog"
"maps"
+ "github.com/apache/arrow-go/v18/arrow"
"github.com/apache/iceberg-go"
iceberginternal "github.com/apache/iceberg-go/internal"
iceio "github.com/apache/iceberg-go/io"
+ tblutils "github.com/apache/iceberg-go/table/internal"
)
// RewriteResult summarizes a completed compaction.
@@ -229,6 +232,7 @@ type compactionGroupConfig struct {
arrowBatchSize int
recordBatchBufferSize int
parquetRowGroupLimit int
+ sortOrder *SortOrder
}
// WithCompactionTargetFileSize sets the size target for output files
@@ -312,6 +316,80 @@ func WithCompactionParquetRowGroupLimit(n int)
CompactionGroupOption {
}
}
+// WithCompactionSortOrder switches the group rewrite from bin-pack to
+// sort: every row of the group is sorted by order before being written,
+// so each output file holds a sorted, non-overlapping range of the group
+// instead of one sorted run per record batch. Rows are split by the
+// table's current partition spec first and each partition is sorted on
+// its own.
+//
+// Each field sorts by its first source column, so a field whose
+// transform does not preserve order (bucket, void) or that has several
+// source columns only clusters rows by that column. Output files record
+// order's id as sort_order_id when order is one of the table's sort
+// orders (for example [Table.SortOrder]), every field is single-source,
+// every field but the last is identity and the last preserves order;
+// otherwise they record none.
+//
+// When every file of the group records order's id as its sort_order_id,
+// every field of order is identity, every file was written under the
+// table's current partition spec and no partition has more files than
+// the scan concurrency ([WithCompactionScanConcurrency], GOMAXPROCS by
+// default), the files of each partition are read as concurrent streams
+// and merged. Memory then holds up to about five batches per file being
+// merged rather than the whole group. A file whose rows turn out not to
+// be sorted fails the group. Otherwise the whole group is held in memory
+// decoded until its files are written; with the batches queued for the
+// writer, that can reach about twice the group's decoded size.
+//
+// Rows are ordered only within a group, so files of different groups
+// overlap on the sort key. Plan groups with
+// [compaction.Config.MaxGroupSizeBytes] large enough to cover a
+// partition, and small enough to fit in memory when its files are not
+// already sorted.
+// [WithCompactionArrowBatchSize] sets the rows per batch read and handed
+// to the writer. An unsorted order is rejected by [ExecuteCompactionGroup].
+func WithCompactionSortOrder(order SortOrder) CompactionGroupOption {
+ return func(c *compactionGroupConfig) {
+ c.sortOrder = &order
+ }
+}
+
+// rewriteStrategy decides how a compaction group's rows are arranged
+// before they are written: which rows are read in what order, and what
+// the writer may assume about that order.
+type rewriteStrategy interface {
+ // records reads tasks with scan, deletes applied, and returns the
+ // rows in the order they are to be written.
+ records(ctx context.Context, scan *Scan, tasks []FileScanTask)
(*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error)
+ // writeOptions are added to the options the group is written with.
+ writeOptions() []WriteRecordOption
+}
+
+// rewriteStrategy returns the strategy the options select: sort when a
+// sort order is set, bin-pack otherwise.
+func (c compactionGroupConfig) rewriteStrategy(tbl *Table) (rewriteStrategy,
error) {
+ if c.sortOrder == nil {
+ return binPackStrategy{}, nil
+ }
+
+ batchRows := int64(c.arrowBatchSize)
+ if batchRows <= 0 {
+ batchRows =
int64(tbl.Properties().GetInt(tblutils.ParquetBatchSizeKey,
tblutils.ParquetBatchSizeDefault))
+ }
+
+ return newSortStrategy(tbl, *c.sortOrder, max(batchRows, 1))
+}
+
+// binPackStrategy writes the group's rows in the order they are read.
+type binPackStrategy struct{}
+
+func (binPackStrategy) records(ctx context.Context, scan *Scan, tasks
[]FileScanTask) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
+ return scan.ReadTasks(ctx, tasks)
+}
+
+func (binPackStrategy) writeOptions() []WriteRecordOption { return nil }
+
// RewriteDataFiles compacts the given groups by reading data with
// deletes applied, writing new consolidated files, and atomically
// replacing the old files. Position delete files that are fully
@@ -459,6 +537,11 @@ func ExecuteCompactionGroup(ctx context.Context, tbl
*Table, group CompactionTas
opt(&cfg)
}
+ strategy, err := cfg.rewriteStrategy(tbl)
+ if err != nil {
+ return CompactionGroupResult{}, fmt.Errorf("compaction group
%q: %w", group.PartitionKey, err)
+ }
+
var scanOpts []ScanOption
if cfg.scanConcurrency > 0 {
scanOpts = append(scanOpts,
WithMaxConcurrency(cfg.scanConcurrency))
@@ -499,7 +582,7 @@ func ExecuteCompactionGroup(ctx context.Context, tbl
*Table, group CompactionTas
}
}
- arrowSchema, records, err := tbl.Scan(scanOpts...).ReadTasks(ctx,
group.Tasks)
+ arrowSchema, records, err := strategy.records(ctx,
tbl.Scan(scanOpts...), group.Tasks)
if err != nil {
return CompactionGroupResult{}, fmt.Errorf("read tasks for
compaction group %q: %w", group.PartitionKey, err)
}
@@ -533,6 +616,7 @@ func ExecuteCompactionGroup(ctx context.Context, tbl
*Table, group CompactionTas
}
writeOpts = append(writeOpts,
WithPreserveRowLineage(projectedSchema))
}
+ writeOpts = append(writeOpts, strategy.writeOptions()...)
var (
newFiles []iceberg.DataFile
diff --git a/table/rolling_data_writer.go b/table/rolling_data_writer.go
index e87d6fc3d..507fde2f3 100644
--- a/table/rolling_data_writer.go
+++ b/table/rolling_data_writer.go
@@ -77,6 +77,11 @@ type writerFactory struct {
equalityFieldIDs []int
sortKeys []compute.SortKey
+ // presorted marks the input as already globally sorted, so batches are
+ // not re-sorted and every file records sortOrderID.
+ presorted bool
+ sortOrderID int
+
// Variant shredding: per-file inference over the first shredBufferRows
rows.
shredEnabled bool
shredBufferRows int
@@ -112,6 +117,18 @@ func withFactoryEqualityFieldIDs(ids []int)
writerFactoryOption {
}
}
+// withFactoryPresorted declares the input records globally sorted by the
+// sort order with id sortOrderID (UnsortedSortOrderID to make no claim).
+// Each rolled file holds a contiguous run of the input, so it is sorted too.
+func withFactoryPresorted(sortOrderID int) writerFactoryOption {
+ return func(w *writerFactory) error {
+ w.presorted = true
+ w.sortOrderID = sortOrderID
+
+ return nil
+ }
+}
+
func withFactoryFileSchema(schema *iceberg.Schema) writerFactoryOption {
return func(w *writerFactory) error {
w.fileSchema = schema
@@ -247,7 +264,7 @@ func newWriterFactory(rootLocation string, args
recordWritingArgs, meta *Metadat
return nil, err
}
- if f.content == iceberg.EntryContentData && meta.defaultSortOrderID !=
UnsortedSortOrderID {
+ if f.content == iceberg.EntryContentData && !f.presorted &&
meta.defaultSortOrderID != UnsortedSortOrderID {
sortOrder, err := meta.GetSortOrderByID(meta.defaultSortOrderID)
if err != nil {
stopCount()
@@ -303,7 +320,8 @@ func (w *writerFactory) openFileWriter(ctx context.Context,
partitionPath string
filePath = w.locProvider.NewDataLocation(fileName)
}
- // No SortOrderID: batches are sorted only individually (see
resolveSortKeys).
+ // SortOrderID is set only for presorted input: otherwise batches are
sorted
+ // individually (see resolveSortKeys), which does not sort the whole
file.
// FileSchema stays logical; only arrowSchema carries the shredded
layout.
return w.format.NewFileWriter(ctx, w.fs, partitionValues,
tblutils.WriteFileInfo{
FileSchema: w.fileSchema,
@@ -316,6 +334,7 @@ func (w *writerFactory) openFileWriter(ctx context.Context,
partitionPath string
Spec: w.currentSpec,
Content: w.content,
EqualityFieldIDs: w.equalityFieldIDs,
+ SortOrderID: w.sortOrderID,
}, arrowSchema)
}
diff --git a/table/scanner.go b/table/scanner.go
index 968c93b47..f9c1f860a 100644
--- a/table/scanner.go
+++ b/table/scanner.go
@@ -2186,40 +2186,100 @@ func bindReadTasksResiduals(
// 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) {
+ as, readTasks, releasePlanIO, err := scan.arrowScanForTasks(ctx, tasks)
+ if err != nil {
+ return nil, nil, err
+ }
+
+ readRecords := scan.readRecords
+ if readRecords == nil {
+ readRecords = (*arrowScan).GetRecords
+ }
+ outSchema, records, err := readRecords(as, ctx, readTasks)
+ if err != nil {
+ // No iterator to drive cleanup on a setup error, so release
here.
+ if releasePlanIO != nil {
+ releasePlanIO()
+ }
+
+ return nil, nil, err
+ }
+
+ if releasePlanIO != nil {
+ records = releasePlanIOAfter(records, releasePlanIO)
+ }
+
+ return outSchema, records, nil
+}
+
+// readTasksPerTask is [Scan.ReadTasks] returning one single-use record
+// stream per task, in task order, instead of one stream for all tasks.
+// The streams share the scan's delete files, so each delete file is read
+// at most once however many streams it applies to, and they may be read
+// concurrently and in any order. Each stream decodes its task on one
+// worker. release must be called exactly once, after the last stream is
+// done being read.
+func (scan *Scan) readTasksPerTask(ctx context.Context, tasks []FileScanTask)
(*arrow.Schema, []iter.Seq2[arrow.RecordBatch, error], func(), error) {
+ as, readTasks, releasePlanIO, err := scan.arrowScanForTasks(ctx, tasks)
+ if err != nil {
+ return nil, nil, nil, err
+ }
+
+ outSchema, streams, releaseDeletes, err := as.GetRecordsPerTask(ctx,
readTasks)
+ if err != nil {
+ if releasePlanIO != nil {
+ releasePlanIO()
+ }
+
+ return nil, nil, nil, err
+ }
+
+ release := func() {
+ releaseDeletes()
+ if releasePlanIO != nil {
+ releasePlanIO()
+ }
+ }
+
+ return outSchema, streams, release, nil
+}
+
+// arrowScanForTasks builds the arrowScan that reads tasks and returns it
+// with the tasks' residuals bound to the scan's schema. A non-nil
+// releasePlanIO frees the plan-scoped FileIO lease and must be called once
+// reading is done.
+func (scan *Scan) arrowScanForTasks(ctx context.Context, tasks []FileScanTask)
(as *arrowScan, readTasks []FileScanTask, releasePlanIO func(), err error) {
if atomic.LoadUint32(&scan.closed) != 0 {
- return nil, nil, fmt.Errorf("%w: scan is closed",
ErrInvalidOperation)
+ return nil, nil, nil, fmt.Errorf("%w: scan is closed",
ErrInvalidOperation)
}
if scan.selectorErr != nil {
- return nil, nil, scan.selectorErr
+ return nil, nil, nil, scan.selectorErr
}
- var (
- boundFilter iceberg.BooleanExpression
- err error
- )
+ var boundFilter iceberg.BooleanExpression
effectiveSchema, err := scan.effectiveSchema()
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
if scan.rowFilter != nil {
boundFilter, err = iceberg.BindExpr(effectiveSchema,
scan.rowFilter, scan.caseSensitive)
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
}
schema, err := scan.Projection()
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
// Bind task residuals against the schema selected by this scan, which
may
// be an older snapshot schema rather than the table's current schema.
- readTasks, err := bindReadTasksResiduals(effectiveSchema, tasks,
scan.caseSensitive)
+ readTasks, err = bindReadTasksResiduals(effectiveSchema, tasks,
scan.caseSensitive)
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
// A plan-scoped FileIO (from remote planning) takes precedence over the
@@ -2227,24 +2287,19 @@ func (scan *Scan) ReadTasks(ctx context.Context, tasks
[]FileScanTask) (*arrow.S
// cannot close the old plan's resources while records are still being
read.
planIO := scan.planIO
var fs io.IO
- var releasePlanIO func()
if planIO != nil {
fs, releasePlanIO, err = planIO.acquire(ctx)
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
} else {
fs, err = scan.ioF(ctx)
}
if err != nil {
- return nil, nil, err
+ return nil, nil, nil, err
}
- readRecords := scan.readRecords
- if readRecords == nil {
- readRecords = (*arrowScan).GetRecords
- }
- outSchema, records, err := readRecords(&arrowScan{
+ return &arrowScan{
metadata: scan.metadata,
fs: fs,
scanSchema: effectiveSchema,
@@ -2257,21 +2312,7 @@ func (scan *Scan) ReadTasks(ctx context.Context, tasks
[]FileScanTask) (*arrow.S
options: scan.options,
concurrency: scan.concurrency,
arrowBatchSize: scan.arrowBatchSize,
- }, ctx, readTasks)
- if err != nil {
- // No iterator to drive cleanup on a setup error, so release
here.
- if releasePlanIO != nil {
- releasePlanIO()
- }
-
- return nil, nil, err
- }
-
- if releasePlanIO != nil {
- records = releasePlanIOAfter(records, releasePlanIO)
- }
-
- return outSchema, records, nil
+ }, readTasks, releasePlanIO, nil
}
// closePlanIO releases the scoped resources associated with the current
diff --git a/table/sort_compaction.go b/table/sort_compaction.go
new file mode 100644
index 000000000..e6ce185b8
--- /dev/null
+++ b/table/sort_compaction.go
@@ -0,0 +1,1009 @@
+// 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 (
+ "bytes"
+ "cmp"
+ "container/heap"
+ "context"
+ "errors"
+ "fmt"
+ "iter"
+ "math"
+ "slices"
+ "sort"
+ "strings"
+
+ "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/decimal128"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+)
+
+// sortStrategy writes a compaction group's rows sorted by order. Groups
+// whose files are already sorted by order are merged; others are sorted
+// in memory.
+type sortStrategy struct {
+ order SortOrder
+ meta Metadata
+ spec iceberg.PartitionSpec
+ schema *iceberg.Schema
+ batchRows int64
+}
+
+func newSortStrategy(tbl *Table, order SortOrder, batchRows int64)
(sortStrategy, error) {
+ if order.IsUnsorted() {
+ return sortStrategy{}, fmt.Errorf("%w: sort compaction requires
a sort order with at least one field", ErrInvalidOperation)
+ }
+ if err := order.CheckCompatibility(tbl.Schema()); err != nil {
+ return sortStrategy{}, fmt.Errorf("sort compaction: %w", err)
+ }
+
+ return sortStrategy{
+ order: order,
+ meta: tbl.Metadata(),
+ spec: tbl.Spec(),
+ schema: tbl.Schema(),
+ batchRows: batchRows,
+ }, nil
+}
+
+func (s sortStrategy) records(ctx context.Context, scan *Scan, tasks
[]FileScanTask) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
+ if parts, ok := presortedPartitions(s.meta, s.spec, s.order, tasks,
scan.concurrency); ok {
+ return mergedCompactionRecords(ctx, scan, s.schema, s.order,
s.batchRows, parts)
+ }
+
+ sc, records, err := scan.ReadTasks(ctx, tasks)
+ if err != nil {
+ return nil, nil, err
+ }
+
+ return sc, sortedCompactionRecords(ctx, s.spec, s.schema, s.order,
s.batchRows, records), nil
+}
+
+func (s sortStrategy) writeOptions() []WriteRecordOption {
+ return
[]WriteRecordOption{withPresortedInput(claimableSortOrderID(s.meta, s.order))}
+}
+
+// sortKeyColumn is one field of a rowComparator.
+type sortKeyColumn struct {
+ col int
+ desc bool
+ nullsFirst bool
+ values valueComparator
+}
+
+// rowComparator orders the rows of record batches that share one schema
+// by the source columns of a sort order, following the Iceberg spec: each
+// field's null order places nulls whatever its direction, floating-point
+// values order as -NaN < -Inf < -x < -0 < 0 < x < Inf < NaN, and strings
+// and binary compare as unsigned bytes. A descending field reverses the
+// order of its non-null values.
+type rowComparator []sortKeyColumn
+
+// newRowComparator resolves order against the scan output schema sc.
+// Sort fields are matched by name through schema because the scan emits
+// the table schema's top-level columns by name.
+func newRowComparator(order SortOrder, schema *iceberg.Schema, sc
*arrow.Schema) (rowComparator, error) {
+ fields := schema.Fields()
+ keys := make(rowComparator, 0, order.Len())
+ for _, field := range order.fields {
+ idx, ok := topLevelFieldIndex(fields, field.FirstSourceID())
+ if !ok {
+ return nil, fmt.Errorf("sort order %d: source id %d is
not a top-level column in schema",
+ order.OrderID(), field.FirstSourceID())
+ }
+ colIdx := sc.FieldIndices(fields[idx].Name)
+ if len(colIdx) != 1 {
+ return nil, fmt.Errorf("sort order %d: column %q not
found in scan output",
+ order.OrderID(), fields[idx].Name)
+ }
+ values, err := newValueComparator(sc.Field(colIdx[0]).Type)
+ if err != nil {
+ return nil, fmt.Errorf("sort order %d: column %q: %w",
order.OrderID(), fields[idx].Name, err)
+ }
+ keys = append(keys, sortKeyColumn{
+ col: colIdx[0],
+ desc: field.Direction == SortDESC,
+ nullsFirst: field.NullOrder == NullsFirst,
+ values: values,
+ })
+ }
+
+ return keys, nil
+}
+
+// compareNulls orders a null against a non-null value. It reports false
+// when neither is null.
+func compareNulls(xNull, yNull, nullsFirst bool) (int, bool) {
+ switch {
+ case !xNull && !yNull:
+ return 0, false
+ case xNull && yNull:
+ return 0, true
+ case xNull == nullsFirst:
+ return -1, true
+ default:
+ return 1, true
+ }
+}
+
+func (c rowComparator) compare(a arrow.RecordBatch, i int, b
arrow.RecordBatch, j int) int {
+ for _, k := range c {
+ x, y := a.Column(k.col), b.Column(k.col)
+ r, isNull := compareNulls(x.IsNull(i), y.IsNull(j),
k.nullsFirst)
+ if !isNull {
+ r = k.values.compare(x, i, y, j)
+ if k.desc {
+ r = -r
+ }
+ }
+ if r != 0 {
+ return r
+ }
+ }
+
+ return 0
+}
+
+// bind returns compare specialized to rows of batches, which resolves each
+// key column's typed values once instead of on every comparison.
+func (c rowComparator) bind(batches []arrow.RecordBatch) func(x, y rowRef) int
{
+ type boundKey struct {
+ cols []arrow.Array
+ hasNulls bool
+ desc bool
+ nullsFirst bool
+ compare func(xb, xi, yb, yi int) int
+ }
+ keys := make([]boundKey, len(c))
+ for n, k := range c {
+ cols := make([]arrow.Array, len(batches))
+ hasNulls := false
+ for b, rec := range batches {
+ cols[b] = rec.Column(k.col)
+ hasNulls = hasNulls || cols[b].NullN() > 0
+ }
+ keys[n] = boundKey{cols, hasNulls, k.desc, k.nullsFirst,
k.values.bind(cols)}
+ }
+
+ return func(x, y rowRef) int {
+ xb, xi, yb, yi := x.batch(), x.row(), y.batch(), y.row()
+ for _, k := range keys {
+ var r int
+ isNull := false
+ if k.hasNulls {
+ r, isNull = compareNulls(k.cols[xb].IsNull(xi),
k.cols[yb].IsNull(yi), k.nullsFirst)
+ }
+ if !isNull {
+ r = k.compare(xb, xi, yb, yi)
+ if k.desc {
+ r = -r
+ }
+ }
+ if r != 0 {
+ return r
+ }
+ }
+
+ return 0
+ }
+}
+
+// valueComparator compares the non-null values of one Arrow type, either
+// at any two positions of any two arrays, or bound to a fixed set of
+// arrays and addressed by (array, row).
+type valueComparator struct {
+ compare func(a arrow.Array, i int, b arrow.Array, j int) int
+ bind func(cols []arrow.Array) func(xb, xi, yb, yi int) int
+}
+
+func orderedValues[T cmp.Ordered, A interface{ Values() []T }]()
valueComparator {
+ return fixedWidthValues[T, A](cmp.Compare[T])
+}
+
+func fixedWidthValues[T any, A interface{ Values() []T }](compare func(T, T)
int) valueComparator {
+ return valueComparator{
+ compare: func(a arrow.Array, i int, b arrow.Array, j int) int {
+ return compare(a.(A).Values()[i], b.(A).Values()[j])
+ },
+ bind: func(cols []arrow.Array) func(xb, xi, yb, yi int) int {
+ values := make([][]T, len(cols))
+ for n, col := range cols {
+ values[n] = col.(A).Values()
+ }
+
+ return func(xb, xi, yb, yi int) int { return
compare(values[xb][xi], values[yb][yi]) }
+ },
+ }
+}
+
+func accessorValues[T any, A interface{ Value(int) T }](compare func(T, T)
int) valueComparator {
+ return valueComparator{
+ compare: func(a arrow.Array, i int, b arrow.Array, j int) int {
+ return compare(a.(A).Value(i), b.(A).Value(j))
+ },
+ bind: func(cols []arrow.Array) func(xb, xi, yb, yi int) int {
+ typed := make([]A, len(cols))
+ for n, col := range cols {
+ typed[n] = col.(A)
+ }
+
+ return func(xb, xi, yb, yi int) int { return
compare(typed[xb].Value(xi), typed[yb].Value(yi)) }
+ },
+ }
+}
+
+func compareFloat[T float32 | float64](x, y T) int {
+ xNaN, yNaN := x != x, y != y
+ switch {
+ case xNaN || yNaN:
+ rank := func(nan bool, v T) int {
+ switch {
+ case !nan:
+ return 0
+ case math.Signbit(float64(v)):
+ return -1
+ default:
+ return 1
+ }
+ }
+
+ return cmp.Compare(rank(xNaN, x), rank(yNaN, y))
+ case x == y:
+ // Only -0 and 0 compare equal while differing in sign.
+ return cmp.Compare(boolRank(!math.Signbit(float64(x))),
boolRank(!math.Signbit(float64(y))))
+ case x < y:
+ return -1
+ default:
+ return 1
+ }
+}
+
+func compareBool(x, y bool) int { return cmp.Compare(boolRank(x), boolRank(y))
}
+
+func boolRank(v bool) int {
+ if v {
+ return 1
+ }
+
+ return 0
+}
+
+func compareDecimal128(x, y decimal128.Num) int { return x.Cmp(y) }
+
+func newValueComparator(dt arrow.DataType) (valueComparator, error) {
+ switch dt.ID() {
+ case arrow.BOOL:
+ return accessorValues[bool, *array.Boolean](compareBool), nil
+ case arrow.INT8:
+ return orderedValues[int8, *array.Int8](), nil
+ case arrow.INT16:
+ return orderedValues[int16, *array.Int16](), nil
+ case arrow.INT32:
+ return orderedValues[int32, *array.Int32](), nil
+ case arrow.INT64:
+ return orderedValues[int64, *array.Int64](), nil
+ case arrow.UINT8:
+ return orderedValues[uint8, *array.Uint8](), nil
+ case arrow.UINT16:
+ return orderedValues[uint16, *array.Uint16](), nil
+ case arrow.UINT32:
+ return orderedValues[uint32, *array.Uint32](), nil
+ case arrow.UINT64:
+ return orderedValues[uint64, *array.Uint64](), nil
+ case arrow.FLOAT32:
+ return fixedWidthValues[float32,
*array.Float32](compareFloat[float32]), nil
+ case arrow.FLOAT64:
+ return fixedWidthValues[float64,
*array.Float64](compareFloat[float64]), nil
+ case arrow.DATE32:
+ return orderedValues[arrow.Date32, *array.Date32](), nil
+ case arrow.DATE64:
+ return orderedValues[arrow.Date64, *array.Date64](), nil
+ case arrow.TIME32:
+ return orderedValues[arrow.Time32, *array.Time32](), nil
+ case arrow.TIME64:
+ return orderedValues[arrow.Time64, *array.Time64](), nil
+ case arrow.TIMESTAMP:
+ return orderedValues[arrow.Timestamp, *array.Timestamp](), nil
+ case arrow.STRING:
+ return accessorValues[string, *array.String](strings.Compare),
nil
+ case arrow.LARGE_STRING:
+ return accessorValues[string,
*array.LargeString](strings.Compare), nil
+ case arrow.STRING_VIEW:
+ return accessorValues[string,
*array.StringView](strings.Compare), nil
+ case arrow.BINARY:
+ return accessorValues[[]byte, *array.Binary](bytes.Compare), nil
+ case arrow.LARGE_BINARY:
+ return accessorValues[[]byte,
*array.LargeBinary](bytes.Compare), nil
+ case arrow.BINARY_VIEW:
+ return accessorValues[[]byte,
*array.BinaryView](bytes.Compare), nil
+ case arrow.FIXED_SIZE_BINARY:
+ return accessorValues[[]byte,
*array.FixedSizeBinary](bytes.Compare), nil
+ case arrow.DECIMAL128:
+ return fixedWidthValues[decimal128.Num,
*array.Decimal128](compareDecimal128), nil
+ case arrow.EXTENSION:
+ storage, err :=
newValueComparator(dt.(arrow.ExtensionType).StorageType())
+ if err != nil {
+ return valueComparator{}, err
+ }
+ unwrap := func(a arrow.Array) arrow.Array { return
a.(array.ExtensionArray).Storage() }
+
+ return valueComparator{
+ compare: func(a arrow.Array, i int, b arrow.Array, j
int) int {
+ return storage.compare(unwrap(a), i, unwrap(b),
j)
+ },
+ bind: func(cols []arrow.Array) func(xb, xi, yb, yi int)
int {
+ stored := make([]arrow.Array, len(cols))
+ for n, col := range cols {
+ stored[n] = unwrap(col)
+ }
+
+ return storage.bind(stored)
+ },
+ }, nil
+ default:
+ return valueComparator{}, fmt.Errorf("%w: cannot sort by values
of type %s", iceberg.ErrNotImplemented, dt)
+ }
+}
+
+// claimableSortOrderID returns the id to record as sort_order_id on files
+// written sorted by order, or UnsortedSortOrderID when the claim would not
+// hold. Rows are ordered by each field's first source column. That matches
+// the spec's ordering only when every field but the last is identity and
+// the last preserves order: ordering by a coarser transform's source
+// column (day(ts)) also orders rows the transform ties on, which breaks
+// the order of the fields after it. The id must also name the same order
+// in the table metadata.
+func claimableSortOrderID(meta Metadata, order SortOrder) int {
+ for i, field := range order.fields {
+ if len(field.SourceIDs) > 1 || field.Transform == nil ||
!field.Transform.PreservesOrder() {
+ return UnsortedSortOrderID
+ }
+ if _, identity := field.Transform.(iceberg.IdentityTransform);
!identity && i < len(order.fields)-1 {
+ return UnsortedSortOrderID
+ }
+ }
+ for _, existing := range meta.SortOrders() {
+ if existing.Equals(order) {
+ return order.OrderID()
+ }
+ }
+
+ return UnsortedSortOrderID
+}
+
+type compactionPartition struct {
+ rec partitionRecord
+ batches []arrow.RecordBatch
+}
+
+func releaseCompactionPartitions(parts []compactionPartition) {
+ for _, p := range parts {
+ for _, b := range p.batches {
+ b.Release()
+ }
+ }
+}
+
+// sortedCompactionRecords materializes every batch of records, splits the
+// rows by the table's current partition spec and yields each partition's
+// rows globally sorted by order, one partition after another, in batches
+// of at most batchRows rows. Partitions are emitted in order of first
+// appearance, so the output stays clustered for the clustered writer even
+// when the group's files were written under an older spec.
+//
+// The whole input is held in memory until the last partition is emitted.
+func sortedCompactionRecords(
+ ctx context.Context,
+ spec iceberg.PartitionSpec,
+ schema *iceberg.Schema,
+ order SortOrder,
+ batchRows int64,
+ records iter.Seq2[arrow.RecordBatch, error],
+) iter.Seq2[arrow.RecordBatch, error] {
+ return func(yield func(arrow.RecordBatch, error) bool) {
+ parts, rows, err := collectCompactionPartitions(ctx, spec,
schema, order, records)
+ defer func() { releaseCompactionPartitions(parts) }()
+ if err != nil {
+ yield(nil, err)
+
+ return
+ }
+
+ for i := range parts {
+ if !yieldSortedPartition(ctx, parts[i].batches, rows,
batchRows, yield) {
+ return
+ }
+ releaseCompactionPartitions(parts[i : i+1])
+ parts[i].batches = nil
+ }
+ }
+}
+
+// collectCompactionPartitions drains records into per-partition batches
+// and resolves order against the scan output schema of the first batch.
+func collectCompactionPartitions(
+ ctx context.Context,
+ spec iceberg.PartitionSpec,
+ schema *iceberg.Schema,
+ order SortOrder,
+ records iter.Seq2[arrow.RecordBatch, error],
+) ([]compactionPartition, rowComparator, error) {
+ var (
+ parts []compactionPartition
+ rows rowComparator
+ plan *partitionExtractionPlan
+ takeFn = partitionBatchByKey(ctx)
+ )
+
+ for rec, err := range records {
+ if err != nil {
+ return parts, nil, err
+ }
+ if rec.NumRows() == 0 {
+ rec.Release()
+
+ continue
+ }
+ if rows == nil {
+ if rows, err = newRowComparator(order, schema,
rec.Schema()); err != nil {
+ rec.Release()
+
+ return parts, nil, err
+ }
+ }
+
+ if spec.IsUnpartitioned() {
+ if len(parts) == 0 {
+ parts = append(parts, compactionPartition{})
+ }
+ parts[0].batches = append(parts[0].batches, rec)
+
+ continue
+ }
+
+ if plan == nil {
+ if plan, err = newPartitionExtractionPlan(spec, schema,
rec.Schema()); err != nil {
+ rec.Release()
+
+ return parts, nil, err
+ }
+ }
+ infos, err := plan.getRecordPartitions(rec)
+ if err != nil {
+ rec.Release()
+
+ return parts, nil, err
+ }
+ slices.SortFunc(infos, func(a, b *partitionInfo) int {
+ return cmp.Compare(a.rows[0], b.rows[0])
+ })
+
+ for _, info := range infos {
+ sub, err := takeFn(rec, info.rows)
+ if err != nil {
+ rec.Release()
+
+ return parts, nil, err
+ }
+ idx := slices.IndexFunc(parts, func(p
compactionPartition) bool {
+ return partitionRecordsEqual(p.rec,
info.partitionRec)
+ })
+ if idx < 0 {
+ parts = append(parts, compactionPartition{rec:
info.partitionRec})
+ idx = len(parts) - 1
+ }
+ parts[idx].batches = append(parts[idx].batches, sub)
+ }
+ rec.Release()
+ }
+
+ return parts, rows, nil
+}
+
+// rowRef addresses one row of a slice of batches: the batch index in the
+// high 32 bits and the row within it in the low 32.
+type rowRef uint64
+
+func newRowRef(batch, row int) rowRef { return rowRef(uint64(batch)<<32 |
uint64(row)) }
+func (r rowRef) batch() int { return int(r >> 32) }
+func (r rowRef) row() int { return int(r & math.MaxUint32) }
+
+// yieldSortedPartition sorts batches as one table and yields the sorted
+// rows. Ownership of each yielded batch passes to the consumer. It returns
+// false when iteration must stop.
+func yieldSortedPartition(
+ ctx context.Context,
+ batches []arrow.RecordBatch,
+ rows rowComparator,
+ batchRows int64,
+ yield func(arrow.RecordBatch, error) bool,
+) bool {
+ if len(batches) == 0 {
+ return true
+ }
+
+ sc := batches[0].Schema()
+ var total int64
+ for _, b := range batches {
+ if !b.Schema().Equal(sc) {
+ yield(nil, errors.New("sort compaction: scan produced
batches with differing schemas"))
+
+ return false
+ }
+ total += b.NumRows()
+ }
+
+ refs := make([]rowRef, 0, total)
+ for b, rec := range batches {
+ for r := range int(rec.NumRows()) {
+ refs = append(refs, newRowRef(b, r))
+ }
+ }
+ slices.SortStableFunc(refs, rows.bind(batches))
+
+ takeFn := partitionBatchByKey(ctx)
+ mem := compute.GetAllocator(ctx)
+ for off := 0; off < len(refs); off += int(batchRows) {
+ if err := ctx.Err(); err != nil {
+ yield(nil, err)
+
+ return false
+ }
+
+ end := min(off+int(batchRows), len(refs))
+ out, err := gatherSortedRows(mem, takeFn, batches,
refs[off:end])
+ if err != nil {
+ yield(nil, fmt.Errorf("sort compaction: gather sorted
rows: %w", err))
+
+ return false
+ }
+ if !yield(out, nil) {
+ return false
+ }
+ }
+
+ return true
+}
+
+// gatherSortedRows builds one batch holding the rows refs of batches, in
+// that order. Each source batch is taken from directly so no column is
+// ever concatenated beyond one output batch.
+func gatherSortedRows(
+ mem memory.Allocator,
+ takeFn partitionBatchFn,
+ batches []arrow.RecordBatch,
+ refs []rowRef,
+) (arrow.RecordBatch, error) {
+ local := make([][]int64, len(batches))
+ pos := make([]int64, len(refs))
+ for i, ref := range refs {
+ b := ref.batch()
+ pos[i] = int64(len(local[b]))
+ local[b] = append(local[b], int64(ref.row()))
+ }
+
+ var pieces []arrow.RecordBatch
+ defer func() {
+ for _, p := range pieces {
+ p.Release()
+ }
+ }()
+ base := make([]int64, len(batches))
+ var n int64
+ for b, idx := range local {
+ if len(idx) == 0 {
+ continue
+ }
+ piece, err := takeFn(batches[b], idx)
+ if err != nil {
+ return nil, err
+ }
+ pieces = append(pieces, piece)
+ base[b] = n
+ n += int64(len(idx))
+ }
+
+ // A single source batch was taken in output order already.
+ if len(pieces) == 1 {
+ pieces[0].Retain()
+
+ return pieces[0], nil
+ }
+
+ merged, err := concatRecordBatches(mem, pieces)
+ if err != nil {
+ return nil, err
+ }
+ defer merged.Release()
+
+ order := make([]int64, len(refs))
+ for i, ref := range refs {
+ order[i] = base[ref.batch()] + pos[i]
+ }
+
+ return takeFn(merged, order)
+}
+
+func concatRecordBatches(mem memory.Allocator, batches []arrow.RecordBatch)
(arrow.RecordBatch, error) {
+ sc := batches[0].Schema()
+ cols := make([]arrow.Array, sc.NumFields())
+ defer func() {
+ for _, c := range cols {
+ if c != nil {
+ c.Release()
+ }
+ }
+ }()
+
+ var rows int64
+ for _, b := range batches {
+ rows += b.NumRows()
+ }
+ chunks := make([]arrow.Array, len(batches))
+ for i := range cols {
+ for j, b := range batches {
+ chunks[j] = b.Column(i)
+ }
+ col, err := array.Concatenate(chunks, mem)
+ if err != nil {
+ return nil, fmt.Errorf("concatenate column %q: %w",
sc.Field(i).Name, err)
+ }
+ cols[i] = col
+ }
+
+ return array.NewRecordBatch(sc, cols, rows), nil
+}
+
+// presortedPartitions reports whether tasks can be merged instead of
+// sorted: every task's file records order's id as its sort_order_id, the
+// id is one claimableSortOrderID accepts, every field of order is
+// identity, every file was written under spec, so each file holds rows of
+// a single partition of spec, and no partition has more than maxStreams
+// files. The tasks are returned grouped by partition, in order of first
+// appearance.
+//
+// The merge compares source values, so a transformed field rules it out:
+// a file sorted by [day(ts)] is validly unsorted by ts within a day.
+// maxStreams bounds the files the merge reads at once.
+func presortedPartitions(meta Metadata, spec iceberg.PartitionSpec, order
SortOrder, tasks []FileScanTask, maxStreams int) ([][]FileScanTask, bool) {
+ id := claimableSortOrderID(meta, order)
+ if id == UnsortedSortOrderID {
+ return nil, false
+ }
+ for _, field := range order.fields {
+ if _, identity := field.Transform.(iceberg.IdentityTransform);
!identity {
+ return nil, false
+ }
+ }
+
+ var (
+ keys []partitionRecord
+ parts [][]FileScanTask
+ )
+ for _, task := range tasks {
+ sortOrderID := task.File.SortOrderID()
+ if sortOrderID == nil || *sortOrderID != id ||
int(task.File.SpecID()) != spec.ID() {
+ return nil, false
+ }
+
+ values := task.File.Partition()
+ key := make(partitionRecord, 0, spec.NumFields())
+ for _, field := range spec.Fields() {
+ key = append(key, values[field.FieldID])
+ }
+ idx := slices.IndexFunc(keys, func(k partitionRecord) bool {
return partitionRecordsEqual(k, key) })
+ if idx < 0 {
+ keys = append(keys, key)
+ parts = append(parts, nil)
+ idx = len(parts) - 1
+ }
+ parts[idx] = append(parts[idx], task)
+ if len(parts[idx]) > maxStreams {
+ return nil, false
+ }
+ }
+
+ return parts, true
+}
+
+// mergedCompactionRecords reads each task of parts (see
+// presortedPartitions) as its own stream and merges the streams of one
+// partition after another into batches of at most batchRows rows sorted
+// by order. Each stream of the partition being merged holds up to about
+// five batches (its scan worker, the channels behind it and the merge
+// cursor), instead of the whole group being held.
+func mergedCompactionRecords(
+ ctx context.Context,
+ scan *Scan,
+ schema *iceberg.Schema,
+ order SortOrder,
+ batchRows int64,
+ parts [][]FileScanTask,
+) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
+ tasks := slices.Concat(parts...)
+ sc, streams, release, err := scan.readTasksPerTask(ctx, tasks)
+ if err != nil {
+ return nil, nil, err
+ }
+ rows, err := newRowComparator(order, schema, sc)
+ if err != nil {
+ release()
+
+ return nil, nil, err
+ }
+
+ paths := make([]string, len(tasks))
+ for i, task := range tasks {
+ paths[i] = task.File.FilePath()
+ }
+
+ records := func(yield func(arrow.RecordBatch, error) bool) {
+ defer release()
+ start := 0
+ for _, part := range parts {
+ end := start + len(part)
+ if !mergeSortedStreams(ctx, rows, paths[start:end],
streams[start:end], batchRows, yield) {
+ return
+ }
+ start = end
+ }
+ }
+
+ return sc, records, nil
+}
+
+// mergeCursor walks one sorted stream of a k-way merge.
+type mergeCursor struct {
+ idx int
+ path string
+ next func() (arrow.RecordBatch, error, bool)
+ stop func()
+ rec arrow.RecordBatch
+ row int
+ // held is rec's index in the pending output's source batches, or -1
+ // while no pending row comes from rec.
+ held int
+}
+
+// load moves the cursor to the first row of the stream's next non-empty
+// batch and reports false at the end of the stream. It checks that the
+// stream is sorted, since a merge of unsorted input would silently yield
+// unsorted output.
+func (c *mergeCursor) load(rows rowComparator) (bool, error) {
+ for {
+ rec, err, ok := c.next()
+ switch {
+ case !ok:
+ c.releaseBatch()
+
+ return false, nil
+ case err != nil:
+ return false, err
+ case rec.NumRows() == 0:
+ rec.Release()
+
+ continue
+ }
+
+ sorted := c.rec == nil || rows.compare(c.rec,
int(c.rec.NumRows())-1, rec, 0) <= 0
+ for r := 1; sorted && r < int(rec.NumRows()); r++ {
+ sorted = rows.compare(rec, r-1, rec, r) <= 0
+ }
+ c.releaseBatch()
+ c.rec, c.row, c.held = rec, 0, -1
+ if !sorted {
+ return false, fmt.Errorf("sort compaction: data file %s
records the sort order being merged but is not sorted by it", c.path)
+ }
+
+ return true, nil
+ }
+}
+
+func (c *mergeCursor) releaseBatch() {
+ if c.rec != nil {
+ c.rec.Release()
+ c.rec = nil
+ }
+}
+
+// mergeHeap orders cursors by their current row, ties broken by stream
+// index so the merge is deterministic.
+type mergeHeap struct {
+ cursors []*mergeCursor
+ rows rowComparator
+}
+
+func (h *mergeHeap) precedes(a *mergeCursor, row int, b *mergeCursor) bool {
+ if c := h.rows.compare(a.rec, row, b.rec, b.row); c != 0 {
+ return c < 0
+ }
+
+ return a.idx < b.idx
+}
+
+func (h *mergeHeap) Len() int { return len(h.cursors) }
+func (h *mergeHeap) Less(i, j int) bool {
+ return h.precedes(h.cursors[i], h.cursors[i].row, h.cursors[j])
+}
+func (h *mergeHeap) Swap(i, j int) { h.cursors[i], h.cursors[j] =
h.cursors[j], h.cursors[i] }
+func (h *mergeHeap) Push(x any) { h.cursors = append(h.cursors,
x.(*mergeCursor)) }
+func (h *mergeHeap) Pop() any {
+ last := h.cursors[len(h.cursors)-1]
+ h.cursors = h.cursors[:len(h.cursors)-1]
+
+ return last
+}
+
+// mergeRun is rows [start, end) of the pending output's source batch
+// held[batch].
+type mergeRun struct{ batch, start, end int }
+
+// mergeSortedStreams k-way merges streams, each sorted by rows, and yields
+// the result in batches of at most batchRows rows. Every step takes the
+// longest run of the minimum stream's current batch that precedes the
+// runner-up's head. An output batch that is one run is a zero-copy slice;
+// otherwise its rows are gathered from their source batches, so
+// interleaved streams cost a gather per output batch rather than a slice
+// per run. Ownership of each yielded batch passes to the consumer. It
+// returns false when iteration must stop.
+func mergeSortedStreams(
+ ctx context.Context,
+ rows rowComparator,
+ paths []string,
+ streams []iter.Seq2[arrow.RecordBatch, error],
+ batchRows int64,
+ yield func(arrow.RecordBatch, error) bool,
+) bool {
+ h := &mergeHeap{rows: rows}
+ cursors := make([]*mergeCursor, 0, len(streams))
+ var (
+ held []arrow.RecordBatch
+ runs []mergeRun
+ pendingRows int64
+ )
+ releaseHeld := func() {
+ for _, b := range held {
+ b.Release()
+ }
+ held, runs, pendingRows = held[:0], runs[:0], 0
+ for _, c := range cursors {
+ c.held = -1
+ }
+ }
+ defer func() {
+ for _, c := range cursors {
+ c.releaseBatch()
+ c.stop()
+ }
+ releaseHeld()
+ }()
+
+ for i, stream := range streams {
+ next, stop := iter.Pull2(stream)
+ c := &mergeCursor{idx: i, path: paths[i], next: next, stop:
stop, held: -1}
+ cursors = append(cursors, c)
+ ok, err := c.load(rows)
+ if err != nil {
+ yield(nil, err)
+
+ return false
+ }
+ if ok {
+ h.cursors = append(h.cursors, c)
+ }
+ }
+ heap.Init(h)
+
+ mem := compute.GetAllocator(ctx)
+ takeFn := partitionBatchByKey(ctx)
+ flush := func() bool {
+ out, err := gatherMergeRuns(mem, takeFn, held, runs,
pendingRows)
+ releaseHeld()
+ if err != nil {
+ yield(nil, fmt.Errorf("sort compaction: merge sorted
rows: %w", err))
+
+ return false
+ }
+
+ return yield(out, nil)
+ }
+
+ for h.Len() > 0 {
+ if err := ctx.Err(); err != nil {
+ yield(nil, err)
+
+ return false
+ }
+
+ c := heap.Pop(h).(*mergeCursor)
+ end := int(c.rec.NumRows())
+ if h.Len() > 0 {
+ next := h.cursors[0]
+ end = c.row + sort.Search(end-c.row, func(k int) bool {
+ return !h.precedes(c, c.row+k, next)
+ })
+ }
+ end = min(end, c.row+int(batchRows-pendingRows))
+
+ if c.held < 0 {
+ c.rec.Retain()
+ held = append(held, c.rec)
+ c.held = len(held) - 1
+ }
+ if last := len(runs) - 1; last >= 0 && runs[last].batch ==
c.held && runs[last].end == c.row {
+ runs[last].end = end
+ } else {
+ runs = append(runs, mergeRun{batch: c.held, start:
c.row, end: end})
+ }
+ pendingRows += int64(end - c.row)
+ c.row = end
+ if pendingRows == batchRows && !flush() {
+ return false
+ }
+
+ if c.row < int(c.rec.NumRows()) {
+ heap.Push(h, c)
+
+ continue
+ }
+ ok, err := c.load(rows)
+ if err != nil {
+ yield(nil, err)
+
+ return false
+ }
+ if ok {
+ heap.Push(h, c)
+ }
+ }
+
+ if pendingRows > 0 {
+ return flush()
+ }
+
+ return true
+}
+
+// gatherMergeRuns builds one batch of the n rows of runs, in order.
+func gatherMergeRuns(
+ mem memory.Allocator,
+ takeFn partitionBatchFn,
+ held []arrow.RecordBatch,
+ runs []mergeRun,
+ n int64,
+) (arrow.RecordBatch, error) {
+ if len(runs) == 1 {
+ r := runs[0]
+
+ return held[r.batch].NewSlice(int64(r.start), int64(r.end)), nil
+ }
+
+ refs := make([]rowRef, 0, n)
+ for _, r := range runs {
+ for row := r.start; row < r.end; row++ {
+ refs = append(refs, newRowRef(r.batch, row))
+ }
+ }
+
+ return gatherSortedRows(mem, takeFn, held, refs)
+}
diff --git a/table/sort_compaction_bench_test.go
b/table/sort_compaction_bench_test.go
new file mode 100644
index 000000000..a82935890
--- /dev/null
+++ b/table/sort_compaction_bench_test.go
@@ -0,0 +1,124 @@
+// 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 (
+ "iter"
+ "strconv"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+)
+
+const (
+ benchMergeStreams = 2
+ benchMergeRows = 131072
+ benchMergeCols = 32
+ benchMergeBatchRows = 8192
+)
+
+// benchMergeInputs builds benchMergeStreams sorted streams of int64 rows,
+// either interleaved key by key or covering disjoint key ranges.
+func benchMergeInputs(b *testing.B, interleaved bool) (rowComparator,
[][]arrow.RecordBatch) {
+ const (
+ streams = benchMergeStreams
+ rows = benchMergeRows
+ cols = benchMergeCols
+ batchRows = benchMergeBatchRows
+ )
+ fields := make([]arrow.Field, cols)
+ ifields := make([]iceberg.NestedField, cols)
+ for c := range cols {
+ fields[c] = arrow.Field{Name: "c" + strconv.Itoa(c), Type:
arrow.PrimitiveTypes.Int64}
+ ifields[c] = iceberg.NestedField{ID: c + 1, Name:
fields[c].Name, Type: iceberg.PrimitiveTypes.Int64, Required: true}
+ }
+ sc := arrow.NewSchema(fields, nil)
+ order, err := NewSortOrder(1, []SortField{{SourceIDs: []int{1},
Transform: iceberg.IdentityTransform{}, Direction: SortASC, NullOrder:
NullsFirst}})
+ if err != nil {
+ b.Fatal(err)
+ }
+ cmp, err := newRowComparator(order, iceberg.NewSchema(0, ifields...),
sc)
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ out := make([][]arrow.RecordBatch, streams)
+ for s := range streams {
+ for off := 0; off < rows; off += batchRows {
+ bldr := array.NewRecordBuilder(memory.DefaultAllocator,
sc)
+ for r := off; r < off+batchRows; r++ {
+ key := int64(s*rows + r)
+ if interleaved {
+ key = int64(r*streams + s)
+ }
+ for c := range cols {
+
bldr.Field(c).(*array.Int64Builder).Append(key)
+ }
+ }
+ out[s] = append(out[s], bldr.NewRecordBatch())
+ bldr.Release()
+ }
+ }
+
+ return cmp, out
+}
+
+func benchMerge(b *testing.B, interleaved bool) {
+ cmp, inputs := benchMergeInputs(b, interleaved)
+ defer func() {
+ for _, batches := range inputs {
+ for _, rec := range batches {
+ rec.Release()
+ }
+ }
+ }()
+ paths := make([]string, len(inputs))
+ b.ResetTimer()
+ for range b.N {
+ streams := make([]iter.Seq2[arrow.RecordBatch, error],
len(inputs))
+ for i, batches := range inputs {
+ streams[i] = func(yield func(arrow.RecordBatch, error)
bool) {
+ for _, rec := range batches {
+ rec.Retain()
+ if !yield(rec, nil) {
+ return
+ }
+ }
+ }
+ }
+ var n int64
+ mergeSortedStreams(b.Context(), cmp, paths, streams,
benchMergeBatchRows, func(rec arrow.RecordBatch, err error) bool {
+ if err != nil {
+ b.Fatal(err)
+ }
+ n += rec.NumRows()
+ rec.Release()
+
+ return true
+ })
+ if n != benchMergeStreams*benchMergeRows {
+ b.Fatalf("got %d rows", n)
+ }
+ }
+}
+
+func BenchmarkMergeInterleaved(b *testing.B) { benchMerge(b, true) }
+func BenchmarkMergeDisjoint(b *testing.B) { benchMerge(b, false) }
diff --git a/table/sort_compaction_internal_test.go
b/table/sort_compaction_internal_test.go
new file mode 100644
index 000000000..0aab3efbe
--- /dev/null
+++ b/table/sort_compaction_internal_test.go
@@ -0,0 +1,460 @@
+// 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 (
+ "encoding/json"
+ "iter"
+ "math"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// TestSortedCompactionRecords_NestedAndNulls sorts rows spread over several
+// batches, with nulls in the sort key and a struct column (which Arrow's
+// chunked Take cannot gather without concatenating) carried along.
+func TestSortedCompactionRecords_NestedAndNulls(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64},
+ iceberg.NestedField{ID: 2, Name: "s", Type:
&iceberg.StructType{FieldList: []iceberg.NestedField{
+ {ID: 3, Name: "x", Type: iceberg.PrimitiveTypes.String},
+ }}},
+ )
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "s", Type: arrow.StructOf(arrow.Field{Name: "x", Type:
arrow.BinaryTypes.String, Nullable: true}), Nullable: true},
+ }, nil)
+
+ inputs := []string{
+ `[{"id": 4, "s": {"x": "d"}}, {"id": null, "s": {"x": "n"}},
{"id": 1, "s": {"x": "a"}}]`,
+ `[{"id": 3, "s": {"x": "c"}}]`,
+ `[{"id": 5, "s": null}, {"id": 2, "s": {"x": "b"}}]`,
+ }
+ records := func(yield func(arrow.RecordBatch, error) bool) {
+ for _, in := range inputs {
+ rec, _, err := array.RecordFromJSON(mem, sc,
strings.NewReader(in))
+ if err != nil {
+ yield(nil, err)
+
+ return
+ }
+ if !yield(rec, nil) {
+ return
+ }
+ }
+ }
+
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform: iceberg.IdentityTransform{},
+ Direction: SortASC, NullOrder: NullsFirst,
+ }})
+ require.NoError(t, err)
+
+ // mem tracks only the input batches: arrow-go's sort_indices kernel
does
+ // not release its scratch buffers, so compute allocations stay
untracked.
+ ctx := t.Context()
+ var got []map[string]any
+ var sizes []int64
+ for rec, err := range sortedCompactionRecords(ctx,
*iceberg.UnpartitionedSpec, schema, order, 4, records) {
+ require.NoError(t, err)
+ sizes = append(sizes, rec.NumRows())
+ b, err := rec.MarshalJSON()
+ rec.Release()
+ require.NoError(t, err)
+ var rows []map[string]any
+ require.NoError(t, json.Unmarshal(b, &rows))
+ got = append(got, rows...)
+ }
+
+ assert.Equal(t, []int64{4, 2}, sizes)
+ want := []map[string]any{
+ {"id": nil, "s": map[string]any{"x": "n"}},
+ {"id": float64(1), "s": map[string]any{"x": "a"}},
+ {"id": float64(2), "s": map[string]any{"x": "b"}},
+ {"id": float64(3), "s": map[string]any{"x": "c"}},
+ {"id": float64(4), "s": map[string]any{"x": "d"}},
+ {"id": float64(5), "s": nil},
+ }
+ assert.Equal(t, want, got)
+}
+
+func TestCompareFloat_IcebergOrder(t *testing.T) {
+ // Ascending order required by the Iceberg spec.
+ want := []float64{
+ math.Copysign(math.NaN(), -1), math.Inf(-1), -1,
math.Copysign(0, -1), 0, 1, math.Inf(1), math.NaN(),
+ }
+ for i, x := range want {
+ for j, y := range want {
+ assert.Equal(t, cmpInts(i, j), compareFloat(x, y),
"compare(%v, %v)", x, y)
+ }
+ }
+}
+
+func cmpInts(a, b int) int {
+ switch {
+ case a < b:
+ return -1
+ case a > b:
+ return 1
+ default:
+ return 0
+ }
+}
+
+// TestSortedCompactionRecords_FloatOrder sorts a float column holding NaN,
+// infinities and a null, which Arrow's sort kernel would order with NaN
+// next to the nulls rather than where the Iceberg spec puts it.
+func TestSortedCompactionRecords_FloatOrder(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := iceberg.NewSchema(0, iceberg.NestedField{ID: 1, Name: "v",
Type: iceberg.PrimitiveTypes.Float64})
+ sc := arrow.NewSchema([]arrow.Field{{Name: "v", Type:
arrow.PrimitiveTypes.Float64, Nullable: true}}, nil)
+
+ tests := []struct {
+ name string
+ direction SortDirection
+ nullOrder NullOrder
+ want []string
+ }{
+ {"asc nulls first", SortASC, NullsFirst, []string{"null",
"-Inf", "-1", "1", "+Inf", "NaN"}},
+ {"desc nulls last", SortDESC, NullsLast, []string{"NaN",
"+Inf", "1", "-1", "-Inf", "null"}},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ records := func(yield func(arrow.RecordBatch, error)
bool) {
+ bldr := array.NewFloat64Builder(mem)
+ defer bldr.Release()
+ bldr.AppendValues([]float64{math.NaN(), 1, 0,
math.Inf(-1), math.Inf(1), -1}, []bool{true, true, false, true, true, true})
+ col := bldr.NewArray()
+ defer col.Release()
+ yield(array.NewRecordBatch(sc,
[]arrow.Array{col}, int64(col.Len())), nil)
+ }
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform:
iceberg.IdentityTransform{},
+ Direction: tt.direction, NullOrder:
tt.nullOrder,
+ }})
+ require.NoError(t, err)
+
+ var got []string
+ for rec, err := range
sortedCompactionRecords(t.Context(), *iceberg.UnpartitionedSpec, schema, order,
100, records) {
+ require.NoError(t, err)
+ col := rec.Column(0).(*array.Float64)
+ for i := range col.Len() {
+ if col.IsNull(i) {
+ got = append(got, "null")
+ } else {
+ got = append(got,
col.ValueStr(i))
+ }
+ }
+ rec.Release()
+ }
+ assert.Equal(t, tt.want, got)
+ })
+ }
+}
+
+func TestClaimableSortOrderID(t *testing.T) {
+ schema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64, Required: true},
+ iceberg.NestedField{ID: 2, Name: "ts", Type:
iceberg.PrimitiveTypes.Timestamp, Required: true},
+ iceberg.NestedField{ID: 3, Name: "a", Type:
iceberg.PrimitiveTypes.Float64},
+ )
+ field := func(id int, transform iceberg.Transform) SortField {
+ return SortField{SourceIDs: []int{id}, Transform: transform,
Direction: SortASC, NullOrder: NullsFirst}
+ }
+
+ tests := []struct {
+ name string
+ fields []SortField
+ claimed bool
+ }{
+ {"identity fields", []SortField{field(1,
iceberg.IdentityTransform{}), field(3, iceberg.IdentityTransform{})}, true},
+ {"order-preserving last field", []SortField{field(1,
iceberg.IdentityTransform{}), field(2, iceberg.DayTransform{})}, true},
+ // Sorting by ts orders rows that share a day, so id is not
sorted within a day.
+ {"order-preserving field before another", []SortField{field(2,
iceberg.DayTransform{}), field(1, iceberg.IdentityTransform{})}, false},
+ {"transform does not preserve order", []SortField{field(1,
iceberg.BucketTransform{NumBuckets: 4})}, false},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ order, err := NewSortOrder(1, tt.fields)
+ require.NoError(t, err)
+ meta, err := NewMetadata(schema,
iceberg.UnpartitionedSpec, order, "s3://bucket/t", nil)
+ require.NoError(t, err)
+
+ want := UnsortedSortOrderID
+ if tt.claimed {
+ want = meta.SortOrder().OrderID()
+ }
+ assert.Equal(t, want, claimableSortOrderID(meta,
meta.SortOrder()))
+ })
+ }
+}
+
+func TestPresortedPartitions(t *testing.T) {
+ schema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64, Required: true},
+ iceberg.NestedField{ID: 2, Name: "p", Type:
iceberg.PrimitiveTypes.String, Required: true},
+ )
+ spec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{2}, FieldID: 1000, Transform:
iceberg.IdentityTransform{}, Name: "p",
+ })
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform: iceberg.IdentityTransform{},
Direction: SortASC, NullOrder: NullsFirst,
+ }})
+ require.NoError(t, err)
+ meta, err := NewMetadata(schema, &spec, order, "s3://bucket/t", nil)
+ require.NoError(t, err)
+ order, spec = meta.SortOrder(), meta.PartitionSpec()
+
+ task := func(path, partition string, sortOrderID *int) FileScanTask {
+ bldr, err := iceberg.NewDataFileBuilder(spec,
iceberg.EntryContentData, path, iceberg.ParquetFile,
+ map[int]any{1000: partition}, nil, nil, 1, 1)
+ require.NoError(t, err)
+ if sortOrderID != nil {
+ bldr.SortOrderID(*sortOrderID)
+ }
+
+ return FileScanTask{File: bldr.Build()}
+ }
+ id, other := order.OrderID(), order.OrderID()+1
+
+ a1, b1, a2 := task("a1", "a", &id), task("b1", "b", &id), task("a2",
"a", &id)
+ parts, ok := presortedPartitions(meta, spec, order, []FileScanTask{a1,
b1, a2}, 2)
+ require.True(t, ok)
+ require.Len(t, parts, 2)
+ assert.Equal(t, []string{"a1", "a2"},
[]string{parts[0][0].File.FilePath(), parts[0][1].File.FilePath()})
+ assert.Equal(t, "b1", parts[1][0].File.FilePath())
+
+ _, ok = presortedPartitions(meta, spec, order, []FileScanTask{a1, b1,
a2}, 1)
+ assert.False(t, ok, "a partition with more files than streams may be
open must be sorted, not merged")
+ _, ok = presortedPartitions(meta, spec, order, []FileScanTask{a1,
task("a3", "a", nil)}, 2)
+ assert.False(t, ok, "a file without sort_order_id must be sorted, not
merged")
+ _, ok = presortedPartitions(meta, spec, order, []FileScanTask{a1,
task("a3", "a", &other)}, 2)
+ assert.False(t, ok, "a file sorted by another order must be sorted, not
merged")
+
+ truncated, err := NewSortOrder(1, []SortField{
+ {SourceIDs: []int{1}, Transform: iceberg.IdentityTransform{},
Direction: SortASC, NullOrder: NullsFirst},
+ {SourceIDs: []int{2}, Transform:
iceberg.TruncateTransform{Width: 1}, Direction: SortASC, NullOrder: NullsFirst},
+ })
+ require.NoError(t, err)
+ truncMeta, err := NewMetadata(schema, &spec, truncated,
"s3://bucket/t", nil)
+ require.NoError(t, err)
+ truncated = truncMeta.SortOrder()
+ truncID := truncated.OrderID()
+ require.Equal(t, truncID, claimableSortOrderID(truncMeta, truncated))
+ _, ok = presortedPartitions(truncMeta, spec, truncated,
+ []FileScanTask{task("t1", "a", &truncID), task("t2", "a",
&truncID)}, 2)
+ assert.False(t, ok, "files sorted by a transform may be unsorted by its
source column, so they must be sorted, not merged")
+}
+
+var mergeTestSchema = arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "src", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+}, nil)
+
+func mergeTestComparator(t *testing.T, direction SortDirection, nullOrder
NullOrder) rowComparator {
+ t.Helper()
+
+ schema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64},
+ iceberg.NestedField{ID: 2, Name: "src", Type:
iceberg.PrimitiveTypes.Int64},
+ )
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform: iceberg.IdentityTransform{},
Direction: direction, NullOrder: nullOrder,
+ }})
+ require.NoError(t, err)
+ rows, err := newRowComparator(order, schema, mergeTestSchema)
+ require.NoError(t, err)
+
+ return rows
+}
+
+// idStream yields one batch per entry of batches, each row tagged with src.
+func idStream(mem memory.Allocator, src int, batches ...[]any)
iter.Seq2[arrow.RecordBatch, error] {
+ return func(yield func(arrow.RecordBatch, error) bool) {
+ for _, ids := range batches {
+ rows := make([]map[string]any, len(ids))
+ for i, id := range ids {
+ rows[i] = map[string]any{"id": id, "src": src}
+ }
+ b, err := json.Marshal(rows)
+ if err != nil {
+ yield(nil, err)
+
+ return
+ }
+ rec, _, err := array.RecordFromJSON(mem,
mergeTestSchema, strings.NewReader(string(b)))
+ if err != nil {
+ yield(nil, err)
+
+ return
+ }
+ if !yield(rec, nil) {
+ return
+ }
+ }
+ }
+}
+
+type mergedRows struct {
+ ids, srcs []any
+ sizes []int64
+}
+
+func runMerge(t *testing.T, rows rowComparator, batchRows int64, streams
...iter.Seq2[arrow.RecordBatch, error]) (mergedRows, error) {
+ t.Helper()
+
+ paths := make([]string, len(streams))
+ for i := range paths {
+ paths[i] = "f" + string(rune('0'+i)) + ".parquet"
+ }
+ var (
+ got mergedRows
+ lastErr error
+ )
+ mergeSortedStreams(t.Context(), rows, paths, streams, batchRows,
func(rec arrow.RecordBatch, err error) bool {
+ if err != nil {
+ lastErr = err
+
+ return false
+ }
+ defer rec.Release()
+ got.sizes = append(got.sizes, rec.NumRows())
+ ids, srcs := rec.Column(0).(*array.Int64),
rec.Column(1).(*array.Int64)
+ for i := range ids.Len() {
+ if ids.IsNull(i) {
+ got.ids = append(got.ids, nil)
+ } else {
+ got.ids = append(got.ids, ids.Value(i))
+ }
+ got.srcs = append(got.srcs, srcs.Value(i))
+ }
+
+ return true
+ })
+
+ return got, lastErr
+}
+
+func TestMergeSortedStreams_Interleaved(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ got, err := runMerge(t, mergeTestComparator(t, SortASC, NullsFirst), 4,
+ idStream(mem, 0, []any{1, 4}, []any{7}),
+ idStream(mem, 1, []any{}, []any{2, 3}, []any{9}),
+ idStream(mem, 2, []any{5, 6, 8}),
+ idStream(mem, 3),
+ )
+ require.NoError(t, err)
+ assert.Equal(t, []any{int64(1), int64(2), int64(3), int64(4), int64(5),
int64(6), int64(7), int64(8), int64(9)}, got.ids)
+ assert.Equal(t, []int64{4, 4, 1}, got.sizes)
+}
+
+func TestMergeSortedStreams_TiesKeepStreamOrder(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ got, err := runMerge(t, mergeTestComparator(t, SortASC, NullsFirst),
100,
+ idStream(mem, 0, []any{1, 2}),
+ idStream(mem, 1, []any{1, 2}),
+ )
+ require.NoError(t, err)
+ assert.Equal(t, []any{int64(1), int64(1), int64(2), int64(2)}, got.ids)
+ assert.Equal(t, []any{int64(0), int64(1), int64(0), int64(1)}, got.srcs)
+}
+
+func TestMergeSortedStreams_DisjointStreamsFillOneBatch(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ got, err := runMerge(t, mergeTestComparator(t, SortASC, NullsFirst),
100,
+ idStream(mem, 1, []any{4, 5, 6}),
+ idStream(mem, 0, []any{1, 2, 3}),
+ )
+ require.NoError(t, err)
+ assert.Equal(t, []any{int64(1), int64(2), int64(3), int64(4), int64(5),
int64(6)}, got.ids)
+ assert.Equal(t, []int64{6}, got.sizes)
+}
+
+func TestMergeSortedStreams_DescendingNullsLast(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ got, err := runMerge(t, mergeTestComparator(t, SortDESC, NullsLast), 2,
+ idStream(mem, 0, []any{5, nil}),
+ idStream(mem, 1, []any{6, 2}, []any{nil}),
+ )
+ require.NoError(t, err)
+ assert.Equal(t, []any{int64(6), int64(5), int64(2), nil, nil}, got.ids)
+}
+
+func TestMergeSortedStreams_RejectsUnsortedInput(t *testing.T) {
+ tests := []struct {
+ name string
+ batches [][]any
+ }{
+ {"within a batch", [][]any{{1, 3, 2}}},
+ {"across batches", [][]any{{1, 5}, {3}}},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ _, err := runMerge(t, mergeTestComparator(t, SortASC,
NullsFirst), 100,
+ idStream(mem, 0, []any{0, 4}),
+ idStream(mem, 1, tt.batches...),
+ )
+ require.Error(t, err)
+ assert.Contains(t, err.Error(), "f1.parquet")
+ })
+ }
+}
+
+func TestMergeSortedStreams_EarlyStopReleases(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ streams := []iter.Seq2[arrow.RecordBatch, error]{
+ idStream(mem, 0, []any{1, 3}, []any{5, 7}),
+ idStream(mem, 1, []any{2, 4}, []any{6, 8}),
+ }
+ batches := 0
+ ok := mergeSortedStreams(t.Context(), mergeTestComparator(t, SortASC,
NullsFirst), []string{"a", "b"}, streams, 3,
+ func(rec arrow.RecordBatch, err error) bool {
+ require.NoError(t, err)
+ rec.Release()
+ batches++
+
+ return false
+ })
+ assert.False(t, ok)
+ assert.Equal(t, 1, batches)
+}
diff --git a/table/sort_compaction_test.go b/table/sort_compaction_test.go
new file mode 100644
index 000000000..8aa3292a5
--- /dev/null
+++ b/table/sort_compaction_test.go
@@ -0,0 +1,308 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package table_test
+
+import (
+ "context"
+ "fmt"
+ "path/filepath"
+ "slices"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet/file"
+ "github.com/apache/arrow-go/v18/parquet/pqarrow"
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/apache/iceberg-go/table"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func idSortOrder(t *testing.T, id int, transform iceberg.Transform)
table.SortOrder {
+ t.Helper()
+
+ order, err := table.NewSortOrder(id, []table.SortField{{
+ SourceIDs: []int{1},
+ Transform: transform,
+ Direction: table.SortASC,
+ NullOrder: table.NullsFirst,
+ }})
+ require.NoError(t, err)
+
+ return order
+}
+
+func newSortCompactionTestTable(t *testing.T, spec *iceberg.PartitionSpec)
*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},
+ )
+ meta, err := table.NewMetadata(schema, spec, idSortOrder(t, 1,
iceberg.IdentityTransform{}), location,
+ iceberg.Properties{table.PropertyFormatVersion: "2"})
+ require.NoError(t, err)
+
+ return table.New(
+ table.Identifier{"db", "sort_compaction_test"},
+ meta, location+"/metadata/v1.metadata.json",
+ func(context.Context) (iceio.IO, error) { return
iceio.LocalFS{}, nil },
+ &rowDeltaCatalog{metadata: meta},
+ )
+}
+
+// addSortTestFile adds one data file with the given ids, all in partition
+// data=partition. The file is written unsorted.
+func addSortTestFile(t *testing.T, tbl *table.Table, name, partition string,
ids ...int64) *table.Table {
+ t.Helper()
+
+ arrowSc, err := table.SchemaToArrowSchema(tbl.Schema(), nil, false,
false)
+ require.NoError(t, err)
+ rows := make([]string, len(ids))
+ for i, id := range ids {
+ rows[i] = fmt.Sprintf(`{"id": %d, "data": %q}`, id, partition)
+ }
+ dataPath := tbl.Location() + "/data/" + name + ".parquet"
+ writeParquetFile(t, dataPath, arrowSc, "["+strings.Join(rows, ",")+"]")
+
+ tx := tbl.NewTransaction()
+ require.NoError(t, tx.AddDataFilesByPath(t.Context(),
[]string{dataPath}, nil, false))
+ tbl, err = tx.Commit(t.Context())
+ require.NoError(t, err)
+
+ return tbl
+}
+
+func sortTestGroup(t *testing.T, tbl *table.Table) table.CompactionTaskGroup {
+ t.Helper()
+
+ tasks, err := tbl.Scan().PlanFiles(t.Context())
+ require.NoError(t, err)
+ slices.SortFunc(tasks, func(a, b table.FileScanTask) int {
+ return strings.Compare(a.File.FilePath(), b.File.FilePath())
+ })
+ group := table.CompactionTaskGroup{PartitionKey: "group", Tasks: tasks}
+ for _, task := range tasks {
+ group.TotalSizeBytes += task.File.FileSizeBytes()
+ }
+
+ return group
+}
+
+func readInt64Column(t *testing.T, path, column string) []int64 {
+ t.Helper()
+
+ pf, err := file.OpenParquetFile(strings.TrimPrefix(path, "file://"),
false)
+ require.NoError(t, err)
+ defer pf.Close()
+
+ rdr, err := pqarrow.NewFileReader(pf, pqarrow.ArrowReadProperties{},
memory.DefaultAllocator)
+ require.NoError(t, err)
+ tbl, err := rdr.ReadTable(t.Context())
+ require.NoError(t, err)
+ defer tbl.Release()
+
+ idx := tbl.Schema().FieldIndices(column)
+ require.Len(t, idx, 1)
+ var out []int64
+ for _, chunk := range tbl.Column(idx[0]).Data().Chunks() {
+ out = append(out, chunk.(*array.Int64).Int64Values()...)
+ }
+
+ return out
+}
+
+func TestExecuteCompactionGroup_SortOrder(t *testing.T) {
+ tbl := newSortCompactionTestTable(t, iceberg.UnpartitionedSpec)
+ tbl = addSortTestFile(t, tbl, "f0", "x", 7, 3, 11)
+ tbl = addSortTestFile(t, tbl, "f1", "x", 2, 9, 5)
+ tbl = addSortTestFile(t, tbl, "f2", "x", 12, 1, 8)
+ tbl = addSortTestFile(t, tbl, "f3", "x", 4, 10, 6)
+ group := sortTestGroup(t, tbl)
+
+ // Five-row output batches span several source files, and a one-byte
+ // target rolls a file after every batch.
+ res, err := table.ExecuteCompactionGroup(t.Context(), tbl, group,
+ table.WithCompactionSortOrder(tbl.SortOrder()),
+ table.WithCompactionArrowBatchSize(5),
+ table.WithCompactionTargetFileSize(1))
+ require.NoError(t, err)
+ require.Len(t, res.NewDataFiles, 3)
+ assert.Len(t, res.OldDataFiles, 4)
+
+ var all []int64
+ for _, df := range res.NewDataFiles {
+ ids := readInt64Column(t, df.FilePath(), "id")
+ assert.True(t, slices.IsSorted(ids), "file %s is not sorted:
%v", df.FilePath(), ids)
+ all = append(all, ids...)
+
+ require.NotNil(t, df.SortOrderID())
+ assert.Equal(t, tbl.SortOrder().OrderID(), *df.SortOrderID())
+ }
+ assert.Equal(t, []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, all,
+ "files must hold consecutive, non-overlapping ranges of the
sorted group")
+}
+
+func TestExecuteCompactionGroup_SortOrderSplitsPartitions(t *testing.T) {
+ spec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{2}, FieldID: 1000, Transform:
iceberg.IdentityTransform{}, Name: "data",
+ })
+ tbl := newSortCompactionTestTable(t, &spec)
+ // Partitions interleave across the group's task order, as they do when
+ // a group spans files written under an older spec.
+ tbl = addSortTestFile(t, tbl, "f0", "a", 5, 1)
+ tbl = addSortTestFile(t, tbl, "f1", "b", 4, 2)
+ tbl = addSortTestFile(t, tbl, "f2", "a", 6, 3)
+ group := sortTestGroup(t, tbl)
+
+ res, err := table.ExecuteCompactionGroup(t.Context(), tbl, group,
+ table.WithCompactionSortOrder(tbl.SortOrder()))
+ require.NoError(t, err)
+ require.Len(t, res.NewDataFiles, 2)
+
+ got := make(map[string][]int64)
+ for _, df := range res.NewDataFiles {
+ part, ok := df.Partition()[1000].(string)
+ require.True(t, ok)
+ got[part] = readInt64Column(t, df.FilePath(), "id")
+ }
+ assert.Equal(t, map[string][]int64{"a": {1, 3, 5, 6}, "b": {2, 4}}, got)
+}
+
+func TestExecuteCompactionGroup_SortOrderIDClaim(t *testing.T) {
+ tbl := newSortCompactionTestTable(t, iceberg.UnpartitionedSpec)
+ tbl = addSortTestFile(t, tbl, "f0", "x", 3, 1)
+ tbl = addSortTestFile(t, tbl, "f1", "x", 2, 4)
+ group := sortTestGroup(t, tbl)
+
+ tests := []struct {
+ name string
+ order table.SortOrder
+ }{
+ {"not a table sort order", idSortOrder(t, 7,
iceberg.IdentityTransform{})},
+ {"transform does not preserve order", idSortOrder(t, 1,
iceberg.BucketTransform{NumBuckets: 4})},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ res, err := table.ExecuteCompactionGroup(t.Context(),
tbl, group,
+ table.WithCompactionSortOrder(tt.order))
+ require.NoError(t, err)
+ require.Len(t, res.NewDataFiles, 1)
+ assert.Equal(t, []int64{1, 2, 3, 4}, readInt64Column(t,
res.NewDataFiles[0].FilePath(), "id"))
+ assert.Nil(t, res.NewDataFiles[0].SortOrderID())
+ })
+ }
+}
+
+func TestExecuteCompactionGroup_UnsortedOrderRejected(t *testing.T) {
+ tbl := newSortCompactionTestTable(t, iceberg.UnpartitionedSpec)
+ tbl = addSortTestFile(t, tbl, "f0", "x", 1)
+
+ _, err := table.ExecuteCompactionGroup(t.Context(), tbl,
sortTestGroup(t, tbl),
+ table.WithCompactionSortOrder(table.UnsortedSortOrder))
+ require.ErrorIs(t, err, table.ErrInvalidOperation)
+}
+
+func TestRewriteDataFiles_SortOrder(t *testing.T) {
+ tbl := newSortCompactionTestTable(t, iceberg.UnpartitionedSpec)
+ tbl = addSortTestFile(t, tbl, "f0", "x", 6, 2)
+ tbl = addSortTestFile(t, tbl, "f1", "x", 5, 1)
+ tbl = addSortTestFile(t, tbl, "f2", "x", 4, 3)
+
+ res, err := tbl.RewriteDataFiles(t.Context(),
[]table.CompactionTaskGroup{sortTestGroup(t, tbl)},
+ table.RewriteDataFilesOptions{
+ GroupOptions:
[]table.CompactionGroupOption{table.WithCompactionSortOrder(tbl.SortOrder())},
+ })
+ require.NoError(t, err)
+ require.NotNil(t, res.Table)
+ assert.Equal(t, 3, res.RemovedDataFiles)
+ assert.Equal(t, 1, res.AddedDataFiles)
+ assertRowCount(t, res.Table, 6)
+
+ tasks, err := res.Table.Scan().PlanFiles(t.Context())
+ require.NoError(t, err)
+ require.Len(t, tasks, 1)
+ require.NotNil(t, tasks[0].File.SortOrderID())
+ assert.Equal(t, tbl.SortOrder().OrderID(), *tasks[0].File.SortOrderID())
+ assert.Equal(t, []int64{1, 2, 3, 4, 5, 6}, readInt64Column(t,
tasks[0].File.FilePath(), "id"))
+}
+
+// TestExecuteCompactionGroup_MergesSortedFiles compacts files that already
+// record the sort order, which takes the streaming merge path, with one
+// position-delete file shared by both inputs.
+func TestExecuteCompactionGroup_MergesSortedFiles(t *testing.T) {
+ tbl := newSortCompactionTestTable(t, iceberg.UnpartitionedSpec)
+ tbl = addSortTestFile(t, tbl, "f0", "x", 7, 3, 11)
+ tbl = addSortTestFile(t, tbl, "f1", "x", 2, 9, 5)
+ tbl = addSortTestFile(t, tbl, "f2", "x", 12, 1, 8)
+ tbl = addSortTestFile(t, tbl, "f3", "x", 4, 10, 6)
+ unsorted := sortTestGroup(t, tbl)
+
+ res, err := tbl.RewriteDataFiles(t.Context(),
[]table.CompactionTaskGroup{
+ {PartitionKey: "a", Tasks: unsorted.Tasks[:2]},
+ {PartitionKey: "b", Tasks: unsorted.Tasks[2:]},
+ }, table.RewriteDataFilesOptions{
+ GroupOptions:
[]table.CompactionGroupOption{table.WithCompactionSortOrder(tbl.SortOrder())},
+ })
+ require.NoError(t, err)
+ tbl = res.Table
+
+ group := sortTestGroup(t, tbl)
+ require.Len(t, group.Tasks, 2)
+ byFirstID := make(map[int64]table.FileScanTask)
+ for _, task := range group.Tasks {
+ require.NotNil(t, task.File.SortOrderID())
+ require.Equal(t, tbl.SortOrder().OrderID(),
*task.File.SortOrderID())
+ byFirstID[readInt64Column(t, task.File.FilePath(), "id")[0]] =
task
+ }
+ // Files hold 2,3,5,7,9,11 and 1,4,6,8,10,12. Delete 2 and 11 from the
+ // first and 4 from the second.
+ first, second := byFirstID[2].File.FilePath(),
byFirstID[1].File.FilePath()
+ delPath := tbl.Location() + "/data/pos-del.parquet"
+ writeParquetFile(t, delPath, table.PositionalDeleteArrowSchema,
fmt.Sprintf(
+ `[{"file_path": %q, "pos": 0}, {"file_path": %q, "pos": 5},
{"file_path": %q, "pos": 1}]`,
+ first, first, second))
+ bldr, err := iceberg.NewDataFileBuilder(*iceberg.UnpartitionedSpec,
iceberg.EntryContentPosDeletes,
+ delPath, iceberg.ParquetFile, nil, nil, nil, 3, 1)
+ require.NoError(t, err)
+ posDel := bldr.Build()
+ for i := range group.Tasks {
+ group.Tasks[i].DeleteFiles = []iceberg.DataFile{posDel}
+ }
+
+ out, err := table.ExecuteCompactionGroup(t.Context(), tbl, group,
+ table.WithCompactionSortOrder(tbl.SortOrder()),
+ table.WithCompactionArrowBatchSize(4),
+ table.WithCompactionTargetFileSize(1))
+ require.NoError(t, err)
+ require.Len(t, out.NewDataFiles, 3)
+
+ var all []int64
+ for _, df := range out.NewDataFiles {
+ ids := readInt64Column(t, df.FilePath(), "id")
+ assert.True(t, slices.IsSorted(ids), "file %s is not sorted:
%v", df.FilePath(), ids)
+ all = append(all, ids...)
+ require.NotNil(t, df.SortOrderID())
+ assert.Equal(t, tbl.SortOrder().OrderID(), *df.SortOrderID())
+ }
+ assert.Equal(t, []int64{1, 3, 5, 6, 7, 8, 9, 10, 12}, all)
+}
diff --git a/table/write_records.go b/table/write_records.go
index 1712dddf2..2f9f94141 100644
--- a/table/write_records.go
+++ b/table/write_records.go
@@ -42,6 +42,17 @@ type writeRecordConfig struct {
preserveRowLineage bool
recordBatchBufferSize int
parquetRowGroupLimit int
+ presorted bool
+ sortOrderID int
+}
+
+// withPresortedInput declares the input records globally sorted by the
+// sort order with id sortOrderID; see withFactoryPresorted.
+func withPresortedInput(sortOrderID int) WriteRecordOption {
+ return func(c *writeRecordConfig) {
+ c.presorted = true
+ c.sortOrderID = sortOrderID
+ }
}
// WithTargetFileSize overrides the table's default target file size.
@@ -170,6 +181,11 @@ func WriteRecords(ctx context.Context, tbl *Table,
errors.New("WithClusteredWrite and WithMaxWriteWorkers
are incompatible: the clustered write path is single-threaded"))
}
+ if cfg.presorted && !cfg.clustered {
+ return internal.SingleErrorIter[iceberg.DataFile](
+ errors.New("presorted input requires the clustered
write path to keep its order"))
+ }
+
if cfg.preserveRowLineage {
if v := tbl.metadata.FormatVersion(); v < 3 {
return internal.SingleErrorIter[iceberg.DataFile](
@@ -258,6 +274,9 @@ func WriteRecords(ctx context.Context, tbl *Table,
if cfg.fileSchema != nil {
args.factoryOpts = append(args.factoryOpts,
withFactoryFileSchema(cfg.fileSchema))
}
+ if cfg.presorted {
+ args.factoryOpts = append(args.factoryOpts,
withFactoryPresorted(cfg.sortOrderID))
+ }
return recordsToDataFiles(ctx, tbl.Location(), meta, args)
}
diff --git a/website/src/api.md b/website/src/api.md
index 162d7a726..6486ff5f4 100644
--- a/website/src/api.md
+++ b/website/src/api.md
@@ -523,6 +523,22 @@ The `table/compaction` subpackage provides bin-packing
planning. The `iceberg co
Compaction tasks represent complete data files, not independent scan splits.
The planner collapses split tasks for the same file into one whole-file task,
rejecting inconsistent delete or lineage metadata and dropping residual filters
because they select files but do not filter rows during compaction. Direct
calls to `ExecuteCompactionGroup` reject partial ranges and residual row
filters. Attached delete files are still applied when rewriting each source
file.
+To sort instead of bin-pack, pass `table.WithCompactionSortOrder` as a group
option. Each group is sorted as a whole, so every output file holds a sorted,
non-overlapping range of it. Output files record the order's `sort_order_id`
when the order is one of the table's sort orders, every field has one source
column, every field but the last is identity and the last preserves order (see
`table.WithCompactionSortOrder`). Files of different groups still overlap, and
by default a group is pac [...]
+
+```go
+cfg := compaction.DefaultConfig()
+cfg.RewriteAll = true
+cfg.MaxGroupSizeBytes = 8 << 30 // 8 GiB per group
+plan, err := compaction.Analyze(ctx, tbl, cfg)
+if err != nil { /* ... */ }
+
+result, err := txn.RewriteDataFiles(ctx, groups /* from plan.Groups */,
table.RewriteDataFilesOptions{
+ GroupOptions:
[]table.CompactionGroupOption{table.WithCompactionSortOrder(tbl.SortOrder())},
+})
+```
+
+When every file of a group already records the order's `sort_order_id` (for
example, files written by an earlier sort rewrite), every sort field is
identity, and no partition of the group has more files than the scan
concurrency (`table.WithCompactionScanConcurrency`, `GOMAXPROCS` by default),
the files are streamed and merged, holding up to about five batches per file
instead of the whole group. Otherwise a sort rewrite holds each group decoded
in memory, and with the batches queued for [...]
+
### Expiring snapshots
```go
