Add partition-local exact semi/anti filter stage

[VibhuJawa]

Summary

Adds a focused PartitionedExactFilterStage for exact bucket-local exclusion or membership after two datasets have been partitioned with the existing ShuffleStage.

• Supports cuDF leftanti and leftsemi joins on one or more exact key columns.
• Matches each left shuffle task to the reference partition with the same partition_index.
• Preserves left-row multiplicity even when reference keys are duplicated.
• Validates complete-partition metadata and schemas before reading data.
• Rejects row-subsetting reads, multi-file writes, destructive path aliases, and missing reference partitions.
• Publishes per-partition input, reference, and output row counts in task metadata.

Why this is the missing primitive

Current Curator already provides the compatible hash shuffle. Campaign-specific grouping can be composed from that shuffle plus workload-specific reducers, but Curator does not currently provide the reusable bucket-local exact semi/anti operation needed for scalable exclusion without broadcasting the complete reference set. This PR deliberately adds only that missing reducer and no scheduler, checkpoint, or workload-specific control plane.

Composition

left ShuffleStage(key_fields, N) ---------+
                                           +--> PartitionedExactFilterStage(mode="leftanti" | "leftsemi")
reference ShuffleStage(key_fields, N) ----+

Both shuffles must use identical keys, hashing, and partition count.

Validation

• 37 focused CPU tests passed; 2 GPU semantics cases are marked for GPU CI.
• git diff --check passed.
• DCO, Ruff, pre-commit, and secrets checks pass on the current head.
• The branch contains three signed, DCO-compliant commits on exact NVIDIA main commit 6fa52e3e597206849acb2ac119413fa6cbfe644d; the latter two are typing-only lint repairs.

