Skip to content

[codex] Add persistent GPU exact-key lookup stage - #9

Draft
VibhuJawa wants to merge 3 commits into
feat/interleaved-lance-materializationfrom
feat/gpu-exact-key-lookup
Draft

[codex] Add persistent GPU exact-key lookup stage#9
VibhuJawa wants to merge 3 commits into
feat/interleaved-lance-materializationfrom
feat/gpu-exact-key-lookup

Conversation

@VibhuJawa

@VibhuJawa VibhuJawa commented Jul 3, 2026

Copy link
Copy Markdown
Owner

Summary

  • add a public GpuExactKeyLookupStage for exact membership enrichment of InterleavedBatch rows
  • build persistent RAPIDS 26.06 FilteredJoin objects once per GPU actor
  • coalesce multiple Curator tasks into one GPU probe while preserving task boundaries, row order, metadata, and passthrough columns
  • support segmented Parquet reference sets, optional storage options, strict type/count validation, and GPU/setup/probe metrics
  • add a complete MINT-1T HTML GPU presence tutorial, including a public Lance-to-Parquet reference builder and runnable Curator pipeline

Why

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:

  1. A row-wise InterleavedBatch task containing input_key_column. In MINT-1T HTML, image rows contain an exact URL in source_ref; text and metadata rows contain null.
  2. One or more immutable Parquet files containing a non-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 InterleavedBatch with exactly the same rows, order, task metadata, and unrelated columns, plus one nullable boolean presence_column:

Input key Presence output
null or empty null; no lookup
exact key in any reference segment True
non-empty key absent from every segment False

Duplicate 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 manifest
  • main.py: runs InterleavedParquetReader -> GpuExactKeyLookupStage -> InterleavedParquetWriterStage
  • explicit normalized MINT input and presence-enriched output schemas
  • manifest-driven segment ordering, byte-size checks, and expected-row validation before Ray starts
  • one-actor-per-GPU scaling guidance and a decision table for GPU membership versus LanceColumnFetchStage

Raw MINT parallel images/texts arrays 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

  • reference segments remain independent to avoid a second full-size string-column allocation during setup
  • process_batch() performs one coalesced probe and restores original task boundaries
  • regular and large Arrow UTF-8 key columns are compatible
  • cuDF is imported lazily
  • the GPU path requires cudf-cu12>=26.6,<26.7; Curator's current deduplication_cuda12 extra pins an older RAPIDS release and should not be installed in the same environment yet

Validation

  • pytest -q tests/stages/interleaved/test_gpu_exact_key_lookup.py tests/stages/interleaved/test_lance_column_fetch.py — 9 passed
  • Ruff checks and git diff --check passed
  • both tutorial CLIs load successfully
  • the reference builder and manifest reader were exercised end-to-end against a real local Lance table with large-string URLs and multiple output segments
  • H100 integration with 355,952,746 exact reference URLs:
    • 47.1 GB persistent GPU state per actor
    • about 37 GB device memory remained free
    • 8.43 seconds to load/build one actor
    • 1.25 million exact probes in 8.8–9.1 ms once warm
    • one real 61,632-row MINT partition in 14.5 ms
    • bit-for-bit identical presence output to the CPU Lance implementation
  • eight H100 actors processed all 4,197 normalized MINT partitions (about 142 GiB) in 7 minutes 50 seconds with no task errors

Stack

This PR targets feat/interleaved-lance-materialization, which contains the generic Lance column-fetch and interleaved Lance reader APIs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant