perf(table): load equality deletes lazily per scan task

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/arrow_scanner.go

@@ -901,16 +901,6 @@ func (as *arrowScan) addTaskProjectedFieldIDs(invariants *arrowScanInvariants, t
 	return nil
 }
 
-func addEqualityDeleteFieldIDs(invariants *arrowScanInvariants, eqDeleteSets map[int][]*equalityDeleteSet) {
-	for _, deleteSets := range eqDeleteSets {
-		for _, deleteSet := range deleteSets {
-			for _, id := range deleteSet.fieldIDs {
-				invariants.projectedIDs[id] = struct{}{}
-			}
-		}
-	}
-}
-
 type enumeratedRecord struct {
 	Record tblutils.Enumerated[arrow.RecordBatch]
 	Task   tblutils.Enumerated[FileScanTask]
@@ -1706,7 +1696,7 @@ func createIterator(ctx context.Context, numWorkers uint, records <-chan enumera
 		return batch.Task.Index < 0
 	}
 
-	sequenced := tblutils.MakeSequencedChan(uint(numWorkers), records,
+	sequenced := tblutils.MakeSequencedChanWithDiscard(uint(numWorkers), records,
 		func(left, right *enumeratedRecord) bool {
 			switch {
 			case isBeforeAny(*left):
@@ -1732,7 +1722,11 @@ func createIterator(ctx context.Context, numWorkers uint, records <-chan enumera
 				return next.Task.Index == prev.Task.Index+1 &&
 					prev.Record.Last && next.Record.Index == 0
 			}
-		}, enumeratedRecord{Task: tblutils.Enumerated[FileScanTask]{Index: -1}})
+		}, enumeratedRecord{Task: tblutils.Enumerated[FileScanTask]{Index: -1}}, func(rec enumeratedRecord) {
+			if rec.Record.Value != nil {
+				rec.Record.Value.Release()
+			}
+		})
 
 	totalRowCount := int64(0)
 
@@ -1799,59 +1793,78 @@ func createIterator(ctx context.Context, numWorkers uint, records <-chan enumera
 	}
 }
 
