perf(table): parallelize non-bulk PurgeFiles deletion

catalog/glue/glue_test.go

@@ -28,6 +28,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -864,16 +865,18 @@ func TestGluePurgeTableSwallowsPurgeFilesError(t *testing.T) {
 	assert := require.New(t)
 	ctx := context.Background()
 	const scheme = "gluepurgefail"
-	dropCalled := false
-	removeBeforeDrop := false
-	removeCalls := 0
+	var (
+		dropCalled       atomic.Bool
+		removeBeforeDrop atomic.Bool
+		removeCalls      atomic.Int64
+	)
 	failingFS := failRemoveIO{
 		MemFS: iceio.NewMemFS(),
 		err:   errGluePurgeRemove,
 		onRemove: func() {
-			removeCalls++
-			if !dropCalled {
-				removeBeforeDrop = true
+			removeCalls.Add(1)
+			if !dropCalled.Load() {
+				removeBeforeDrop.Store(true)
 			}
 		},
 	}
@@ -899,15 +902,15 @@ func TestGluePurgeTableSwallowsPurgeFilesError(t *testing.T) {
 		DatabaseName: aws.String("test_database"),
 		Name:         aws.String("test_table"),
 	}, mock.Anything).Run(func(mock.Arguments) {
-		dropCalled = true
+		dropCalled.Store(true)
 	}).Return(&glue.DeleteTableOutput{}, nil).Once()
 
 	glueCatalog := &Catalog{glueSvc: mockGlueSvc}
 
 	assert.NoError(glueCatalog.PurgeTable(ctx, TableIdentifier("test_database", "test_table")))
-	assert.True(dropCalled)
-	assert.Positive(removeCalls)
-	assert.False(removeBeforeDrop, "PurgeTable should drop the catalog entry before removing files")
+	assert.True(dropCalled.Load())
+	assert.Positive(removeCalls.Load())
+	assert.False(removeBeforeDrop.Load(), "PurgeTable should drop the catalog entry before removing files")
 	file, err := failingFS.Open(dataFile)
 	assert.NoError(err, "data file should remain when FileIO remove fails")
 	assert.NotNil(file)

catalog/hive/hive_test.go

@@ -25,6 +25,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/apache/iceberg-go"
@@ -916,16 +917,18 @@ func TestHivePurgeTableSwallowsPurgeFilesError(t *testing.T) {
 	assert := require.New(t)
 	ctx := context.Background()
 	const scheme = "hivepurgefail"
-	dropCalled := false
-	removeBeforeDrop := false
-	removeCalls := 0
+	var (
+		dropCalled       atomic.Bool
+		removeBeforeDrop atomic.Bool
+		removeCalls      atomic.Int64
+	)
 	failingFS := failRemoveIO{
 		MemFS: iceio.NewMemFS(),
 		err:   errHivePurgeRemove,
 		onRemove: func() {
-			removeCalls++
-			if !dropCalled {
-				removeBeforeDrop = true
+			removeCalls.Add(1)
+			if !dropCalled.Load() {
+				removeBeforeDrop.Store(true)
 			}
 		},
 	}
@@ -948,15 +951,15 @@ func TestHivePurgeTableSwallowsPurgeFilesError(t *testing.T) {
 		Return(hiveTable, nil).Twice()
 	mockClient.On("DropTable", mock.Anything, "test_database", "test_table", false).
 		Run(func(mock.Arguments) {
-			dropCalled = true
+			dropCalled.Store(true)
 		}).Return(nil).Once()
 
 	hiveCatalog := NewCatalogWithClient(mockClient, iceberg.Properties{})
 
 	assert.NoError(hiveCatalog.PurgeTable(ctx, TableIdentifier("test_database", "test_table")))
-	assert.True(dropCalled)
-	assert.Positive(removeCalls)
-	assert.False(removeBeforeDrop, "PurgeTable should drop the catalog entry before removing files")
+	assert.True(dropCalled.Load())
+	assert.Positive(removeCalls.Load())
+	assert.False(removeBeforeDrop.Load(), "PurgeTable should drop the catalog entry before removing files")
 	file, err := failingFS.Open(dataFile)
 	assert.NoError(err, "data file should remain when FileIO remove fails")
 	assert.NotNil(file)

schema.go

@@ -346,8 +346,7 @@ func (s *Schema) MarshalJSON() ([]byte, error) {
 
 	type Alias Schema
 
-	aliasCopy := *(*Alias)(s)
-	aliasCopy.IdentifierFieldIDs = ids
+	aliasCopy := Alias{ID: s.ID, IdentifierFieldIDs: ids}
 
 	return json.Marshal(struct {
 		Type   string        `json:"type"`

schema_test.go

@@ -24,6 +24,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/apache/iceberg-go"
@@ -2293,3 +2294,48 @@ func TestVisitGeoSchemaWithSchemaVisitorPerPrimitiveType(t *testing.T) {
 	assert.Equal(t, 1, v.geometryCalls)
 	assert.Equal(t, 1, v.geographyCalls)
 }
+
+func TestSchemaMarshalJSONConcurrentLazyLookups(t *testing.T) {
+	for range 32 {
+		schema := iceberg.NewSchemaWithIdentifiers(17, nil,
+			iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true},
+			iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String},
+		)
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		for range 8 {
+			wg.Go(func() {
+				<-start
+				for range 8 {
+					_, err := json.Marshal(schema)
+					assert.NoError(t, err)
+				}
+			})
+			wg.Go(func() {
+				<-start
+				_, found := schema.FindFieldByID(1)
+				assert.True(t, found)
+				_, found = schema.FindFieldByName("data")
+				assert.True(t, found)
+				_, found = schema.FindFieldByNameCaseInsensitive("DATA")
+				assert.True(t, found)
+				name, found := schema.FindColumnName(2)
+				assert.True(t, found)
+				assert.Equal(t, "data", name)
+			})
+		}
+		close(start)
+		wg.Wait()
+
+		data, err := json.Marshal(schema)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{
+			"type": "struct", "schema-id": 17, "identifier-field-ids": [],
+			"fields": [
+				{"id": 1, "name": "id", "type": "long", "required": true},
+				{"id": 2, "name": "data", "type": "string", "required": false}
+			]
+		}`, string(data))
+		assert.Nil(t, schema.IdentifierFieldIDs)
+	}
+}

table/orphan_cleanup.go

@@ -688,14 +688,27 @@ func deleteFiles(ctx context.Context, fs iceio.IO, orphanFiles []string, cfg *or
 		}
 	}
 
-	if cfg.maxConcurrency == 1 {
-		return deleteFilesSequential(fs, orphanFiles, cfg)
+	if cfg.maxConcurrency <= 1 {
+		return deleteFilesSequential(ctx, fs, orphanFiles, cfg)
 	}
 
-	return deleteFilesParallel(fs, orphanFiles, cfg)
+	deleteFunc := fs.Remove
+	if cfg.deleteFunc != nil {
+		deleteFunc = cfg.deleteFunc
+	}
+
+	return deleteFilesParallel(
+		ctx,
+		orphanFiles,
+		cfg.maxConcurrency,
+		deleteFunc,
+		func(file string, err error) error {
+			return fmt.Errorf("failed to delete orphan file %s: %w", file, err)
+		},
+	)
 }
 
-func deleteFilesSequential(fs iceio.IO, orphanFiles []string, cfg *orphanCleanupConfig) ([]string, error) {
+func deleteFilesSequential(ctx context.Context, fs iceio.IO, orphanFiles []string, cfg *orphanCleanupConfig) ([]string, error) {
 	var deletedFiles []string
 
 	deleteFunc := fs.Remove
@@ -705,6 +718,11 @@ func deleteFilesSequential(fs iceio.IO, orphanFiles []string, cfg *orphanCleanup
 
 	var result error
 	for _, file := range orphanFiles {
+		if err := ctx.Err(); err != nil {
+			result = errors.Join(result, err)
+
+			break
+		}
 		if err := deleteFunc(file); err != nil {
 			result = errors.Join(result, fmt.Errorf("failed to delete orphan file %s: %w", file, err))
 
@@ -716,55 +734,85 @@ func deleteFilesSequential(fs iceio.IO, orphanFiles []string, cfg *orphanCleanup
 	return deletedFiles, result
 }
 
-func deleteFilesParallel(fs iceio.IO, orphanFiles []string, cfg *orphanCleanupConfig) ([]string, error) {
-	deleteFunc := fs.Remove
-	if cfg.deleteFunc != nil {
-		deleteFunc = cfg.deleteFunc
+func deleteFilesParallel(
+	ctx context.Context,
+	files []string,
+	maxConcurrency int,
+	deleteFunc func(string) error,
+	wrapError func(string, error) error,
+) ([]string, error) {
+	workers := min(max(maxConcurrency, 1), len(files))
+	jobs := make(chan int)
+	deleted := make([]bool, len(files))
+	deleteErrors := make([]error, len(files))
+
+	var cancellationErr error
+	var cancellationOnce sync.Once
+	recordCancellation := func(err error) {
+		cancellationOnce.Do(func() {
+			cancellationErr = err
+		})
 	}
 
-	in := make(chan string, cfg.maxConcurrency)
-	out := make(chan string, cfg.maxConcurrency)
-	errList := make([][]error, cfg.maxConcurrency)
-
-	go func() {
-		defer close(in)
-		for _, file := range orphanFiles {
-			in <- file
-		}
-	}()
-
 	var wg sync.WaitGroup
-	wg.Add(cfg.maxConcurrency)
-	for i := range cfg.maxConcurrency {
-		go func(workerID int) {
+	wg.Add(workers)
+	for range workers {
+		go func() {
 			defer wg.Done()
-			for file := range in {
-				if err := deleteFunc(file); err != nil {
-					errList[workerID] = append(errList[workerID], fmt.Errorf("failed to delete orphan file %s: %w", file, err))
-				} else {
-					out <- file
+			for {
+				select {
+				case <-ctx.Done():
+					recordCancellation(ctx.Err())
+
+					return
+				case index, ok := <-jobs:
+					if !ok {
+						return
+					}
+					if err := ctx.Err(); err != nil {
+						recordCancellation(err)
+
+						return
+					}
+
+					if err := deleteFunc(files[index]); err != nil {
+						deleteErrors[index] = wrapError(files[index], err)
+					} else {
+						deleted[index] = true
+					}
 				}
 			}
-		}(i)
+		}()
 	}
 
-	go func() {
-		wg.Wait()
-		close(out)
-	}()
+send:
+	for index := range files {
+		select {
+		case jobs <- index:
+		case <-ctx.Done():
+			recordCancellation(ctx.Err())
 
-	deletedFiles := make([]string, 0, len(orphanFiles))
-	for file := range out {
-		deletedFiles = append(deletedFiles, file)
+			break send
+		}
 	}
+	close(jobs)
+	wg.Wait()
 
-	var allErrors []error
-	for _, workerErrors := range errList {
-		allErrors = append(allErrors, workerErrors...)
+	deletedFiles := make([]string, 0, len(files))
+	allErrors := make([]error, 0)
+	for index, file := range files {
+		if deleted[index] {
+			deletedFiles = append(deletedFiles, file)
+		}
+		if deleteErrors[index] != nil {
+			allErrors = append(allErrors, deleteErrors[index])
+		}
+	}
+	if cancellationErr != nil {
+		allErrors = append(allErrors, cancellationErr)
 	}
-	err := errors.Join(allErrors...)
 
-	return deletedFiles, err
+	return deletedFiles, errors.Join(allErrors...)
 }
 
 // normalizeFilePath normalizes file paths for comparison by handling different
@@ -1310,15 +1358,23 @@ func (t Table) PurgeFiles(ctx context.Context) error {
 				errs = append(errs, fmt.Errorf("bulk deletion failed: %w", bulkErr))
 			}
 		} else {
-			for _, file := range files {
-				if err := ctx.Err(); err != nil {
-					errs = append(errs, err)
+			_, removeErr := deleteFilesParallel(
+				ctx,
+				files,
+				runtime.GOMAXPROCS(0),
+				func(file string) error {
+					if err := fs.Remove(file); err != nil && !os.IsNotExist(err) {
+						return err
+					}
 
-					break
-				}
-				if rmErr := fs.Remove(file); rmErr != nil && !os.IsNotExist(rmErr) {
-					errs = append(errs, fmt.Errorf("failed to remove %s: %w", file, rmErr))
-				}
+					return nil
+				},
+				func(file string, err error) error {
+					return fmt.Errorf("failed to remove %s: %w", file, err)
+				},
+			)
+			if removeErr != nil {
+				errs = append(errs, removeErr)
 			}
 		}
 	}