Refactor multimodal column validation and shared schema casting

Author: VibhuJawa

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

@@ -21,6 +21,7 @@
 from nemo_curator.tasks import MultimodalBatch, _EmptyTask
 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.multimodal_utils import cast_required_fields
 
 _DEFAULT_PARQUET_EXTENSIONS: Final[list[str]] = [".parquet"]
 
@@ -43,8 +44,8 @@ class ParquetMultimodalReaderStage(BaseMultimodalReaderStage):
 
     def __post_init__(self) -> None:
         super().__post_init__()
-        self.columns = self._validate_columns(self.columns)
-        self.metadata_columns = self._validate_metadata_columns(self.metadata_columns)
+        self.columns = self._validate_column_selection(self.columns, field_name="data.columns")
+        self.metadata_columns = self._validate_column_selection(self.metadata_columns, field_name="metadata.columns")
         if self.columns is not None:
             missing_required = [name for name in MULTIMODAL_SCHEMA.names if name not in self.columns]
             if missing_required:
@@ -55,37 +56,18 @@ def __post_init__(self) -> None:
                 raise ValueError(msg)
 
     @staticmethod
-    def _validate_columns(columns: list[str] | None) -> list[str] | None:
-        """Validate optional data column selection."""
+    def _validate_column_selection(columns: list[str] | None, *, field_name: str) -> list[str] | None:
+        """Validate optional column selection for data and metadata inputs."""
         if columns is None:
             return None
         if len(columns) == 0:
-            msg = "columns must be a non-empty list when provided"
+            msg = f"{field_name} 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."""
-        if columns is None:
-            return None
-        if len(columns) == 0:
-            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 = "metadata_columns entries must be non-empty strings"
+                msg = f"{field_name} entries must be non-empty strings"
                 raise ValueError(msg)
             if column not in seen:
                 seen.add(column)
@@ -114,7 +96,7 @@ def _read_metadata_table(self, metadata_path: str) -> pa.Table:
     def _normalize_data_table(self, table: pa.Table) -> pa.Table:
         source = self._ensure_optional_string_column(table, "element_metadata_json")
         self._ensure_required_columns(source, MULTIMODAL_SCHEMA.names)
-        return self._cast_required_fields(source, MULTIMODAL_SCHEMA)
+        return cast_required_fields(source, MULTIMODAL_SCHEMA)
 
     def _normalize_metadata_table(self, table: pa.Table) -> pa.Table:
         self._ensure_required_columns(
@@ -124,19 +106,7 @@ def _normalize_metadata_table(self, table: pa.Table) -> pa.Table:
         )
         source = self._ensure_optional_string_column(table, "sample_type")
         source = self._ensure_optional_string_column(source, "metadata_json")
-        return self._cast_required_fields(source, METADATA_SCHEMA)
-
-    @staticmethod
-    def _cast_required_fields(table: pa.Table, required_schema: pa.Schema) -> pa.Table:
-        """Cast required fields in-place while preserving any extra columns."""
-        out = table
-        for required_field in required_schema:
-            col_idx = out.schema.get_field_index(required_field.name)
-            if col_idx >= 0:
-                col = out[required_field.name]
-                if not col.type.equals(required_field.type):
-                    out = out.set_column(col_idx, required_field.name, col.cast(required_field.type))
-        return out
+        return cast_required_fields(source, METADATA_SCHEMA)
 
     @staticmethod
     def _ensure_optional_string_column(table: pa.Table, column_name: str) -> pa.Table:

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

@@ -25,7 +25,7 @@
     resolve_sidecar_output_path,
     resolve_task_scoped_output_path,
 )
-from nemo_curator.utils.multimodal_utils import sort_multimodal_table
+from nemo_curator.utils.multimodal_utils import cast_required_fields, sort_multimodal_table
 
 _SUPPORTED_WRITE_MODES: set[str] = {"overwrite", "error", "ignore"}
 
@@ -137,19 +137,7 @@ def _build_output_table(self, task: MultimodalBatch) -> pa.Table:
         if missing:
             msg = f"{self.__class__.__name__} requires columns: {missing}"
             raise ValueError(msg)
-        return self._cast_required_fields(table, MULTIMODAL_SCHEMA)
-
-    @staticmethod
-    def _cast_required_fields(table: pa.Table, required_schema: pa.Schema) -> pa.Table:
-        """Cast required fields in-place while preserving any extra columns."""
-        out = table
-        for required_field in required_schema:
-            col_idx = out.schema.get_field_index(required_field.name)
-            if col_idx >= 0:
-                col = out[required_field.name]
-                if not col.type.equals(required_field.type):
-                    out = out.set_column(col_idx, required_field.name, col.cast(required_field.type))
-        return out
+        return cast_required_fields(table, MULTIMODAL_SCHEMA)
 
     def _write_tabular(self, table: pa.Table, output_path: str, format_name: Literal["parquet", "arrow"]) -> None:
         """Write Arrow table to parquet or arrow artifact path."""
@@ -201,7 +189,7 @@ def _build_metadata_table(task: MultimodalBatch) -> pa.Table:
                     column_name,
                     pa.nulls(metadata_table.num_rows, type=pa.string()),
                 )
-        metadata_table = BaseMultimodalWriterStage._cast_required_fields(metadata_table, METADATA_SCHEMA)
+        metadata_table = cast_required_fields(metadata_table, METADATA_SCHEMA)
         return metadata_table.sort_by([("sample_id", "ascending")])
 
     @staticmethod

nemo_curator/utils/multimodal_utils.py

@@ -17,6 +17,18 @@
     import pyarrow as pa
 
 
+def cast_required_fields(table: pa.Table, required_schema: pa.Schema) -> pa.Table:
+    """Cast required fields in-place while preserving any extra columns."""
+    out = table
+    for required_field in required_schema:
+        col_idx = out.schema.get_field_index(required_field.name)
+        if col_idx >= 0:
+            col = out[required_field.name]
+            if not col.type.equals(required_field.type):
+                out = out.set_column(col_idx, required_field.name, col.cast(required_field.type))
+    return out
+
+
 def validate_content_path_loading_mode(*, content_path: str, row_indices: list[int], content_keys: list[object | None]) -> None:
     """Ensure one content_path group uses a single loading mode."""
     for idx in row_indices: