iceberg: sort compaction bins on disk instead of in memory (#11112)

A sorted rewrite collected every row of a bin into one slice and sorted it
there, so a bin larger than the worker's heap could not be sorted at all.
sort_max_input_mb existed for that reason and skipped the bins it capped.

parquet-go's SortingWriter buffers sort_buffer_rows rows, encodes each buffer
as a sorted run, and merges the runs at close; backing those runs with a
FileBufferPool keeps them in files rather than on the heap. sort_spill_dir says
where, defaulting to the system temp directory — NewFileBufferPool resolves an
empty path to the working directory, which is not what an unset setting means.

The output now also declares its sorting columns, which the plain writer the
sorted path used never did.

Claude-Session: https://claude.ai/code/session_015SZkLTUvd1svDu4xdr6Q3y

weed/worker/tasks/iceberg/compact.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"path"
 	"sort"
 	"strings"
@@ -1149,19 +1150,48 @@ func mergeParquetFilesSorted(
 		return nil, 0, fmt.Errorf("resolve equality columns: %w", err)
 	}
 
-	comparator := parquetSchema.Comparator(sortingColumns...)
+	// Sorted rewrites stream rows through parquet-go's sorting writer rather
-	var allRows []parquet.Row
+	// than holding the whole bin: rows accumulate until sortBufferRows, then a
-
+	// sorted run is encoded into a temporary row group, and Close merges the
-	collectRows := func(reader *parquet.Reader, source string) (int64, error) {
+	// runs into the output. The runs live in files from a buffer pool instead
-		return visitFilteredParquetRows(ctx, reader, source, bucketName, dataPath, positionDeletes, resolvedEqGroups, func(filtered []parquet.Row) error {
+	// of on the heap, so a bin larger than memory sorts rather than failing.
-			for _, row := range filtered {
+	spillDir := rewritePlan.spillDir
-				allRows = append(allRows, row.Clone())
+	if spillDir == "" {
+		// NewFileBufferPool takes filepath.Abs of what it is given, and
+		// filepath.Abs("") is the working directory, not the temp directory.
+		spillDir = os.TempDir()
 	}
-			return nil
+	sortBufferRows := rewritePlan.bufferRows
+	if sortBufferRows < minSortBufferRows {
+		sortBufferRows = defaultSortBufferRows
+	}
+
+	var outputBuf bytes.Buffer
+	writer := parquet.NewSortingWriter[any](&outputBuf, sortBufferRows, parquetSchema,
+		parquet.SortingWriterConfig(
+			parquet.SortingColumns(sortingColumns...),
+			parquet.SortingBuffers(parquet.NewFileBufferPool(spillDir, "seaweedfs-iceberg-sort-*")),
+		),
+	)
+	// Closing returns the run files to the pool, which deletes them. A failure
+	// before that has to release them too, or the job leaves temp files behind.
+	closed := false
+	defer func() {
+		if !closed {
+			writer.Reset(io.Discard)
+		}
+	}()
+
+	writeRows := func(reader *parquet.Reader, source string) (int64, error) {
+		return visitFilteredParquetRows(ctx, reader, source, bucketName, dataPath, positionDeletes, resolvedEqGroups, func(filtered []parquet.Row) error {
+			// The writer's row buffer copies the values it is handed, so these
+			// rows need no cloning out of the reader's own buffer first.
+			_, writeErr := writer.WriteRows(filtered)
+			return writeErr
 		})
 	}
 
-	totalRows, err := collectRows(firstReader, entries[0].DataFile().FilePath())
+	totalRows, err := writeRows(firstReader, entries[0].DataFile().FilePath())
 	if err != nil {
 		return nil, 0, err
 	}
@@ -1184,28 +1214,17 @@ func mergeParquetFilesSorted(
 			return nil, 0, fmt.Errorf("schema mismatch in %s: cannot merge files with different schemas", entry.DataFile().FilePath())
 		}
 
-		rowsCollected, err := collectRows(reader, entry.DataFile().FilePath())
+		rowsWritten, err := writeRows(reader, entry.DataFile().FilePath())
 		if err != nil {
 			return nil, 0, err
 		}
-		totalRows += rowsCollected
+		totalRows += rowsWritten
 	}
 
-	sort.SliceStable(allRows, func(i, j int) bool {
-		return comparator(allRows[i], allRows[j]) < 0
-	})
-
-	var outputBuf bytes.Buffer
-	writer := parquet.NewWriter(&outputBuf, parquetSchema)
-	if len(allRows) > 0 {
-		if _, err := writer.WriteRows(allRows); err != nil {
-			writer.Close()
-			return nil, 0, fmt.Errorf("write sorted rows: %w", err)
-		}
-	}
 	if err := writer.Close(); err != nil {
-		return nil, 0, fmt.Errorf("close writer: %w", err)
+		return nil, 0, fmt.Errorf("close sorting writer: %w", err)
 	}
+	closed = true
 
 	return outputBuf.Bytes(), totalRows, nil
 }

weed/worker/tasks/iceberg/config.go

@@ -26,6 +26,8 @@ const (
 	defaultDeleteMaxOutputFiles   = 8
 	defaultRewriteStrategy        = "binpack"
 	rewriteStrategyAuto           = "auto"
+	defaultSortBufferRows         = 262144
+	minSortBufferRows             = 1024
 	defaultMinManifestsToRewrite  = 5
 	minManifestsToRewrite         = 2
 	defaultOperations             = "all"
@@ -126,6 +128,8 @@ type Config struct {
 	Where                       string
 	RewriteStrategy             string
 	SortMaxInputBytes           int64
+	SortBufferRows              int64
+	SortSpillDir                string
 }
 
 // ParseConfig extracts an iceberg maintenance Config from plugin config values.
@@ -149,6 +153,8 @@ func ParseConfig(values map[string]*plugin_pb.ConfigValue) Config {
 		Where:                       strings.TrimSpace(readStringConfig(values, "where", "")),
 		RewriteStrategy:             strings.TrimSpace(strings.ToLower(readStringConfig(values, "rewrite_strategy", defaultRewriteStrategy))),
 		SortMaxInputBytes:           readSizeMBConfig(values, "sort_max_input_mb", 0),
+		SortBufferRows:              readInt64Config(values, "sort_buffer_rows", defaultSortBufferRows),
+		SortSpillDir:                strings.TrimSpace(readStringConfig(values, "sort_spill_dir", "")),
 	}
 
 	// Clamp the fields that are always defaulted by worker config parsing.
@@ -199,6 +205,11 @@ func applyThresholdDefaults(cfg Config) Config {
 	if cfg.SortMaxInputBytes < 0 {
 		cfg.SortMaxInputBytes = 0
 	}
+	// A run smaller than this buys no memory back worth the extra runs, and
+	// the writer stops at 32K of them, so a tiny value would cap the output.
+	if cfg.SortBufferRows < minSortBufferRows {
+		cfg.SortBufferRows = defaultSortBufferRows
+	}
 	if cfg.MinManifestsToRewrite < minManifestsToRewrite {
 		cfg.MinManifestsToRewrite = minManifestsToRewrite
 	}

weed/worker/tasks/iceberg/handler.go

@@ -192,6 +192,21 @@ func (h *Handler) Descriptor() *plugin_pb.JobTypeDescriptor {
 							Widget:      plugin_pb.ConfigWidget_CONFIG_WIDGET_TEXT,
 							Placeholder: "binpack or sort",
 						},
+						{
+							Name:        "sort_buffer_rows",
+							Label:       "Sort Buffer Rows",
+							Description: "Rows a sorted rewrite holds in memory before spilling a sorted run to disk. The writer merges at most 32K runs, so this also bounds how many rows one output file can hold.",
+							FieldType:   plugin_pb.ConfigFieldType_CONFIG_FIELD_TYPE_INT64,
+							Widget:      plugin_pb.ConfigWidget_CONFIG_WIDGET_NUMBER,
+							MinValue:    &plugin_pb.ConfigValue{Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: minSortBufferRows}},
+						},
+						{
+							Name:        "sort_spill_dir",
+							Label:       "Sort Spill Directory",
+							Description: "Directory holding a sorted rewrite's temporary runs. Empty uses the system temp directory.",
+							FieldType:   plugin_pb.ConfigFieldType_CONFIG_FIELD_TYPE_STRING,
+							Widget:      plugin_pb.ConfigWidget_CONFIG_WIDGET_TEXT,
+						},
 						{
 							Name:        "sort_max_input_mb",
 							Label:       "Sort Max Input (MB)",
@@ -327,6 +342,8 @@ func (h *Handler) Descriptor() *plugin_pb.JobTypeDescriptor {
 				"table_properties_override":     {Kind: &plugin_pb.ConfigValue_BoolValue{BoolValue: true}},
 				"rewrite_strategy":              {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: defaultRewriteStrategy}},
 				"sort_max_input_mb":             {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 0}},
+				"sort_buffer_rows":              {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: defaultSortBufferRows}},
+				"sort_spill_dir":                {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: ""}},
 				"where":                         {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: ""}},
 			},
 		},
