Skip to content
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 163 additions & 23 deletions table/arrow_scanner.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ func releasePerFilePosDeletes(deletesPerFile perFilePosDeletes) {
}
}

// readAllDeleteFiles is retained for the eager-path benchmark and regression
// tests; scans use lazyPositionDeleteLoader instead.
func readAllDeleteFiles(ctx context.Context, fs iceio.IO, tasks []FileScanTask, concurrency int) (perFilePosDeletes, error) {
deletesPerFile := make(perFilePosDeletes)
uniqueDeletes := make(map[string]iceberg.DataFile)
Expand Down Expand Up @@ -138,6 +140,125 @@ 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.
// The loader has the lifetime of exactly one scan. Each file's first load also
// locks in its result, including context errors, for every caller; a loader
// must not be reused for a retry with a different context.
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() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

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
cached.err = fmt.Errorf("read position deletes from %s: %w",
cached.dataFile.FilePath(), cached.err)
}
})
if cached.err != nil {
return nil, cached.err

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

}

if chunk := cached.deletes[targetPath]; chunk != nil {
deletes = append(deletes, chunk)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor — Lazy loader narrows positional-delete scope from global to task-scoped, changing scan results

The eager path built perFilePosDeletes keyed by every data-file path found inside each delete file (arrow_scanner.go:134-138), so a delete file referenced by one task also applied to any other task's data file it happened to mention. load() instead consults only the delete files listed on the task itself and looks up cached.deletes[targetPath] (:252-254). The new semantics match Java's task-scoped DeleteFileIndex and are, I believe, the correct ones - but this is a result-set change in a PR whose description says only error timing moves, and no test pins it. Planner-built tasks should not hit it; Scan.ReadTasks accepts caller-supplied tasks and can. Suggest a regression test asserting a delete file is applied only to tasks that reference it, plus a line in the PR description. Flagging for your judgement rather than blocking.

}

return deletes, nil
}

func (l *lazyPositionDeleteLoader) release() {
if l == nil {
return
}

l.releaseOnce.Do(func() {
for _, cached := range l.files {
releasePosDeletes(cached.deletes)
cached.deletes = nil
}
})
}

// perFileDVBitmaps maps each data-file path to the deletion-vector bitmap
// that applies to it. Kept separate from perFilePosDeletes so the row-filter
// pipeline can use compute.Filter on a Boolean mask built from Contains()
Expand Down Expand Up @@ -1956,6 +2077,10 @@ func (as *arrowScan) producePosDeletesFromTask(ctx context.Context, task tblutil
}

func createIterator(ctx context.Context, numWorkers uint, records <-chan enumeratedRecord, deletesPerFile perFilePosDeletes, cancel context.CancelCauseFunc, rowLimit int64) iter.Seq2[arrow.RecordBatch, error] {
return createIteratorWithCleanup(ctx, numWorkers, records, deletesPerFile, cancel, rowLimit, nil)
}

func createIteratorWithCleanup(ctx context.Context, numWorkers uint, records <-chan enumeratedRecord, deletesPerFile perFilePosDeletes, cancel context.CancelCauseFunc, rowLimit int64, cleanup func()) iter.Seq2[arrow.RecordBatch, error] {
isBeforeAny := func(batch enumeratedRecord) bool {
return batch.Task.Index < 0
}
Expand Down Expand Up @@ -2003,6 +2128,9 @@ func createIterator(ctx context.Context, numWorkers uint, records <-chan enumera
}

releasePerFilePosDeletes(deletesPerFile)
if cleanup != nil {
cleanup()
}
}()

defer cancel(nil)
Expand Down Expand Up @@ -2057,15 +2185,12 @@ func createIterator(ctx context.Context, numWorkers uint, records <-chan enumera
}
}

func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks []FileScanTask, deletesPerFile perFilePosDeletes, dvBitmaps perFileDVBitmaps, equalityDeleteLoader *lazyEqualityDeleteLoader, invariants *arrowScanInvariants) iter.Seq2[arrow.RecordBatch, error] {
func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks []FileScanTask, positionDeleteLoader *lazyPositionDeleteLoader, 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))
taskChan := make(chan tblutils.Enumerated[FileScanTask], len(tasks))
records := make(chan enumeratedRecord, numWorkers)

var wg sync.WaitGroup
Expand All @@ -2086,16 +2211,34 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
}

filePath := task.Value.File.FilePath()
var positionalDeletes positionDeletes
if positionDeleteLoader != nil {
var err error
positionalDeletes, err = positionDeleteLoader.load(scanCtx, task.Value)
if err != nil {
select {
case records <- enumeratedRecord{Task: task, Err: err}:
case <-scanCtx.Done():
}
cancel(err)

return
}
}

eqDeleteSets, err := equalityDeleteLoader.load(scanCtx, task.Value)
if err != nil {
records <- enumeratedRecord{Task: task, Err: err}
select {
case records <- enumeratedRecord{Task: task, Err: err}:
case <-scanCtx.Done():
}
cancel(err)

return
}

if err := as.recordsFromTask(scanCtx, task, records,
deletesPerFile[filePath],
positionalDeletes,
dvBitmaps[filePath],
eqDeleteSets,
invariants); err != nil {
Expand Down Expand Up @@ -2126,11 +2269,19 @@ func (as *arrowScan) recordBatchesFromTasksAndDeletes(ctx context.Context, tasks
}
}()

createIterator(scanCtx, uint(numWorkers), records, deletesPerFile,
cancel, as.rowLimit)(yield)
var cleanup func()
if positionDeleteLoader != nil {
cleanup = positionDeleteLoader.release
}
createIteratorWithCleanup(scanCtx, uint(numWorkers), records, nil,
cancel, as.rowLimit, cleanup)(yield)
}
}

// GetRecords prepares the projected Arrow schema and a single-use record
// iterator. Positional- and equality-delete files are opened and read during
// iteration, so errors from those files are returned by the iterator;
// deletion-vector errors are returned before the iterator is created.
func (as *arrowScan) GetRecords(ctx context.Context, tasks []FileScanTask) (*arrow.Schema, iter.Seq2[arrow.RecordBatch, error], error) {
var err error
as.useLargeTypes, err = strconv.ParseBool(as.options.Get(ScanOptionArrowUseLargeTypes, "false"))
Expand Down Expand Up @@ -2168,36 +2319,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
}

equalityDeleteLoader, err := newLazyEqualityDeleteLoader(
as.fs, invariants.tableSchema, invariants.nameMapping, tasks)
if err != nil {
releasePerFilePosDeletes(deletesPerFile)

return nil, nil, err
}
equalityDeleteLoader.addFieldIDs(invariants.projectedIDs)

positionDeleteLoader := newLazyPositionDeleteLoader(as.fs, tasks)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


return resultSchema, as.recordBatchesFromTasksAndDeletes(ctx, tasks,
deletesPerFile, dvBitmaps, equalityDeleteLoader, invariants), nil
positionDeleteLoader, dvBitmaps, equalityDeleteLoader, invariants), nil
}
Loading
Loading