Align interleaved source locators with Parquet FILE

nemo_curator/stages/interleaved/README.md

@@ -47,7 +47,9 @@ These are set and managed by pipeline stages. Users should not write to them dir
 | `content_type` | string | Content | MIME type (e.g. `text/plain`, `image/jpeg`) |
 | `text_content` | string | Content | Text payload for text rows |
 | `binary_content` | large_binary | Content | Image bytes (populated by materialization) |
-| `source_ref` | string | Internal | JSON locator `{path, member, byte_offset, byte_size, frame_index}`. `path` alone = direct/remote read; + `member` = tar extract; + `byte_offset/size` = range read (fastest). `path` accepts local or remote (`s3://`) URIs. |
+| `source_ref` | FILE | Internal | Parquet FILE-compatible reference: `uri`, `offset`, `size`, `content_type`, `checksum`, `inline` |
+| `source_member` | string | Internal | Archive member metadata adjacent to FILE |
+| `source_frame_index` | int32 | Internal | Multi-frame index metadata adjacent to FILE |
 | `materialize_error` | string | Internal | Error message if materialization failed |
 
 ### User columns (passthrough)
@@ -73,41 +75,27 @@ Class attributes:
 - `REQUIRED_COLUMNS` -- frozenset of columns that must always be present (non-nullable schema fields)
 
 Key methods:
-- `build_source_ref(path, member, byte_offset, byte_size, frame_index)` -- build a JSON locator string
-- `parse_source_ref(value)` -- parse back with soft migration for older formats
+- `build_source_ref(uri, offset, size)` -- build a FILE-compatible external reference
 - `with_parsed_source_ref_columns(prefix)` -- expand source_ref into DataFrame columns
 - `to_pyarrow()` / `to_pandas()` -- conversion between formats
 
 ### source_ref
 
-A JSON string embedded in each row that tracks where the original content lives:
-
-```json
-{
-  "path": "/data/shard-00000.tar",
-  "member": "abc123.jpg",
-  "byte_offset": 1024,
-  "byte_size": 45678,
-  "frame_index": null
-}
-```
-
-- `path` + `member` -- tar archive path and member name
-- `path` alone (no member) -- direct file path
-- `byte_offset` + `byte_size` -- enables range reads without opening the tar
-- `frame_index` (optional) -- selects a single frame from a multi-frame TIFF during materialization
+`source_ref` uses all six fields from the closed Parquet FILE specification.
+PyArrow cannot yet emit the new FILE footer annotation, so current output uses
+the compatible physical group. Archive member and TIFF frame metadata remain
+in the adjacent `source_member` and `source_frame_index` columns.
 
 ### Materialization
 
-Binary content (images) can be loaded lazily. Three I/O strategies dispatch automatically based on `source_ref` content (`utils/materialization.py`):
+Binary content (images) can be loaded lazily. Two I/O strategies dispatch automatically based on `source_ref` content (`utils/materialization.py`):
 
 | Strategy | When | How |
 |----------|------|-----|
-| **Range read** | `byte_offset` + `byte_size` present | `fs.cat_ranges()` -- batched HTTP range requests per path |
-| **Tar extract** | `member` present, no byte range | Open tar once, `extractfile()` per member |
-| **Direct read** | No `member` | Read entire file via `fsspec.open()` |
+| **FILE read** | External FILE reference | Deduplicated `fs.cat_ranges()` calls, batched per filesystem |
+| **Tar extract** | `source_member` present, no byte range | Open tar once, `extractfile()` per member |
 
-When `frame_index` is set in the `source_ref`, materialization extracts a single frame from a multi-frame TIFF and returns it as a standalone TIFF. Non-TIFF content is returned unchanged regardless of `frame_index`.
+When `source_frame_index` is set, materialization extracts a single frame from a multi-frame TIFF and returns it as a standalone TIFF. Non-TIFF content is returned unchanged regardless of `frame_index`.
 
 Materialization can happen at read time (`materialize_on_read=True`) or write time (`materialize_on_write=True`).
 

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

