Add partitioned exact filter stage

[VibhuJawa]

Description

Adds a small GPU stage for exact bucket-local left-semi and left-anti filtering after two datasets have been partitioned by compatible ShuffleStage runs.

The stage reads only the matching reference partition, preserves every matching or non-matching left row and its multiplicity, and avoids broadcasting a large reference set. It validates matching partition metadata, rejects row-subsetting reads and multi-file writes that would violate the complete-partition contract, and emits deterministic part.N.parquet outputs.

Callers must use the same key fields, hash behavior, and partition count for both upstream shuffles. ShuffleStage emits every partition, including empty partitions, so a missing matching reference file is treated as an incomplete shuffle and fails closed.

Usage

pipeline.add_stage(
    PartitionedExactFilterStage(
        key_fields=["url_id"],
        reference_path="/data/downloaded-shuffle",
        output_path="/data/cleaned",
        total_partitions=65536,
        mode="leftanti",
    )
)

Use mode="leftsemi" to emit the exact overlap instead.

Validation

• 33 focused CPU tests passed.
• Two true-cuDF tests cover duplicate reference keys and left multiplicity for both modes; they are registered for normal GPU CI.
• GPU coverage configuration passes for all 44 GPU test files.
• Public import remains lazy and does not import cuDF on CPU-only code paths.

This stage is intentionally separate from persistent membership filtering in interleaved batches: it reduces compatible disk-shuffled FileGroupTask partitions.

Checklist

• I am familiar with the Contributing Guide.
• New tests cover these changes.
• The stage API documentation is up to date.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds PartitionedExactFilterStage, a GPU stage that performs a cuDF left-semi or left-anti join between a shuffled left partition and its matching reference partition (by part.N.parquet index), avoiding any broadcast of the full reference set.

• The stage validates partition metadata, blocks row-subsetting read kwargs and multi-file write kwargs, pre-checks the reference schema with PyArrow before issuing the cuDF read, and uses os.path.realpath-based path identity to guard against output overwriting input — all three concerns raised in the previous review round are addressed in the current implementation.
• All 33 CPU tests pass, two true-cuDF GPU tests cover duplicate-key and multiplicity invariants for both modes, and the public import remains lazy so cuDF is never loaded on CPU-only code paths.

Confidence Score: 5/5

Safe to merge — the new stage is self-contained, all three concerns from the previous review round have been addressed, and the test suite covers both the happy path and all major failure modes.

The implementation is clean and defensive: the reference schema is pre-validated via PyArrow before any cuDF read, path identity uses os.path.realpath to catch symlink aliases, and construction-time validation blocks every category of invalid input. No correctness or data-integrity issues were found on the changed code paths.

No files require special attention.

Important Files Changed

