laskoviymishka commented on code in PR #1938:
URL: https://github.com/apache/iceberg-go/pull/1938#discussion_r3892542507
##########
table/arrow_scanner.go:
##########
@@ -135,6 +135,120 @@ func readAllDeleteFiles(ctx context.Context, fs iceio.IO,
tasks []FileScanTask,
return deletesPerFile, nil
}
+// lazyPositionDeleteLoader indexes positional-delete metadata for a scan, but
+// waits to open each delete file until a worker reaches a task that references
+// it. A delete file can apply to more than one data file, so the cache keeps
+// the complete grouped result for the delete file rather than caching only one
+// task's positions.
+//
+// The grouped Arrow chunks are owned by the loader until release. The iterator
+// calls release after all workers have stopped, which keeps shared chunks
alive
+// while multiple tasks use them and also covers early iterator termination.
+type lazyPositionDeleteLoader struct {
+ fs iceio.IO
+ files map[string]*lazyPositionDeleteFile
+
+ releaseOnce sync.Once
+}
+
+type lazyPositionDeleteFile struct {
+ dataFile iceberg.DataFile
+
+ once sync.Once
+ deletes map[string]*arrow.Chunked
+ err error
+}
+
+func newLazyPositionDeleteLoader(fs iceio.IO, tasks []FileScanTask)
*lazyPositionDeleteLoader {
+ loader := &lazyPositionDeleteLoader{
+ fs: fs,
+ files: make(map[string]*lazyPositionDeleteFile),
+ }
+
+ for _, task := range tasks {
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() !=
iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if _, ok := loader.files[path]; !ok {
+ loader.files[path] =
&lazyPositionDeleteFile{dataFile: deleteFile}
+ }
+ }
+ }
+
+ return loader
+}
+
+func (l *lazyPositionDeleteLoader) load(ctx context.Context, task
FileScanTask) (positionDeletes, error) {
+ if len(task.DeleteFiles) == 0 {
+ return nil, nil
+ }
+
+ targetPath := task.File.FilePath()
+ deletes := make(positionDeletes, 0, len(task.DeleteFiles))
+ // Most scan tasks carry one positional delete file. Avoid allocating a
+ // deduplication map unless there can actually be duplicate entries.
+ var seen map[string]struct{}
+ if len(task.DeleteFiles) > 1 {
+ seen = make(map[string]struct{}, len(task.DeleteFiles))
+ }
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() != iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if seen != nil {
+ if _, ok := seen[path]; ok {
+ continue
+ }
+ seen[path] = struct{}{}
+ }
+
+ cached, ok := l.files[path]
+ if !ok {
+ // The loader is normally built from the same task
slice supplied to
+ // this method. Keep this guard so a malformed caller
cannot panic a
+ // scan if it changes a task after loader construction.
+ continue
+ }
+
+ cached.once.Do(func() {
Review Comment:
This caches whatever `readDeletes` returns, including `context.Canceled`,
for the life of the loader, so a later `load()` with a fresh context still gets
the stale error. That's fine today because the loader is built per `GetRecords`
and every worker shares `scanCtx`, but nothing in the type says so.
I'd add a sentence to `lazyPositionDeleteFile` (or `load`) spelling out that
the loader lives for exactly one scan and that any error, including transient
context errors, is locked in for all callers regardless of their own context.
If we ever reuse a loader across retries the `once.Do` would need to become
cancellation-aware, and I'd rather that be written down before someone hits it.
wdyt?
##########
table/arrow_scanner.go:
##########
@@ -135,6 +135,120 @@ func readAllDeleteFiles(ctx context.Context, fs iceio.IO,
tasks []FileScanTask,
return deletesPerFile, nil
}
+// lazyPositionDeleteLoader indexes positional-delete metadata for a scan, but
+// waits to open each delete file until a worker reaches a task that references
+// it. A delete file can apply to more than one data file, so the cache keeps
+// the complete grouped result for the delete file rather than caching only one
+// task's positions.
+//
+// The grouped Arrow chunks are owned by the loader until release. The iterator
+// calls release after all workers have stopped, which keeps shared chunks
alive
+// while multiple tasks use them and also covers early iterator termination.
+type lazyPositionDeleteLoader struct {
+ fs iceio.IO
+ files map[string]*lazyPositionDeleteFile
+
+ releaseOnce sync.Once
+}
+
+type lazyPositionDeleteFile struct {
+ dataFile iceberg.DataFile
+
+ once sync.Once
+ deletes map[string]*arrow.Chunked
+ err error
+}
+
+func newLazyPositionDeleteLoader(fs iceio.IO, tasks []FileScanTask)
*lazyPositionDeleteLoader {
+ loader := &lazyPositionDeleteLoader{
+ fs: fs,
+ files: make(map[string]*lazyPositionDeleteFile),
+ }
+
+ for _, task := range tasks {
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() !=
iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if _, ok := loader.files[path]; !ok {
+ loader.files[path] =
&lazyPositionDeleteFile{dataFile: deleteFile}
+ }
+ }
+ }
+
+ return loader
+}
+
+func (l *lazyPositionDeleteLoader) load(ctx context.Context, task
FileScanTask) (positionDeletes, error) {
+ if len(task.DeleteFiles) == 0 {
+ return nil, nil
+ }
+
+ targetPath := task.File.FilePath()
+ deletes := make(positionDeletes, 0, len(task.DeleteFiles))
+ // Most scan tasks carry one positional delete file. Avoid allocating a
+ // deduplication map unless there can actually be duplicate entries.
+ var seen map[string]struct{}
+ if len(task.DeleteFiles) > 1 {
+ seen = make(map[string]struct{}, len(task.DeleteFiles))
+ }
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() != iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if seen != nil {
+ if _, ok := seen[path]; ok {
+ continue
+ }
+ seen[path] = struct{}{}
+ }
+
+ cached, ok := l.files[path]
+ if !ok {
+ // The loader is normally built from the same task
slice supplied to
+ // this method. Keep this guard so a malformed caller
cannot panic a
+ // scan if it changes a task after loader construction.
+ continue
+ }
+
+ cached.once.Do(func() {
+ cached.deletes, cached.err = readDeletes(ctx, l.fs,
cached.dataFile)
+ if cached.err != nil {
+ // readDeletes currently returns nil on errors.
Release defensively
+ // in case a future reader returns partial
Arrow ownership.
+ releasePosDeletes(cached.deletes)
+ cached.deletes = nil
+ }
+ })
+ if cached.err != nil {
+ return nil, cached.err
Review Comment:
When `readDeletes` fails we return the raw error without the delete-file
path, so the caller sees something like `file not found` with no clue which
file. `readAllDeletionVectors` already wraps with the puffin path; I'd match it
here, something like `fmt.Errorf("read position deletes from %s: %w",
cached.dataFile.FilePath(), cached.err)` inside the `once.Do` so the path
travels with the cached error.
##########
table/arrow_scanner.go:
##########
@@ -1799,59 +1924,83 @@ 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, positionDeleteLoader *lazyPositionDeleteLoader, dvBitmaps
perFileDVBitmaps, eqDeleteSets map[int][]*equalityDeleteSet, 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))
+ numWorkers := min(as.concurrency, len(tasks))
+ taskChan := make(chan tblutils.Enumerated[FileScanTask],
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()
+ var positionalDeletes
positionDeletes
+ if positionDeleteLoader != nil {
+ var err error
+ positionalDeletes, err
= positionDeleteLoader.load(scanCtx, task.Value)
+ if err != nil {
+ records <-
enumeratedRecord{Task: task, Err: err}
Review Comment:
This send is unconditional while the receive side can stall.
`MakeSequencedChanWithDiscard` drains `records`, but if `sequenced` fills first
(buffer is `numWorkers`) the helper blocks on `out <- *previous` until the
consumer reads, which delays it reading `records`, which delays this error
send, which delays this worker returning, which delays `wg.Wait()` and the
`records` close.
It's bounded backpressure rather than a deadlock, but under a slow consumer
and high concurrency it stretches shutdown, and that's exactly what the 500ms
deadline in `TestArrowScanPreCancelledIteratorTearsDownProducer` is up against,
so it's a plausible CI flake. I'd make it a `select { case records <- ...: case
<-scanCtx.Done(): return }`, matching the feeder. wdyt?
##########
table/arrow_scanner.go:
##########
@@ -1884,36 +2033,25 @@ func (as *arrowScan) GetRecords(ctx context.Context,
tasks []FileScanTask) (*arr
return nil, nil, err
}
- deletesPerFile, err := readAllDeleteFiles(ctx, as.fs, tasks,
as.concurrency)
- if err != nil {
- // readAllDeleteFiles can return a partially-populated map
alongside
- // the error if some goroutines completed before the failure.
- releasePerFilePosDeletes(deletesPerFile)
-
- return nil, nil, err
- }
-
// DV bitmaps stay in their native form rather than being materialized
// into int64 positions and merged with the Parquet pos-delete map.
// filterByDeletionVector applies the bitmap to each batch via a Boolean
// keep-mask + compute.Filter — O(1) Contains lookups, vectorized
Filter,
// no intermediate position set.
dvBitmaps, err := readAllDeletionVectors(ctx, as.fs, tasks,
as.concurrency)
if err != nil {
- releasePerFilePosDeletes(deletesPerFile)
-
return nil, nil, err
}
eqDeleteSets, err := readAllEqualityDeleteFiles(ctx, as.fs,
invariants.tableSchema, invariants.nameMapping, tasks,
as.concurrency)
if err != nil {
- // Positional deletes were fully loaded; release them before
aborting.
- releasePerFilePosDeletes(deletesPerFile)
-
return nil, nil, err
}
addEqualityDeleteFieldIDs(invariants, eqDeleteSets)
- return resultSchema, as.recordBatchesFromTasksAndDeletes(ctx, tasks,
deletesPerFile, dvBitmaps, eqDeleteSets, invariants), nil
+ positionDeleteLoader := newLazyPositionDeleteLoader(as.fs, tasks)
Review Comment:
With the lazy loader, an unreadable positional delete file no longer fails
`GetRecords`; the error now surfaces mid-iteration as `enumeratedRecord.Err`. A
caller doing `schema, iter, err := GetRecords(...); if err != nil { return }`
and then consuming will miss delete-file errors unless they also check the
per-item error.
DV and equality-delete errors still surface eagerly from `GetRecords`, so
the two now behave differently. I'd add a godoc line on `GetRecords` noting
that positional-delete read errors are delivered through the iterator while DV
and equality errors surface before it returns, and probably a CHANGELOG note
since it's a caller-visible behavioral change.
`TestArrowScanDefersPositionDeleteReadsUntilIteration` already pins the new
behavior, so this is just documenting it.
##########
table/arrow_scanner.go:
##########
@@ -135,6 +135,120 @@ func readAllDeleteFiles(ctx context.Context, fs iceio.IO,
tasks []FileScanTask,
return deletesPerFile, nil
}
+// lazyPositionDeleteLoader indexes positional-delete metadata for a scan, but
+// waits to open each delete file until a worker reaches a task that references
+// it. A delete file can apply to more than one data file, so the cache keeps
+// the complete grouped result for the delete file rather than caching only one
+// task's positions.
+//
+// The grouped Arrow chunks are owned by the loader until release. The iterator
+// calls release after all workers have stopped, which keeps shared chunks
alive
+// while multiple tasks use them and also covers early iterator termination.
+type lazyPositionDeleteLoader struct {
+ fs iceio.IO
+ files map[string]*lazyPositionDeleteFile
+
+ releaseOnce sync.Once
+}
+
+type lazyPositionDeleteFile struct {
+ dataFile iceberg.DataFile
+
+ once sync.Once
+ deletes map[string]*arrow.Chunked
+ err error
+}
+
+func newLazyPositionDeleteLoader(fs iceio.IO, tasks []FileScanTask)
*lazyPositionDeleteLoader {
+ loader := &lazyPositionDeleteLoader{
+ fs: fs,
+ files: make(map[string]*lazyPositionDeleteFile),
+ }
+
+ for _, task := range tasks {
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() !=
iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if _, ok := loader.files[path]; !ok {
+ loader.files[path] =
&lazyPositionDeleteFile{dataFile: deleteFile}
+ }
+ }
+ }
+
+ return loader
+}
+
+func (l *lazyPositionDeleteLoader) load(ctx context.Context, task
FileScanTask) (positionDeletes, error) {
+ if len(task.DeleteFiles) == 0 {
+ return nil, nil
+ }
+
+ targetPath := task.File.FilePath()
+ deletes := make(positionDeletes, 0, len(task.DeleteFiles))
+ // Most scan tasks carry one positional delete file. Avoid allocating a
+ // deduplication map unless there can actually be duplicate entries.
+ var seen map[string]struct{}
+ if len(task.DeleteFiles) > 1 {
+ seen = make(map[string]struct{}, len(task.DeleteFiles))
+ }
+ for _, deleteFile := range task.DeleteFiles {
+ if deleteFile.ContentType() != iceberg.EntryContentPosDeletes {
+ continue
+ }
+
+ path := deleteFile.FilePath()
+ if seen != nil {
+ if _, ok := seen[path]; ok {
+ continue
+ }
+ seen[path] = struct{}{}
+ }
+
+ cached, ok := l.files[path]
+ if !ok {
+ // The loader is normally built from the same task
slice supplied to
+ // this method. Keep this guard so a malformed caller
cannot panic a
+ // scan if it changes a task after loader construction.
+ continue
+ }
+
+ cached.once.Do(func() {
+ cached.deletes, cached.err = readDeletes(ctx, l.fs,
cached.dataFile)
+ if cached.err != nil {
+ // readDeletes currently returns nil on errors.
Release defensively
+ // in case a future reader returns partial
Arrow ownership.
+ releasePosDeletes(cached.deletes)
+ cached.deletes = nil
+ }
+ })
+ if cached.err != nil {
+ return nil, cached.err
+ }
+
+ if chunk := cached.deletes[targetPath]; chunk != nil {
+ deletes = append(deletes, chunk)
Review Comment:
These chunks are borrowed from the loader without a `Retain()`, so multiple
workers hold the same `*arrow.Chunked` while the loader still owns it. It's
safe only because `release()` runs strictly after `wg.Wait()`, so nothing reads
a freed chunk, but that ordering is the entire thing keeping it correct and it
isn't visible from the types.
`release()` also nils `cached.deletes`, so a second range over the returned
`iter.Seq2` would hit the done `once`, read `nil`, and silently yield zero
positional deletes. I'd keep the nil-write (dropping it turns a second range
into a use-after-free on released chunks, which is worse) and add a one-line
comment that the iterator is single-use. If we'd rather not lean on the
ordering invariant at all, `Retain()` on append plus `Release()` after
`collectPosDeletePositions` makes it self-contained. Either way, I'd make it
explicit.
##########
table/arrow_scanner_lazy_delete_bench_test.go:
##########
@@ -0,0 +1,165 @@
+// 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 (
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/pqarrow"
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+)
+
+const (
+ lazyPositionDeleteBenchmarkTaskCount = 10_000
+ lazyPositionDeleteBenchmarkFileCount = 1_000
+ lazyPositionDeleteBenchmarkTasksPerFile =
lazyPositionDeleteBenchmarkTaskCount / lazyPositionDeleteBenchmarkFileCount
+ lazyPositionDeleteBenchmarkDataFileBytes = 128
+)
+
+type lazyPositionDeleteBenchmarkFixture struct {
+ fs *iceio.MemFS
+ tasks []FileScanTask
+}
+
+func newLazyPositionDeleteBenchmarkFixture(b *testing.B)
lazyPositionDeleteBenchmarkFixture {
+ b.Helper()
+
+ fs := iceio.NewMemFS()
+ deleteFiles := make([]iceberg.DataFile,
lazyPositionDeleteBenchmarkFileCount)
+ for deleteIndex := range deleteFiles {
+ deletePath :=
fmt.Sprintf("mem://benchmark/deletes/delete-%04d.parquet", deleteIndex)
+ var content strings.Builder
+ content.WriteByte('[')
+ for taskOffset := range lazyPositionDeleteBenchmarkTasksPerFile
{
+ if taskOffset > 0 {
+ content.WriteByte(',')
+ }
+ taskIndex := deleteIndex +
taskOffset*lazyPositionDeleteBenchmarkFileCount
+ fmt.Fprintf(&content,
+
`{"file_path":"mem://benchmark/data/data-%05d.parquet","pos":0}`,
+ taskIndex)
+ }
+ content.WriteByte(']')
+
+ benchmarkWritePosDeleteParquet(b, fs, deletePath,
content.String())
+ builder, err := iceberg.NewDataFileBuilder(
+ *iceberg.UnpartitionedSpec,
iceberg.EntryContentPosDeletes,
+ deletePath, iceberg.ParquetFile, nil, nil, nil,
+ lazyPositionDeleteBenchmarkTasksPerFile,
lazyPositionDeleteBenchmarkDataFileBytes)
+ if err != nil {
+ b.Fatal(err)
+ }
+ deleteFiles[deleteIndex] = builder.Build()
+ }
+
+ tasks := make([]FileScanTask, lazyPositionDeleteBenchmarkTaskCount)
+ for taskIndex := range tasks {
+ dataPath :=
fmt.Sprintf("mem://benchmark/data/data-%05d.parquet", taskIndex)
+ builder, err := iceberg.NewDataFileBuilder(
+ *iceberg.UnpartitionedSpec, iceberg.EntryContentData,
+ dataPath, iceberg.ParquetFile, nil, nil, nil, 1,
lazyPositionDeleteBenchmarkDataFileBytes)
+ if err != nil {
+ b.Fatal(err)
+ }
+ tasks[taskIndex] = FileScanTask{
+ File: builder.Build(),
+ DeleteFiles:
[]iceberg.DataFile{deleteFiles[taskIndex%lazyPositionDeleteBenchmarkFileCount]},
+ }
+ }
+
+ return lazyPositionDeleteBenchmarkFixture{fs: fs, tasks: tasks}
+}
+
+func benchmarkWritePosDeleteParquet(b *testing.B, fs *iceio.MemFS, path,
content string) {
+ b.Helper()
+
+ record := mustLoadRecordBatchFromJSON(PositionalDeleteArrowSchema,
content)
+ defer record.Release()
+ tbl := array.NewTableFromRecords(PositionalDeleteArrowSchema,
[]arrow.RecordBatch{record})
+ defer tbl.Release()
+
+ file, err := fs.Create(path)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if err := pqarrow.WriteTable(tbl, file, record.NumRows(),
+ parquet.NewWriterProperties(parquet.WithStats(true)),
+ pqarrow.DefaultWriterProps()); err != nil {
+ b.Fatal(err)
+ }
+ if err := file.Close(); err != nil {
+ b.Fatal(err)
+ }
+}
+
+func BenchmarkLazyPositionDeleteLoading(b *testing.B) {
+ fixture := newLazyPositionDeleteBenchmarkFixture(b)
+
+ b.Run("eager_all_delete_files", func(b *testing.B) {
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ deletes, err := readAllDeleteFiles(b.Context(),
fixture.fs, fixture.tasks, 16)
Review Comment:
After this change `readAllDeleteFiles` is only reachable from this
benchmark; the production path goes through the lazy loader now. A short
"retained for benchmarking the eager path" comment on the function would keep
someone from deleting it as dead code or wiring it back into a scan by mistake.
##########
table/arrow_scanner_lazy_delete_regression_test.go:
##########
@@ -0,0 +1,380 @@
+// 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"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/apache/iceberg-go/table/internal"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+type countingOpenMemFS struct {
+ *iceio.MemFS
+ opens atomic.Int64
+}
+
+func (f *countingOpenMemFS) Open(name string) (iceio.File, error) {
+ f.opens.Add(1)
+
+ return f.MemFS.Open(name)
+}
+
+func newLazyDataFile(t *testing.T, path string) iceberg.DataFile {
+ t.Helper()
+
+ builder, err := iceberg.NewDataFileBuilder(
+ *iceberg.UnpartitionedSpec, iceberg.EntryContentData,
+ path, iceberg.ParquetFile, nil, nil, nil, 1, 128)
+ require.NoError(t, err)
+
+ return builder.Build()
+}
+
+func TestLazyPositionDeleteLoaderDefersReadsAndSharesResults(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ const (
+ deletePath = "mem://bucket/deletes/shared.parquet"
+ dataPathA = "mem://bucket/data/a.parquet"
+ dataPathB = "mem://bucket/data/b.parquet"
+ )
+ memFS := &countingOpenMemFS{MemFS: iceio.NewMemFS()}
+ writePosDeleteParquetToMemFS(t, memFS.MemFS, deletePath, `[
+ {"file_path": "`+dataPathA+`", "pos": 1},
+ {"file_path": "`+dataPathB+`", "pos": 3}
+ ]`)
+
+ deleteFile := newPosDeleteFile(t, deletePath, 2, 128)
+ tasks := []FileScanTask{
+ {
+ File: newLazyDataFile(t, dataPathA),
+ DeleteFiles: []iceberg.DataFile{deleteFile, deleteFile},
+ },
+ {
+ File: newLazyDataFile(t, dataPathB),
+ DeleteFiles: []iceberg.DataFile{deleteFile},
+ },
+ }
+ loader := newLazyPositionDeleteLoader(memFS, tasks)
+
+ assert.Zero(t, memFS.opens.Load(), "constructing the scan loader must
not open delete files")
+
+ gotA, err := loader.load(ctx, tasks[0])
+ require.NoError(t, err)
+ require.Len(t, gotA, 1, "duplicate delete references must be read once
per task")
+ assert.Equal(t, []int64{1}, int64Values(gotA[0]))
+ assert.Equal(t, int64(1), memFS.opens.Load())
+
+ gotB, err := loader.load(ctx, tasks[1])
+ require.NoError(t, err)
+ require.Len(t, gotB, 1)
+ assert.Equal(t, []int64{3}, int64Values(gotB[0]))
+ assert.Equal(t, int64(1), memFS.opens.Load(), "shared delete files must
use one read")
+
+ loader.release()
+}
+
+func TestLazyPositionDeleteLoaderSingleflightsConcurrentLoads(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ const (
+ deletePath = "mem://bucket/deletes/concurrent.parquet"
+ dataPath = "mem://bucket/data/a.parquet"
+ )
+ fs := &countingOpenMemFS{MemFS: iceio.NewMemFS()}
+ writePosDeleteParquetToMemFS(t, fs.MemFS, deletePath,
+ `[{"file_path":"`+dataPath+`","pos":7}]`)
+ deleteFile := newPosDeleteFile(t, deletePath, 1, 128)
+ task := FileScanTask{
+ File: newLazyDataFile(t, dataPath),
+ DeleteFiles: []iceberg.DataFile{deleteFile},
+ }
+ loader := newLazyPositionDeleteLoader(fs, []FileScanTask{task})
+
+ const callers = 8
+ results := make(chan positionDeletes, callers)
+ errs := make(chan error, callers)
+ var wg sync.WaitGroup
+ wg.Add(callers)
+ for range callers {
+ go func() {
+ defer wg.Done()
+ deletes, err := loader.load(ctx, task)
+ results <- deletes
+ errs <- err
+ }()
+ }
+ wg.Wait()
+ close(results)
+ close(errs)
+
+ for err := range errs {
+ require.NoError(t, err)
+ }
+ for deletes := range results {
+ require.Len(t, deletes, 1)
+ assert.Equal(t, []int64{7}, int64Values(deletes[0]))
+ }
+ assert.Equal(t, int64(1), fs.opens.Load(), "concurrent users must share
one delete-file read")
+
+ loader.release()
+}
+
+func TestLazyPositionDeleteLoaderCachesErrors(t *testing.T) {
+ fs := &countingOpenMemFS{MemFS: iceio.NewMemFS()}
+ deleteFile := newPosDeleteFile(t,
"mem://bucket/deletes/missing.parquet", 1, 128)
+ task := FileScanTask{
+ File: newLazyDataFile(t, "mem://bucket/data/a.parquet"),
+ DeleteFiles: []iceberg.DataFile{deleteFile},
+ }
+ loader := newLazyPositionDeleteLoader(fs, []FileScanTask{task})
+
+ first, err := loader.load(context.Background(), task)
+ require.Error(t, err)
+ assert.Nil(t, first)
+ assert.Equal(t, int64(1), fs.opens.Load())
+
+ second, secondErr := loader.load(context.Background(), task)
+ require.Error(t, secondErr)
+ assert.Nil(t, second)
+ assert.ErrorIs(t, secondErr, err)
+ assert.Equal(t, int64(1), fs.opens.Load(), "a failed delete file must
not be retried by other tasks")
+
+ loader.release()
+}
+
+func TestLazyPositionDeleteLoaderCachesCancellation(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx, cancel := context.WithCancel(compute.WithAllocator(t.Context(),
mem))
+ cancel()
+
+ fs := &countingOpenMemFS{MemFS: iceio.NewMemFS()}
+ deletePath := "mem://bucket/deletes/cancelled.parquet"
+ writePosDeleteParquetToMemFS(t, fs.MemFS, deletePath,
`[{"file_path":"mem://bucket/data/a.parquet","pos":1}]`)
+ deleteFile := newPosDeleteFile(t, deletePath, 1, 128)
+ task := FileScanTask{
+ File: newLazyDataFile(t, "mem://bucket/data/a.parquet"),
+ DeleteFiles: []iceberg.DataFile{deleteFile},
+ }
+ loader := newLazyPositionDeleteLoader(fs, []FileScanTask{task})
+
+ _, firstErr := loader.load(ctx, task)
+ require.ErrorIs(t, firstErr, context.Canceled)
+ _, secondErr := loader.load(context.Background(), task)
+ require.ErrorIs(t, secondErr, context.Canceled)
+ assert.Equal(t, int64(1), fs.opens.Load(), "cancellation must not cause
a second read")
+
+ loader.release()
+}
+
+func TestArrowScanDefersPositionDeleteReadsUntilIteration(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(t.Context(), mem)
+
+ schema := iceberg.NewSchema(1, iceberg.NestedField{
+ ID: 1, Name: "value", Type: iceberg.PrimitiveTypes.Int64,
+ })
+ metadata, err := NewMetadata(schema, iceberg.UnpartitionedSpec,
+ UnsortedSortOrder, "mem://bucket/table", nil)
+ require.NoError(t, err)
+
+ const deletePath = "mem://bucket/deletes/one.parquet"
+ memFS := &countingOpenMemFS{MemFS: iceio.NewMemFS()}
+ writePosDeleteParquetToMemFS(t, memFS.MemFS, deletePath,
+ `[{"file_path":"mem://bucket/data/missing.parquet","pos":1}]`)
+ deleteFile := newPosDeleteFile(t, deletePath, 1, 128)
+ task := FileScanTask{
+ File: newLazyDataFile(t,
"mem://bucket/data/missing.parquet"),
+ DeleteFiles: []iceberg.DataFile{deleteFile},
+ }
+ scan := &arrowScan{
+ metadata: metadata,
+ fs: memFS,
+ scanSchema: schema,
+ projectedSchema: schema,
+ boundRowFilter: iceberg.AlwaysTrue{},
+ rowLimit: -1,
+ concurrency: 1,
+ }
+
+ _, records, err := scan.GetRecords(ctx, []FileScanTask{task})
+ require.NoError(t, err)
+ assert.Zero(t, memFS.opens.Load(), "GetRecords must not read position
deletes")
+
+ var iterErr error
+ for record, err := range records {
+ if record != nil {
+ record.Release()
+ }
+ iterErr = err
+
+ break
+ }
+ require.Error(t, iterErr, "iteration should reach the missing data file
after loading its delete")
+ assert.Equal(t, int64(2), memFS.opens.Load(), "the first task should
open one delete and one data file")
+}
+
+func TestLazyPositionDeleteLoaderReleasesChunksWhenIteratorStops(t *testing.T)
{
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ loader := &lazyPositionDeleteLoader{files:
map[string]*lazyPositionDeleteFile{
+ "mem://bucket/deletes/one.parquet": {
+ deletes: map[string]*arrow.Chunked{
+ "mem://bucket/data/a.parquet":
chunkedPosDelete(t, mem, []int64{1}),
+ },
+ },
+ }}
+
+ batch := checkedInt64RecordBatch(mem, 1)
+ records := make(chan enumeratedRecord, 1)
+ records <- enumeratedRecord{
+ Record: internal.Enumerated[arrow.RecordBatch]{
+ Value: batch,
+ Index: 0,
+ Last: true,
+ },
+ Task: internal.Enumerated[FileScanTask]{Index: 0, Last: true},
+ }
+ close(records)
+
+ ctx, cancel := context.WithCancelCause(context.Background())
+ itr := createIteratorWithCleanup(ctx, 1, records, nil, cancel, 0,
loader.release)
+ for record, err := range itr {
+ require.NoError(t, err)
+ record.Release()
+
+ break
+ }
+}
+
+func TestArrowScanPreCancelledIteratorTearsDownProducer(t *testing.T) {
+ ctx, cancel := context.WithCancelCause(context.Background())
+ cancel(context.Canceled)
+
+ scan := &arrowScan{concurrency: 1, rowLimit: -1}
+ tasks := []FileScanTask{{
+ File: newLazyDataFile(t,
"mem://bucket/data/pre-cancelled.parquet"),
+ }}
+ records := scan.recordBatchesFromTasksAndDeletes(ctx, tasks, nil, nil,
nil, nil)
+
+ done := make(chan error, 1)
+ go func() {
+ var iterErr error
+ for _, err := range records {
+ if err != nil {
+ iterErr = err
+ }
+ }
+ done <- iterErr
+ }()
+
+ select {
+ case err := <-done:
+ require.ErrorIs(t, err, context.Canceled)
+ case <-time.After(500 * time.Millisecond):
+ t.Fatal("pre-cancelled scan iterator did not terminate")
+ }
+}
+
+func TestCreateIteratorReleasesOutOfOrderBatchAfterError(t *testing.T) {
Review Comment:
This proves the discard callback releases a queued out-of-order batch, but
it does it with a hand-built channel and no live workers, so it doesn't
exercise the actual scan path that produced the leak concern in the earlier
review.
None of the new tests run `arrowScan.GetRecords` with `concurrency >= 2`,
several tasks sharing a delete file, an early consumer `break`, and a
`CheckedAllocator` asserting `AssertSize(t, 0)` at the end. That end-to-end
combination is the one that was actually broken, and it's the thing I'd most
want locked down before merge: a real `arrowScan` with concurrency 4, ~8 tasks
over a few shared delete files, real Parquet, consume one batch and break, then
assert zero bytes outstanding. A failing-delete-file variant with concurrency >
1 would be a nice bonus.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