@@ -358,6 +375,8 @@ func (h *Handler) Descriptor() *plugin_pb.JobTypeDescriptor {
 			"table_properties_override":     {Kind: &plugin_pb.ConfigValue_BoolValue{BoolValue: true}},
 			"rewrite_strategy":              {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: defaultRewriteStrategy}},
 			"sort_max_input_mb":             {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: 0}},
+			"sort_buffer_rows":              {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: defaultSortBufferRows}},
+			"sort_spill_dir":                {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: ""}},
 			"where":                         {Kind: &plugin_pb.ConfigValue_StringValue{StringValue: ""}},
 		},
 	}

weed/worker/tasks/iceberg/sort_merge_test.go

@@ -0,0 +1,161 @@
+package iceberg
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/apache/iceberg-go"
+	"github.com/apache/iceberg-go/table"
+	"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
+	"github.com/seaweedfs/seaweedfs/weed/s3api/s3tables"
+)
+
+// seedSortableTable creates a table sorted ascending on id whose data files are
+// deliberately out of order: ids interleave across the files and descend within
+// each one, so neither the file order nor the row order is already sorted.
+func seedSortableTable(t *testing.T, files, rowsPerFile int) (*fakeFilerServer, filer_pb.SeaweedFilerClient, tableSetup) {
+	t.Helper()
+
+	fs, client := startFakeFiler(t)
+	sortOrder, err := table.NewSortOrder(1, []table.SortField{{
+		SourceIDs: []int{1},
+		Transform: iceberg.IdentityTransform{},
+		Direction: table.SortASC,
+		NullOrder: table.NullsFirst,
+	}})
+	if err != nil {
+		t.Fatalf("new sort order: %v", err)
+	}
+
+	type dataRow = struct {
+		ID   int64
+		Name string
+	}
+	dataFiles := make([]struct {
+		Name string
+		Rows []dataRow
+	}, 0, files)
+	for f := 0; f < files; f++ {
+		rows := make([]dataRow, 0, rowsPerFile)
+		for i := 0; i < rowsPerFile; i++ {
+			id := int64((rowsPerFile-1-i)*files + f)
+			rows = append(rows, dataRow{ID: id, Name: fmt.Sprintf("row-%d", id)})
+		}
+		dataFiles = append(dataFiles, struct {
+			Name string
+			Rows []dataRow
+		}{Name: fmt.Sprintf("d%d.parquet", f), Rows: rows})
+	}
+
+	setup := tableSetup{BucketName: "tb", Namespace: "ns", TableName: "tbl"}
+	populateTableWithDeleteFilesAndSortOrder(t, fs, setup, dataFiles, nil, nil, sortOrder)
+	return fs, client, setup
+}
+
+func sortCompactionConfig(spillDir string) Config {
+	return Config{
+		TargetFileSizeBytes: 256 * 1024 * 1024,
+		MinInputFiles:       2,
+		MaxCommitRetries:    3,
+		ApplyDeletes:        true,
+		RewriteStrategy:     "sort",
+		SortBufferRows:      minSortBufferRows,
+		SortSpillDir:        spillDir,
+	}
+}
+
+// compactedFileNames lists the outputs a compaction wrote, which is empty when
+// every bin was skipped.
+func compactedFileNames(fs *fakeFilerServer, setup tableSetup) []string {
+	dataDir := path.Join(s3tables.TablesPath, setup.BucketName, setup.tablePath(), "data")
+	var names []string
+	for _, e := range fs.listDir(dataDir) {
+		if strings.HasPrefix(e.Name, "compact-") {
+			names = append(names, e.Name)
+		}
+	}
+	return names
+}
+
+// A sorted rewrite whose bin does not fit in one buffer has to spill: each
+// buffer is encoded as a sorted run and the runs are merged at close. This is
+// the case that used to hold every row of the bin in memory, so it is worth
+// proving both that the output is globally ordered and that the runs do not
+// outlive the job.
+func TestCompactDataFilesSortSpillsRunsAndCleansUp(t *testing.T) {
+	const files, rowsPerFile = 3, 1200 // 3600 rows against a 1024-row buffer
+	fs, client, setup := seedSortableTable(t, files, rowsPerFile)
+
+	spillDir := t.TempDir()
+	handler := NewHandler(nil)
+	result, _, err := handler.compactDataFiles(context.Background(), client, setup.BucketName, setup.tablePath(), sortCompactionConfig(spillDir), nil)
+	if err != nil {
+		t.Fatalf("compactDataFiles: %v", err)
+	}
+	if !strings.Contains(result, "using sort") {
+		t.Fatalf("expected sorted compaction result, got %q", result)
+	}
+
+	rows := readCompactedRows(t, fs, setup)
+	if len(rows) != files*rowsPerFile {
+		t.Fatalf("expected %d compacted rows, got %d", files*rowsPerFile, len(rows))
+	}
+	for i := 1; i < len(rows); i++ {
+		if rows[i-1].ID > rows[i].ID {
+			t.Fatalf("rows are not sorted by id at %d: %d then %d", i, rows[i-1].ID, rows[i].ID)
+		}
+	}
+
+	leftovers, err := os.ReadDir(spillDir)
+	if err != nil {
+		t.Fatalf("read spill dir: %v", err)
+	}
+	if len(leftovers) != 0 {
+		t.Fatalf("expected the sorted runs to be removed, found %d file(s)", len(leftovers))
+	}
+}
+
+// An empty spill directory at the end of a job cannot tell a run written to
+// disk from one held in memory. A directory that cannot hold the runs can: the
+// same data compacts with a usable one and does not with an unusable one, which
+// it would not do if the runs never reached the pool. A bin whose merge fails
+// is logged and skipped rather than failing the job, so the difference shows up
+// as a missing output file rather than an error.
+func TestCompactDataFilesSortSpillDirDecidesOutcome(t *testing.T) {
+	// Two files that together cross the 1024-row buffer, so a run is spilled.
+	const files, rowsPerFile = 2, 700
+
+	for _, tc := range []struct {
+		name         string
+		spillDir     func(t *testing.T) string
+		wantOutcomes int
+	}{
+		{
+			name:         "usable",
+			spillDir:     func(t *testing.T) string { return t.TempDir() },
+			wantOutcomes: 1,
+		},
+		{
+			name:         "unusable",
+			spillDir:     func(t *testing.T) string { return filepath.Join(t.TempDir(), "not-created") },
+			wantOutcomes: 0,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			fs, client, setup := seedSortableTable(t, files, rowsPerFile)
+			handler := NewHandler(nil)
+
+			if _, _, err := handler.compactDataFiles(context.Background(), client, setup.BucketName, setup.tablePath(), sortCompactionConfig(tc.spillDir(t)), nil); err != nil {
+				t.Fatalf("compactDataFiles: %v", err)
+			}
+			if names := compactedFileNames(fs, setup); len(names) != tc.wantOutcomes {
+				t.Fatalf("expected %d compacted file(s), got %v", tc.wantOutcomes, names)
+			}
+		})
+	}
+}

