Tighten multimodal reader/writer contracts and reduce reader complexity

Author: VibhuJawa

nemo_curator/stages/multimodal/io/readers/base.py

@@ -10,7 +10,6 @@
 
 import json
 from abc import ABC, abstractmethod
-from collections import OrderedDict
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -23,14 +22,17 @@
 from nemo_curator.tasks.multimodal import METADATA_SCHEMA, MULTIMODAL_SCHEMA
 from nemo_curator.utils.file_utils import resolve_fs_and_path
 from nemo_curator.utils.grouping import split_by_chunk_size
-from nemo_curator.utils.multimodal_utils import sort_multimodal_table
+from nemo_curator.utils.multimodal_utils import (
+    metadata_map_from_tables,
+    metadata_rows_for_table,
+    sort_multimodal_table,
+)
 from nemo_curator.utils.webdataset_utils import content_type_from_name
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
 ReaderTask = FileGroupTask | tuple[FileGroupTask, FileGroupTask | None]
-_PAIR_ELEMENT_COUNT = 2
 
 
 @dataclass
@@ -128,7 +130,7 @@ def _build_batches_from_tables(
     ) -> MultimodalBatch | list[MultimodalBatch]:
         table = self._concat_data_tables_or_empty(data_tables)
         table = sort_multimodal_table(table)
