feat(interleaved): Foundation IO improvements — schema utilities, base-class fixes, and bug fixes

.gitignore

@@ -157,3 +157,4 @@ data/
 .DS_Store
 AGENTS.md
 alm_output/
+benchmark_results/

nemo_curator/stages/interleaved/README.md

@@ -9,7 +9,7 @@ WebDataset tar shards
         |
         v
 ┌─────────────────────────┐
-│  WebdatasetReader       │  CompositeStage: FilePartitioning + WebdatasetReaderStage
+│  WebdatasetReader       │  CompositeStage: FilePartitioning + InterleavedWebdatasetReaderStage
 │  (io/reader.py)         │  Parses tar members -> normalized interleaved rows
 └────────┬────────────────┘
          |  InterleavedBatch (Arrow/Pandas)
@@ -141,7 +141,7 @@ stages/interleaved/
 │   ├── reader.py                   # WebdatasetReader (CompositeStage)
 │   ├── readers/
 │   │   ├── base.py                 # BaseInterleavedReader
-│   │   └── webdataset.py           # WebdatasetReaderStage (ProcessingStage)
+│   │   └── webdataset.py           # InterleavedWebdatasetReaderStage (ProcessingStage)
 │   └── writers/
 │       ├── base.py                 # BaseInterleavedWriter (filesystem + materialization + process)
 │       └── tabular.py              # InterleavedParquetWriterStage

nemo_curator/stages/interleaved/io/__init__.py

@@ -15,4 +15,7 @@
 from nemo_curator.stages.interleaved.io.reader import WebdatasetReader
 from nemo_curator.stages.interleaved.io.writers.tabular import InterleavedParquetWriterStage
 
-__all__ = ["InterleavedParquetWriterStage", "WebdatasetReader"]
+__all__ = [
+    "InterleavedParquetWriterStage",
+    "WebdatasetReader",
+]

nemo_curator/stages/interleaved/io/reader.py

@@ -12,17 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
 from typing import Any
 
 from nemo_curator.stages.base import CompositeStage
 from nemo_curator.stages.file_partitioning import FilePartitioningStage
-from nemo_curator.stages.interleaved.io.readers.webdataset import WebdatasetReaderStage
+from nemo_curator.stages.interleaved.io.readers.webdataset import InterleavedWebdatasetReaderStage
 from nemo_curator.stages.interleaved.utils import (
     DEFAULT_IMAGE_EXTENSIONS,
     DEFAULT_JSON_EXTENSIONS,
     DEFAULT_WEBDATASET_EXTENSIONS,
-    require_source_id_field,
     resolve_storage_options,
 )
 from nemo_curator.tasks import InterleavedBatch, _EmptyTask
@@ -41,7 +42,7 @@ class WebdatasetReader(CompositeStage[_EmptyTask, InterleavedBatch]):
     file_extensions: list[str] = field(default_factory=lambda: list(DEFAULT_WEBDATASET_EXTENSIONS))
     json_extensions: list[str] = field(default_factory=lambda: list(DEFAULT_JSON_EXTENSIONS))
     image_extensions: list[str] = field(default_factory=lambda: list(DEFAULT_IMAGE_EXTENSIONS))
-    source_id_field: str = ""
+    source_id_field: str | None = None
     sample_id_field: str | None = None
     texts_field: str = "texts"
     images_field: str = "images"
@@ -53,7 +54,6 @@ class WebdatasetReader(CompositeStage[_EmptyTask, InterleavedBatch]):
 
     def __post_init__(self):
         super().__init__()
-        self.source_id_field = require_source_id_field(self.source_id_field)
         self.storage_options = resolve_storage_options(io_kwargs=self.read_kwargs)
 
     def decompose(self) -> list:
@@ -65,7 +65,7 @@ def decompose(self) -> list:
                 file_extensions=self.file_extensions,
                 storage_options=self.storage_options,
             ),