[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 new bucket-local reduction primitive that performs a cuDF left-semi or left-anti join between a shuffled left partition and its matching reference partition (by partition_index), enabling scalable exact exclusion/membership filtering without broadcasting the full reference set.

• The core process() method validates partition metadata, guards against output-overwrites and symlink aliases, schema-checks the reference partition cheaply via PyArrow before any GPU I/O, performs the join, and publishes per-partition row count metrics in the output task metadata.
• A __init__.py exports the new class from the shuffle_utils subpackage, and 37 focused CPU tests (with 2 GPU-semantics tests marked for GPU CI) cover constructor validation, metadata enforcement, overwrite guards, schema errors, and _stage_perf propagation.

Confidence Score: 4/5

The new stage is well-guarded and test coverage is thorough; the findings are quality-of-life improvements rather than correctness issues.

The join logic, partition-index matching, overwrite guards, and schema validation are all correct. The left partition schema check runs after the full GPU load (unlike the cheaper pre-load reference check), and the reference file is opened twice for remote filesystems — both are efficiency gaps rather than data-correctness bugs. The path-stripping on Windows is theoretical for this GPU-only workload.

The core implementation file partitioned_exact_filter.py deserves a second look around the ordering of schema validation and the double remote file-open for the reference partition.

Important Files Changed

Filename	Overview
nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py	New stage implementing bucket-local semi/anti join via cuDF merge; correctness looks solid but the left partition schema is validated after GPU load (unlike the reference schema which is checked via PyArrow first), and the reference file is opened twice for remote filesystems.
nemo_curator/stages/deduplication/shuffle_utils/init.py	Exports PartitionedExactFilterStage from the new submodule; straightforward and correct.
tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py	37 focused CPU tests cover constructor validation, metadata checks, overwrite guards, schema errors, and stage-perf propagation; GPU semantics tests (@pytest.mark.gpu) are correctly marked for GPU CI.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Stage as PartitionedExactFilterStage
    participant RefFS as reference_fs (fsspec)
    participant CuDF as cudf
    participant OutFS as output_fs (fsspec)

    Caller->>Stage: process(task)
    Stage->>Stage: _partition_index(task) -- validate metadata
    Stage->>Stage: _output_partition_path(idx) -- compute output path
    Stage->>Stage: "check output != left input (overwrite guard)"
    Stage->>Stage: _reference_partition_path(idx) -- compute ref path
    Stage->>RefFS: isfile(reference_file) -- existence check
    Stage->>RefFS: open(reference_file, rb) -- schema check
    RefFS-->>Stage: stream
    Stage->>Stage: pq.read_schema(stream) -- validate ref key columns
    Stage->>CuDF: read_parquet(task.data) -- load left partition
    CuDF-->>Stage: left DataFrame
    Stage->>CuDF: "read_parquet(reference_file, columns=key_fields)"
    CuDF-->>Stage: reference DataFrame
    Stage->>Stage: _require_key_fields(left) -- validate left schema
    Stage->>Stage: "left.merge(reference, how=mode) -- semi/anti join"
    Stage->>Stage: "guard len(output) <= len(left)"
    Stage->>OutFS: "output.to_parquet(output_file, index=False)"
    Stage->>Stage: _log_metrics(row counts)
    Stage-->>Caller: FileGroupTask(output_file, enriched 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 CuDF as cudf
    participant OutFS as output_fs (fsspec)

    Caller->>Stage: process(task)
    Stage->>Stage: _partition_index(task) -- validate metadata
    Stage->>Stage: _output_partition_path(idx) -- compute output path
    Stage->>Stage: "check output != left input (overwrite guard)"
    Stage->>Stage: _reference_partition_path(idx) -- compute ref path
    Stage->>RefFS: isfile(reference_file) -- existence check
    Stage->>RefFS: open(reference_file, rb) -- schema check
    RefFS-->>Stage: stream
    Stage->>Stage: pq.read_schema(stream) -- validate ref key columns
    Stage->>CuDF: read_parquet(task.data) -- load left partition
    CuDF-->>Stage: left DataFrame
    Stage->>CuDF: "read_parquet(reference_file, columns=key_fields)"
    CuDF-->>Stage: reference DataFrame
    Stage->>Stage: _require_key_fields(left) -- validate left schema
    Stage->>Stage: "left.merge(reference, how=mode) -- semi/anti join"
    Stage->>Stage: "guard len(output) <= len(left)"
    Stage->>OutFS: "output.to_parquet(output_file, index=False)"
    Stage->>Stage: _log_metrics(row counts)
    Stage-->>Caller: FileGroupTask(output_file, enriched metadata)

Loading

_{Reviews (1): Last reviewed commit: "Keep typing-only container import lazy" | Re-trigger Greptile}

[greptile-apps]

Left schema validated after full GPU load

The reference partition schema is validated via PyArrow before any cuDF read (lines 202–204), so a missing key column on the reference side is caught cheaply. The left partition, however, is fully loaded into GPU memory on line 205 before its schema is checked on line 211. For large left partitions that are missing a key column, all the data lands on the GPU before the ValueError is raised, wasting GPU memory and time. Pre-checking the left schema via pq.read_schema() on task.data[0] (or iterating task.data) before calling cudf.read_parquet would mirror the reference-side guard and catch the error cheaply.

[greptile-apps]

Reference file opened twice for remote filesystems

reference_file is opened once via self.reference_fs.open(reference_file, "rb") for the PyArrow schema check (line 202), then opened again by cudf.read_parquet(reference_file, ...) (line 206). For remote filesystems (S3, GCS, HDFS) this causes two separate I/O round-trips per partition. The schema could be extracted from the already-read cuDF DataFrame after the cudf.read_parquet call (reference.dtypes), or the file could be downloaded once into a local buffer that both PyArrow and cuDF consume, to avoid the duplicate remote I/O.

[greptile-apps]

rstrip("/") may not strip Windows-style separators on non-POSIX paths

reference_path.rstrip("/") only strips forward slashes. If a caller passes a Windows-style path like C:\output\ or a mixed-separator path, the trailing backslash is not removed. The emptiness check on line 130 and the _path_identity comparison on line 133 would still work for the comparison, but self.reference_path and self.output_path would retain trailing backslashes, causing malformed file paths in _reference_partition_path and _output_partition_path. This is mostly academic for GPU workloads (Linux), but using path.rstrip("/\\") or os.path.normpath would be more robust.