From 538b0dccbf0971ff576640b675e3303319569b90 Mon Sep 17 00:00:00 2001 From: Vibhu Jawa Date: Mon, 13 Jul 2026 23:18:40 -0700 Subject: [PATCH 1/2] Add partitioned exact filter stage Signed-off-by: Vibhu Jawa --- .../deduplication/shuffle_utils/__init__.py | 17 ++ .../shuffle_utils/partitioned_exact_filter.py | 226 ++++++++++++++++++ .../test_partitioned_exact_filter.py | 226 ++++++++++++++++++ 3 files changed, 469 insertions(+) create mode 100644 nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py create mode 100644 tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py diff --git a/nemo_curator/stages/deduplication/shuffle_utils/__init__.py b/nemo_curator/stages/deduplication/shuffle_utils/__init__.py index e69de29bb2..4b5796f531 100644 --- a/nemo_curator/stages/deduplication/shuffle_utils/__init__.py +++ b/nemo_curator/stages/deduplication/shuffle_utils/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .partitioned_exact_filter import PartitionedExactFilterStage + +__all__ = ["PartitionedExactFilterStage"] diff --git a/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py b/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py new file mode 100644 index 0000000000..cff90da5c6 --- /dev/null +++ b/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.stages.resources import Resources +from nemo_curator.tasks import FileGroupTask +from nemo_curator.utils.file_utils import check_disallowed_kwargs, get_fs + +if TYPE_CHECKING: + import cudf + + +ExactFilterMode = Literal["leftsemi", "leftanti"] + +_ROW_SUBSETTING_READ_KWARGS = ["columns", "filters", "row_groups", "nrows", "skip_rows"] +_NON_SINGLE_FILE_WRITE_KWARGS = [ + "index", + "metadata_file_path", + "partition_cols", + "partition_file_name", + "partition_offsets", +] + + +class PartitionedExactFilterStage(ProcessingStage[FileGroupTask, FileGroupTask]): + """Filter complete shuffled partitions against matching reference partitions. + + Both inputs must have been partitioned with the same key columns, hash + function, and partition count. Each input task is matched to + ``part..parquet`` below ``reference_path``. The stage + performs a cuDF left-semi or left-anti join and therefore preserves the + multiplicity of rows on the left; duplicate reference keys cannot multiply + the output. + + This stage deliberately does not perform a shuffle. It is the bucket-local + reduction step to compose after two compatible :class:`ShuffleStage` + operations. + + Parameters + ---------- + key_fields + Exact key column or columns shared by both shuffled datasets. + reference_path + Directory containing reference files named + ``part..parquet``. + output_path + Directory for the filtered partition files. + total_partitions + Expected partition count for both shuffled datasets. + mode + ``"leftsemi"`` keeps matching left rows and ``"leftanti"`` keeps + non-matching left rows. + read_kwargs + Keyword arguments for reading complete left-side Parquet files with + cuDF. Row projection, filtering, and row-subsetting options are not + accepted. + reference_read_kwargs + Keyword arguments for reading complete reference Parquet files with + cuDF. Row projection, filtering, and row-subsetting options are not + accepted. + write_kwargs + Keyword arguments for writing one result Parquet file with cuDF. + Dataset-partitioning and auxiliary-metadata output options are not + accepted. + """ + + name = "PartitionedExactFilter" + resources = Resources(gpus=1.0) + + def __init__( # noqa: PLR0913 + self, + key_fields: str | list[str], + reference_path: str, + output_path: str, + total_partitions: int, + mode: ExactFilterMode = "leftanti", + read_kwargs: dict[str, Any] | None = None, + reference_read_kwargs: dict[str, Any] | None = None, + write_kwargs: dict[str, Any] | None = None, + ) -> None: + super().__init__() + fields = [key_fields] if isinstance(key_fields, str) else list(key_fields) + if not fields or any(not isinstance(field, str) or not field for field in fields): + msg = "key_fields must contain at least one non-empty column name" + raise ValueError(msg) + if len(set(fields)) != len(fields): + msg = "key_fields must not contain duplicate column names" + raise ValueError(msg) + if mode not in ("leftsemi", "leftanti"): + msg = "mode must be 'leftsemi' or 'leftanti'" + raise ValueError(msg) + if isinstance(total_partitions, bool) or not isinstance(total_partitions, int) or total_partitions <= 0: + msg = "total_partitions must be a positive integer" + raise ValueError(msg) + + reference_root = reference_path.rstrip("/") + output_root = output_path.rstrip("/") + if not reference_root or not output_root: + msg = "reference_path and output_path must be non-root directory paths" + raise ValueError(msg) + if reference_root == output_root: + msg = "output_path must be distinct from reference_path" + raise ValueError(msg) + + self.key_fields = fields + self.reference_path = reference_root + self.output_path = output_root + self.total_partitions = total_partitions + self.mode = mode + self.read_kwargs = dict(read_kwargs or {}) + self.reference_read_kwargs = dict(reference_read_kwargs or {}) + self.write_kwargs = dict(write_kwargs or {}) + + check_disallowed_kwargs(self.read_kwargs, _ROW_SUBSETTING_READ_KWARGS) + check_disallowed_kwargs(self.reference_read_kwargs, _ROW_SUBSETTING_READ_KWARGS) + check_disallowed_kwargs(self.write_kwargs, _NON_SINGLE_FILE_WRITE_KWARGS) + + self.reference_fs = get_fs( + self.reference_path, + storage_options=self.reference_read_kwargs.get("storage_options", {}), + ) + self.output_fs = get_fs(self.output_path, storage_options=self.write_kwargs.get("storage_options", {})) + self.output_fs.mkdirs(self.output_path, exist_ok=True) + + def inputs(self) -> tuple[list[str], list[str]]: + return (["data", "_metadata"], []) + + def outputs(self) -> tuple[list[str], list[str]]: + return (["data", "_metadata"], []) + + def _partition_index(self, task: FileGroupTask) -> int: + partition_index = task._metadata.get("partition_index") + task_total = task._metadata.get("total_partitions") + if isinstance(partition_index, bool) or not isinstance(partition_index, int): + msg = "input task metadata must contain an integer partition_index" + raise TypeError(msg) + if partition_index < 0 or partition_index >= self.total_partitions: + msg = f"partition_index {partition_index} is outside [0, {self.total_partitions})" + raise ValueError(msg) + if task_total != self.total_partitions: + msg = ( + f"input task total_partitions {task_total!r} does not match the configured " + f"value {self.total_partitions}" + ) + raise ValueError(msg) + return partition_index + + def _reference_partition_path(self, partition_index: int) -> str: + return self.reference_fs.sep.join([self.reference_path, f"part.{partition_index}.parquet"]) + + def _output_partition_path(self, partition_index: int) -> str: + return self.output_fs.sep.join([self.output_path, f"part.{partition_index}.parquet"]) + + def process(self, task: FileGroupTask) -> FileGroupTask: + partition_index = self._partition_index(task) + output_file = self._output_partition_path(partition_index) + if output_file in task.data: + msg = "output partition would overwrite its left input partition" + raise ValueError(msg) + reference_file = self._reference_partition_path(partition_index) + if not self.reference_fs.isfile(reference_file): + msg = f"matching reference partition is missing: {reference_file}" + raise FileNotFoundError(msg) + + import cudf + + left = cudf.read_parquet(task.data, **self.read_kwargs) + reference = cudf.read_parquet( + reference_file, + columns=self.key_fields, + **self.reference_read_kwargs, + ) + self._require_key_fields(left, "left") + self._require_key_fields(reference, "reference") + + reference_rows = len(reference) + output = left.merge(reference, how=self.mode, on=self.key_fields) + if len(output) > len(left): + msg = "exact semi/anti join unexpectedly increased left-row multiplicity" + raise RuntimeError(msg) + + output.to_parquet(output_file, index=False, **self.write_kwargs) + + self._log_metrics( + { + "input_rows": len(left), + "reference_rows": reference_rows, + "output_rows": len(output), + } + ) + return FileGroupTask( + dataset_name=f"{task.dataset_name}_{self.name}", + data=[output_file], + _metadata={ + **task._metadata, + "partition_index": partition_index, + "total_partitions": self.total_partitions, + "key_fields": self.key_fields, + "filter_mode": self.mode, + "input_rows": len(left), + "reference_rows": reference_rows, + "output_rows": len(output), + }, + _stage_perf=task._stage_perf, + ) + + def _require_key_fields(self, frame: cudf.DataFrame, side: str) -> None: + missing = [field for field in self.key_fields if field not in frame.columns] + if missing: + msg = f"{side} partition is missing exact key columns: {missing}" + raise ValueError(msg) diff --git a/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py b/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py new file mode 100644 index 0000000000..2e1f6ef2d5 --- /dev/null +++ b/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar + +import pytest + +from nemo_curator.stages.deduplication.shuffle_utils.partitioned_exact_filter import ( + PartitionedExactFilterStage, +) +from nemo_curator.tasks import FileGroupTask + + +def make_stage(tmp_path: Path, **kwargs) -> PartitionedExactFilterStage: + reference_path = tmp_path / "reference" + reference_path.mkdir(exist_ok=True) + return PartitionedExactFilterStage( + key_fields=kwargs.pop("key_fields", "key"), + reference_path=str(reference_path), + output_path=str(tmp_path / "output"), + total_partitions=kwargs.pop("total_partitions", 2), + **kwargs, + ) + + +@pytest.mark.parametrize("key_fields", [[], "", ["key", "key"], ["key", ""]]) +def test_rejects_invalid_key_fields(tmp_path: Path, key_fields: str | list[str]) -> None: + with pytest.raises(ValueError, match="key_fields"): + make_stage(tmp_path, key_fields=key_fields) + + +@pytest.mark.parametrize("total_partitions", [0, -1, True, 1.5]) +def test_rejects_invalid_partition_count(tmp_path: Path, total_partitions: int) -> None: + with pytest.raises(ValueError, match="total_partitions"): + make_stage(tmp_path, total_partitions=total_partitions) + + +@pytest.mark.parametrize("kwarg", ["columns", "filters", "row_groups", "nrows", "skip_rows"]) +@pytest.mark.parametrize("side", ["left", "reference"]) +def test_rejects_row_subsetting_read_kwargs(tmp_path: Path, kwarg: str, side: str) -> None: + kwargs = {"read_kwargs" if side == "left" else "reference_read_kwargs": {kwarg: object()}} + with pytest.raises(ValueError, match=kwarg): + make_stage(tmp_path, **kwargs) + + +@pytest.mark.parametrize( + "kwarg", + ["index", "metadata_file_path", "partition_cols", "partition_file_name", "partition_offsets"], +) +def test_rejects_non_single_file_write_kwargs(tmp_path: Path, kwarg: str) -> None: + with pytest.raises(ValueError, match=kwarg): + make_stage(tmp_path, write_kwargs={kwarg: object()}) + + +def test_rejects_invalid_filter_mode(tmp_path: Path) -> None: + with pytest.raises(ValueError, match=r"leftsemi.*leftanti"): + make_stage(tmp_path, mode="inner") + + +def test_rejects_destructive_output_root(tmp_path: Path) -> None: + reference_path = tmp_path / "partitions" + reference_path.mkdir() + with pytest.raises(ValueError, match="distinct from reference_path"): + PartitionedExactFilterStage( + key_fields="key", + reference_path=str(reference_path), + output_path=str(reference_path), + total_partitions=1, + ) + + +@pytest.mark.parametrize( + ("metadata", "exception", "message"), + [ + ({"total_partitions": 2}, TypeError, "partition_index"), + ({"partition_index": True, "total_partitions": 2}, TypeError, "partition_index"), + ({"partition_index": -1, "total_partitions": 2}, ValueError, "outside"), + ({"partition_index": 2, "total_partitions": 2}, ValueError, "outside"), + ({"partition_index": 0, "total_partitions": 3}, ValueError, "does not match"), + ], +) +def test_requires_complete_shuffle_partition_metadata( + tmp_path: Path, + metadata: dict[str, int], + exception: type[Exception], + message: str, +) -> None: + stage = make_stage(tmp_path) + task = FileGroupTask(dataset_name="left", data=["left.parquet"], _metadata=metadata) + with pytest.raises(exception, match=message): + stage.process(task) + + +def test_requires_matching_reference_partition(tmp_path: Path) -> None: + stage = make_stage(tmp_path) + task = FileGroupTask( + dataset_name="left", + data=["left.parquet"], + _metadata={"partition_index": 1, "total_partitions": 2}, + ) + with pytest.raises(FileNotFoundError, match=r"part\.1\.parquet"): + stage.process(task) + + +def test_rejects_left_input_overwrite(tmp_path: Path) -> None: + stage = make_stage(tmp_path) + task = FileGroupTask( + dataset_name="left", + data=[str(tmp_path / "output" / "part.0.parquet")], + _metadata={"partition_index": 0, "total_partitions": 2}, + ) + with pytest.raises(ValueError, match="overwrite its left input"): + stage.process(task) + + +def test_process_preserves_stage_perf_without_gpu(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + class FakeFrame: + columns: ClassVar[list[str]] = ["key", "value"] + + def __init__(self, rows: int) -> None: + self.rows = rows + self.write: tuple[str, bool] | None = None + + def __len__(self) -> int: + return self.rows + + def merge(self, other: "FakeFrame", *, how: str, on: list[str]) -> "FakeFrame": + assert len(other) == 2 + assert how == "leftanti" + assert on == ["key"] + return FakeFrame(3) + + def to_parquet(self, path: str, *, index: bool, **kwargs) -> None: + assert not kwargs + self.write = (path, index) + + left = FakeFrame(5) + reference = FakeFrame(2) + + def read_parquet(path: str | list[str], **kwargs) -> FakeFrame: + if isinstance(path, list): + assert kwargs == {} + return left + assert kwargs == {"columns": ["key"]} + return reference + + monkeypatch.setitem(__import__("sys").modules, "cudf", SimpleNamespace(read_parquet=read_parquet)) + stage = make_stage(tmp_path, total_partitions=1) + (tmp_path / "reference" / "part.0.parquet").touch() + task = FileGroupTask( + dataset_name="left", + data=[str(tmp_path / "left.parquet")], + _metadata={"partition_index": 0, "total_partitions": 1}, + _stage_perf=["upstream-stage"], + ) + + result = stage.process(task) + + assert result._stage_perf == task._stage_perf + assert result._metadata["input_rows"] == 5 + assert result._metadata["reference_rows"] == 2 + assert result._metadata["output_rows"] == 3 + + +@pytest.mark.gpu +@pytest.mark.parametrize( + ("mode", "expected_ids"), + [ + ("leftsemi", [1, 2, 4]), + ("leftanti", [0, 3]), + ], +) +def test_exact_filter_preserves_left_multiplicity( + tmp_path: Path, + mode: str, + expected_ids: list[int], +) -> None: + import cudf + + reference_path = tmp_path / "reference" + reference_path.mkdir() + left_path = tmp_path / "left.parquet" + cudf.DataFrame( + { + "id": [0, 1, 2, 3, 4], + "key": ["a", "b", "b", "c", "d"], + "value": [10, 20, 21, 30, 40], + } + ).to_parquet(left_path) + cudf.DataFrame({"key": ["b", "d", "d"]}).to_parquet(reference_path / "part.0.parquet") + + stage = PartitionedExactFilterStage( + key_fields="key", + reference_path=str(reference_path), + output_path=str(tmp_path / "output"), + total_partitions=1, + mode=mode, + ) + task = FileGroupTask( + dataset_name="left", + data=[str(left_path)], + _metadata={"partition_index": 0, "total_partitions": 1}, + _stage_perf=["upstream-stage"], + ) + result = stage.process(task) + + output = cudf.read_parquet(result.data[0]).sort_values("id") + assert output["id"].to_arrow().to_pylist() == expected_ids + assert list(output.columns) == ["id", "key", "value"] + assert result._metadata["input_rows"] == 5 + assert result._metadata["reference_rows"] == 3 + assert result._metadata["output_rows"] == len(expected_ids) + assert result._stage_perf == task._stage_perf From f2d17a539714496e82453eb18ac8724d7c68aaa2 Mon Sep 17 00:00:00 2001 From: Vibhu Jawa Date: Wed, 15 Jul 2026 00:34:28 -0700 Subject: [PATCH 2/2] Harden partitioned exact filter validation --- .../shuffle_utils/partitioned_exact_filter.py | 37 +++++++++-- .../test_partitioned_exact_filter.py | 64 +++++++++++++++++-- 2 files changed, 88 insertions(+), 13 deletions(-) diff --git a/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py b/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py index cff90da5c6..ef6df95301 100644 --- a/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py +++ b/nemo_curator/stages/deduplication/shuffle_utils/partitioned_exact_filter.py @@ -14,8 +14,12 @@ from __future__ import annotations +import os +import posixpath from typing import TYPE_CHECKING, Any, Literal +from fsspec.core import split_protocol + from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import FileGroupTask @@ -28,8 +32,8 @@ ExactFilterMode = Literal["leftsemi", "leftanti"] _ROW_SUBSETTING_READ_KWARGS = ["columns", "filters", "row_groups", "nrows", "skip_rows"] +_CONTROLLED_WRITE_KWARGS = ["index"] _NON_SINGLE_FILE_WRITE_KWARGS = [ - "index", "metadata_file_path", "partition_cols", "partition_file_name", @@ -37,6 +41,17 @@ ] +def _path_identity(path: str) -> tuple[str, str]: + """Return a comparison identity without changing the path used for I/O.""" + protocol, stripped_path = split_protocol(path) + normalized_protocol = (protocol or "file").casefold() + if normalized_protocol == "file": + normalized_path = os.path.realpath(stripped_path) + else: + normalized_path = posixpath.normpath(stripped_path) + return normalized_protocol, normalized_path + + class PartitionedExactFilterStage(ProcessingStage[FileGroupTask, FileGroupTask]): """Filter complete shuffled partitions against matching reference partitions. @@ -75,8 +90,8 @@ class PartitionedExactFilterStage(ProcessingStage[FileGroupTask, FileGroupTask]) accepted. write_kwargs Keyword arguments for writing one result Parquet file with cuDF. - Dataset-partitioning and auxiliary-metadata output options are not - accepted. + The stage controls ``index``. Dataset-partitioning and + auxiliary-metadata output options are not accepted. """ name = "PartitionedExactFilter" @@ -113,7 +128,7 @@ def __init__( # noqa: PLR0913 if not reference_root or not output_root: msg = "reference_path and output_path must be non-root directory paths" raise ValueError(msg) - if reference_root == output_root: + if _path_identity(reference_root) == _path_identity(output_root): msg = "output_path must be distinct from reference_path" raise ValueError(msg) @@ -128,6 +143,7 @@ def __init__( # noqa: PLR0913 check_disallowed_kwargs(self.read_kwargs, _ROW_SUBSETTING_READ_KWARGS) check_disallowed_kwargs(self.reference_read_kwargs, _ROW_SUBSETTING_READ_KWARGS) + check_disallowed_kwargs(self.write_kwargs, _CONTROLLED_WRITE_KWARGS) check_disallowed_kwargs(self.write_kwargs, _NON_SINGLE_FILE_WRITE_KWARGS) self.reference_fs = get_fs( @@ -169,7 +185,8 @@ def _output_partition_path(self, partition_index: int) -> str: def process(self, task: FileGroupTask) -> FileGroupTask: partition_index = self._partition_index(task) output_file = self._output_partition_path(partition_index) - if output_file in task.data: + output_identity = _path_identity(output_file) + if any(_path_identity(input_file) == output_identity for input_file in task.data): msg = "output partition would overwrite its left input partition" raise ValueError(msg) reference_file = self._reference_partition_path(partition_index) @@ -178,7 +195,11 @@ def process(self, task: FileGroupTask) -> FileGroupTask: raise FileNotFoundError(msg) import cudf + import pyarrow.parquet as pq + with self.reference_fs.open(reference_file, "rb") as reference_stream: + reference_schema = pq.read_schema(reference_stream) + self._require_column_names(reference_schema.names, "reference") left = cudf.read_parquet(task.data, **self.read_kwargs) reference = cudf.read_parquet( reference_file, @@ -186,7 +207,6 @@ def process(self, task: FileGroupTask) -> FileGroupTask: **self.reference_read_kwargs, ) self._require_key_fields(left, "left") - self._require_key_fields(reference, "reference") reference_rows = len(reference) output = left.merge(reference, how=self.mode, on=self.key_fields) @@ -220,7 +240,10 @@ def process(self, task: FileGroupTask) -> FileGroupTask: ) def _require_key_fields(self, frame: cudf.DataFrame, side: str) -> None: - missing = [field for field in self.key_fields if field not in frame.columns] + self._require_column_names(frame.columns, side) + + def _require_column_names(self, column_names: Any, side: str) -> None: + missing = [field for field in self.key_fields if field not in column_names] if missing: msg = f"{side} partition is missing exact key columns: {missing}" raise ValueError(msg) diff --git a/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py b/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py index 2e1f6ef2d5..ad29f27172 100644 --- a/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py +++ b/tests/stages/deduplication/shuffle_utils/test_partitioned_exact_filter.py @@ -16,6 +16,8 @@ from types import SimpleNamespace from typing import ClassVar +import pyarrow as pa +import pyarrow.parquet as pq import pytest from nemo_curator.stages.deduplication.shuffle_utils.partitioned_exact_filter import ( @@ -82,6 +84,20 @@ def test_rejects_destructive_output_root(tmp_path: Path) -> None: ) +def test_rejects_destructive_output_root_alias(tmp_path: Path) -> None: + reference_path = tmp_path / "partitions" + reference_path.mkdir() + alias_path = tmp_path / "partition-alias" + alias_path.symlink_to(reference_path, target_is_directory=True) + with pytest.raises(ValueError, match="distinct from reference_path"): + PartitionedExactFilterStage( + key_fields="key", + reference_path=str(reference_path), + output_path=str(alias_path), + total_partitions=1, + ) + + @pytest.mark.parametrize( ("metadata", "exception", "message"), [ @@ -126,7 +142,26 @@ def test_rejects_left_input_overwrite(tmp_path: Path) -> None: stage.process(task) -def test_process_preserves_stage_perf_without_gpu(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_left_input_overwrite_alias(tmp_path: Path) -> None: + stage = make_stage(tmp_path) + output_alias = tmp_path / "output-alias" + output_alias.symlink_to(tmp_path / "output", target_is_directory=True) + task = FileGroupTask( + dataset_name="left", + data=[str(output_alias / "." / "part.0.parquet")], + _metadata={"partition_index": 0, "total_partitions": 2}, + ) + with pytest.raises(ValueError, match="overwrite its left input"): + stage.process(task) + + +@pytest.mark.parametrize(("mode", "output_rows"), [("leftanti", 3), ("leftsemi", 2)]) +def test_process_preserves_stage_perf_without_gpu( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mode: str, + output_rows: int, +) -> None: class FakeFrame: columns: ClassVar[list[str]] = ["key", "value"] @@ -139,9 +174,9 @@ def __len__(self) -> int: def merge(self, other: "FakeFrame", *, how: str, on: list[str]) -> "FakeFrame": assert len(other) == 2 - assert how == "leftanti" + assert how == mode assert on == ["key"] - return FakeFrame(3) + return FakeFrame(output_rows) def to_parquet(self, path: str, *, index: bool, **kwargs) -> None: assert not kwargs @@ -158,8 +193,8 @@ def read_parquet(path: str | list[str], **kwargs) -> FakeFrame: return reference monkeypatch.setitem(__import__("sys").modules, "cudf", SimpleNamespace(read_parquet=read_parquet)) - stage = make_stage(tmp_path, total_partitions=1) - (tmp_path / "reference" / "part.0.parquet").touch() + stage = make_stage(tmp_path, total_partitions=1, mode=mode) + pq.write_table(pa.table({"key": pa.array([], type=pa.string())}), tmp_path / "reference" / "part.0.parquet") task = FileGroupTask( dataset_name="left", data=[str(tmp_path / "left.parquet")], @@ -172,7 +207,24 @@ def read_parquet(path: str | list[str], **kwargs) -> FakeFrame: assert result._stage_perf == task._stage_perf assert result._metadata["input_rows"] == 5 assert result._metadata["reference_rows"] == 2 - assert result._metadata["output_rows"] == 3 + assert result._metadata["output_rows"] == output_rows + + +def test_reference_schema_error_precedes_cudf_read(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def read_parquet(_path: str | list[str], **_kwargs) -> None: + pytest.fail("cuDF data read should not run after schema validation fails") + + monkeypatch.setitem(__import__("sys").modules, "cudf", SimpleNamespace(read_parquet=read_parquet)) + stage = make_stage(tmp_path, total_partitions=1) + pq.write_table(pa.table({"other": pa.array([], type=pa.string())}), tmp_path / "reference" / "part.0.parquet") + task = FileGroupTask( + dataset_name="left", + data=[str(tmp_path / "left.parquet")], + _metadata={"partition_index": 0, "total_partitions": 1}, + ) + + with pytest.raises(ValueError, match="reference partition is missing exact key columns"): + stage.process(task) @pytest.mark.gpu