Filename	Overview
nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py	New stage performing GPU left-semi/anti join on matching shuffled partitions; well-validated construction, correct schema pre-check before cuDF read, and robust path-identity overwrite guard.
nemo_curator/stages/deduplication/shuffle_utils/init.py	Minimal package init that re-exports PartitionedExactFilterStage; import is lazy and does not trigger cuDF on CPU paths.
tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py	33 CPU tests cover construction validation, partition metadata checks, overwrite guards (including symlink aliases), and a monkeypatched integration path; two GPU tests cover duplicate reference keys and left-row multiplicity for both modes.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Stage as PartitionedExactFilterStage
    participant RefFS as reference_fs (fsspec)
    participant PyArrow as pq.read_schema
    participant cuDF

    Caller->>Stage: process(task: FileGroupTask)
    Stage->>Stage: _partition_index(task) — validates metadata
    Stage->>Stage: _output_partition_path(idx)
    Stage->>Stage: overwrite guard (_path_identity check)
    Stage->>Stage: _reference_partition_path(idx)
    Stage->>RefFS: isfile(reference_file)
    RefFS-->>Stage: True / raise FileNotFoundError
    Stage->>RefFS: open(reference_file, "rb")
    RefFS-->>Stage: stream
    Stage->>PyArrow: read_schema(stream)
    PyArrow-->>Stage: reference_schema
    Stage->>Stage: _require_column_names(schema.names, "reference")
    Stage->>cuDF: "read_parquet(task.data, **read_kwargs)"
    cuDF-->>Stage: left DataFrame (all columns)
    Stage->>cuDF: "read_parquet(reference_file, columns=key_fields, **reference_read_kwargs)"
    cuDF-->>Stage: reference DataFrame (key columns only)
    Stage->>Stage: _require_key_fields(left, "left")
    Stage->>cuDF: "left.merge(reference, how=mode, on=key_fields)"
    cuDF-->>Stage: output DataFrame
    Stage->>Stage: "len(output) > len(left) guard"
    Stage->>cuDF: "output.to_parquet(output_file, index=False, **write_kwargs)"
    Stage->>Stage: _log_metrics(...)
    Stage-->>Caller: FileGroupTask (output_file, updated metadata)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant Stage as PartitionedExactFilterStage
    participant RefFS as reference_fs (fsspec)
    participant PyArrow as pq.read_schema
    participant cuDF

    Caller->>Stage: process(task: FileGroupTask)
    Stage->>Stage: _partition_index(task) — validates metadata
    Stage->>Stage: _output_partition_path(idx)
    Stage->>Stage: overwrite guard (_path_identity check)
    Stage->>Stage: _reference_partition_path(idx)
    Stage->>RefFS: isfile(reference_file)
    RefFS-->>Stage: True / raise FileNotFoundError
    Stage->>RefFS: open(reference_file, "rb")
    RefFS-->>Stage: stream
    Stage->>PyArrow: read_schema(stream)
    PyArrow-->>Stage: reference_schema
    Stage->>Stage: _require_column_names(schema.names, "reference")
    Stage->>cuDF: "read_parquet(task.data, **read_kwargs)"
    cuDF-->>Stage: left DataFrame (all columns)
    Stage->>cuDF: "read_parquet(reference_file, columns=key_fields, **reference_read_kwargs)"
    cuDF-->>Stage: reference DataFrame (key columns only)
    Stage->>Stage: _require_key_fields(left, "left")
    Stage->>cuDF: "left.merge(reference, how=mode, on=key_fields)"
    cuDF-->>Stage: output DataFrame
    Stage->>Stage: "len(output) > len(left) guard"
    Stage->>cuDF: "output.to_parquet(output_file, index=False, **write_kwargs)"
    Stage->>Stage: _log_metrics(...)
    Stage-->>Caller: FileGroupTask (output_file, updated metadata)

Loading

_{Reviews (2): Last reviewed commit: "Harden partitioned exact filter validati..." | Re-trigger Greptile}

[greptile-apps]

_require_key_fields(reference, ...) is unreachable in the error case

Because cudf.read_parquet is called with columns=self.key_fields, cuDF will raise its own error (e.g. RuntimeError) if any requested column is absent from the Parquet schema — before execution ever reaches _require_key_fields(reference, "reference"). When the read succeeds, the returned frame has exactly self.key_fields as its columns, so the check trivially passes. The user therefore always gets a raw cuDF error instead of the more descriptive "reference partition is missing exact key columns" message. Moving the schema check to before the read (or reading without the columns constraint and selecting afterwards) would let the custom error fire correctly.

[greptile-apps]

"index" is semantically misplaced in _NON_SINGLE_FILE_WRITE_KWARGS

partition_cols, metadata_file_path, partition_offsets, and partition_file_name are genuinely multi-file/partitioning kwargs. "index" is unrelated to that — it controls whether the DataFrame index is written as a column. Grouping it here means the docstring claim ("Dataset-partitioning and auxiliary-metadata output options are not accepted") does not describe index. A caller blocked by this list gets a confusing error, and the list name would mislead future maintainers. Consider separating it into its own constant (e.g. _CONTROLLED_WRITE_KWARGS = ["index"]) or at least adding a comment.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Overwrite guard uses naive string equality on paths

output_file in task.data is an exact string comparison. If task.data paths come from the shuffler with a different representation than self.output_path — for example, with a double slash, a resolved symlink, or a differently-cased URI — the check silently passes and the stage overwrites its own input. Since ShuffleStage derives its output paths from an external C++ library, path normalisation is not guaranteed to be consistent with rstrip("/"). Using os.path.realpath (or fsspec's _strip_protocol + normpath) on both sides before comparing would make the guard robust.