diff --git a/nemo_curator/stages/interleaved/README.md b/nemo_curator/stages/interleaved/README.md index d84298c20c..57e5d7e067 100644 --- a/nemo_curator/stages/interleaved/README.md +++ b/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`). diff --git a/nemo_curator/stages/interleaved/io/readers/webdataset.py b/nemo_curator/stages/interleaved/io/readers/webdataset.py index 5f4d8df6d5..ed296cd71d 100644 --- a/nemo_curator/stages/interleaved/io/readers/webdataset.py +++ b/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: diff --git a/nemo_curator/stages/interleaved/io/writers/base.py b/nemo_curator/stages/interleaved/io/writers/base.py index ff7ecfafa5..52e152f253 100644 --- a/nemo_curator/stages/interleaved/io/writers/base.py +++ b/nemo_curator/stages/interleaved/io/writers/base.py @@ -17,21 +17,24 @@ 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): @@ -39,7 +42,7 @@ class BaseInterleavedWriter(ProcessingStage[InterleavedBatch, FileGroupTask], AB 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"): diff --git a/nemo_curator/stages/interleaved/io/writers/tabular.py b/nemo_curator/stages/interleaved/io/writers/tabular.py index 754176ef3b..da6dcd127d 100644 --- a/nemo_curator/stages/interleaved/io/writers/tabular.py +++ b/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) diff --git a/nemo_curator/stages/interleaved/io/writers/webdataset.py b/nemo_curator/stages/interleaved/io/writers/webdataset.py index efc59ef6e8..5ff03a6742 100644 --- a/nemo_curator/stages/interleaved/io/writers/webdataset.py +++ b/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)}" diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py index 0d47eec33d..2f19c13ff7 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/utils.py +++ b/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, } diff --git a/nemo_curator/stages/interleaved/utils/materialization.py b/nemo_curator/stages/interleaved/utils/materialization.py index d4df54cecf..61a57ef999 100644 --- a/nemo_curator/stages/interleaved/utils/materialization.py +++ b/nemo_curator/stages/interleaved/utils/materialization.py @@ -31,13 +31,12 @@ _TAR_EXTENSIONS = (".tar", ".tar.gz", ".tgz") # Exact columns added by InterleavedBatch.with_parsed_source_ref_columns(prefix="_src_"). # Drop only these — not everything starting with "_src_" — to preserve user passthrough columns. -_SRC_PARSE_COLS = ("_src_path", "_src_member", "_src_byte_offset", "_src_byte_size", "_src_frame_index") +_SRC_PARSE_COLS = ("_src_uri", "_src_member", "_src_offset", "_src_size", "_src_frame_index") class _ClassifiedRows(NamedTuple): tar_extract: dict[str, list[tuple[int, str, int | None]]] - range_read: dict[str, list[tuple[int, str, int, int, int | None]]] - direct_read: dict[str, list[int]] + range_read: dict[str, list[tuple[int, str, int, int | None, int | None]]] missing: list[int] @@ -45,7 +44,7 @@ def _get_frame_index(df: pd.DataFrame, idx: int) -> int | None: if "_src_frame_index" not in df.columns: return None val = df.loc[idx, "_src_frame_index"] - if val is None or (isinstance(val, float) and pd.isna(val)): + if pd.isna(val): return None return int(val) @@ -54,44 +53,38 @@ def _classify_rows( df: pd.DataFrame, image_mask: pd.Series, ) -> _ClassifiedRows: - """Partition pending image rows into three I/O strategy groups. + """Partition pending image rows into two I/O strategy groups. - - tar_extract: has member name but no byte_offset (must open tar and extractfile) - - range_read: has member + byte_offset + byte_size (can use fs.cat_ranges) - - direct_read: no member (path is the file itself) - - missing: path is None/NaN + - tar_extract: has adjacent member metadata but no FILE range + - range_read: external FILE reference (whole file or byte range) + - missing: FILE uri is absent """ tar_extract: dict[str, list[tuple[int, str, int | None]]] = {} - range_read: dict[str, list[tuple[int, str, int, int, int | None]]] = {} - direct_read: dict[str, list[int]] = {} + range_read: dict[str, list[tuple[int, str, int, int | None, int | None]]] = {} missing: list[int] = [] for idx in df[image_mask].index: - path = df.loc[idx, "_src_path"] - if path is None or (isinstance(path, float) and pd.isna(path)) or path == "": + uri = df.loc[idx, "_src_uri"] + if uri is None or (isinstance(uri, float) and pd.isna(uri)) or uri == "": missing.append(idx) continue - path_str = str(path) + path_str = str(uri) raw_member = df.loc[idx, "_src_member"] - has_member = raw_member not in (None, "") and pd.notna(raw_member) - - if not has_member: - direct_read.setdefault(path_str, []).append(idx) - continue - - member_str = str(raw_member) + has_member = pd.notna(raw_member) and raw_member != "" frame_idx = _get_frame_index(df, idx) - raw_offset = df.loc[idx, "_src_byte_offset"] - raw_size = df.loc[idx, "_src_byte_size"] - has_range = raw_offset is not None and raw_size is not None and pd.notna(raw_offset) and pd.notna(raw_size) + raw_offset = df.loc[idx, "_src_offset"] + raw_size = df.loc[idx, "_src_size"] + offset = 0 if pd.isna(raw_offset) else int(raw_offset) + size = None if pd.isna(raw_size) else int(raw_size) - if has_range and int(raw_size) > 0: - range_read.setdefault(path_str, []).append((idx, member_str, int(raw_offset), int(raw_size), frame_idx)) - else: - tar_extract.setdefault(path_str, []).append((idx, member_str, frame_idx)) + if size is not None or not has_member: + label = str(raw_member) if has_member else path_str + range_read.setdefault(path_str, []).append((idx, label, offset, size, frame_idx)) + elif has_member: + tar_extract.setdefault(path_str, []).append((idx, str(raw_member), frame_idx)) - return _ClassifiedRows(tar_extract=tar_extract, range_read=range_read, direct_read=direct_read, missing=missing) + return _ClassifiedRows(tar_extract=tar_extract, range_read=range_read, missing=missing) def _extract_tiff_frame(tiff_bytes: bytes, frame_index: int) -> bytes | None: @@ -156,8 +149,8 @@ def _fill_tar_extract_rows( def _scatter_range_blobs( blobs: list[object], - range_keys: list[tuple[str, int, int]], - unique_ranges: dict[tuple[str, int, int], list[tuple[int, str, int | None]]], + range_keys: list[tuple[str, int, int | None]], + unique_ranges: dict[tuple[str, int, int | None], list[tuple[int, str, int | None]]], binary_values: list[object], error_values: list[str | None], ) -> None: @@ -166,7 +159,7 @@ def _scatter_range_blobs( if isinstance(blob, Exception): for idx, member, _fi in unique_ranges[key]: error_values[idx] = f"range read error for member '{member}'" - elif blob is None or len(blob) == 0: + elif blob is None: for idx, member, _fi in unique_ranges[key]: error_values[idx] = f"empty range read for member '{member}'" else: @@ -181,10 +174,10 @@ def _scatter_range_blobs( def _build_global_range_index( - groups: dict[str, list[tuple[int, str, int, int, int | None]]], + groups: dict[str, list[tuple[int, str, int, int | None, int | None]]], storage_options: dict[str, object], error_values: list[str | None], -) -> list[tuple[object, dict[tuple[str, int, int], list[tuple[int, str, int | None]]]]]: +) -> list[tuple[object, dict[tuple[str, int, int | None], list[tuple[int, str, int | None]]]]]: """Resolve filesystem paths and build per-filesystem deduplicated range indices. Returns a list of ``(fs, unique_ranges)`` pairs — one entry per distinct @@ -201,12 +194,12 @@ def _build_global_range_index( # id(fs) -> (fs, unique_ranges); fsspec caches fs instances so same-backend # paths naturally share the same id. - fs_groups: dict[int, tuple[object, dict[tuple[str, int, int], list[tuple[int, str, int | None]]]]] = {} + fs_groups: dict[int, tuple[object, dict[tuple[str, int, int | None], list[tuple[int, str, int | None]]]]] = {} for path, entries in groups.items(): try: fs, fs_path = url_to_fs(path, **storage_options) - except (ValueError, OSError) as exc: + except (OSError, RuntimeError, ValueError) as exc: logger.warning("Failed to resolve filesystem for path {!r}: {}", path, exc) for idx, *_ in entries: error_values[idx] = "failed to resolve filesystem" @@ -224,7 +217,7 @@ def _build_global_range_index( def _fill_range_read_rows( - groups: dict[str, list[tuple[int, str, int, int, int | None]]], + groups: dict[str, list[tuple[int, str, int, int | None, int | None]]], storage_options: dict[str, object], binary_values: list[object], error_values: list[str | None], @@ -243,7 +236,7 @@ def _fill_range_read_rows( range_keys = list(unique_ranges.keys()) cat_paths = [fp for fp, _, _ in range_keys] cat_starts = [off for _, off, _ in range_keys] - cat_ends = [off + sz for _, off, sz in range_keys] + cat_ends = [None if sz is None else off + sz for _, off, sz in range_keys] try: blobs = fs.cat_ranges(cat_paths, cat_starts, cat_ends) @@ -257,31 +250,6 @@ def _fill_range_read_rows( _scatter_range_blobs(blobs, range_keys, unique_ranges, binary_values, error_values) -def _fill_direct_read_rows( - groups: dict[str, list[int]], - storage_options: dict[str, object], - binary_values: list[object], - error_values: list[str | None], -) -> None: - """Read each direct file once, share bytes across all rows referencing it.""" - for path, row_idxs in groups.items(): - payload = _read_direct_file(path, storage_options) - for idx in row_idxs: - if payload is not None: - binary_values[idx] = payload - error_values[idx] = None - else: - error_values[idx] = "failed to read path" - - -def _read_direct_file(path: str, storage_options: dict[str, object]) -> bytes | None: - try: - with fsspec.open(path, mode="rb", **storage_options) as fobj: - return fobj.read() - except (OSError, RuntimeError, ValueError): - return None - - def _fill_materialized_bytes( df: pd.DataFrame, image_mask: pd.Series, @@ -297,7 +265,6 @@ def _fill_materialized_bytes( _fill_tar_extract_rows(classified.tar_extract, storage_options, binary_values, error_values) _fill_range_read_rows(classified.range_read, storage_options, binary_values, error_values) - _fill_direct_read_rows(classified.direct_read, storage_options, binary_values, error_values) def _init_materialization_buffers(df: pd.DataFrame) -> tuple[list[object], list[str | None]]: @@ -344,10 +311,9 @@ def materialize_task_binary_content( ) -> InterleavedBatch: """Return a task with image-row binary content materialized from source_ref. - Dispatches to three I/O strategies based on source_ref contents: - - range_read: byte_offset + byte_size present -> batched fs.cat_ranges() - - tar_extract: member present, no byte range -> open tar + extractfile - - direct_read: no member -> read file directly + Dispatches to two I/O strategies based on source_ref contents: + - range_read: external FILE reference -> deduplicated fs.cat_ranges() + - tar_extract: adjacent member present, no byte range -> open tar + extractfile """ df = task.with_parsed_source_ref_columns(prefix="_src_").reset_index(drop=True) if df.empty: diff --git a/nemo_curator/stages/interleaved/utils/schema.py b/nemo_curator/stages/interleaved/utils/schema.py index 08b2981d5e..e50492ea55 100644 --- a/nemo_curator/stages/interleaved/utils/schema.py +++ b/nemo_curator/stages/interleaved/utils/schema.py @@ -20,10 +20,13 @@ from __future__ import annotations +import json + import pyarrow as pa from loguru import logger from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA, RESERVED_COLUMNS +from nemo_curator.utils.storage_utils import FILE_REFERENCE_TYPE _LARGE_COMPAT: dict[tuple[pa.DataType, pa.DataType], pa.DataType] = { (pa.large_string(), pa.string()): pa.large_string(), @@ -31,6 +34,10 @@ } +def _parse_legacy_json(column: pa.ChunkedArray) -> list[dict[str, object]]: + return [json.loads(value) if value else {} for value in column.to_pylist()] + + def reconcile_schema(inferred: pa.Schema) -> pa.Schema: """Build a schema with canonical types for reserved columns and inferred types for passthrough. @@ -117,6 +124,31 @@ def align_interleaved_table(table: pa.Table, schema: pa.Schema | None = None) -> types while passthrough columns are preserved. With an explicit schema, the table is padded, reordered, and cast exactly to that schema. """ + source_ref_index = table.schema.get_field_index("source_ref") + source_ref = table.column(source_ref_index) if source_ref_index >= 0 else None + if source_ref is not None and source_ref.type in (pa.string(), pa.large_string()): + refs = _parse_legacy_json(source_ref) + for ref in refs: + if (path := ref.get("path")) is not None: + ref.update( + uri=str(path), + offset=int(ref["byte_offset"]) if ref.get("byte_offset") is not None else None, + size=int(ref["byte_size"]) if ref.get("byte_size") is not None else None, + ) + file_refs = (ref if ref.get("path") is not None else None for ref in refs) + table = table.set_column(source_ref_index, "source_ref", pa.array(file_refs, type=FILE_REFERENCE_TYPE)) + for name, key, dtype in ( + ("source_member", "member", pa.string()), + ("source_frame_index", "frame_index", pa.int32()), + ("page_number", "page", pa.int64()), + ("source_bbox", "bbox", pa.list_(pa.float64())), + ): + if name not in table.column_names: + table = table.append_column(name, pa.array((ref.get(key) for ref in refs), type=dtype)) + elif source_ref is not None and pa.types.is_struct(source_ref.type) and source_ref.type != FILE_REFERENCE_TYPE: + file_refs = pa.array(source_ref.to_pylist(), type=FILE_REFERENCE_TYPE) + table = table.set_column(source_ref_index, "source_ref", file_refs) if schema is not None: return align_table(table, schema) - return table.cast(reconcile_schema(table.schema)) + reconciled = reconcile_schema(table.schema) + return table if table.schema == reconciled else table.cast(reconciled) diff --git a/nemo_curator/tasks/interleaved.py b/nemo_curator/tasks/interleaved.py index bea6f93b80..c173c47870 100644 --- a/nemo_curator/tasks/interleaved.py +++ b/nemo_curator/tasks/interleaved.py @@ -29,12 +29,9 @@ ``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 + ``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 ================== ============= =========== =============================================== @@ -42,7 +39,6 @@ ``fields`` parameter on the reader. These flow through the pipeline untouched. """ -import json from dataclasses import dataclass, field import pandas as pd @@ -50,6 +46,8 @@ import pyarrow.compute as pc from loguru import logger +from nemo_curator.utils.storage_utils import FILE_REFERENCE_TYPE + from .tasks import Task INTERLEAVED_SCHEMA = pa.schema( @@ -60,7 +58,9 @@ pa.field("content_type", pa.string(), nullable=True), pa.field("text_content", pa.string(), nullable=True), pa.field("binary_content", pa.large_binary(), nullable=True), - pa.field("source_ref", pa.string(), nullable=True), + pa.field("source_ref", FILE_REFERENCE_TYPE, nullable=True), + pa.field("source_member", pa.string(), nullable=True), + pa.field("source_frame_index", pa.int32(), nullable=True), pa.field("materialize_error", pa.string(), nullable=True), ] ) @@ -174,59 +174,36 @@ def delete_rows(self, mask: pd.Series) -> "InterleavedBatch": @staticmethod def build_source_ref( - path: str | None, - member: str | None, - byte_offset: int | None = None, - byte_size: int | None = None, - frame_index: int | None = None, - ) -> str: - """Build a ``source_ref`` JSON locator string.""" - ref: dict[str, object] = { - "path": path, - "member": member, - "byte_offset": byte_offset, - "byte_size": byte_size, - } - if frame_index is not None: - ref["frame_index"] = frame_index - return json.dumps(ref, ensure_ascii=True) - - @staticmethod - def parse_source_ref(source_value: str | None) -> dict[str, str | int | None]: - """Parse a ``source_ref`` JSON string into a locator dict.""" - if source_value is None or pd.isna(source_value) or source_value == "": - return {"path": None, "member": None, "byte_offset": None, "byte_size": None, "frame_index": None} - parsed = json.loads(source_value) - if not isinstance(parsed, dict): - msg = "source_ref must decode to a JSON object" - raise TypeError(msg) - - path = parsed.get("path") - member = parsed.get("member") - byte_offset = parsed.get("byte_offset") - byte_size = parsed.get("byte_size") - frame_index = parsed.get("frame_index") - - return { - "path": path if path is None else str(path), - "member": member if member is None else str(member), - "byte_offset": int(byte_offset) if byte_offset is not None else None, - "byte_size": int(byte_size) if byte_size is not None else None, - "frame_index": int(frame_index) if frame_index is not None else None, - } + uri: str | None, + offset: int | None = None, + size: int | None = None, + ) -> dict[str, object] | None: + """Build a Parquet FILE-compatible external reference.""" + if (offset is not None and size is None) or any(value is not None and value < 0 for value in (offset, size)): + msg = "source_ref offset and size must be non-negative, with size set for offset" + raise ValueError(msg) + return ( + { + "uri": uri, + "offset": offset, + "size": size, + "content_type": None, + "checksum": None, + "inline": None, + } + if uri + else None + ) def with_parsed_source_ref_columns(self, prefix: str = "_src_") -> pd.DataFrame: - """Return a DataFrame copy with parsed ``source_ref`` columns added. - - Columns: ``{prefix}path``, ``{prefix}member``, ``{prefix}byte_offset``, - ``{prefix}byte_size``, ``{prefix}frame_index``. - """ + """Return a DataFrame with FILE locator and adjacent source columns expanded.""" df = self.to_pandas().copy() - parsed = [self.parse_source_ref(value) for value in df["source_ref"].tolist()] - parsed_df = pd.DataFrame.from_records( - parsed, - columns=["path", "member", "byte_offset", "byte_size", "frame_index"], + parsed = pd.DataFrame.from_records( + (value if isinstance(value, dict) else {} for value in df["source_ref"]), + columns=["uri", "offset", "size"], ) - for col in parsed_df.columns: - df[f"{prefix}{col}"] = parsed_df[col].to_numpy(copy=False) + for col in parsed: + df[f"{prefix}{col}"] = parsed[col].to_numpy(copy=False) + df[f"{prefix}member"] = df.get("source_member") + df[f"{prefix}frame_index"] = df.get("source_frame_index") return df diff --git a/nemo_curator/utils/storage_utils.py b/nemo_curator/utils/storage_utils.py index b03e7ad15c..192eb0a430 100644 --- a/nemo_curator/utils/storage_utils.py +++ b/nemo_curator/utils/storage_utils.py @@ -14,6 +14,19 @@ import pathlib +import pyarrow as pa + +FILE_REFERENCE_TYPE = pa.struct( + [ + pa.field("uri", pa.string()), + pa.field("offset", pa.int64()), + pa.field("size", pa.int64()), + pa.field("content_type", pa.string()), + pa.field("checksum", pa.string()), + pa.field("inline", pa.binary()), + ] +) + def _get_local_path(localpath: pathlib.Path, *args: str) -> pathlib.Path: """Construct a full local path from a base path and additional components. diff --git a/tests/stages/interleaved/conftest.py b/tests/stages/interleaved/conftest.py index 8ac0d14d4a..d4ebc84322 100644 --- a/tests/stages/interleaved/conftest.py +++ b/tests/stages/interleaved/conftest.py @@ -206,12 +206,8 @@ def make_image_row( "content_type": content_type, "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref( - path=path, - member=member, - byte_offset=byte_offset, - byte_size=byte_size, - ), + "source_ref": InterleavedBatch.build_source_ref(path, byte_offset, byte_size), + "source_member": member, "materialize_error": None, } @@ -311,9 +307,8 @@ def single_row_table() -> pa.Table: position=0, modality="text", text_content="hello", - source_ref=json.dumps( - {"path": "/dataset/shard.tar", "member": "s1.json", "byte_offset": 10, "byte_size": 20} - ), + source_ref=InterleavedBatch.build_source_ref("/dataset/shard.tar", 10, 20), + source_member="s1.json", ) ], schema=INTERLEAVED_SCHEMA, diff --git a/tests/stages/interleaved/pdf/nemotron_parse/test_utils.py b/tests/stages/interleaved/pdf/nemotron_parse/test_utils.py index 0f8efbf5d2..2e1ebea04d 100644 --- a/tests/stages/interleaved/pdf/nemotron_parse/test_utils.py +++ b/tests/stages/interleaved/pdf/nemotron_parse/test_utils.py @@ -158,6 +158,9 @@ def test_basic_output(self): text_rows = [r for r in rows if r["modality"] == "text"] assert len(text_rows) == 1 assert text_rows[0]["text_content"] == "Hello" + assert text_rows[0]["source_ref"] is None + assert text_rows[0]["page_number"] == 0 + assert text_rows[0]["source_bbox"] == [0.0, 0.0, 1.0, 1.0] def test_empty_pages(self): rows = build_interleaved_rows("s1", "http://example.com", "test.pdf", [], []) diff --git a/tests/stages/interleaved/test_base_writer.py b/tests/stages/interleaved/test_base_writer.py index 0e2f96dd79..d23e9ae6f9 100644 --- a/tests/stages/interleaved/test_base_writer.py +++ b/tests/stages/interleaved/test_base_writer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for BaseInterleavedWriter: on_materialize_error modes, _write_dataframe, +"""Tests for BaseInterleavedWriter: on_materialize_error modes, _write_table, and _align_output schema alignment.""" from dataclasses import dataclass @@ -120,12 +120,12 @@ def test_all_null_materialize_error_returns_all_rows(tmp_path: Path) -> None: # --------------------------------------------------------------------------- -# _write_dataframe: NotImplementedError on abstract base +# _write_table: NotImplementedError on abstract base # --------------------------------------------------------------------------- -def test_base_write_dataframe_raises_not_implemented(tmp_path: Path) -> None: - """BaseInterleavedWriter._write_dataframe() raises NotImplementedError.""" +def test_base_write_table_raises_not_implemented(tmp_path: Path) -> None: + """BaseInterleavedWriter._write_table() raises NotImplementedError.""" @dataclass class _StubWriter(BaseInterleavedWriter): @@ -133,8 +133,8 @@ class _StubWriter(BaseInterleavedWriter): name: str = "stub_writer" writer = _StubWriter(path=str(tmp_path / "out"), mode="overwrite") - with pytest.raises(NotImplementedError, match="_write_dataframe"): - writer._write_dataframe(pd.DataFrame(), "dummy.stub", {}) + with pytest.raises(NotImplementedError, match="_write_table"): + writer._write_table(pa.table({}), "dummy.stub", {}) # --------------------------------------------------------------------------- @@ -159,8 +159,8 @@ def test_align_output_with_schema_drops_extra_columns(tmp_path: Path) -> None: ) df = pd.DataFrame([{"sample_id": "s1", "position": 0, "modality": "text", "content_type": "text/plain"}]) result = writer._align_output(df) - assert "content_type" not in result.columns - assert "sample_id" in result.columns + assert "content_type" not in result.column_names + assert "sample_id" in result.column_names def test_align_output_without_schema_preserves_extra_columns(tmp_path: Path) -> None: @@ -173,8 +173,8 @@ def test_align_output_without_schema_preserves_extra_columns(tmp_path: Path) -> ) df = pd.DataFrame([{"sample_id": "s1", "position": 0, "modality": "text", "my_custom_column": "keep_me"}]) result = writer._align_output(df) - assert "my_custom_column" in result.columns - assert result["my_custom_column"].iloc[0] == "keep_me" + assert "my_custom_column" in result.column_names + assert result["my_custom_column"][0].as_py() == "keep_me" # --------------------------------------------------------------------------- diff --git a/tests/stages/interleaved/test_interleaved_task.py b/tests/stages/interleaved/test_interleaved_task.py index 90b5e37e58..425bcd49c4 100644 --- a/tests/stages/interleaved/test_interleaved_task.py +++ b/tests/stages/interleaved/test_interleaved_task.py @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json - import pandas as pd import pyarrow as pa import pytest from nemo_curator.tasks import InterleavedBatch from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA +from nemo_curator.utils.storage_utils import FILE_REFERENCE_TYPE _SAMPLE_ROW = { "sample_id": "s1", @@ -128,43 +127,27 @@ def test_add_rows_and_delete_rows_not_implemented() -> None: task.delete_rows(pd.Series([True])) -# --- parse_source_ref edge cases --- - +# --- source_ref --- -def test_parse_source_ref_non_dict_raises() -> None: - with pytest.raises(TypeError, match="source_ref must decode to a JSON object"): - InterleavedBatch.parse_source_ref("[1, 2]") +def test_source_ref_uses_exact_file_children() -> None: + ref = InterleavedBatch.build_source_ref("/a.tar", offset=10, size=20) + assert FILE_REFERENCE_TYPE.names == ["uri", "offset", "size", "content_type", "checksum", "inline"] + assert ref == {**dict.fromkeys(FILE_REFERENCE_TYPE.names), "uri": "/a.tar", "offset": 10, "size": 20} -def test_parse_source_ref_with_frame_index() -> None: - ref = json.dumps({"path": "/a.tar", "member": "m.jpg", "byte_offset": 10, "byte_size": 20, "frame_index": 5}) - parsed = InterleavedBatch.parse_source_ref(ref) - assert parsed["path"] == "/a.tar" - assert parsed["member"] == "m.jpg" - assert parsed["byte_offset"] == 10 - assert parsed["byte_size"] == 20 - assert parsed["frame_index"] == 5 +@pytest.mark.parametrize(("offset", "size"), [(-1, 1), (1, -1), (1, None)]) +def test_build_source_ref_rejects_invalid_ranges(offset: int, size: int | None) -> None: + with pytest.raises(ValueError, match="source_ref"): + InterleavedBatch.build_source_ref(uri="/a", offset=offset, size=size) -# --- build_source_ref --- - -@pytest.mark.parametrize( - ("frame_index", "key_present"), - [ - pytest.param(3, True, id="with_frame_index"), - pytest.param(None, False, id="without_frame_index"), - ], -) -def test_build_source_ref_frame_index(frame_index: int | None, key_present: bool) -> None: - ref_str = InterleavedBatch.build_source_ref( - path="/a.tar", - member="m.jpg", - byte_offset=10, - byte_size=20, - frame_index=frame_index, - ) - parsed = json.loads(ref_str) - assert ("frame_index" in parsed) is key_present - if key_present: - assert parsed["frame_index"] == frame_index +def test_source_metadata_is_adjacent_to_file() -> None: + row = { + **_SAMPLE_ROW, + "source_ref": InterleavedBatch.build_source_ref(uri="/a.tar", offset=10, size=20), + "source_member": "m.tiff", + "source_frame_index": 5, + } + parsed = _make_batch(pd.DataFrame([row])).with_parsed_source_ref_columns() + assert (parsed.loc[0, "_src_member"], parsed.loc[0, "_src_frame_index"]) == ("m.tiff", 5) diff --git a/tests/stages/interleaved/test_materialization.py b/tests/stages/interleaved/test_materialization.py index 420172a1e8..b7d50833ed 100644 --- a/tests/stages/interleaved/test_materialization.py +++ b/tests/stages/interleaved/test_materialization.py @@ -47,6 +47,7 @@ [ pytest.param(None, None, id="none_value"), pytest.param(float("nan"), None, id="nan_value"), + pytest.param(pd.NA, None, id="pd_na_value"), ], ) def test_get_frame_index_returns_none_for_missing_values(val: object, expected: None) -> None: @@ -64,10 +65,10 @@ def test_get_frame_index_returns_none_for_missing_values(val: object, expected: def test_classify_rows_missing_path_variants(path_val: object) -> None: df = pd.DataFrame( { - "_src_path": [path_val], + "_src_uri": [path_val], "_src_member": [None], - "_src_byte_offset": [None], - "_src_byte_size": [None], + "_src_offset": [None], + "_src_size": [None], } ) result = _classify_rows(df, pd.Series([True])) @@ -77,15 +78,25 @@ def test_classify_rows_missing_path_variants(path_val: object) -> None: def test_classify_rows_range_with_zero_size() -> None: df = pd.DataFrame( { - "_src_path": ["/shard.tar"], + "_src_uri": ["/shard.tar"], "_src_member": ["img.jpg"], - "_src_byte_offset": [100], - "_src_byte_size": [0], + "_src_offset": [100], + "_src_size": [0], } ) result = _classify_rows(df, pd.Series([True])) - assert "/shard.tar" in result.tar_extract - assert not result.range_read + assert "/shard.tar" in result.range_read + assert not result.tar_extract + + +def test_materialize_size_without_offset_reads_prefix(tmp_path: Path) -> None: + path = tmp_path / "image.bin" + path.write_bytes(b"abcTRAILING") + task = make_image_task([make_image_row(path=str(path), byte_size=3)]) + + result = materialize_task_binary_content(task).to_pandas() + + assert result.loc[0, "binary_content"] == b"abc" # --- _extract_tiff_frame --- @@ -194,7 +205,6 @@ def _make_range_setup( [ pytest.param(RuntimeError("fail"), "range read error", id="exception_blob"), pytest.param(None, "empty range read", id="none_blob"), - pytest.param(b"", "empty range read", id="empty_blob"), ], ) def test_scatter_range_blobs_error_cases(blob: object, expected_error_substr: str) -> None: @@ -204,6 +214,13 @@ def test_scatter_range_blobs_error_cases(blob: object, expected_error_substr: st assert expected_error_substr in error_values[0] +def test_scatter_range_blobs_accepts_empty_range() -> None: + range_keys, unique_ranges, binary_values, error_values = _make_range_setup("empty", 0, 0) + _scatter_range_blobs([b""], range_keys, unique_ranges, binary_values, error_values) + assert binary_values[0] == b"" + assert error_values[0] is None + + def test_scatter_range_blobs_bytearray_conversion() -> None: range_keys, unique_ranges, binary_values, error_values = _make_range_setup("img.jpg", 0, 10) _scatter_range_blobs([bytearray(b"image-data")], range_keys, unique_ranges, binary_values, error_values) @@ -404,7 +421,7 @@ def test_materialize_with_only_missing_binary_false(tmp_path: Path) -> None: "content_type": "image/jpeg", "text_content": None, "binary_content": b"old-bytes", - "source_ref": InterleavedBatch.build_source_ref(path=str(img_path), member=None), + "source_ref": InterleavedBatch.build_source_ref(str(img_path)), "materialize_error": None, } ] @@ -434,7 +451,7 @@ def test_materialize_preserves_passthrough_columns_with_src_prefix(tmp_path: Pat "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path=str(img_path), member=None), + "source_ref": InterleavedBatch.build_source_ref(str(img_path)), "materialize_error": None, "_src_html": "keep-me", "_src_metadata": "also-keep-me", diff --git a/tests/stages/interleaved/test_multimodal_core.py b/tests/stages/interleaved/test_multimodal_core.py index 340b4f78aa..29942acf80 100644 --- a/tests/stages/interleaved/test_multimodal_core.py +++ b/tests/stages/interleaved/test_multimodal_core.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json +import tarfile from io import BytesIO from pathlib import Path from unittest.mock import MagicMock, patch @@ -30,7 +30,6 @@ ) from nemo_curator.stages.interleaved.utils.materialization import ( _classify_rows, - _read_direct_file, materialize_task_binary_content, ) from nemo_curator.tasks import InterleavedBatch @@ -41,24 +40,10 @@ def test_with_parsed_source_ref_columns(single_row_task: InterleavedBatch) -> None: df = single_row_task.with_parsed_source_ref_columns() - assert df.loc[0, "_src_path"] == "/dataset/shard.tar" + assert df.loc[0, "_src_uri"] == "/dataset/shard.tar" assert df.loc[0, "_src_member"] == "s1.json" - assert df.loc[0, "_src_byte_offset"] == 10 - assert df.loc[0, "_src_byte_size"] == 20 - - -def test_parse_source_ref_ignores_legacy_keys() -> None: - legacy_format = json.dumps({"content_path": "/old/path.tar", "content_key": "old.json"}) - parsed = InterleavedBatch.parse_source_ref(legacy_format) - assert parsed["path"] is None - assert parsed["member"] is None - assert parsed["byte_offset"] is None - assert parsed["byte_size"] is None - - -def test_parse_source_ref_empty_values() -> None: - assert InterleavedBatch.parse_source_ref(None)["path"] is None - assert InterleavedBatch.parse_source_ref("")["path"] is None + assert df.loc[0, "_src_offset"] == 10 + assert df.loc[0, "_src_size"] == 20 # --- classify_rows tests --- @@ -67,9 +52,10 @@ def test_parse_source_ref_empty_values() -> None: @pytest.mark.parametrize( ("src_path", "src_member", "byte_range", "expected_bucket", "expected_missing"), [ - pytest.param("/img.jpg", None, None, "direct_read", [], id="direct_read"), + pytest.param("/img.jpg", None, None, "range_read", [], id="whole_file"), pytest.param("/shard.tar", "img.jpg", None, "tar_extract", [], id="tar_extract"), pytest.param("/shard.tar", "img.jpg", (512, 1024), "range_read", [], id="range_read"), + pytest.param("/shard.tar", None, (512, 1024), "range_read", [], id="range_read_without_member"), pytest.param(None, None, None, None, [0], id="missing_path"), ], ) @@ -83,40 +69,46 @@ def test_classify_rows( byte_offset, byte_size = byte_range if byte_range else (None, None) df = pd.DataFrame( { - "_src_path": [src_path], + "_src_uri": [src_path], "_src_member": [src_member], - "_src_byte_offset": [byte_offset], - "_src_byte_size": [byte_size], + "_src_offset": [byte_offset], + "_src_size": [byte_size], } ) result = _classify_rows(df, pd.Series([True])) assert result.missing == expected_missing if expected_bucket is not None: assert src_path in getattr(result, expected_bucket) - for other in {"direct_read", "tar_extract", "range_read"} - {expected_bucket}: + for other in {"tar_extract", "range_read"} - {expected_bucket}: assert not getattr(result, other) if expected_bucket == "range_read": - assert result.range_read[src_path][0] == (0, src_member, byte_offset, byte_size, None) + assert result.range_read[src_path][0] == ( + 0, + src_member or src_path, + 0 if byte_offset is None else byte_offset, + byte_size, + None, + ) def test_classify_rows_mixed_batch() -> None: df = pd.DataFrame( { - "_src_path": ["/img.jpg", "/shard.tar", "/shard.tar", None], + "_src_uri": ["/img.jpg", "/shard.tar", "/shard.tar", None], "_src_member": [None, "a.jpg", "b.jpg", None], - "_src_byte_offset": [None, None, 100, None], - "_src_byte_size": [None, None, 200, None], + "_src_offset": [None, None, 100, None], + "_src_size": [None, None, 200, None], } ) mask = pd.Series([True, True, True, True]) result = _classify_rows(df, mask) - assert len(result.direct_read["/img.jpg"]) == 1 + assert len(result.range_read["/img.jpg"]) == 1 assert len(result.tar_extract["/shard.tar"]) == 1 assert len(result.range_read["/shard.tar"]) == 1 assert result.missing == [3] -# --- materialize: direct read --- +# --- materialize: whole-file read --- def test_materialize_fills_binary_from_direct_path(tmp_path: Path) -> None: @@ -158,14 +150,13 @@ def test_materialize_tar_extract_missing_member(tmp_path: Path) -> None: # --- materialize: range read (with byte_offset/byte_size) --- -def test_materialize_fills_binary_from_range_read(tmp_path: Path) -> None: +def test_materialize_fills_binary_from_tar_range_without_member(tmp_path: Path) -> None: payload = b"range-read-image-bytes" - raw_file = tmp_path / "data.bin" - raw_file.write_bytes(b"HEADER" + payload + b"FOOTER") + tar_path = write_tar(tmp_path / "shard.tar", {"image.jpg": payload}) + with tarfile.open(tar_path) as tf: + info = tf.getmember("image.jpg") - task = make_image_task( - [make_image_row(path=str(raw_file), member="data.bin", byte_offset=6, byte_size=len(payload))] - ) + task = make_image_task([make_image_row(path=tar_path, byte_offset=info.offset_data, byte_size=info.size)]) result = materialize_task_binary_content(task) df = result.to_pandas() assert df.loc[0, "binary_content"] == payload @@ -756,23 +747,9 @@ def test_webdataset_reader_composite_decompose(tmp_path: Path) -> None: # --- exception broadening in materialization --- -def test_read_direct_file_handles_non_oserror_exceptions() -> None: - """_read_direct_file must gracefully return None for non-OSError exceptions - (e.g. RuntimeError from fsspec plugins) instead of crashing. - """ - with patch( - "nemo_curator.stages.interleaved.utils.materialization.fsspec.open", side_effect=RuntimeError("plugin error") - ): - result = _read_direct_file("/some/path.jpg", {}) - assert result is None - - -def test_materialize_records_error_for_non_oserror_on_direct_read() -> None: - """Non-OSError exceptions during direct-read materialization must be - recorded as materialize_error, not crash the pipeline. - """ +def test_materialize_records_error_for_non_oserror_on_file_read() -> None: task = make_image_task([make_image_row(path="/fake/path.jpg")]) - with patch("nemo_curator.stages.interleaved.utils.materialization.fsspec.open", side_effect=RuntimeError("boom")): + with patch("nemo_curator.stages.interleaved.utils.materialization.url_to_fs", side_effect=RuntimeError("boom")): result = materialize_task_binary_content(task) df = result.to_pandas() assert isinstance(df.loc[0, "materialize_error"], str) @@ -820,7 +797,7 @@ def test_iter_materialized_bytes_only_yields_masked_rows(tmp_path: Path) -> None "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path=str(file_a), member=None), + "source_ref": InterleavedBatch.build_source_ref(str(file_a)), "materialize_error": None, }, { @@ -830,7 +807,7 @@ def test_iter_materialized_bytes_only_yields_masked_rows(tmp_path: Path) -> None "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path=str(file_b), member=None), + "source_ref": InterleavedBatch.build_source_ref(str(file_b)), "materialize_error": None, }, ] @@ -868,7 +845,7 @@ def test_iter_materialized_bytes_preserves_original_indices(tmp_path: Path) -> N "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path=str(img_path), member=None), + "source_ref": InterleavedBatch.build_source_ref(str(img_path)), "materialize_error": None, }, ] @@ -905,11 +882,9 @@ def test_materialize_extracts_individual_tiff_frames(tmp_path: Path) -> None: "content_type": "image/tiff", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref( - path=tar_path, - member="doc.tiff", - frame_index=i, - ), + "source_ref": InterleavedBatch.build_source_ref(tar_path), + "source_member": "doc.tiff", + "source_frame_index": i, "materialize_error": None, } ) diff --git a/tests/stages/interleaved/test_multimodal_reader.py b/tests/stages/interleaved/test_multimodal_reader.py index 4484d9e7d9..bc5f75ae95 100644 --- a/tests/stages/interleaved/test_multimodal_reader.py +++ b/tests/stages/interleaved/test_multimodal_reader.py @@ -255,13 +255,8 @@ def test_reader_image_tokens_with_frame_index(tmp_path: Path) -> None: assert image_rows.iloc[0]["position"] == 1, "First non-None image at interleaved position 1" assert image_rows.iloc[1]["position"] == 2, "Second non-None image at interleaved position 2" - refs = [InterleavedBatch.parse_source_ref(v) for v in image_rows["source_ref"].tolist()] - - assert refs[0]["member"] == "doc.pdf.tiff", "Non-matching string should resolve to default TIFF" - assert refs[0]["frame_index"] == 0, "First non-None token gets frame_index=0" - - assert refs[1]["member"] == "doc.pdf.tiff" - assert refs[1]["frame_index"] == 1, "Second non-None token gets frame_index=1" + assert image_rows["source_member"].tolist() == ["doc.pdf.tiff", "doc.pdf.tiff"] + assert image_rows["source_frame_index"].tolist() == [0, 1] text_rows = df[df["modality"] == "text"] assert len(text_rows) == 3 @@ -598,15 +593,8 @@ def test_reader_frame_counter_resets_per_content_key(tmp_path: Path) -> None: image_rows = df[df["modality"] == "image"].sort_values("position") assert len(image_rows) == 4 - refs = [InterleavedBatch.parse_source_ref(v) for v in image_rows["source_ref"].tolist()] - assert refs[0]["member"] == "a.tiff" - assert refs[0]["frame_index"] == 0 - assert refs[1]["member"] == "a.tiff" - assert refs[1]["frame_index"] == 1 - assert refs[2]["member"] == "b.tiff" - assert refs[2]["frame_index"] == 0, "frame_index must reset to 0 for a different TIFF file" - assert refs[3]["member"] == "b.tiff" - assert refs[3]["frame_index"] == 1 + assert image_rows["source_member"].tolist() == ["a.tiff", "a.tiff", "b.tiff", "b.tiff"] + assert image_rows["source_frame_index"].tolist() == [0, 1, 0, 1] def test_reader_materialize_preserves_raw_bytes_on_frame_extraction_failure(tmp_path: Path) -> None: @@ -700,8 +688,7 @@ def test_reader_materialize_on_read_jpeg_png_bytes_preserved( assert row["binary_content"] == image_bytes, f"{fmt} bytes must be preserved verbatim" assert pd.isna(row["materialize_error"]) or row["materialize_error"] is None - ref = InterleavedBatch.parse_source_ref(row["source_ref"]) - assert ref["frame_index"] is None, f"{fmt} must not have a frame_index in source_ref" + assert pd.isna(row["source_frame_index"]), f"{fmt} must not have a source_frame_index" # Confirm PIL can decode the round-tripped bytes decoded = Image.open(BytesIO(row["binary_content"])) diff --git a/tests/stages/interleaved/test_multimodal_writer.py b/tests/stages/interleaved/test_multimodal_writer.py index 9c549fd987..4d39583de0 100644 --- a/tests/stages/interleaved/test_multimodal_writer.py +++ b/tests/stages/interleaved/test_multimodal_writer.py @@ -43,24 +43,15 @@ def _read_batch(input_task: FileGroupTask) -> InterleavedBatch: return batch -def _source_ref(content_path: str, content_key: str | None) -> str: - return json.dumps( - { - "path": content_path, - "member": content_key, - "byte_offset": None, - "byte_size": None, - } - ) - - def test_writer_marks_materialize_error_on_bad_source_path(tmp_path: Path, input_task: FileGroupTask) -> None: batch = _read_batch(input_task) df = batch.to_pandas().copy() image_mask = df["modality"] == "image" assert image_mask.any() first_image_idx = df[image_mask].index[0] - df.loc[first_image_idx, "source_ref"] = _source_ref("/definitely/missing/path.tar", "abc123.tiff") + df["source_ref"] = df["source_ref"].astype(object) + df.loc[[first_image_idx], "source_ref"] = [InterleavedBatch.build_source_ref("/definitely/missing/path.tar")] + df.loc[first_image_idx, "source_member"] = "abc123.tiff" bad_batch = InterleavedBatch( dataset_name=batch.dataset_name, data=df, @@ -96,7 +87,7 @@ def test_writer_materializes_direct_content_path_without_key(tmp_path: Path) -> "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": _source_ref(str(raw_path), None), + "source_ref": InterleavedBatch.build_source_ref(str(raw_path)), "materialize_error": None, } ], @@ -142,6 +133,18 @@ def test_writer_does_not_persist_dataframe_index(tmp_path: Path) -> None: assert "__index_level_0__" not in written.columns +def test_parquet_writer_uses_configured_fsspec_filesystem() -> None: + task = InterleavedBatch(dataset_name="test", data=pa.Table.from_pylist([make_row()], schema=INTERLEAVED_SCHEMA)) + writer = InterleavedParquetWriterStage( + path="memory://interleaved-writer", materialize_on_write=False, mode="overwrite" + ) + + output = writer.process(task) + + with writer.fs.open(output.data[0], "rb") as fobj: + assert pq.read_table(fobj).num_rows == 1 + + def test_interleaved_ordering_preserved_through_filter_and_write(tmp_path: Path) -> None: """End-to-end: interleaved text+image rows survive filtering and parquet roundtrip.""" @@ -359,7 +362,7 @@ def test_writer_no_materialize_preserves_null_binary(tmp_path: Path) -> None: "content_type": "image/jpeg", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path="/fake/img.jpg", member=None), + "source_ref": InterleavedBatch.build_source_ref("/fake/img.jpg"), "materialize_error": None, } ], @@ -658,7 +661,7 @@ def test_wds_writer_null_binary_skips_member(tmp_path: Path) -> None: "content_type": "image/png", "text_content": None, "binary_content": None, - "source_ref": InterleavedBatch.build_source_ref(path="/fake/img.png", member=None), + "source_ref": InterleavedBatch.build_source_ref("/fake/img.png"), "materialize_error": None, }, ], diff --git a/tests/stages/interleaved/test_schema_utils.py b/tests/stages/interleaved/test_schema_utils.py index b966407069..b08f0075be 100644 --- a/tests/stages/interleaved/test_schema_utils.py +++ b/tests/stages/interleaved/test_schema_utils.py @@ -18,7 +18,12 @@ import pyarrow as pa -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, + align_table, + reconcile_schema, + resolve_schema, +) from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA, RESERVED_COLUMNS _BASE = [pa.field("sample_id", pa.string()), pa.field("position", pa.int32()), pa.field("modality", pa.string())] @@ -85,3 +90,25 @@ def test_align_table_drops_extra_columns_and_casts_passthrough() -> None: result = align_table(table, target) assert result.schema.names == ["sample_id", "score"] assert result.schema.field("score").type == pa.float64() + + +def test_align_interleaved_table_migrates_legacy_source_refs() -> None: + table = pa.table( + {"source_ref": ['{"path":"/a.tar","member":"a.jpg","byte_offset":10,"byte_size":20,"frame_index":2}']} + ) + result = align_interleaved_table(table) + + assert result.column("source_ref")[0].as_py() == { + "uri": "/a.tar", + "offset": 10, + "size": 20, + "content_type": None, + "checksum": None, + "inline": None, + } + assert result.column("source_member").to_pylist() == ["a.jpg"] + assert result.column("source_frame_index").to_pylist() == [2] + pdf = align_interleaved_table(pa.table({"source_ref": ['{"page":2,"bbox":[0.1,0.2,0.8,0.9]}']})) + assert pdf.select(["source_ref", "page_number", "source_bbox"]).to_pylist() == [ + {"source_ref": None, "page_number": 2, "source_bbox": [0.1, 0.2, 0.8, 0.9]} + ]