[codex] Add persistent GPU exact-key lookup stage - #9
Draft
VibhuJawa wants to merge 3 commits into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
GpuExactKeyLookupStagefor exact membership enrichment ofInterleavedBatchrowsFilteredJoinobjects once per GPU actorWhy
Lance scalar-index scans return matching row IDs but do not expose the matched input key through the public Python API. Presence-only workflows therefore had to reread the key column from remote Lance data to associate row IDs with exact URLs. On the MINT image table, those random remote reads dominated runtime even after moving the B-tree metadata and index segments to local storage.
This stage uses an immutable exact-key Parquet sidecar and a persistent GPU hash set. It is intentionally separate from
LanceColumnFetchStage: the GPU stage answers bulk membership only, while the Lance stage remains responsible for stable-row-ID payload projection.Input and output
The stage has two inputs:
InterleavedBatchtask containinginput_key_column. In MINT-1T HTML, image rows contain an exact URL insource_ref; text and metadata rows contain null.reference_key_column. Their union is the exact reference set. The tutorial builds these files from a required, pinned Lance version using only Lance's public scanner API.It returns an
InterleavedBatchwith exactly the same rows, order, task metadata, and unrelated columns, plus one nullable booleanpresence_column:TrueFalseDuplicate input keys preserve all input rows. Existing destination columns fail instead of being overwritten. The stage does not fetch image bytes or other Lance columns, mutate Lance, persist row IDs, or implement Lance B-tree traversal on the GPU.
MINT-1T tutorial
The new tutorial provides:
build_reference.py: streams only the exact URL column from a pinned Lance version into segmented Zstd Parquet, rejects null/non-string keys, and writes a deterministic manifestmain.py: runsInterleavedParquetReader -> GpuExactKeyLookupStage -> InterleavedParquetWriterStageLanceColumnFetchStageRaw MINT parallel
images/textsarrays are normalized once before this pipeline. The tutorial begins at the persisted row-wise interleaved boundary, making the expensive presence pass independently resumable and repeatable.Behavior and dependencies
process_batch()performs one coalesced probe and restores original task boundariescudf-cu12>=26.6,<26.7; Curator's currentdeduplication_cuda12extra pins an older RAPIDS release and should not be installed in the same environment yetValidation
pytest -q tests/stages/interleaved/test_gpu_exact_key_lookup.py tests/stages/interleaved/test_lance_column_fetch.py— 9 passedgit diff --checkpassedStack
This PR targets
feat/interleaved-lance-materialization, which contains the generic Lance column-fetch and interleaved Lance reader APIs.