diff --git a/schema.go b/schema.go index 24b0ff4b1..4330103dc 100644 --- a/schema.go +++ b/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"` diff --git a/schema_test.go b/schema_test.go index 7fa16e345..6d794ab71 100644 --- a/schema_test.go +++ b/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) + } +} diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go index 73a3cf8ff..8bca0d15f 100644 --- a/table/arrow_scanner.go +++ b/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 } diff --git a/table/arrow_scanner_delete_error_test.go b/table/arrow_scanner_delete_error_test.go new file mode 100644 index 000000000..3f13fe721 --- /dev/null +++ b/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) +} diff --git a/table/equality_delete_reader.go b/table/equality_delete_reader.go index 7338864ed..4c1ee9f87 100644 --- a/table/equality_delete_reader.go +++ b/table/equality_delete_reader.go @@ -26,6 +26,7 @@ import ( "fmt" "math" "slices" + "sync" "unsafe" "github.com/apache/arrow-go/v18/arrow" @@ -198,6 +199,177 @@ func newEqualityDeleteFileSet(id int, deleteSet *equalityDeleteSet) *equalityDel } } +type lazyEqualityDeleteLoader struct { + fs iceio.IO + tableSchema *iceberg.Schema + nameMapping iceberg.NameMapping + files map[string]*lazyEqualityDeleteFile + combinations sync.Map +} + +type lazyEqualityDeleteFile struct { + id int + dataFile iceberg.DataFile + fieldIDs []int + + once sync.Once + set *equalityDeleteFileSet + err error +} + +type lazyEqualityDeleteCombination struct { + once sync.Once + set *equalityDeleteSet +} + +func newLazyEqualityDeleteLoader( + fs iceio.IO, + tableSchema *iceberg.Schema, + nameMapping iceberg.NameMapping, + tasks []FileScanTask, +) (*lazyEqualityDeleteLoader, error) { + loader := &lazyEqualityDeleteLoader{ + fs: fs, + tableSchema: tableSchema, + nameMapping: nameMapping, + files: make(map[string]*lazyEqualityDeleteFile), + } + + for _, task := range tasks { + for _, dataFile := range task.EqualityDeleteFiles { + if dataFile.ContentType() != iceberg.EntryContentEqDeletes { + continue + } + + fieldIDs := dataFile.EqualityFieldIDs() + if len(fieldIDs) == 0 { + return nil, fmt.Errorf("%w: equality delete file %s", ErrEmptyEqualityFieldIDs, dataFile.FilePath()) + } + + path := dataFile.FilePath() + if _, ok := loader.files[path]; ok { + continue + } + + loader.files[path] = &lazyEqualityDeleteFile{ + id: len(loader.files), + dataFile: dataFile, + fieldIDs: fieldIDs, + } + } + } + + if len(loader.files) == 0 { + return nil, nil + } + + return loader, nil +} + +func (l *lazyEqualityDeleteLoader) addFieldIDs(idset set[int]) { + if l == nil { + return + } + + for _, file := range l.files { + for _, fieldID := range file.fieldIDs { + idset[fieldID] = struct{}{} + } + } +} + +func (l *lazyEqualityDeleteLoader) loadFile(ctx context.Context, file *lazyEqualityDeleteFile) (*equalityDeleteFileSet, error) { + file.once.Do(func() { + keys, colNames, err := readEqualityDeleteFile( + ctx, l.fs, l.tableSchema, l.nameMapping, file.dataFile, file.fieldIDs) + if err != nil { + file.err = err + + return + } + + file.set = newEqualityDeleteFileSet(file.id, &equalityDeleteSet{ + fieldIDs: file.fieldIDs, + colNames: colNames, + keys: keys, + }) + }) + + return file.set, file.err +} + +func (l *lazyEqualityDeleteLoader) combine(files []*equalityDeleteFileSet) *equalityDeleteSet { + files = normalizeEqualityDeleteFiles(files) + if len(files) == 1 { + return files[0].equalityDeleteSet + } + + key := equalityDeleteSetCombinationKey(files) + entryValue, _ := l.combinations.LoadOrStore(key, &lazyEqualityDeleteCombination{}) + entry := entryValue.(*lazyEqualityDeleteCombination) + entry.once.Do(func() { + entry.set = mergeEqualityDeleteSets(files) + }) + + return entry.set +} + +func (l *lazyEqualityDeleteLoader) load(ctx context.Context, task FileScanTask) ([]*equalityDeleteSet, error) { + if l == nil || len(task.EqualityDeleteFiles) == 0 { + return nil, nil + } + if len(task.EqualityDeleteFiles) == 1 { + dataFile := task.EqualityDeleteFiles[0] + if dataFile.ContentType() != iceberg.EntryContentEqDeletes { + return nil, nil + } + + file, ok := l.files[dataFile.FilePath()] + if !ok { + return nil, nil + } + + fileSet, err := l.loadFile(ctx, file) + if err != nil { + return nil, err + } + if len(fileSet.keys) == 0 { + return nil, nil + } + + return []*equalityDeleteSet{fileSet.equalityDeleteSet}, nil + } + + perFile := make(map[string]*equalityDeleteFileSet, len(task.EqualityDeleteFiles)) + for _, dataFile := range task.EqualityDeleteFiles { + if dataFile.ContentType() != iceberg.EntryContentEqDeletes { + continue + } + + path := dataFile.FilePath() + if _, seen := perFile[path]; seen { + continue + } + + file, ok := l.files[path] + if !ok { + continue + } + + fileSet, err := l.loadFile(ctx, file) + if err != nil { + return nil, err + } + perFile[path] = fileSet + } + + if len(perFile) == 0 { + return nil, nil + } + + return buildEqualityDeleteSetsForTask(task, perFile, l.combine), nil +} + // readAllEqualityDeleteFiles reads all unique equality delete files from // the tasks and builds per-task delete key sets. Returns nil if there are // no equality deletes. Delete files with different equality field IDs are @@ -301,91 +473,120 @@ func buildEqualityDeleteSetsPerTask( // File IDs are sufficient as the cache key because each ID identifies one // immutable delete set with a fixed equality-field group for this call. sharedSets := make(map[string]*equalityDeleteSet) + combine := func(files []*equalityDeleteFileSet) *equalityDeleteSet { + return equalityDeleteSetForFiles(files, sharedSets) + } + for i, t := range tasks { - if len(t.EqualityDeleteFiles) == 0 { - continue + sets := buildEqualityDeleteSetsForTask(t, perFile, combine) + if len(sets) > 0 { + perTask[i] = sets } + } - var ( - groupKey string - groupFiles []*equalityDeleteFileSet - groups map[string][]*equalityDeleteFileSet - ) + return perTask +} - for _, d := range t.EqualityDeleteFiles { - dk, ok := perFile[d.FilePath()] - if !ok { - continue - } +func buildEqualityDeleteSetsForTask( + task FileScanTask, + perFile map[string]*equalityDeleteFileSet, + combine func([]*equalityDeleteFileSet) *equalityDeleteSet, +) []*equalityDeleteSet { + if len(task.EqualityDeleteFiles) == 0 { + return nil + } - if groups != nil { - groups[dk.groupKey] = append(groups[dk.groupKey], dk) - } else if len(groupFiles) == 0 { - groupKey = dk.groupKey - groupFiles = append(groupFiles, dk) - } else if dk.groupKey != groupKey { - groups = make(map[string][]*equalityDeleteFileSet, 2) - groups[groupKey] = groupFiles - groupFiles = nil - groups[dk.groupKey] = append(groups[dk.groupKey], dk) - } else { - groupFiles = append(groupFiles, dk) - } - } + var ( + groupKey string + groupFiles []*equalityDeleteFileSet + groups map[string][]*equalityDeleteFileSet + ) - if groups == nil { - if len(groupFiles) == 0 { - continue - } + for _, dataFile := range task.EqualityDeleteFiles { + fileSet, ok := perFile[dataFile.FilePath()] + if !ok { + continue + } - deleteSet := equalityDeleteSetForFiles(groupFiles, sharedSets) - if len(deleteSet.keys) > 0 { - perTask[i] = []*equalityDeleteSet{deleteSet} - } + if groups != nil { + groups[fileSet.groupKey] = append(groups[fileSet.groupKey], fileSet) + } else if len(groupFiles) == 0 { + groupKey = fileSet.groupKey + groupFiles = append(groupFiles, fileSet) + } else if fileSet.groupKey != groupKey { + groups = make(map[string][]*equalityDeleteFileSet, 2) + groups[groupKey] = groupFiles + groupFiles = nil + groups[fileSet.groupKey] = append(groups[fileSet.groupKey], fileSet) + } else { + groupFiles = append(groupFiles, fileSet) + } + } - continue + if groups == nil { + if len(groupFiles) == 0 { + return nil } - sets := make([]*equalityDeleteSet, 0, len(groups)) - for _, files := range groups { - deleteSet := equalityDeleteSetForFiles(files, sharedSets) - if len(deleteSet.keys) > 0 { - sets = append(sets, deleteSet) - } + deleteSet := combine(groupFiles) + if len(deleteSet.keys) == 0 { + return nil } - if len(sets) > 0 { - perTask[i] = sets + return []*equalityDeleteSet{deleteSet} + } + + sets := make([]*equalityDeleteSet, 0, len(groups)) + for _, files := range groups { + deleteSet := combine(files) + if len(deleteSet.keys) > 0 { + sets = append(sets, deleteSet) } } - return perTask + return sets } func equalityDeleteSetForFiles( files []*equalityDeleteFileSet, sharedSets map[string]*equalityDeleteSet, ) *equalityDeleteSet { + files = normalizeEqualityDeleteFiles(files) + if len(files) == 1 { + return files[0].equalityDeleteSet + } + + key := equalityDeleteSetCombinationKey(files) + if deleteSet, ok := sharedSets[key]; ok { + return deleteSet + } + + deleteSet := mergeEqualityDeleteSets(files) + sharedSets[key] = deleteSet + + return deleteSet +} + +func normalizeEqualityDeleteFiles(files []*equalityDeleteFileSet) []*equalityDeleteFileSet { slices.SortFunc(files, func(a, b *equalityDeleteFileSet) int { return cmp.Compare(a.id, b.id) }) - files = slices.CompactFunc(files, func(a, b *equalityDeleteFileSet) bool { + + return slices.CompactFunc(files, func(a, b *equalityDeleteFileSet) bool { return a.id == b.id }) +} - if len(files) == 1 { - return files[0].equalityDeleteSet - } - +func equalityDeleteSetCombinationKey(files []*equalityDeleteFileSet) string { combinationKey := make([]byte, 0, len(files)*8) for _, file := range files { combinationKey = binary.LittleEndian.AppendUint64(combinationKey, uint64(file.id)) } - key := string(combinationKey) - if deleteSet, ok := sharedSets[key]; ok { - return deleteSet - } + return string(combinationKey) +} + +func mergeEqualityDeleteSets(files []*equalityDeleteFileSet) *equalityDeleteSet { deleteSet := &equalityDeleteSet{ keys: make(set[string]), fieldIDs: files[0].fieldIDs, @@ -397,8 +598,6 @@ func equalityDeleteSetForFiles( } } - sharedSets[key] = deleteSet - return deleteSet } diff --git a/table/equality_delete_reader_bench_test.go b/table/equality_delete_reader_bench_test.go index ee0ba913a..f30eeaf8c 100644 --- a/table/equality_delete_reader_bench_test.go +++ b/table/equality_delete_reader_bench_test.go @@ -22,9 +22,11 @@ import ( "context" "encoding/binary" "fmt" + "io" "os" "path/filepath" "strings" + "sync" "testing" "github.com/apache/arrow-go/v18/arrow" @@ -115,6 +117,186 @@ func BenchmarkReadEqualityDeleteFile(b *testing.B) { } } +var equalityDeleteLoadingBenchmarkSink int + +type equalityDeleteLoadingBenchmarkInput struct { + fs *countingEqualityDeleteOpenFS + tableSchema *iceberg.Schema + tasks []FileScanTask + deleteFiles int +} + +func BenchmarkLazyEqualityDeleteLoading(b *testing.B) { + input := newEqualityDeleteLoadingBenchmarkInput(b) + + b.Run("eager/all_tasks", func(b *testing.B) { + input.fs.opens.Store(0) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + perTask, err := readAllEqualityDeleteFiles( + b.Context(), input.fs, input.tableSchema, nil, input.tasks, 16) + if err != nil { + b.Fatal(err) + } + + totalKeys := 0 + for _, sets := range perTask { + for _, deleteSet := range sets { + totalKeys += len(deleteSet.keys) + } + } + equalityDeleteLoadingBenchmarkSink = totalKeys + } + b.StopTimer() + b.ReportMetric(float64(input.deleteFiles), "delete_files/op") + b.ReportMetric(float64(input.fs.opens.Load())/float64(b.N), "opens/op") + }) + + for _, benchmark := range []struct { + name string + taskCount int + }{ + {name: "lazy/unread", taskCount: 0}, + {name: "lazy/first_task", taskCount: 1}, + {name: "lazy/ten_tasks", taskCount: 10}, + {name: "lazy/full_scan", taskCount: len(input.tasks)}, + } { + b.Run(benchmark.name, func(b *testing.B) { + input.fs.opens.Store(0) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + loader, err := newLazyEqualityDeleteLoader( + input.fs, input.tableSchema, nil, input.tasks) + if err != nil { + b.Fatal(err) + } + + totalKeys, err := loadEqualityDeleteTasksConcurrently( + b.Context(), loader, input.tasks[:benchmark.taskCount], 16) + if err != nil { + b.Fatal(err) + } + equalityDeleteLoadingBenchmarkSink = totalKeys + } + b.StopTimer() + b.ReportMetric(float64(input.deleteFiles), "delete_files/op") + b.ReportMetric(float64(input.fs.opens.Load())/float64(b.N), "opens/op") + }) + } +} + +func loadEqualityDeleteTasksConcurrently( + ctx context.Context, + loader *lazyEqualityDeleteLoader, + tasks []FileScanTask, + concurrency int, +) (int, error) { + if len(tasks) == 0 { + return 0, nil + } + + taskChan := make(chan FileScanTask, len(tasks)) + for _, task := range tasks { + taskChan <- task + } + close(taskChan) + totals := make(chan int, len(tasks)) + errs := make(chan error, 1) + var wg sync.WaitGroup + workerCount := min(concurrency, len(tasks)) + wg.Add(workerCount) + for range workerCount { + go func() { + defer wg.Done() + for task := range taskChan { + sets, err := loader.load(ctx, task) + if err != nil { + select { + case errs <- err: + default: + } + + return + } + + total := 0 + for _, deleteSet := range sets { + total += len(deleteSet.keys) + } + totals <- total + } + }() + } + wg.Wait() + close(totals) + + total := 0 + for count := range totals { + total += count + } + + select { + case err := <-errs: + return total, err + default: + return total, nil + } +} + +func newEqualityDeleteLoadingBenchmarkInput(b *testing.B) equalityDeleteLoadingBenchmarkInput { + b.Helper() + + const ( + deleteFileCount = 1_000 + taskCount = 10_000 + ) + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + firstPath := "mem://benchmark-lazy-equality/delete-0000.parquet" + writeEqualityDeleteParquetToMemFS(b, fs.MemFS, firstPath, + `[{"id": 0}, {"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}, {"id": 5}, {"id": 6}, {"id": 7}, {"id": 8}, {"id": 9}]`) + file, err := fs.MemFS.Open(firstPath) + if err != nil { + b.Fatal(err) + } + contents, err := io.ReadAll(file) + if closeErr := file.Close(); err == nil { + err = closeErr + } + if err != nil { + b.Fatal(err) + } + + deleteFiles := make([]iceberg.DataFile, deleteFileCount) + for i := range deleteFileCount { + path := fmt.Sprintf("mem://benchmark-lazy-equality/delete-%04d.parquet", i) + if i > 0 { + if err := fs.WriteFile(path, contents); err != nil { + b.Fatal(err) + } + } + deleteFiles[i] = newEqualityDeleteSetAssemblyTestFile(b, path, []int{1}) + } + + tasks := make([]FileScanTask, taskCount) + for i := range tasks { + tasks[i] = FileScanTask{ + EqualityDeleteFiles: []iceberg.DataFile{deleteFiles[i%deleteFileCount]}, + } + } + + return equalityDeleteLoadingBenchmarkInput{ + fs: fs, + tableSchema: tableSchema, + tasks: tasks, + deleteFiles: deleteFileCount, + } +} + func benchEqDeletes( b *testing.B, buildRec func(memory.Allocator, int) arrow.RecordBatch, diff --git a/table/equality_delete_reader_internal_test.go b/table/equality_delete_reader_internal_test.go index 838f179b1..2944d71a1 100644 --- a/table/equality_delete_reader_internal_test.go +++ b/table/equality_delete_reader_internal_test.go @@ -25,6 +25,8 @@ import ( "os" "path/filepath" "strings" + "sync" + "sync/atomic" "testing" "github.com/apache/arrow-go/v18/arrow" @@ -41,6 +43,22 @@ import ( "github.com/stretchr/testify/require" ) +type countingEqualityDeleteOpenFS struct { + *iceio.MemFS + attempts atomic.Int64 + opens atomic.Int64 +} + +func (f *countingEqualityDeleteOpenFS) Open(name string) (iceio.File, error) { + f.attempts.Add(1) + file, err := f.MemFS.Open(name) + if err == nil { + f.opens.Add(1) + } + + return file, err +} + func TestMakeColEncoderMatchesGenericForNullFastPathTypes(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0) @@ -239,6 +257,224 @@ func TestReadAllEqualityDeleteFilesRejectsEmptyEqualityFieldIDs(t *testing.T) { ) require.ErrorIs(t, err, ErrEmptyEqualityFieldIDs) require.ErrorContains(t, err, "empty-equality-fields.parquet") + + _, err = newLazyEqualityDeleteLoader( + iceio.NewMemFS(), schema, nil, + []FileScanTask{{EqualityDeleteFiles: []iceberg.DataFile{deleteFile}}}) + require.ErrorIs(t, err, ErrEmptyEqualityFieldIDs) + require.ErrorContains(t, err, "empty-equality-fields.parquet") +} + +func TestLazyEqualityDeleteLoaderLoadsFilesOnDemand(t *testing.T) { + t.Parallel() + + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + deleteAPath := "mem://lazy-equality/a.parquet" + deleteBPath := "mem://lazy-equality/b.parquet" + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, deleteAPath, `[{"id": 1}, {"id": 2}]`) + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, deleteBPath, `[{"id": 3}, {"id": 4}]`) + + deleteA := newEqualityDeleteSetAssemblyTestFile(t, deleteAPath, []int{1}) + deleteB := newEqualityDeleteSetAssemblyTestFile(t, deleteBPath, []int{1}) + tasks := []FileScanTask{ + {EqualityDeleteFiles: []iceberg.DataFile{deleteA}}, + {EqualityDeleteFiles: []iceberg.DataFile{deleteA, deleteB}}, + {EqualityDeleteFiles: []iceberg.DataFile{deleteB, deleteA}}, + } + + loader, err := newLazyEqualityDeleteLoader(fs, tableSchema, nil, tasks) + require.NoError(t, err) + assert.Equal(t, int64(0), fs.opens.Load()) + + projectedIDs := set[int]{} + loader.addFieldIDs(projectedIDs) + assert.Equal(t, set[int]{1: {}}, projectedIDs) + + first, err := loader.load(t.Context(), tasks[0]) + require.NoError(t, err) + require.Len(t, first, 1) + assert.Len(t, first[0].keys, 2) + assert.Equal(t, int64(1), fs.opens.Load()) + assert.Equal(t, int64(1), fs.attempts.Load()) + + combined, err := loader.load(t.Context(), tasks[1]) + require.NoError(t, err) + require.Len(t, combined, 1) + assert.Len(t, combined[0].keys, 4) + assert.Equal(t, int64(2), fs.opens.Load()) + assert.Equal(t, int64(2), fs.attempts.Load()) + + reversed, err := loader.load(t.Context(), tasks[2]) + require.NoError(t, err) + require.Len(t, reversed, 1) + assert.Same(t, combined[0], reversed[0]) + assert.Equal(t, int64(2), fs.opens.Load()) + assert.Equal(t, int64(2), fs.attempts.Load()) +} + +func TestLazyEqualityDeleteLoaderReadsSharedFileOnceConcurrently(t *testing.T) { + t.Parallel() + + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + deletePath := "mem://lazy-equality/concurrent.parquet" + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, deletePath, `[{"id": 1}]`) + deleteFile := newEqualityDeleteSetAssemblyTestFile(t, deletePath, []int{1}) + tasks := []FileScanTask{{EqualityDeleteFiles: []iceberg.DataFile{deleteFile}}} + + loader, err := newLazyEqualityDeleteLoader(fs, tableSchema, nil, tasks) + require.NoError(t, err) + + const workers = 16 + start := make(chan struct{}) + errs := make(chan error, workers) + var wg sync.WaitGroup + wg.Add(workers) + for range workers { + go func() { + defer wg.Done() + <-start + sets, err := loader.load(t.Context(), tasks[0]) + if err == nil && len(sets) != 1 { + err = fmt.Errorf("got %d equality delete sets, want 1", len(sets)) + } + errs <- err + }() + } + close(start) + wg.Wait() + close(errs) + + for err := range errs { + require.NoError(t, err) + } + assert.Equal(t, int64(1), fs.opens.Load()) + assert.Equal(t, int64(1), fs.attempts.Load()) +} + +func TestLazyEqualityDeleteLoaderCachesReadErrors(t *testing.T) { + t.Parallel() + + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + deleteFile := newEqualityDeleteSetAssemblyTestFile( + t, "mem://lazy-equality/missing.parquet", []int{1}) + tasks := []FileScanTask{{EqualityDeleteFiles: []iceberg.DataFile{deleteFile}}} + + loader, err := newLazyEqualityDeleteLoader(fs, tableSchema, nil, tasks) + require.NoError(t, err) + _, firstErr := loader.load(t.Context(), tasks[0]) + _, secondErr := loader.load(t.Context(), tasks[0]) + require.Error(t, firstErr) + require.Error(t, secondErr) + assert.Equal(t, int64(0), fs.opens.Load()) + assert.Equal(t, int64(1), fs.attempts.Load()) +} + +func TestLazyEqualityDeleteLoaderHonorsCanceledContext(t *testing.T) { + t.Parallel() + + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + deletePath := "mem://lazy-equality/canceled.parquet" + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, deletePath, `[{"id": 1}]`) + deleteFile := newEqualityDeleteSetAssemblyTestFile(t, deletePath, []int{1}) + tasks := []FileScanTask{{EqualityDeleteFiles: []iceberg.DataFile{deleteFile}}} + + loader, err := newLazyEqualityDeleteLoader(fs, tableSchema, nil, tasks) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = loader.load(ctx, tasks[0]) + require.ErrorIs(t, err, context.Canceled) +} + +func TestArrowScanDefersEqualityDeleteIOUntilIteration(t *testing.T) { + t.Parallel() + + tableSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + metadata, err := NewMetadata(tableSchema, iceberg.UnpartitionedSpec, + UnsortedSortOrder, "mem://lazy-equality/table", iceberg.Properties{PropertyFormatVersion: "2"}) + require.NoError(t, err) + tableSchema = metadata.CurrentSchema() + + fs := &countingEqualityDeleteOpenFS{MemFS: iceio.NewMemFS()} + dataPath := "mem://lazy-equality/data.parquet" + deletePath := "mem://lazy-equality/delete.parquet" + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, dataPath, `[{"id": 1}, {"id": 2}]`) + writeEqualityDeleteParquetToMemFS(t, fs.MemFS, deletePath, `[{"id": 2}]`) + + dataBuilder, err := iceberg.NewDataFileBuilder( + *iceberg.UnpartitionedSpec, iceberg.EntryContentData, + dataPath, iceberg.ParquetFile, nil, nil, nil, 2, 128) + require.NoError(t, err) + deleteFile := newEqualityDeleteSetAssemblyTestFile(t, deletePath, []int{1}) + + scanner := &arrowScan{ + fs: fs, + metadata: metadata, + scanSchema: tableSchema, + projectedSchema: tableSchema, + boundRowFilter: iceberg.AlwaysTrue{}, + filterSchema: tableSchema, + caseSensitive: true, + rowLimit: ScanNoLimit, + concurrency: 1, + } + _, records, err := scanner.GetRecords(t.Context(), []FileScanTask{{ + File: dataBuilder.Build(), + EqualityDeleteFiles: []iceberg.DataFile{deleteFile}, + }}) + require.NoError(t, err) + assert.Equal(t, int64(0), fs.opens.Load(), "unread iterators must not open equality deletes") + assert.Equal(t, int64(0), fs.attempts.Load(), "unread iterators must not attempt equality-delete I/O") + + var ids []int64 + for record, err := range records { + require.NoError(t, err) + column := record.Column(0).(*array.Int64) + for i := range column.Len() { + ids = append(ids, column.Value(i)) + } + record.Release() + } + + assert.Equal(t, []int64{1}, ids) + assert.Equal(t, int64(2), fs.opens.Load(), "iteration should open the data and equality-delete files") + assert.Equal(t, int64(2), fs.attempts.Load()) +} + +func writeEqualityDeleteParquetToMemFS(t testing.TB, fs *iceio.MemFS, path, content string) { + t.Helper() + + iceSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + arrowSchema, err := SchemaToArrowSchema(iceSchema, nil, true, false) + require.NoError(t, err) + record := mustLoadRecordBatchFromJSON(arrowSchema, content) + defer record.Release() + + tbl := array.NewTableFromRecords(arrowSchema, []arrow.RecordBatch{record}) + defer tbl.Release() + + file, err := fs.Create(path) + require.NoError(t, err) + require.NoError(t, pqarrow.WriteTable(tbl, file, record.NumRows(), + parquet.NewWriterProperties(parquet.WithStats(true)), pqarrow.DefaultWriterProps())) + require.NoError(t, file.Close()) } func TestReadEqualityDeleteFileMatchesMaterializedRead(t *testing.T) { @@ -777,7 +1013,7 @@ func TestBuildEqualityDeleteSetsPerTaskSkipsMissingDeleteFiles(t *testing.T) { } func newEqualityDeleteSetAssemblyTestFile( - t *testing.T, + t testing.TB, path string, fieldIDs []int, ) iceberg.DataFile { diff --git a/table/equality_delete_reader_test.go b/table/equality_delete_reader_test.go index ee772c2a9..34e9f647b 100644 --- a/table/equality_delete_reader_test.go +++ b/table/equality_delete_reader_test.go @@ -465,8 +465,21 @@ func TestEqualityDeleteReadRejectsAmbiguousColumns(t *testing.T) { tbl, err = tx.Commit(t.Context()) require.NoError(t, err) - _, _, err = tbl.Scan().ToArrowRecords(t.Context()) - require.ErrorIs(t, err, table.ErrAmbiguousEqualityColumn) + _, records, err := tbl.Scan().ToArrowRecords(t.Context()) + require.NoError(t, err) + + var scanErr error + for record, err := range records { + if record != nil { + record.Release() + } + if err != nil { + scanErr = err + + break + } + } + require.ErrorIs(t, scanErr, table.ErrAmbiguousEqualityColumn) } func TestEqualityDeleteMatchingAcrossPartitionSpecEvolution(t *testing.T) { diff --git a/table/internal/utils.go b/table/internal/utils.go index 5fb7e6164..743985db8 100644 --- a/table/internal/utils.go +++ b/table/internal/utils.go @@ -83,11 +83,27 @@ func (pq *pqueue[T]) Pop() any { // based on the comesAfter and isNext functions. The values are read in from // the provided source and then re-ordered before being sent to the output. func MakeSequencedChan[T any](bufferSize uint, source <-chan T, comesAfter, isNext func(a, b *T) bool, initial T) <-chan T { + return MakeSequencedChanWithDiscard(bufferSize, source, comesAfter, isNext, initial, nil) +} + +// MakeSequencedChanWithDiscard is MakeSequencedChan with a callback for values +// that remain in the reorder queue when the source is closed. Values already +// sent to the output are not passed to discard. +func MakeSequencedChanWithDiscard[T any](bufferSize uint, source <-chan T, comesAfter, isNext func(a, b *T) bool, initial T, discard func(T)) <-chan T { pq := pqueue[T]{queue: make([]*T, 0), compare: comesAfter} heap.Init(&pq) previous, out := &initial, make(chan T, bufferSize) go func() { defer close(out) + defer func() { + if discard == nil { + return + } + + for _, val := range pq.queue { + discard(*val) + } + }() for val := range source { heap.Push(&pq, &val) for pq.Len() > 0 && isNext(previous, pq.queue[0]) {