-func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks []FileScanTask, deletesPerFile perFilePosDeletes, dvBitmaps perFileDVBitmaps, eqDeleteSets map[int][]*equalityDeleteSet, invariants *arrowScanInvariants) iter.Seq2[arrow.RecordBatch, error] {
-	extSet := substrait.NewExtensionSet()
-
-	ctx, cancel := context.WithCancelCause(exprs.WithExtensionIDSet(ctx, extSet))
-	taskChan := make(chan tblutils.Enumerated[FileScanTask], len(tasks))
-
-	// numWorkers := 1
-	numWorkers := min(as.concurrency, len(tasks))
-	records := make(chan enumeratedRecord, numWorkers)
-
-	var wg sync.WaitGroup
-	wg.Add(numWorkers)
-	for range numWorkers {
-		go func() {
-			defer wg.Done()
-			for {
-				select {
-				case <-ctx.Done():
-					return
-				case task, ok := <-taskChan:
-					if !ok {
+func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks []FileScanTask, deletesPerFile perFilePosDeletes, dvBitmaps perFileDVBitmaps, equalityDeleteLoader *lazyEqualityDeleteLoader, invariants *arrowScanInvariants) iter.Seq2[arrow.RecordBatch, error] {
+	return func(yield func(arrow.RecordBatch, error) bool) {
+		extSet := substrait.NewExtensionSet()
+
+		scanCtx, cancel := context.WithCancelCause(exprs.WithExtensionIDSet(ctx, extSet))
+		taskChan := make(chan tblutils.Enumerated[FileScanTask], len(tasks))
+
+		// numWorkers := 1
+		numWorkers := min(as.concurrency, len(tasks))
+		records := make(chan enumeratedRecord, numWorkers)
+
+		var wg sync.WaitGroup
+		wg.Add(numWorkers)
+		for range numWorkers {
+			go func() {
+				defer wg.Done()
+				for {
+					select {
+					case <-scanCtx.Done():
 						return
+					case task, ok := <-taskChan:
+						if !ok {
+							return
+						}
+						if scanCtx.Err() != nil {
+							return
+						}
+
+						filePath := task.Value.File.FilePath()
+						eqDeleteSets, err := equalityDeleteLoader.load(scanCtx, task.Value)
+						if err != nil {
+							records <- enumeratedRecord{Task: task, Err: err}
+							cancel(err)
+
+							return
+						}
+
+						if err := as.recordsFromTask(scanCtx, task, records,
+							deletesPerFile[filePath],
+							dvBitmaps[filePath],
+							eqDeleteSets,
+							invariants); err != nil {
+							cancel(err)
+
+							return
+						}
 					}
+				}
+			}()
+		}
 
-					filePath := task.Value.File.FilePath()
-					if err := as.recordsFromTask(ctx, task, records,
-						deletesPerFile[filePath],
-						dvBitmaps[filePath],
-						eqDeleteSets[task.Index],
-						invariants); err != nil {
-						cancel(err)
+		go func() {
+			defer func() {
+				close(taskChan)
+				wg.Wait()
+				close(records)
+			}()
 
-						return
-					}
+			for i, t := range tasks {
+				select {
+				case <-scanCtx.Done():
+					return
+				case taskChan <- tblutils.Enumerated[FileScanTask]{
+					Value: t, Index: i, Last: i == len(tasks)-1,
+				}:
 				}
 			}
 		}()
-	}
 
-	go func() {
-		for i, t := range tasks {
-			taskChan <- tblutils.Enumerated[FileScanTask]{
-				Value: t, Index: i, Last: i == len(tasks)-1,
-			}
-		}
-		close(taskChan)
-
-		wg.Wait()
-		close(records)
-	}()
-
-	return createIterator(ctx, uint(numWorkers), records, deletesPerFile,
-		cancel, as.rowLimit)
+		createIterator(scanCtx, uint(numWorkers), records, deletesPerFile,
+			cancel, as.rowLimit)(yield)
+	}
 }
 
 func (as *arrowScan) GetRecords(ctx context.Context, tasks []FileScanTask) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
@@ -1905,15 +1918,15 @@ func (as *arrowScan) GetRecords(ctx context.Context, tasks []FileScanTask) (*arr
 		return nil, nil, err
 	}
 
-	eqDeleteSets, err := readAllEqualityDeleteFiles(ctx, as.fs,
-		invariants.tableSchema, invariants.nameMapping, tasks, as.concurrency)
+	equalityDeleteLoader, err := newLazyEqualityDeleteLoader(
+		as.fs, invariants.tableSchema, invariants.nameMapping, tasks)
 	if err != nil {
-		// Positional deletes were fully loaded; release them before aborting.
 		releasePerFilePosDeletes(deletesPerFile)
 
 		return nil, nil, err
 	}
-	addEqualityDeleteFieldIDs(invariants, eqDeleteSets)
+	equalityDeleteLoader.addFieldIDs(invariants.projectedIDs)
 
-	return resultSchema, as.recordBatchesFromTasksAndDeletes(ctx, tasks, deletesPerFile, dvBitmaps, eqDeleteSets, invariants), nil
+	return resultSchema, as.recordBatchesFromTasksAndDeletes(ctx, tasks,
+		deletesPerFile, dvBitmaps, equalityDeleteLoader, invariants), nil
 }

table/arrow_scanner_delete_error_test.go

@@ -0,0 +1,58 @@
+// 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"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/memory"
+	"github.com/apache/iceberg-go/table/internal"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCreateIteratorReleasesQueuedBatchesAfterDeleteError(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer mem.AssertSize(t, 0)
+
+	expectedErr := errors.New("delete load failed")
+	records := make(chan enumeratedRecord, 2)
+	records <- enumeratedRecord{
+		Record: internal.Enumerated[arrow.RecordBatch]{
+			Value: checkedInt64RecordBatch(mem, 1),
+			Index: 0,
+			Last:  true,
+		},
+		Task: internal.Enumerated[FileScanTask]{Index: 1, Last: true},
+	}
+	records <- enumeratedRecord{
+		Task: internal.Enumerated[FileScanTask]{Index: 0},
+		Err:  expectedErr,
+	}
+	close(records)
+
+	ctx, cancel := context.WithCancelCause(context.Background())
+	var gotErr error
+	for _, err := range createIterator(ctx, 2, records, nil, cancel, 0) {
+		gotErr = err
+	}
+
+	require.ErrorIs(t, gotErr, expectedErr)
+}