weed/worker/tasks/iceberg/sort_strategy.go

@@ -13,6 +13,9 @@ import (
 type compactionRewritePlan struct {
 	strategy   string
 	sortFields []compactionSortField
+	// Carried on the plan so the merge does not need the whole Config.
+	bufferRows int64
+	spillDir   string
 }
 
 type compactionSortField struct {
@@ -37,7 +40,7 @@ func resolveCompactionRewritePlan(config Config, meta table.Metadata) (*compacti
 			glog.V(2).Infof("iceberg compact: auto strategy falling back to binpack: %v", err)
 			return &compactionRewritePlan{strategy: defaultRewriteStrategy}, nil
 		}
-		return &compactionRewritePlan{strategy: "sort", sortFields: sortFields}, nil
+		return newSortPlan(config, sortFields), nil
 	}
 	if strategy != "sort" {
 		return nil, fmt.Errorf("unsupported rewrite strategy %q", config.RewriteStrategy)
@@ -52,10 +55,18 @@ func resolveCompactionRewritePlan(config Config, meta table.Metadata) (*compacti
 		return nil, err
 	}
 
+	return newSortPlan(config, sortFields), nil
+}
+
+// newSortPlan pairs the resolved sort fields with the settings the sorted
+// merge needs while writing.
+func newSortPlan(config Config, sortFields []compactionSortField) *compactionRewritePlan {
 	return &compactionRewritePlan{
-		strategy:   strategy,
+		strategy:   "sort",
 		sortFields: sortFields,
-	}, nil
+		bufferRows: config.SortBufferRows,
+		spillDir:   config.SortSpillDir,
+	}
 }
 
 var errUnsupportedTableSortOrder = fmt.Errorf("unsupported table sort order")