@@ -35,6 +35,7 @@
 from nemo_curator.stages.interleaved.utils.materialization import _extract_tiff_frame
 from nemo_curator.tasks import FileGroupTask, InterleavedBatch
 from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA, RESERVED_COLUMNS
+from nemo_curator.utils.storage_utils import FILE_REFERENCE_TYPE
 
 from .base import BaseInterleavedReader
 
@@ -92,24 +93,19 @@ def _build_source_ref(
         self,
         ctx: _SampleContext,
         content_key: str | None,
-        *,
-        frame_index: int | None = None,
-    ) -> str:
+    ) -> pa.StructScalar | None:
         if content_key is None:
-            return InterleavedBatch.build_source_ref(path=None, member=None)
-        byte_offset = None
-        byte_size = None
-        if ctx.member_info and content_key in ctx.member_info:
+            return None
+        offset = size = None
+        if (
+            ctx.member_info
+            and content_key in ctx.member_info
+            and not ctx.tar_path.lower().split("?", 1)[0].endswith((".tar.gz", ".tgz"))
+        ):
             info = ctx.member_info[content_key]
-            byte_offset = info.offset_data
-            byte_size = info.size
-        return InterleavedBatch.build_source_ref(
-            path=ctx.tar_path,
-            member=content_key,
-            byte_offset=byte_offset,
-            byte_size=byte_size,
-            frame_index=frame_index,
-        )
+            offset, size = info.offset_data, info.size
+        value = InterleavedBatch.build_source_ref(ctx.tar_path, offset, size)
+        return pa.scalar(value, type=FILE_REFERENCE_TYPE)
 
     # -- row builders (override in subclasses for custom formats) --
 
@@ -123,6 +119,8 @@ def _build_row(ctx: _SampleContext, row_fields: dict[str, Any]) -> dict[str, Any
             "text_content": row_fields.get("text_content"),
             "binary_content": row_fields.get("binary_content"),
             "source_ref": row_fields.get("source_ref"),
+            "source_member": row_fields.get("source_member"),
+            "source_frame_index": row_fields.get("source_frame_index"),
             "materialize_error": None,
         }
 
@@ -135,6 +133,7 @@ def _metadata_row(self, ctx: _SampleContext) -> dict[str, Any]:
                     "modality": "metadata",
                     "content_type": "application/json",
                     "source_ref": self._build_source_ref(ctx, ctx.json_member_name),
+                    "source_member": ctx.json_member_name,
                 },
             ),
             **ctx.passthrough,
@@ -188,6 +187,7 @@ def _text_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
                     "content_type": "text/plain",
                     "text_content": str(text_value),
                     "source_ref": source_ref,
+                    "source_member": ctx.json_member_name,
                 },
             )
             self._apply_per_modality_fields(row, ctx.per_text_passthrough, non_none_counter)
@@ -225,7 +225,9 @@ def _image_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
                     "position": idx,
                     "modality": "image",
                     "content_type": content_type or ("application/octet-stream" if image_member_name else None),
-                    "source_ref": self._build_source_ref(ctx, content_key, frame_index=frame_index),
+                    "source_ref": self._build_source_ref(ctx, content_key),
+                    "source_member": content_key,
+                    "source_frame_index": frame_index,
                 },
             )
             self._apply_per_modality_fields(row, ctx.per_image_passthrough, non_none_counter)
