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

.gitignore

@@ -156,6 +156,8 @@ data/
 # macOS Files
 .DS_Store
 AGENTS.md
+alm_output/
+benchmark_results/
 
 # Fern: generated API reference from `libraries` in docs.yml (`fern docs md generate`).
 # Regenerated in CI (fern-docs-ci.yml, publish-fern-docs.yml); do not commit.

benchmarking/scripts/multimodal_mint1t_benchmark.py

@@ -42,7 +42,6 @@ def create_pipeline(args: argparse.Namespace) -> Pipeline:
     )
     pipeline.add_stage(
         WebdatasetReader(
-            source_id_field="pdf_name",
             file_paths=args.input_path,
             files_per_partition=args.files_per_partition,
             blocksize=args.input_blocksize,

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)
@@ -52,7 +52,6 @@ Extra fields from the source data flow through the pipeline as additional column
 
 ```python
 reader = WebdatasetReader(
-    source_id_field="pdf_name",
     file_paths="/data/shards/",
     fields=("p_hash", "score", "aux"),  # These become extra columns
 )
@@ -117,7 +116,6 @@ from nemo_curator.stages.interleaved.stages import InterleavedAspectRatioFilterS
 
 pipeline = Pipeline(name="mint1t_pipeline")
 pipeline.add_stage(WebdatasetReader(
-    source_id_field="pdf_name",
     file_paths="/data/mint1t/shards/",
 ))
 pipeline.add_stage(InterleavedAspectRatioFilterStage(drop_invalid_rows=True))
@@ -141,7 +139,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,6 @@ 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 = ""
     sample_id_field: str | None = None
     texts_field: str = "texts"
     images_field: str = "images"
@@ -53,7 +53,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,13 +64,12 @@ 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,
                 json_extensions=tuple(self.json_extensions),
                 image_extensions=tuple(self.image_extensions),
-                source_id_field=self.source_id_field,
                 sample_id_field=self.sample_id_field,
                 texts_field=self.texts_field,
                 images_field=self.images_field,

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,51 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+import pyarrow as pa
+
 from nemo_curator.stages.base import ProcessingStage
+from nemo_curator.stages.interleaved.utils.schema import align_table, reconcile_schema, resolve_schema
 from nemo_curator.tasks import FileGroupTask, InterleavedBatch
 
 
 @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 is not None or 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))