-        metadata_by_sample = self._metadata_map_from_tables(metadata_tables)
+        metadata_by_sample = metadata_map_from_tables(metadata_tables)
         table_splits = self.split_table(table)
         batches = [
             self._build_batch(
@@ -181,7 +183,7 @@ def split_table_by_sample_max_bytes(self, table: pa.Table, max_batch_bytes: int)
         """Split table by sample groups while preserving sample row locality."""
         if table.num_rows == 0:
             return [table]
-        row_indices_by_sample: OrderedDict[str, list[int]] = OrderedDict()
+        row_indices_by_sample: dict[str, list[int]] = {}
         for idx, sample_id in enumerate(table["sample_id"].to_pylist()):
             sid = str(sample_id)
             row_indices_by_sample.setdefault(sid, [])
@@ -205,6 +207,7 @@ def _text_row(  # noqa: PLR0913
         content_type: str,
         text_content: str,
         element_metadata_json: str | None = None,
+        source_id: str | None = None,
     ) -> dict[str, object]:
         """Build one normalized text row payload."""
         return {
@@ -215,7 +218,7 @@ def _text_row(  # noqa: PLR0913
             "text_content": text_content,
             "binary_content": None,
             "element_metadata_json": element_metadata_json,
-            "source_id": sid,
+            "source_id": source_id or sid,
             "source_shard": source_shard,
             "content_path": None,
             "content_key": None,
@@ -284,24 +287,6 @@ def _task_metadata(self, task: FileGroupTask) -> dict[str, Any]:
         """Propagate task metadata and attach storage options used for reads."""
         return {**task._metadata, "storage_options": dict(self.storage_options)}
 
-    @staticmethod
-    def _metadata_map_from_tables(metadata_tables: list[pa.Table]) -> dict[str, str]:
-        """Build first-wins sample->metadata_json map from metadata tables."""
-        metadata_by_sample: dict[str, str] = {}
-        for metadata_table in metadata_tables:
-            has_rows = metadata_table.num_rows > 0
-            has_sample_id = "sample_id" in metadata_table.column_names
-            if has_rows and has_sample_id:
-                sample_ids = metadata_table["sample_id"].to_pylist()
-                if "metadata_json" in metadata_table.column_names:
-                    metadata_json_values = metadata_table["metadata_json"].to_pylist()
-                else:
-                    metadata_json_values = [None] * len(sample_ids)
-                for sample_id, metadata_json in zip(sample_ids, metadata_json_values, strict=True):
-                    if isinstance(metadata_json, str):
-                        metadata_by_sample.setdefault(str(sample_id), metadata_json)
-        return metadata_by_sample
-
     def _build_batch(
         self,
         task: FileGroupTask,
@@ -311,7 +296,7 @@ def _build_batch(
         split_output: bool,
     ) -> MultimodalBatch:
         """Assemble one ``MultimodalBatch`` from normalized data and metadata."""
-        metadata_rows = self._metadata_rows_for_table(table, metadata_by_sample)
+        metadata_rows = metadata_rows_for_table(table, metadata_by_sample)
         metadata_table = (
             pa.Table.from_pylist(metadata_rows, schema=METADATA_SCHEMA)
             if metadata_rows
@@ -326,45 +311,3 @@ def _build_batch(
             _metadata=self._task_metadata(task),
             _stage_perf=task._stage_perf,
         )
-
-    @staticmethod
-    def _metadata_rows_for_table(
-        table: pa.Table,
-        metadata_by_sample: dict[str, str],
-    ) -> list[dict[str, object]]:
-        """Build metadata rows with single-pass sample type inference."""
-        if table.num_rows == 0:
-            return []
-
-        sample_stats: OrderedDict[str, tuple[int, bool, bool]] = OrderedDict()
-        for sample_id, modality in zip(table["sample_id"].to_pylist(), table["modality"].to_pylist(), strict=True):
-            sid = str(sample_id)
-            modality_name = str(modality)
-            count, has_image, has_text = sample_stats.get(sid, (0, False, False))
-            sample_stats[sid] = (
-                count + 1,
-                has_image or modality_name == "image",
-                has_text or modality_name == "text",
-            )
-
-        return [
-            {
-                "sample_id": sid,
-                "sample_type": BaseMultimodalReaderStage._sample_type_from_summary(
-                    num_rows=num_rows,
-                    has_image=has_image,
-                    has_text=has_text,
-                ),
-                "metadata_json": metadata_by_sample.get(sid),
-            }
-            for sid, (num_rows, has_image, has_text) in sample_stats.items()
-        ]
-
-    @staticmethod
-    def _sample_type_from_summary(num_rows: int, has_image: bool, has_text: bool) -> str:
-        """Infer sample type from in-sample modality ordering."""
-        if num_rows == 1:
-            return "single"
-        if num_rows == _PAIR_ELEMENT_COUNT and has_image and has_text:
-            return "pair"
-        return "interleaved"

nemo_curator/stages/multimodal/io/readers/parquet.py

@@ -57,26 +57,35 @@ def __post_init__(self) -> None:
     @staticmethod
     def _validate_columns(columns: list[str] | None) -> list[str] | None:
         """Validate optional data column selection."""
-        return ParquetMultimodalReaderStage._validate_column_selection(columns, option_name="columns")
+        if columns is None:
+            return None
+        if len(columns) == 0:
+            msg = "columns must be a non-empty list when provided"
+            raise ValueError(msg)
+        seen: set[str] = set()
+        normalized: list[str] = []
+        for column in columns:
+            if not isinstance(column, str) or not column:
+                msg = "columns entries must be non-empty strings"
+                raise ValueError(msg)
+            if column not in seen:
+                seen.add(column)
+                normalized.append(column)
+        return normalized
 
     @staticmethod
     def _validate_metadata_columns(columns: list[str] | None) -> list[str] | None:
         """Validate optional metadata sidecar column selection."""
-        return ParquetMultimodalReaderStage._validate_column_selection(columns, option_name="metadata_columns")
-
-    @staticmethod
-    def _validate_column_selection(columns: list[str] | None, option_name: str) -> list[str] | None:
-        """Validate and de-duplicate a selected column list."""
         if columns is None:
             return None
         if len(columns) == 0:
-            msg = f"{option_name} must be a non-empty list when provided"
+            msg = "metadata_columns must be a non-empty list when provided"
             raise ValueError(msg)
         seen: set[str] = set()
         normalized: list[str] = []
         for column in columns:
             if not isinstance(column, str) or not column:
-                msg = f"{option_name} entries must be non-empty strings"
+                msg = "metadata_columns entries must be non-empty strings"
                 raise ValueError(msg)
             if column not in seen:
                 seen.add(column)
@@ -165,12 +174,6 @@ class ParquetMultimodalReader(CompositeStage[_EmptyTask, MultimodalBatch]):
 
     def __post_init__(self) -> None:
         super().__init__()
-        if isinstance(self.file_paths, str) and not self.file_paths.endswith(".parquet"):
-            msg = (
-                "When file_paths is a string, it must point to a .parquet file. "
-                "Use an explicit list of parquet file paths when reading multiple files."
-            )
-            raise ValueError(msg)
 
     def decompose(self) -> list[ProcessingStage]:
         return [

nemo_curator/stages/multimodal/io/readers/webdataset.py

@@ -49,6 +49,10 @@
 InterleavedSegment = dict[str, object]
 
 
+class WebDatasetMemberParseError(ValueError):
+    """Expected parse/validation failure for one WebDataset member."""
+
+
 @dataclass
 class RowBuildState:
     """Per-shard mutable parse state.
@@ -82,7 +86,7 @@ def _required_segment_str(segment: InterleavedSegment, field: str) -> str:
     value = segment.get(field)
     if not isinstance(value, str) or not value:
         msg = f"Interleaved segment must include non-empty string '{field}'"
-        raise ValueError(msg)
+        raise WebDatasetMemberParseError(msg)
     return value
 
 
@@ -92,25 +96,25 @@ def _validate_interleaved_payload(
 ) -> tuple[str, list[InterleavedSegment]]:
     if not isinstance(decoded, dict):
         msg = "Interleaved JSON payload must decode to an object"
-        raise TypeError(msg)
+        raise WebDatasetMemberParseError(msg)
 
     sample_id_field = field_map["sample_id"]
     segments_field = field_map["segments"]
     sample_id = decoded.get(sample_id_field)
     if not isinstance(sample_id, str) or not sample_id:
         msg = f"Interleaved JSON payload must include non-empty string '{sample_id_field}'"
-        raise ValueError(msg)
+        raise WebDatasetMemberParseError(msg)
 
     segments = decoded.get(segments_field)
     if not isinstance(segments, list):
         msg = f"Interleaved JSON payload must include list field '{segments_field}'"
-        raise TypeError(msg)
+        raise WebDatasetMemberParseError(msg)
 
     typed_segments: list[InterleavedSegment] = []
     for idx, segment in enumerate(segments):
         if not isinstance(segment, dict):
             msg = f"Interleaved segment at index={idx} for sample_id='{sample_id}' must be an object"
-            raise TypeError(msg)
+            raise WebDatasetMemberParseError(msg)
         typed_segments.append(segment)
     return sample_id, typed_segments
 
@@ -195,7 +199,7 @@ def read_data(self, data_path: str, metadata_path: str | None) -> tuple[pa.Table
                     try:
                         payload = self._member_payload(tf, member_name, member)
                         rows.extend(self._rows_from_member(state, member_name, payload, source))
-                    except Exception as err:  # noqa: BLE001
+                    except (OSError, UnicodeDecodeError, json.JSONDecodeError, WebDatasetMemberParseError) as err:
                         self._handle_member_error(member_name, err)
         return self._rows_to_table(rows), pa.Table.from_pylist(state.metadata_rows, schema=METADATA_SCHEMA)
 
@@ -283,7 +287,7 @@ def _rows_from_interleaved_json(
                     f"Unsupported interleaved modality='{modality}' for sample_id='{sample_id}' "
                     "in WebDatasetReaderStage (supported: text, image)"
                 )
-                raise ValueError(msg)
+                raise WebDatasetMemberParseError(msg)
             if self._loads_modality(modality):
                 if modality == "text":
                     rows.append(
@@ -317,20 +321,25 @@ def _maybe_rows_from_interleaved_json_member(
     ) -> list[dict[str, object]] | None:
         if payload is None:
             msg = f"JSON member '{member_name}' missing payload bytes"
-            raise ValueError(msg)
+            raise WebDatasetMemberParseError(msg)
         try:
             parsed = self._rows_from_interleaved_json(payload, source, state)
-        except (KeyError, TypeError, ValueError):
+        except WebDatasetMemberParseError:
             if self.sample_format == "interleaved":
                 raise
             return None
+        except KeyError as err:
+            if self.sample_format == "interleaved":
+                msg = f"Interleaved JSON missing required field: {err}"
+                raise WebDatasetMemberParseError(msg) from err
+            return None
         return parsed
 
     @staticmethod
     def _decode_text_payload(payload: bytes | None, member_name: str) -> str:
         if payload is None:
             msg = f"Text member '{member_name}' missing payload bytes"
-            raise ValueError(msg)
+            raise WebDatasetMemberParseError(msg)
         return payload.decode("utf-8") if payload else ""
 
     @staticmethod
@@ -339,13 +348,13 @@ def _binary_modality_for_member(member_name: str) -> str:
         modality = modality_from_content_type(content_type)
         if modality == "unknown":
             msg = f"Unsupported content_type='{content_type}' for member '{member_name}' in WebDatasetReaderStage"
-            raise ValueError(msg)
+            raise WebDatasetMemberParseError(msg)
         if modality != "image":
             msg = (
                 f"Unsupported binary modality='{modality}' for member '{member_name}' "
                 "in WebDatasetReaderStage (supported: image)"
             )
-            raise ValueError(msg)
+            raise WebDatasetMemberParseError(msg)
         return modality
 
     def _rows_from_binary_member(