-            WebdatasetReaderStage(
+            InterleavedWebdatasetReaderStage(
                 read_kwargs=self.read_kwargs,
                 materialize_on_read=self.materialize_on_read,
                 max_batch_bytes=self.max_batch_bytes,

nemo_curator/stages/interleaved/io/readers/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemo_curator.stages.interleaved.io.readers.webdataset import WebdatasetReaderStage
+from nemo_curator.stages.interleaved.io.readers.webdataset import InterleavedWebdatasetReaderStage
 
-__all__ = ["WebdatasetReaderStage"]
+__all__ = ["InterleavedWebdatasetReaderStage"]

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

@@ -15,19 +15,78 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+import pyarrow as pa
+from loguru import logger
+
 from nemo_curator.stages.base import ProcessingStage
+from nemo_curator.stages.interleaved.utils.schema import align_table, reconcile_schema
 from nemo_curator.tasks import FileGroupTask, InterleavedBatch
+from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA
 
 
 @dataclass
 class BaseInterleavedReader(ProcessingStage[FileGroupTask, InterleavedBatch]):
-    """Base contract for interleaved readers."""
+    """Base contract for interleaved readers.
+
+    By default (``schema=None``) user-added passthrough columns are preserved
+    and only reserved-column types are reconciled via ``reconcile_schema``.
+
+    If *schema* is set explicitly, every output table is strictly aligned to it
+    (missing columns become typed nulls, extra columns are dropped).
+
+    Use *schema_overrides* to add or override individual field types relative to
+    ``INTERLEAVED_SCHEMA`` while keeping strict alignment:
+
+    .. code-block:: python
+
+        reader = InterleavedParquetReader(
+            "data.parquet",
+            schema_overrides={"url": pa.string(), "timestamp": pa.int64()},
+        )
+    """
 
     read_kwargs: dict[str, Any] = field(default_factory=dict)
+    schema: pa.Schema | None = None
+    schema_overrides: dict[str, pa.DataType] | None = None
     name: str = "base_interleaved_reader"
 
+    def __post_init__(self) -> None:
+        if self.schema_overrides is not None:
+            self.schema = _resolve_schema(self.schema, self.schema_overrides)
+
     def inputs(self) -> tuple[list[str], list[str]]:
         return ["data"], []
 
     def outputs(self) -> tuple[list[str], list[str]]:
         return ["data"], ["sample_id", "position", "modality"]
+
+    def _align_output(self, table: pa.Table) -> pa.Table:
+        """Reconcile or align *table* to the declared schema."""
+        if self.schema is not None:
+            return align_table(table, self.schema)
+        return table.cast(reconcile_schema(table.schema))
+
+
+def _resolve_schema(
+    schema: pa.Schema | None,
+    overrides: dict[str, pa.DataType] | None,
+) -> pa.Schema | None:
+    """Return the effective schema from user-supplied *schema* or *overrides*.
+
+    Priority: *schema* > *schema_overrides* merged on top of ``INTERLEAVED_SCHEMA``.
+    Raises ``ValueError`` if both inputs are ``None``.
+    """
+    if schema is not None:
+        if overrides:
+            logger.warning("schema_overrides ignored because schema= is already set; use one or the other, not both")
+        return schema
+    if overrides:
+        fields = {f.name: f for f in INTERLEAVED_SCHEMA}
+        for name, dtype in overrides.items():
+            orig = fields.get(name)
+            nullable = orig.nullable if orig is not None else True
+            metadata = orig.metadata if orig is not None else None
+            fields[name] = pa.field(name, dtype, nullable=nullable, metadata=metadata)
+        return pa.schema(list(fields.values()))
+    msg = "At least one of schema= or schema_overrides= must be provided"
+    raise ValueError(msg)

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

@@ -29,7 +29,6 @@
 from nemo_curator.stages.interleaved.utils import (
     DEFAULT_IMAGE_EXTENSIONS,
     DEFAULT_JSON_EXTENSIONS,
-    require_source_id_field,
     resolve_storage_options,
     validate_and_project_source_fields,
 )
@@ -67,14 +66,14 @@ class _SampleContext:
 
 
 @dataclass
-class WebdatasetReaderStage(BaseInterleavedReader):
+class InterleavedWebdatasetReaderStage(BaseInterleavedReader):
     """Read MINT1T-style WebDataset shards into a row-wise multimodal task."""
 
     materialize_on_read: bool = False
     max_batch_bytes: int | None = None
     json_extensions: tuple[str, ...] = DEFAULT_JSON_EXTENSIONS
     image_extensions: tuple[str, ...] = field(default_factory=lambda: DEFAULT_IMAGE_EXTENSIONS)
-    source_id_field: str = ""
+    source_id_field: str | None = None
     sample_id_field: str | None = None
     texts_field: str = "texts"
     images_field: str = "images"
@@ -85,7 +84,7 @@ class WebdatasetReaderStage(BaseInterleavedReader):
     name: str = "webdataset_reader"
 
     def __post_init__(self) -> None:
-        self.source_id_field = require_source_id_field(self.source_id_field)
+        super().__post_init__()
 
     # -- source_ref construction --
 
@@ -254,7 +253,7 @@ def _rows_from_sample(self, ctx: _SampleContext) -> list[dict[str, Any]]:
 
     def _build_passthrough_row(self, sample: dict[str, Any]) -> dict[str, Any]:
         excluded = RESERVED_COLUMNS | {
-            self.source_id_field,
+            *([self.source_id_field] if self.source_id_field else []),
             *([self.sample_id_field] if self.sample_id_field else []),
             self.texts_field,
             self.images_field,
@@ -283,30 +282,24 @@ def _extract_per_modality_fields(
         return result
 
     def _empty_output_schema(self) -> pa.Schema:
-        schema = INTERLEAVED_SCHEMA
+        # Use explicit schema if set; otherwise fall back to INTERLEAVED_SCHEMA as base
+        base = self.schema if self.schema is not None else INTERLEAVED_SCHEMA
         seen = set(self.fields or ())
         all_extra = list(self.fields or ())
         for f in (*self.per_image_fields, *self.per_text_fields):
             if f not in seen:
                 all_extra.append(f)
                 seen.add(f)
         if not all_extra:
-            return schema
-        existing = set(schema.names)
-        extra_fields = [pa.field(name, pa.null()) for name in all_extra if name not in existing]
-        return pa.schema([*schema, *extra_fields]) if extra_fields else schema
+            return base
+        existing = set(base.names)
+        extra_fields = []
+        for name in all_extra:
+            if name not in existing:
+                extra_fields.append(pa.field(name, pa.null()))
+        return pa.schema([*base, *extra_fields]) if extra_fields else base
 
-    @staticmethod
-    def _reconcile_schema(inferred: pa.Schema) -> pa.Schema:
-        """Build a schema with canonical types for reserved columns and inferred types for passthrough."""
-        canonical = {f.name: f for f in INTERLEAVED_SCHEMA}
-        fields = []
-        for f in inferred:
-            if f.name in canonical:
-                fields.append(canonical[f.name])
-            else:
-                fields.append(f)
-        return pa.schema(fields)
+    # _align_output is inherited from BaseInterleavedReader
 
     # -- image member resolution --
 
@@ -398,11 +391,11 @@ def _rows_from_member(
                 else:
                     frame_index = parsed_ref.get("frame_index")
                     if frame_index is not None:
-                        extracted = _extract_tiff_frame(raw_bytes, frame_index)
-                        if extracted is None:
+                        tiff_frame = _extract_tiff_frame(raw_bytes, frame_index)
+                        if tiff_frame is None:
                             row["materialize_error"] = f"failed to extract frame {frame_index} from '{content_key}'"
                         else:
-                            raw_bytes = extracted
+                            raw_bytes = tiff_frame
                 row["binary_content"] = raw_bytes
             read_ctx.byte_cache.clear()
         return sample_rows
@@ -434,7 +427,7 @@ def process(self, task: FileGroupTask) -> InterleavedBatch | list[InterleavedBat
 
         if rows:
             table = pa.Table.from_pylist(rows)
-            table = table.cast(self._reconcile_schema(table.schema))
+            table = self._align_output(table)
         else:
             # Empty tables use _empty_output_schema(); passthrough columns get
             # pa.null() type which is intentional (no data to infer from).

nemo_curator/stages/interleaved/io/writers/__init__.py

@@ -14,4 +14,6 @@
 
 from nemo_curator.stages.interleaved.io.writers.tabular import InterleavedParquetWriterStage
 
-__all__ = ["InterleavedParquetWriterStage"]
+__all__ = [
+    "InterleavedParquetWriterStage",
+]