This is an automated email from the ASF dual-hosted git repository.
laskoviymishka pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new d883b2c68 perf(table): avoid gathering already ordered compaction
batches (#2373)
d883b2c68 is described below
commit d883b2c6873386babaa6ea2c5cbea18a30c2459c
Author: Minh Vu <[email protected]>
AuthorDate: Sat Oct 31 11:09:24 2026 +0100
perf(table): avoid gathering already ordered compaction batches (#2373)
---
table/sort_compaction.go | 52 +++
table/sort_compaction_ordered_fastpath_test.go | 510 +++++++++++++++++++++++++
2 files changed, 562 insertions(+)
diff --git a/table/sort_compaction.go b/table/sort_compaction.go
index e6ce185b8..ebebfaad2 100644
--- a/table/sort_compaction.go
+++ b/table/sort_compaction.go
@@ -561,6 +561,33 @@ func yieldSortedPartition(
}
slices.SortStableFunc(refs, rows.bind(batches))
+ if sortedRefsAreIdentity(refs, batches) {
+ // Keep the scan's natural order. Every output owns its
reference to
+ // the original batch; no Arrow Take arrays are needed.
+ for _, batch := range batches {
+ for off := int64(0); off < batch.NumRows(); off +=
batchRows {
+ if err := ctx.Err(); err != nil {
+ yield(nil, err)
+
+ return false
+ }
+ end := min(off+batchRows, batch.NumRows())
+ var out arrow.RecordBatch
+ if off == 0 && end == batch.NumRows() {
+ batch.Retain()
+ out = batch
+ } else {
+ out = batch.NewSlice(off, end)
+ }
+ if !yield(out, nil) {
+ return false
+ }
+ }
+ }
+
+ return true
+ }
+
takeFn := partitionBatchByKey(ctx)
mem := compute.GetAllocator(ctx)
for off := 0; off < len(refs); off += int(batchRows) {
@@ -585,6 +612,31 @@ func yieldSortedPartition(
return true
}
+// sortedRefsAreIdentity reports whether the stable sort preserved input order.
+func sortedRefsAreIdentity(refs []rowRef, batches []arrow.RecordBatch) bool {
+ // Check the final row first to reject a common late inversion quickly.
+ if len(refs) > 0 {
+ lastBatch := len(batches) - 1
+ for batches[lastBatch].NumRows() == 0 {
+ lastBatch--
+ }
+ if refs[len(refs)-1] != newRowRef(lastBatch,
int(batches[lastBatch].NumRows())-1) {
+ return false
+ }
+ }
+ index := 0
+ for batch, rec := range batches {
+ for row := range int(rec.NumRows()) {
+ if refs[index] != newRowRef(batch, row) {
+ return false
+ }
+ index++
+ }
+ }
+
+ 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.
diff --git a/table/sort_compaction_ordered_fastpath_test.go
b/table/sort_compaction_ordered_fastpath_test.go
new file mode 100644
index 000000000..55c104758
--- /dev/null
+++ b/table/sort_compaction_ordered_fastpath_test.go
@@ -0,0 +1,510 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package table
+
+import (
+ "context"
+ "errors"
+ "slices"
+ "strconv"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+)
+
+type orderedPartitionTestRow struct {
+ key int64
+ valid bool
+ label string
+}
+
+func orderedPartitionTestSetup(t testing.TB) (*arrow.Schema, rowComparator) {
+ t.Helper()
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "key", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "label", Type: arrow.BinaryTypes.String},
+ }, nil)
+ schema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "key", Type:
iceberg.PrimitiveTypes.Int64},
+ iceberg.NestedField{ID: 2, Name: "label", Type:
iceberg.PrimitiveTypes.String},
+ )
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform: iceberg.IdentityTransform{},
+ Direction: SortASC, NullOrder: NullsFirst,
+ }})
+ if err != nil {
+ t.Fatal(err)
+ }
+ cmp, err := newRowComparator(order, schema, sc)
+ if err != nil {
+ t.Fatal(err)
+ }
+
+ return sc, cmp
+}
+
+func makeOrderedPartitionTestBatch(mem memory.Allocator, sc *arrow.Schema,
rows []orderedPartitionTestRow) arrow.RecordBatch {
+ builder := array.NewRecordBuilder(mem, sc)
+ defer builder.Release()
+ keys := builder.Field(0).(*array.Int64Builder)
+ labels := builder.Field(1).(*array.StringBuilder)
+ for _, row := range rows {
+ if row.valid {
+ keys.Append(row.key)
+ } else {
+ keys.AppendNull()
+ }
+ labels.Append(row.label)
+ }
+
+ return builder.NewRecordBatch()
+}
+
+func TestYieldSortedPartitionFastPathPreservesTiesAndNulls(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ sc, cmp := orderedPartitionTestSetup(t)
+ batches := []arrow.RecordBatch{
+ makeOrderedPartitionTestBatch(mem, sc,
[]orderedPartitionTestRow{
+ {label: "null"}, {key: 1, valid: true, label: "a"},
{key: 1, valid: true, label: "b"},
+ }),
+ makeOrderedPartitionTestBatch(mem, sc,
[]orderedPartitionTestRow{
+ {key: 1, valid: true, label: "c"},
+ {key: 2, valid: true, label: "d"},
+ {key: 3, valid: true, label: "e"},
+ {key: 3, valid: true, label: "f"},
+ {key: 4, valid: true, label: "g"},
+ }),
+ }
+ defer func() {
+ for _, batch := range batches {
+ batch.Release()
+ }
+ }()
+ var got []string
+ var sizes []int64
+ calls := 0
+ ok := yieldSortedPartition(t.Context(), batches, cmp, 3, func(rec
arrow.RecordBatch, err error) bool {
+ if err != nil {
+ t.Fatalf("yield: %v", err)
+ }
+ defer rec.Release()
+ if calls == 0 && rec != batches[0] {
+ t.Fatal("full input batch was copied")
+ }
+ calls++
+ sizes = append(sizes, rec.NumRows())
+ col := rec.Column(1).(*array.String)
+ for i := range col.Len() {
+ got = append(got, col.Value(i))
+ }
+
+ return true
+ })
+ if !ok {
+ t.Fatal("iteration stopped early")
+ }
+ if want := []int64{3, 3, 2}; !slices.Equal(sizes, want) {
+ t.Fatalf("batch sizes %v; want %v", sizes, want)
+ }
+ if want := []string{"null", "a", "b", "c", "d", "e", "f", "g"};
!slices.Equal(got, want) {
+ t.Fatalf("row order %v; want %v", got, want)
+ }
+}
+
+func TestYieldSortedPartitionFastPathMatchesLegacyAcrossBatchBoundaries(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ sc, cmp := orderedPartitionTestSetup(t)
+
+ // Source batches intentionally do not align with the four-row output
+ // limit. Ties cross the batch boundary, and the first key is null.
+ input := [][]orderedPartitionTestRow{
+ {
+ {label: "null"},
+ {key: 1, valid: true, label: "a"},
+ {key: 1, valid: true, label: "b"},
+ },
+ {
+ {key: 1, valid: true, label: "c"},
+ {key: 2, valid: true, label: "d"},
+ {key: 2, valid: true, label: "e"},
+ {key: 3, valid: true, label: "f"},
+ {key: 4, valid: true, label: "g"},
+ },
+ }
+ batches := make([]arrow.RecordBatch, len(input))
+ var want []orderedPartitionTestRow
+ for i, rows := range input {
+ batches[i] = makeOrderedPartitionTestBatch(mem, sc, rows)
+ want = append(want, rows...)
+ }
+ defer func() {
+ for _, batch := range batches {
+ batch.Release()
+ }
+ }()
+
+ var legacyRows []orderedPartitionTestRow
+ for _, method := range []struct {
+ name string
+ fn func(context.Context, []arrow.RecordBatch,
rowComparator, int64, func(arrow.RecordBatch, error) bool) bool
+ wantSizes []int64
+ }{
+ {"legacy", legacyYieldSortedPartition, []int64{4, 4}},
+ {"fastpath", yieldSortedPartition, []int64{3, 4, 1}},
+ } {
+ t.Run(method.name, func(t *testing.T) {
+ var got []orderedPartitionTestRow
+ var sizes []int64
+ ok := method.fn(t.Context(), batches, cmp, 4, func(rec
arrow.RecordBatch, err error) bool {
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer rec.Release()
+ sizes = append(sizes, rec.NumRows())
+ keys := rec.Column(0).(*array.Int64)
+ labels := rec.Column(1).(*array.String)
+ for i := range int(rec.NumRows()) {
+ row := orderedPartitionTestRow{
+ valid: !keys.IsNull(i),
+ label: labels.Value(i),
+ }
+ if row.valid {
+ row.key = keys.Value(i)
+ }
+ got = append(got, row)
+ }
+
+ return true
+ })
+ if !ok {
+ t.Fatal("iteration stopped early")
+ }
+ if !slices.Equal(got, want) {
+ t.Fatalf("rows = %v, want %v", got, want)
+ }
+ if !slices.Equal(sizes, method.wantSizes) {
+ t.Fatalf("batch sizes = %v, want %v", sizes,
method.wantSizes)
+ }
+ if method.name == "legacy" {
+ legacyRows = slices.Clone(got)
+ } else if !slices.Equal(got, legacyRows) {
+ t.Fatalf("fast-path rows = %v, legacy rows =
%v", got, legacyRows)
+ }
+ })
+ }
+}
+
+func TestYieldSortedPartitionFastPathFallsBack(t *testing.T) {
+ tests := []struct {
+ name string
+ input [][]orderedPartitionTestRow
+ want []string
+ }{
+ {name: "overlapping batches", input:
[][]orderedPartitionTestRow{
+ {{key: 1, valid: true, label: "a"}, {key: 4, valid:
true, label: "b"}, {key: 4, valid: true, label: "c"}},
+ {{key: 0, valid: true, label: "d"}, {key: 2, valid:
true, label: "e"}, {key: 4, valid: true, label: "f"}},
+ }, want: []string{"d", "a", "e", "b", "c", "f"}},
+ {name: "unordered within batch", input:
[][]orderedPartitionTestRow{
+ {{key: 5, valid: true, label: "a"}, {key: 1, valid:
true, label: "b"}, {key: 5, valid: true, label: "c"}},
+ {{key: 2, valid: true, label: "d"}, {key: 3, valid:
true, label: "e"}, {key: 4, valid: true, label: "f"}},
+ }, want: []string{"b", "d", "e", "f", "a", "c"}},
+ }
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ sc, cmp := orderedPartitionTestSetup(t)
+ batches := make([]arrow.RecordBatch, len(tc.input))
+ for i, rows := range tc.input {
+ batches[i] = makeOrderedPartitionTestBatch(mem,
sc, rows)
+ }
+ defer func() {
+ for _, batch := range batches {
+ batch.Release()
+ }
+ }()
+ var got []string
+ ok := yieldSortedPartition(t.Context(), batches, cmp,
2, func(rec arrow.RecordBatch, err error) bool {
+ if err != nil {
+ t.Fatalf("yield: %v", err)
+ }
+ defer rec.Release()
+ col := rec.Column(1).(*array.String)
+ for i := range col.Len() {
+ got = append(got, col.Value(i))
+ }
+
+ return true
+ })
+ if !ok || !slices.Equal(got, tc.want) {
+ t.Fatalf("got %v (ok=%v), want %v", got, ok,
tc.want)
+ }
+ })
+ }
+}
+
+func TestYieldSortedPartitionFastPathCancellation(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ sc, cmp := orderedPartitionTestSetup(t)
+ batch := makeOrderedPartitionTestBatch(mem, sc,
[]orderedPartitionTestRow{{key: 1, valid: true, label: "a"}})
+ defer batch.Release()
+ ctx, cancel := context.WithCancel(t.Context())
+ cancel()
+ calls := 0
+ ok := yieldSortedPartition(ctx, []arrow.RecordBatch{batch}, cmp, 1,
func(rec arrow.RecordBatch, err error) bool {
+ calls++
+ if rec != nil {
+ rec.Release()
+ t.Fatal("expected cancellation, got record")
+ }
+ if err != context.Canceled {
+ t.Fatalf("got %v, want context.Canceled", err)
+ }
+
+ return true
+ })
+ if ok || calls != 1 {
+ t.Fatalf("ok=%v calls=%d", ok, calls)
+ }
+}
+
+// legacyYieldSortedPartition retains the pre-fast-path implementation to
+// measure before and after on the same Go toolchain and Arrow version.
+func legacyYieldSortedPartition(
+ 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 _, rec := range batches {
+ if !rec.Schema().Equal(sc) {
+ yield(nil, errors.New("different schemas"))
+
+ return false
+ }
+ total += rec.NumRows()
+ }
+ refs := make([]rowRef, 0, total)
+ for batch, rec := range batches {
+ for row := range int(rec.NumRows()) {
+ refs = append(refs, newRowRef(batch, row))
+ }
+ }
+ 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
+ }
+ out, err := gatherSortedRows(mem, takeFn, batches,
refs[off:min(off+int(batchRows), len(refs))])
+ if err != nil {
+ yield(nil, err)
+
+ return false
+ }
+ if !yield(out, nil) {
+ return false
+ }
+ }
+
+ return true
+}
+
+func BenchmarkYieldSortedPartitionFastPath(b *testing.B) {
+ const (
+ batchCount = 8
+ rowsPerBatch = 4096
+ columnCount = 16
+ outputRows = 8192
+ )
+ for _, mode := range []string{"globally_sorted", "nearly_sorted",
"near_middle", "near_center", "center_swap", "tail_swap", "interleaved",
"shuffled"} {
+ b.Run(mode, func(b *testing.B) {
+ fields := make([]arrow.Field, columnCount)
+ for i := range fields {
+ fields[i] = arrow.Field{Name: "c" +
strconv.Itoa(i), Type: arrow.PrimitiveTypes.Int64}
+ }
+ sc := arrow.NewSchema(fields, nil)
+ schemaFields := make([]iceberg.NestedField, columnCount)
+ for i := range fields {
+ schemaFields[i] = iceberg.NestedField{ID: i +
1, Name: fields[i].Name, Type: iceberg.PrimitiveTypes.Int64, Required: true}
+ }
+ order, err := NewSortOrder(1, []SortField{{
+ SourceIDs: []int{1}, Transform:
iceberg.IdentityTransform{},
+ Direction: SortASC, NullOrder: NullsFirst,
+ }})
+ if err != nil {
+ b.Fatal(err)
+ }
+ comparator, err := newRowComparator(order,
iceberg.NewSchema(0, schemaFields...), sc)
+ if err != nil {
+ b.Fatal(err)
+ }
+ batches := make([]arrow.RecordBatch, batchCount)
+ for batch := range batches {
+ builder :=
array.NewRecordBuilder(memory.DefaultAllocator, sc)
+ for row := range rowsPerBatch {
+ var key int64
+ switch mode {
+ case "globally_sorted",
"nearly_sorted", "near_middle", "near_center", "center_swap", "tail_swap":
+ key = int64(batch*rowsPerBatch
+ row)
+ if mode == "center_swap" &&
batch == batchCount/2-1 {
+ switch row {
+ case rowsPerBatch / 2:
+ key++
+ case rowsPerBatch/2 + 1:
+ key--
+ }
+ }
+ if mode == "tail_swap" && batch
== batchCount-1 {
+ switch row {
+ case rowsPerBatch - 2:
+ key++
+ case rowsPerBatch - 1:
+ key--
+ }
+ }
+ if (mode == "nearly_sorted" &&
batch == batchCount-1 && row == rowsPerBatch-1) ||
+ (mode == "near_middle"
&& batch == batchCount/2 && row == 0) ||
+ (mode == "near_center"
&& batch == batchCount/2-1 && row == rowsPerBatch/2) {
+ key = 0
+ }
+ case "interleaved":
+ key = int64(row*batchCount +
batch)
+ case "shuffled":
+ key =
int64((row*37%rowsPerBatch)*batchCount + batch)
+ }
+ for col := range columnCount {
+
builder.Field(col).(*array.Int64Builder).Append(key + int64(col))
+ }
+ }
+ batches[batch] = builder.NewRecordBatch()
+ builder.Release()
+ }
+ b.Cleanup(func() {
+ for _, batch := range batches {
+ batch.Release()
+ }
+ })
+ for _, method := range []struct {
+ name string
+ fn func(context.Context, []arrow.RecordBatch,
rowComparator, int64, func(arrow.RecordBatch, error) bool) bool
+ }{
+ {"legacy", legacyYieldSortedPartition},
+ {"fastpath", yieldSortedPartition},
+ } {
+ b.Run(method.name, func(b *testing.B) {
+ b.ReportAllocs()
+ b.ReportMetric(batchCount*rowsPerBatch,
"rows/op")
+ b.ResetTimer()
+ for range b.N {
+ var count int64
+ ok := method.fn(b.Context(),
batches, comparator, outputRows, func(rec arrow.RecordBatch, err error) bool {
+ if err != nil {
+ b.Fatal(err)
+ }
+ count += rec.NumRows()
+ rec.Release()
+
+ return true
+ })
+ if !ok || count !=
batchCount*rowsPerBatch {
+ b.Fatalf("got %d rows
(ok=%v)", count, ok)
+ }
+ }
+ })
+ }
+ })
+ }
+}
+
+func TestSortedRefsAreIdentityChecksEveryPair(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ schema, comparator := orderedPartitionTestSetup(t)
+ for _, inversion := range []int{-1, 1, 2, 3, 4, 5, 6, 7} {
+ t.Run(strconv.Itoa(inversion), func(t *testing.T) {
+ var batches []arrow.RecordBatch
+ for batch := range 4 {
+ var rows []orderedPartitionTestRow
+ for row := range 2 {
+ i := batch*2 + row
+ key := int64(i)
+ if i == inversion {
+ key = -1
+ }
+ rows = append(rows,
orderedPartitionTestRow{key: key, valid: true})
+ }
+ batches = append(batches,
makeOrderedPartitionTestBatch(mem, schema, rows))
+ if batch == 1 {
+ batches = append(batches,
makeOrderedPartitionTestBatch(mem, schema, nil))
+ }
+ }
+ defer func() {
+ for _, batch := range batches {
+ batch.Release()
+ }
+ }()
+ var refs []rowRef
+ for batch, rec := range batches {
+ for row := range int(rec.NumRows()) {
+ refs = append(refs, newRowRef(batch,
row))
+ }
+ }
+ slices.SortStableFunc(refs, comparator.bind(batches))
+ got := sortedRefsAreIdentity(refs, batches)
+ if got != (inversion < 0) {
+ t.Fatalf("inversion at %d: got sorted=%t",
inversion, got)
+ }
+ })
+ }
+}
+
+func TestSortedRefsAreIdentityEmptyBatches(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ t.Cleanup(func() { mem.AssertSize(t, 0) })
+ schema, comparator := orderedPartitionTestSetup(t)
+ empty := makeOrderedPartitionTestBatch(mem, schema, nil)
+ defer empty.Release()
+ batch := makeOrderedPartitionTestBatch(mem, schema,
[]orderedPartitionTestRow{{key: 1, valid: true}})
+ defer batch.Release()
+ for _, batches := range [][]arrow.RecordBatch{{empty}, {empty, batch,
empty}} {
+ var refs []rowRef
+ for batch, rec := range batches {
+ for row := range int(rec.NumRows()) {
+ refs = append(refs, newRowRef(batch, row))
+ }
+ }
+ slices.SortStableFunc(refs, comparator.bind(batches))
+ if !sortedRefsAreIdentity(refs, batches) {
+ t.Fatal("expected unchanged row order")
+ }
+ }
+}
