Fix Pairwise IO and IdentifyDuplicates in SemDedup for Cloud I/O (#1253)

Author: praateekmahajan

nemo_curator/stages/deduplication/semantic/identify_duplicates.py

@@ -88,14 +88,23 @@ def process_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]:
             return []
 
         all_files = [file for task in tasks for file in task.data]
-        # Read using filters
-        df: pd.DataFrame = pd.read_parquet(
-            all_files,
-            storage_options=self.input_storage_options,
-            **self.read_kwargs,
-            filters=[("cosine_sim_score", ">=", 1.0 - self.eps)],
-            engine="pyarrow",
-        )[["id"]]  # TODO: If we want we can add other columns
+        # We read using filters
+        # We read file by file since list[files] when files are remote urls can fail
+        # See https://github.com/pandas-dev/pandas/issues/62922
+        df: pd.DataFrame = pd.concat(
+            (
+                pd.read_parquet(
+                    f,
+                    storage_options=self.input_storage_options,
+                    **self.read_kwargs,
+                    filters=[("cosine_sim_score", ">=", 1.0 - self.eps)],
+                    columns=["id"],
+                    engine="pyarrow",
+                )
+                for f in all_files
+            ),
+            ignore_index=True,
+        )
         # Write out sorted and with multiple row groups
         df.sort_values("id", inplace=True)  # noqa: PD002
 

nemo_curator/stages/deduplication/semantic/pairwise_io.py

@@ -21,6 +21,7 @@
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 from nemo_curator.tasks import FileGroupTask, _EmptyTask
+from nemo_curator.utils.client_utils import is_remote_url
 from nemo_curator.utils.file_utils import get_all_file_paths_under, get_fs, infer_dataset_name_from_path
 
 if TYPE_CHECKING:
@@ -52,6 +53,7 @@ def __init__(
         self.name = "pairwise_file_partitioning"
         self.resources = Resources(cpus=0.5)
         self.fs: AbstractFileSystem | None = None
+        self.path_normalizer = lambda x: x
 
     def inputs(self) -> tuple[list[str], list[str]]:
         return ["data"], []
@@ -61,6 +63,7 @@ def outputs(self) -> tuple[list[str], list[str]]:
 
     def setup(self, _: WorkerMetadata | None = None) -> None:
         self.fs = get_fs(self.input_path, storage_options=self.storage_options)
+        self.path_normalizer = self.fs.unstrip_protocol if is_remote_url(self.input_path) else (lambda x: x)
 
     def ray_stage_spec(self) -> dict[str, Any]:
         """Ray stage specification for this stage."""
@@ -83,7 +86,7 @@ def process(self, _: _EmptyTask) -> list[FileGroupTask]:
             # Extract centroid ID from directory name (e.g., "centroid=0" -> 0)
             if "centroid=" in entry:
                 centroid_id = int(entry.split("centroid=")[-1])
-                centroid_dirs[centroid_id] = entry
+                centroid_dirs[centroid_id] = self.path_normalizer(entry)
 
         logger.debug(
             f"Found {len(centroid_dirs)} centroid directories e.g. {next(iter(centroid_dirs.values())) if centroid_dirs else None}"