@@ -380,15 +382,14 @@ def _rows_from_member(
             for row in sample_rows:
                 if row["modality"] != "image" or row["position"] < 0:
                     continue
-                parsed_ref = InterleavedBatch.parse_source_ref(row["source_ref"])
-                content_key = parsed_ref.get("member")
+                content_key = row.get("source_member")
                 if not content_key:
                     continue
                 raw_bytes = self._extract_tar_member(tf, content_key, read_ctx.byte_cache)
                 if raw_bytes is None:
                     row["materialize_error"] = f"missing member '{content_key}'"
                 else:
-                    frame_index = parsed_ref.get("frame_index")
+                    frame_index = row.get("source_frame_index")
                     if frame_index is not None:
                         tiff_frame = _extract_tiff_frame(raw_bytes, frame_index)
                         if tiff_frame is None:

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

@@ -17,29 +17,32 @@
 import uuid
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
-import pandas as pd
 import pyarrow as pa
+import pyarrow.compute as pc
 from fsspec.core import url_to_fs
 from loguru import logger
 
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.interleaved.utils import materialize_task_binary_content
-from nemo_curator.stages.interleaved.utils.schema import align_table, reconcile_schema, resolve_schema
+from nemo_curator.stages.interleaved.utils.schema import align_interleaved_table, resolve_schema
 from nemo_curator.tasks import FileGroupTask, InterleavedBatch
 from nemo_curator.utils.client_utils import is_remote_url
 from nemo_curator.utils.file_utils import check_output_mode
 from nemo_curator.utils.hash_utils import get_deterministic_hash
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 
 @dataclass
 class BaseInterleavedWriter(ProcessingStage[InterleavedBatch, FileGroupTask], ABC):
     """Base class for interleaved writers.
 
     Handles filesystem setup, deterministic file naming, optional binary
     materialization, schema alignment, and process() orchestration.
-    Subclasses implement ``_write_dataframe`` for format-specific output.
+    Subclasses implement ``_write_table`` for format-specific output.
 
     If *schema* is set, every output table is aligned to it (missing columns
     become typed nulls, extra columns are dropped, types are reconciled).
@@ -111,33 +114,44 @@ def _materialize_dataframe(self, task: InterleavedBatch) -> pd.DataFrame:
             logger.info("materialize: dropped {} samples with errors", len(bad_samples))
         return out
 
-    def _align_output(self, df: pd.DataFrame) -> pd.DataFrame:
+    def _align_output(self, df: pd.DataFrame) -> pa.Table:
         """Reconcile or align *df* to the declared schema."""
         table = pa.Table.from_pandas(df, preserve_index=False)
-        if self.schema is not None:
-            table = align_table(table, self.schema)
-        else:
-            table = table.cast(reconcile_schema(table.schema))
-        return table.to_pandas(types_mapper=pd.ArrowDtype)
+        return align_interleaved_table(table, self.schema)
+
+    def _prepare_table(self, task: InterleavedBatch) -> pa.Table:
+        table = task.data
+        error_index = table.schema.get_field_index("materialize_error") if isinstance(table, pa.Table) else -1
+        if (
+            isinstance(table, pa.Table)
+            and not self.materialize_on_write
+            and (error_index < 0 or table.column(error_index).null_count == table.num_rows)
+        ):
+            table = align_interleaved_table(table, self.schema)
+            image_rows = pc.sum(pc.equal(table.column("modality"), "image")).as_py() or 0
+            self._log_metrics({"rows_out": float(table.num_rows), "image_rows": float(image_rows)})
+            if error_index >= 0:
+                self._log_metric("materialize_errors", 0.0)
+            return table
 
-    def _write_dataframe(self, df: pd.DataFrame, file_path: str, write_kwargs: dict[str, Any]) -> None:
-        """Format-specific DataFrame writer. Subclasses must implement this.
+        with self._time_metric("materialize_dataframe_total_s"):
+            return self._align_output(self._materialize_dataframe(task))
+
+    def _write_table(self, table: pa.Table, file_path: str, write_kwargs: dict[str, Any]) -> None:
+        """Format-specific Arrow table writer. Subclasses must implement this.
 
         Subclasses that override ``write_data()`` or ``process()`` directly
         (e.g. writers that do not follow the one-file-per-task pattern) may
         override this method as a no-op instead.
         """
         msg = (
-            f"{type(self).__name__} must override `_write_dataframe()`, or override "
+            f"{type(self).__name__} must override `_write_table()`, or override "
             "`write_data()` / `process()` so that output data is actually written."
         )
         raise NotImplementedError(msg)
 
     def write_data(self, task: InterleavedBatch, file_path: str) -> None:
-        with self._time_metric("materialize_dataframe_total_s"):
-            df = self._materialize_dataframe(task)
-        df = self._align_output(df)
-        self._write_dataframe(df, file_path, self._effective_write_kwargs)
+        self._write_table(self._prepare_table(task), file_path, self._effective_write_kwargs)
 
     def process(self, task: InterleavedBatch) -> FileGroupTask:
         if source_files := task._metadata.get("source_files"):

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

@@ -15,12 +15,21 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
 
 from .base import BaseInterleavedWriter
 
-if TYPE_CHECKING:
-    import pandas as pd
+_DICTIONARY_COLUMNS = [
+    "position",
+    "modality",
+    "content_type",
+    "source_ref.uri",
+    "source_ref.content_type",
+    "source_frame_index",
+]
 
 
 @dataclass
@@ -30,8 +39,10 @@ class InterleavedParquetWriterStage(BaseInterleavedWriter):
     file_extension: str = "parquet"
     name: str = "interleaved_parquet_writer"
 
-    def _write_dataframe(self, df: pd.DataFrame, file_path: str, write_kwargs: dict[str, Any]) -> None:
+    def _write_table(self, table: pa.Table, file_path: str, write_kwargs: dict[str, Any]) -> None:
         write_kwargs.setdefault("compression", "snappy")
         write_kwargs.setdefault("row_group_size", 128_000)
-        with self._time_metric("parquet_write_s"):
-            df.to_parquet(file_path, **write_kwargs)
+        write_kwargs.setdefault("use_dictionary", _DICTIONARY_COLUMNS)
+        write_kwargs.pop("index", None)
+        with self.fs.open(file_path, "wb") as fobj, self._time_metric("parquet_write_s"):
+            pq.write_table(table.replace_schema_metadata(), fobj, **write_kwargs)

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

@@ -20,7 +20,7 @@
 import urllib.parse
 from dataclasses import dataclass
 from io import BytesIO
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import fsspec
 import pandas as pd
@@ -29,6 +29,9 @@
 
 from .base import BaseInterleavedWriter
 
+if TYPE_CHECKING:
+    import pyarrow as pa
+
 # ---------------------------------------------------------------------------
 # Module-level helpers (importable in tests)
 # ---------------------------------------------------------------------------
@@ -197,7 +200,8 @@ class InterleavedWebdatasetWriterStage(BaseInterleavedWriter):
 
     _SUPPORTED_MODALITIES: ClassVar[frozenset[str]] = frozenset({"metadata", "text", "image"})
 
-    def _write_dataframe(self, df: pd.DataFrame, file_path: str, _write_kwargs: dict[str, Any]) -> None:
+    def _write_table(self, table: pa.Table, file_path: str, _write_kwargs: dict[str, Any]) -> None:
+        df = table.to_pandas(types_mapper=pd.ArrowDtype)
         unsupported = set(df["modality"].dropna().unique()) - self._SUPPORTED_MODALITIES
         if unsupported:
             msg = f"Unsupported modality {sorted(unsupported)!r}. Supported: {sorted(self._SUPPORTED_MODALITIES)}"

nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py

@@ -369,6 +369,7 @@ def build_interleaved_rows(  # noqa: PLR0913
             "source_ref": None,
             "url": url,
             "page_number": None,
+            "source_bbox": None,
             "pdf_name": pdf_name,
             "element_class": None,
         }
@@ -389,7 +390,6 @@ def build_interleaved_rows(  # noqa: PLR0913
         for elem in ordered:
             cls = elem["class"]
             bbox = elem.get("bbox")
-            source_ref = json.dumps({"page": page_num, "bbox": bbox})
 
             if cls == "Picture":
                 modality, content_type = "image", "image/png"
@@ -412,9 +412,10 @@ def build_interleaved_rows(  # noqa: PLR0913
                     "content_type": content_type,
                     "text_content": text,
                     "binary_content": binary,
-                    "source_ref": source_ref,
+                    "source_ref": None,
                     "url": url,
                     "page_number": page_num,
+                    "source_bbox": bbox,
                     "pdf_name": pdf_name,
                     "element_class": cls,
                 }